♻️ 优化取网卡流量及硬盘容量

Co-authored-by: CHN-STUDENT <21209416+CHN-STUDENT@users.noreply.github.com>

cmd/agent/main.go

@@ -55,17 +55,17 @@ func doSelfUpdate() {
 		updateCh <- struct{}{}
 	}()
 	v := semver.MustParse(version)
-	log.Println("Check update", v)
+	println("Check update", v)
 	latest, err := selfupdate.UpdateSelf(v, "naiba/nezha")
 	if err != nil {
-		log.Println("Binary update failed:", err)
+		println("Binary update failed:", err)
 		return
 	}
 	if latest.Version.Equals(v) {
 		// latest version is the same as current version. It means current binary is up to date.
-		log.Println("Current binary is the latest version", version)
+		println("Current binary is the latest version", version)
 	} else {
-		log.Println("Successfully updated to version", latest.Version)
+		println("Successfully updated to version", latest.Version)
 		os.Exit(1)
 	}
 }
@@ -120,18 +120,18 @@ func run() {
 	var conn *grpc.ClientConn
 
 	retry := func() {
-		log.Println("Error to close connection ...")
+		println("Error to close connection ...")
 		if conn != nil {
 			conn.Close()
 		}
 		time.Sleep(delayWhenError)
-		log.Println("Try to reconnect ...")
+		println("Try to reconnect ...")
 	}
 
 	for {
 		conn, err = grpc.Dial(server, grpc.WithInsecure(), grpc.WithPerRPCCredentials(&auth))
 		if err != nil {
-			log.Printf("grpc.Dial err: %v", err)
+			println("grpc.Dial err: ", err)
 			retry()
 			continue
 		}
@@ -139,26 +139,26 @@ func run() {
 		// 第一步注册
 		_, err = client.ReportSystemInfo(ctx, monitor.GetHost().PB())
 		if err != nil {
-			log.Printf("client.ReportSystemInfo err: %v", err)
+			println("client.ReportSystemInfo err: ", err)
 			retry()
 			continue
 		}
 		// 执行 Task
 		tasks, err := client.RequestTask(ctx, monitor.GetHost().PB())
 		if err != nil {
-			log.Printf("client.RequestTask err: %v", err)
+			println("client.RequestTask err: ", err)
 			retry()
 			continue
 		}
 		err = receiveTasks(tasks)
-		log.Printf("receiveTasks exit to main: %v", err)
+		println("receiveTasks exit to main: ", err)
 		retry()
 	}
 }
 
 func receiveTasks(tasks pb.NezhaService_RequestTaskClient) error {
 	var err error
-	defer log.Printf("receiveTasks exit %v => %v", time.Now(), err)
+	defer println("receiveTasks exit", time.Now(), "=>", err)
 	for {
 		var task *pb.Task
 		task, err = tasks.Recv()
@@ -264,7 +264,7 @@ func doTask(task *pb.Task) {
 		}
 		result.Delay = float32(time.Since(startedAt).Seconds())
 	default:
-		log.Printf("Unknown action: %v", task)
+		println("Unknown action: ", task)
 	}
 	client.ReportTask(ctx, &result)
 }
@@ -272,13 +272,13 @@ func doTask(task *pb.Task) {
 func reportState() {
 	var lastReportHostInfo time.Time
 	var err error
-	defer log.Printf("reportState exit %v => %v", time.Now(), err)
+	defer println("reportState exit", time.Now(), "=>", err)
 	for {
 		if client != nil {
 			monitor.TrackNetworkSpeed()
 			_, err = client.ReportSystemState(ctx, monitor.GetState(dao.ReportDelay).PB())
 			if err != nil {
-				log.Printf("reportState error %v", err)
+				println("reportState error", err)
 				time.Sleep(delayWhenError)
 			}
 			if lastReportHostInfo.Before(time.Now().Add(-10 * time.Minute)) {
@@ -288,3 +288,9 @@ func reportState() {
 		}
 	}
 }
+
+func println(v ...interface{}) {
+	if dao.Conf.Debug {
+		log.Println(v...)
+	}
+}

cmd/agent/monitor/monitor.go

@@ -22,7 +22,7 @@ func GetHost() *model.Host {
 	hi, _ := host.Info()
 	var cpuType string
 	if hi.VirtualizationSystem != "" {
-		cpuType = "Virtual"
+		cpuType = "Vrtual"
 	} else {
 		cpuType = "Physical"
 	}
@@ -37,15 +37,15 @@ func GetHost() *model.Host {
 	}
 	mv, _ := mem.VirtualMemory()
 	ms, _ := mem.SwapMemory()
-	u, _ := disk.Usage("/")
+	diskTotal, _ := getDiskTotalAndUsed()
 
 	return &model.Host{
 		Platform:        hi.OS,
 		PlatformVersion: hi.PlatformVersion,
 		CPU:             cpus,
 		MemTotal:        mv.Total,
-		DiskTotal:       u.Total,
 		SwapTotal:       ms.Total,
+		DiskTotal:       diskTotal,
 		Arch:            hi.KernelArch,
 		Virtualization:  hi.VirtualizationSystem,
 		BootTime:        hi.BootTime,
@@ -57,23 +57,19 @@ func GetHost() *model.Host {
 
 func GetState(delay int64) *model.HostState {
 	hi, _ := host.Info()
-	// Memory
 	mv, _ := mem.VirtualMemory()
 	ms, _ := mem.SwapMemory()
-	// CPU
 	var cpuPercent float64
 	cp, err := cpu.Percent(time.Second*time.Duration(delay), false)
 	if err == nil {
 		cpuPercent = cp[0]
 	}
-	// Disk
+	_, diskUsed := getDiskTotalAndUsed()
-	u, _ := disk.Usage("/")
-
 	return &model.HostState{
 		CPU:            cpuPercent,
 		MemUsed:        mv.Used,
 		SwapUsed:       ms.Used,
-		DiskUsed:       u.Used,
+		DiskUsed:       diskUsed,
 		NetInTransfer:  atomic.LoadUint64(&netInTransfer),
 		NetOutTransfer: atomic.LoadUint64(&netOutTransfer),
 		NetInSpeed:     atomic.LoadUint64(&netInSpeed),
@@ -84,10 +80,22 @@ func GetState(delay int64) *model.HostState {
 
 func TrackNetworkSpeed() {
 	var innerNetInTransfer, innerNetOutTransfer uint64
-	nc, err := net.IOCounters(false)
+	nc, err := net.IOCounters(true)
 	if err == nil {
-		innerNetInTransfer += nc[0].BytesRecv
+		for _, v := range nc {
-		innerNetOutTransfer += nc[0].BytesSent
+			if strings.Contains(v.Name, "lo") ||
+				strings.Contains(v.Name, "tun") ||
+				strings.Contains(v.Name, "docker") ||
+				strings.Contains(v.Name, "veth") ||
+				strings.Contains(v.Name, "br-") ||
+				strings.Contains(v.Name, "vmbr") ||
+				strings.Contains(v.Name, "vnet") ||
+				strings.Contains(v.Name, "kube") {
+				continue
+			}
+			innerNetInTransfer += v.BytesRecv
+			innerNetOutTransfer += v.BytesSent
+		}
 		now := uint64(time.Now().Unix())
 		diff := now - atomic.LoadUint64(&lastUpdate)
 		if diff > 0 {
@@ -99,3 +107,32 @@ func TrackNetworkSpeed() {
 		atomic.StoreUint64(&lastUpdate, now)
 	}
 }
+
+func getDiskTotalAndUsed() (total uint64, used uint64) {
+	diskList, _ := disk.Partitions(false)
+	for _, d := range diskList {
+		fsType := strings.ToLower(d.Fstype)
+		if strings.Contains(fsType, "ext4") ||
+			strings.Contains(fsType, "ext3") ||
+			strings.Contains(fsType, "ext2") ||
+			strings.Contains(fsType, "reiserfs") ||
+			strings.Contains(fsType, "jfs") ||
+			strings.Contains(fsType, "btrfs") ||
+			strings.Contains(fsType, "fuseblk") ||
+			strings.Contains(fsType, "zfs") ||
+			strings.Contains(fsType, "simfs") ||
+			strings.Contains(fsType, "ntfs") ||
+			strings.Contains(fsType, "fat32") ||
+			strings.Contains(fsType, "exfat") ||
+			strings.Contains(fsType, "xfs") {
+			diskUsageOf, _ := disk.Usage(d.Mountpoint)
+			path := diskUsageOf.Path
+			// 不统计 K8s 的虚拟挂载点，see here：https://github.com/shirou/gopsutil/issues/1007
+			if !strings.Contains(path, "/var/lib/kubelet") {
+				total += diskUsageOf.Total
+				used += diskUsageOf.Used
+			}
+		}
+	}
+	return
+}