modularize monitor, reduce init usage (#81)

* kill process tree using syscall on windows & cleanup (#80) * kill process tree using syscall on windows & cleanup * use job api * add error check for cmd.Start * modularize monitor, reduce init usage * replace slices with sort * update gopsutil & other dependencies
2024-11-03 21:53:09 +08:00
parent 0cba96bae1
commit af41e4d843
23 changed files with 665 additions and 381 deletions
@@ -0,0 +1,122 @@
+package disk
+
+import (
+	"context"
+	"os/exec"
+	"runtime"
+	"strconv"
+	"strings"
+
+	psDisk "github.com/shirou/gopsutil/v4/disk"
+
+	"github.com/nezhahq/agent/pkg/util"
+)
+
+type DiskKeyType string
+
+const DiskKey DiskKeyType = "disk"
+
+var expectDiskFsTypes = []string{
+	"apfs", "ext4", "ext3", "ext2", "f2fs", "reiserfs", "jfs", "btrfs",
+	"fuseblk", "zfs", "simfs", "ntfs", "fat32", "exfat", "xfs", "fuse.rclone",
+}
+
+func GetHost(ctx context.Context) (uint64, error) {
+	devices, err := getDevices(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	var total uint64
+	for _, mountPath := range devices {
+		diskUsageOf, err := psDisk.Usage(mountPath)
+		if err == nil {
+			total += diskUsageOf.Total
+		}
+	}
+
+	// Fallback 到这个方法,仅统计根路径,适用于OpenVZ之类的.
+	if runtime.GOOS == "linux" && total == 0 {
+		cmd := exec.Command("df")
+		out, err := cmd.CombinedOutput()
+		if err == nil {
+			s := strings.Split(string(out), "\n")
+			for _, c := range s {
+				info := strings.Fields(c)
+				if len(info) == 6 {
+					if info[5] == "/" {
+						total, _ = strconv.ParseUint(info[1], 0, 64)
+						// 默认获取的是1K块为单位的.
+						total = total * 1024
+					}
+				}
+			}
+		}
+	}
+
+	return total, nil
+}
+
+func GetState(ctx context.Context) (uint64, error) {
+	devices, err := getDevices(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	var used uint64
+	for _, mountPath := range devices {
+		diskUsageOf, err := psDisk.Usage(mountPath)
+		if err == nil {
+			used += diskUsageOf.Used
+		}
+	}
+
+	// Fallback 到这个方法,仅统计根路径,适用于OpenVZ之类的.
+	if runtime.GOOS == "linux" && used == 0 {
+		cmd := exec.Command("df")
+		out, err := cmd.CombinedOutput()
+		if err == nil {
+			s := strings.Split(string(out), "\n")
+			for _, c := range s {
+				info := strings.Fields(c)
+				if len(info) == 6 {
+					if info[5] == "/" {
+						used, _ = strconv.ParseUint(info[2], 0, 64)
+						// 默认获取的是1K块为单位的.
+						used = used * 1024
+					}
+				}
+			}
+		}
+	}
+
+	return used, nil
+}
+
+func getDevices(ctx context.Context) (map[string]string, error) {
+	devices := make(map[string]string)
+
+	// 如果配置了白名单，使用白名单的列表
+	if s, ok := ctx.Value(DiskKey).([]string); ok && len(s) > 0 {
+		for i, v := range s {
+			devices[strconv.Itoa(i)] = v
+		}
+		return devices, nil
+	}
+
+	// 否则使用默认过滤规则
+	diskList, err := psDisk.Partitions(false)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, d := range diskList {
+		fsType := strings.ToLower(d.Fstype)
+		// 不统计 K8s 的虚拟挂载点：https://github.com/shirou/gopsutil/issues/1007
+		if devices[d.Device] == "" && util.ContainsStr(expectDiskFsTypes, fsType) && !strings.Contains(d.Mountpoint, "/var/lib/kubelet") {
+			devices[d.Device] = d.Mountpoint
+		}
+	}
+
+	return devices, nil
+}