From 09c4ef17645baec7c9afec4aa21977cd1e8e942f Mon Sep 17 00:00:00 2001
From: UUBulb <35923940+uubulb@users.noreply.github.com>
Date: Sun, 7 Jul 2024 10:18:59 +0800
Subject: [PATCH] refactor(pty): switch to using interface (#34)

* refactor(pty): switch to using interface

for better flexibility (although may not be useful)

* chore: reduce error messages

* delete redundant messages
---
 pkg/gpu/gpu.go         |  11 ++--
 pkg/gpu/gpu_darwin.go  |   8 ++-
 pkg/monitor/monitor.go | 120 ++++++++++++++++++++++++++---------------
 pkg/pty/pty_windows.go |  89 ++++++++++++++++--------------
 4 files changed, 134 insertions(+), 94 deletions(-)

diff --git a/pkg/gpu/gpu.go b/pkg/gpu/gpu.go
index 65780a8..c1d9355 100644
--- a/pkg/gpu/gpu.go
+++ b/pkg/gpu/gpu.go
@@ -4,25 +4,24 @@
 package gpu
 
 import (
-	"fmt"
+	"errors"
 
 	"github.com/jaypipes/ghw"
 )
 
-func GetGPUModel() []string {
+func GetGPUModel() ([]string, error) {
 	var gpuModel []string
 	gi, err := ghw.GPU(ghw.WithDisableWarnings())
 	if err != nil {
-		fmt.Printf("Error getting GPU info: %v", err)
-		return nil
+		return nil, err
 	}
 
 	for _, card := range gi.GraphicsCards {
 		if card.DeviceInfo == nil {
-			return nil
+			return nil, errors.New("Cannot find device info")
 		}
 		gpuModel = append(gpuModel, card.DeviceInfo.Product.Name)
 	}
 
-	return gpuModel
+	return gpuModel, nil
 }
diff --git a/pkg/gpu/gpu_darwin.go b/pkg/gpu/gpu_darwin.go
index 7768e77..bebfa2f 100644
--- a/pkg/gpu/gpu_darwin.go
+++ b/pkg/gpu/gpu_darwin.go
@@ -3,7 +3,6 @@
 package gpu
 
 import (
-	"fmt"
 	"os/exec"
 	"regexp"
 	"strings"
@@ -26,7 +25,7 @@ func extractGPUInfo(cmd *exec.Cmd) ([]string, error) {
 	return modelNames, nil
 }
 
-func GetGPUModel() []string {
+func GetGPUModel() ([]string, error) {
 	vendorNames := []string{
 		"AMD", "Intel", "Nvidia", "Apple",
 	}
@@ -37,8 +36,7 @@ func GetGPUModel() []string {
 		ioreg = exec.Command("ioreg", "-rd1", "-c", "IOPCIDevice")
 		gi, err = extractGPUInfo(ioreg)
 		if err != nil {
-			fmt.Println("Error executing ioreg:", err)
-			return nil
+			return nil, err
 		}
 	}
 
@@ -51,5 +49,5 @@ func GetGPUModel() []string {
 			}
 		}
 	}
-	return gpuModel
+	return gpuModel, nil
 }
diff --git a/pkg/monitor/monitor.go b/pkg/monitor/monitor.go
index d1ceb7d..28f244f 100644
--- a/pkg/monitor/monitor.go
+++ b/pkg/monitor/monitor.go
@@ -49,8 +49,17 @@ var (
 // 获取设备数据的最大尝试次数
 const maxDeviceDataFetchAttempts = 3
 
+// 获取主机数据的尝试次数，Key 为 Host 的属性名
+var hostDataFetchAttempts = map[string]int{
+	"CPU": 0,
+	"GPU": 0,
+}
+
 // 获取状态数据的尝试次数，Key 为 HostState 的属性名
-var deviceDataFetchAttempts = map[string]int{
+var statDataFetchAttempts = map[string]int{
+	"CPU":          0,
+	"Load":         0,
+	"GPU":          0,
 	"Temperatures": 0,
 }
 
@@ -67,7 +76,7 @@ func GetHost(agentConfig *model.AgentConfig) *model.Host {
 	var cpuType string
 	hi, err := host.Info()
 	if err != nil {
-		util.Println("host.Info error:", err)
+		util.Println("host.Info error: ", err)
 	} else {
 		if hi.VirtualizationRole == "guest" {
 			cpuType = "Virtual"
@@ -86,31 +95,43 @@ func GetHost(agentConfig *model.AgentConfig) *model.Host {
 	}
 
 	cpuModelCount := make(map[string]int)
-	ci, err := cpu.Info()
-	if err != nil {
-		util.Println("cpu.Info error:", err)
-	} else {
-		for i := 0; i < len(ci); i++ {
-			cpuModelCount[ci[i].ModelName]++
-		}
-		for model, count := range cpuModelCount {
-			if len(ci) > 1 {
-				ret.CPU = append(ret.CPU, fmt.Sprintf("%s %d %s Core", model, count, cpuType))
-			} else {
-				ret.CPU = append(ret.CPU, fmt.Sprintf("%s %d %s Core", model, ci[0].Cores, cpuType))
+	if hostDataFetchAttempts["CPU"] < maxDeviceDataFetchAttempts {
+		ci, err := cpu.Info()
+		if err != nil {
+			hostDataFetchAttempts["CPU"]++
+			util.Println("cpu.Info error: ", err, ", attempt: ", hostDataFetchAttempts["CPU"])
+		} else {
+			hostDataFetchAttempts["CPU"] = 0
+			for i := 0; i < len(ci); i++ {
+				cpuModelCount[ci[i].ModelName]++
+			}
+			for model, count := range cpuModelCount {
+				if len(ci) > 1 {
+					ret.CPU = append(ret.CPU, fmt.Sprintf("%s %d %s Core", model, count, cpuType))
+				} else {
+					ret.CPU = append(ret.CPU, fmt.Sprintf("%s %d %s Core", model, ci[0].Cores, cpuType))
+				}
 			}
 		}
 	}
 
 	if agentConfig.GPU {
-		ret.GPU = gpu.GetGPUModel()
+		if hostDataFetchAttempts["GPU"] < maxDeviceDataFetchAttempts {
+			ret.GPU, err = gpu.GetGPUModel()
+			if err != nil {
+				hostDataFetchAttempts["GPU"]++
+				util.Println("gpu.GetGPUModel error: ", err, ", attempt: ", hostDataFetchAttempts["GPU"])
+			} else {
+				hostDataFetchAttempts["GPU"] = 0
+			}
+		}
 	}
 
 	ret.DiskTotal, _ = getDiskTotalAndUsed(agentConfig)
 
 	mv, err := mem.VirtualMemory()
 	if err != nil {
-		util.Println("mem.VirtualMemory error:", err)
+		util.Println("mem.VirtualMemory error: ", err)
 	} else {
 		ret.MemTotal = mv.Total
 		if runtime.GOOS != "windows" {
@@ -121,7 +142,7 @@ func GetHost(agentConfig *model.AgentConfig) *model.Host {
 	if runtime.GOOS == "windows" {
 		ms, err := mem.SwapMemory()
 		if err != nil {
-			util.Println("mem.SwapMemory error:", err)
+			util.Println("mem.SwapMemory error: ", err)
 		} else {
 			ret.SwapTotal = ms.Total
 		}
@@ -139,16 +160,20 @@ func GetHost(agentConfig *model.AgentConfig) *model.Host {
 func GetState(agentConfig *model.AgentConfig, skipConnectionCount bool, skipProcsCount bool) *model.HostState {
 	var ret model.HostState
 
-	cp, err := cpu.Percent(0, false)
-	if err != nil || len(cp) == 0 {
-		util.Println("cpu.Percent error:", err)
-	} else {
-		ret.CPU = cp[0]
+	if statDataFetchAttempts["CPU"] < maxDeviceDataFetchAttempts {
+		cp, err := cpu.Percent(0, false)
+		if err != nil || len(cp) == 0 {
+			statDataFetchAttempts["CPU"]++
+			util.Println("cpu.Percent error: ", err, ", attempt: ", statDataFetchAttempts["CPU"])
+		} else {
+			statDataFetchAttempts["CPU"] = 0
+			ret.CPU = cp[0]
+		}
 	}
 
 	vm, err := mem.VirtualMemory()
 	if err != nil {
-		util.Println("mem.VirtualMemory error:", err)
+		util.Println("mem.VirtualMemory error: ", err)
 	} else {
 		ret.MemUsed = vm.Total - vm.Available
 		if runtime.GOOS != "windows" {
@@ -159,7 +184,7 @@ func GetState(agentConfig *model.AgentConfig, skipConnectionCount bool, skipProc
 		// gopsutil 在 Windows 下不能正确取 swap
 		ms, err := mem.SwapMemory()
 		if err != nil {
-			util.Println("mem.SwapMemory error:", err)
+			util.Println("mem.SwapMemory error: ", err)
 		} else {
 			ret.SwapUsed = ms.Used
 		}
@@ -167,20 +192,24 @@ func GetState(agentConfig *model.AgentConfig, skipConnectionCount bool, skipProc
 
 	_, ret.DiskUsed = getDiskTotalAndUsed(agentConfig)
 
-	loadStat, err := load.Avg()
-	if err != nil {
-		util.Println("load.Avg error:", err)
-	} else {
-		ret.Load1 = loadStat.Load1
-		ret.Load5 = loadStat.Load5
-		ret.Load15 = loadStat.Load15
+	if statDataFetchAttempts["Load"] < maxDeviceDataFetchAttempts {
+		loadStat, err := load.Avg()
+		if err != nil {
+			statDataFetchAttempts["Load"]++
+			util.Println("load.Avg error: ", err, ", attempt: ", statDataFetchAttempts["Load"])
+		} else {
+			statDataFetchAttempts["Load"] = 0
+			ret.Load1 = loadStat.Load1
+			ret.Load5 = loadStat.Load5
+			ret.Load15 = loadStat.Load15
+		}
 	}
 
 	var procs []int32
 	if !skipProcsCount {
 		procs, err = process.Pids()
 		if err != nil {
-			util.Println("process.Pids error:", err)
+			util.Println("process.Pids error: ", err)
 		} else {
 			ret.ProcessCount = uint64(len(procs))
 		}
@@ -324,16 +353,19 @@ func updateGPUStat(agentConfig *model.AgentConfig, gpuStat *uint64) {
 		return
 	}
 	defer atomic.StoreInt32(&updateGPUStatus, 0)
+
 	if agentConfig.GPU {
-		gs, err := gpustat.GetGPUStat()
-		if err != nil {
-			util.Println("gpustat.GetGPUStat error:", err)
-			atomicStoreFloat64(gpuStat, gs)
-		} else {
-			atomicStoreFloat64(gpuStat, gs)
+		if statDataFetchAttempts["GPU"] < maxDeviceDataFetchAttempts {
+			gs, err := gpustat.GetGPUStat()
+			if err != nil {
+				statDataFetchAttempts["GPU"]++
+				util.Println("gpustat.GetGPUStat error: ", err, ", attempt: ", statDataFetchAttempts["GPU"])
+				atomicStoreFloat64(gpuStat, gs)
+			} else {
+				statDataFetchAttempts["GPU"] = 0
+				atomicStoreFloat64(gpuStat, gs)
+			}
 		}
-	} else {
-		atomicStoreFloat64(gpuStat, 0)
 	}
 }
 
@@ -343,13 +375,13 @@ func updateTemperatureStat() {
 	}
 	defer atomic.StoreInt32(&updateTempStatus, 0)
 
-	if deviceDataFetchAttempts["Temperatures"] <= maxDeviceDataFetchAttempts {
+	if statDataFetchAttempts["Temperatures"] < maxDeviceDataFetchAttempts {
 		temperatures, err := sensors.SensorsTemperatures()
 		if err != nil {
-			deviceDataFetchAttempts["Temperatures"]++
-			util.Println("host.SensorsTemperatures error:", err, "attempt:", deviceDataFetchAttempts["Temperatures"])
+			statDataFetchAttempts["Temperatures"]++
+			util.Println("host.SensorsTemperatures error: ", err, ", attempt: ", statDataFetchAttempts["Temperatures"])
 		} else {
-			deviceDataFetchAttempts["Temperatures"] = 0
+			statDataFetchAttempts["Temperatures"] = 0
 			tempStat := []model.SensorTemperature{}
 			for _, t := range temperatures {
 				if t.Temperature > 0 {
diff --git a/pkg/pty/pty_windows.go b/pkg/pty/pty_windows.go
index 9f9f712..6515228 100644
--- a/pkg/pty/pty_windows.go
+++ b/pkg/pty/pty_windows.go
@@ -22,8 +22,19 @@ import (
 
 var isWin10 bool
 
-type Pty struct {
-	tty interface{}
+type Pty interface {
+	Write(p []byte) (n int, err error)
+	Read(p []byte) (n int, err error)
+	Setsize(cols, rows uint32) error
+	Close() error
+}
+
+type winPTY struct {
+	tty *winpty.WinPTY
+}
+
+type conPty struct {
+	tty *conpty.ConPty
 }
 
 func init() {
@@ -108,9 +119,7 @@ func getExecutableFilePath() (string, error) {
 	return filepath.Dir(ex), nil
 }
 
-func Start() (*Pty, error) {
-	var tty interface{}
-
+func Start() (Pty, error) {
 	shellPath, err := exec.LookPath("powershell.exe")
 	if err != nil || shellPath == "" {
 		shellPath = "cmd.exe"
@@ -120,46 +129,48 @@ func Start() (*Pty, error) {
 		return nil, err
 	}
 	if !isWin10 {
-		tty, err = winpty.OpenDefault(path, shellPath)
+		tty, err := winpty.OpenDefault(path, shellPath)
+		return &winPTY{tty: tty}, err
 	} else {
-		tty, err = conpty.Start(shellPath, conpty.ConPtyWorkDir(path))
-	}
-	return &Pty{tty: tty}, err
-}
-
-func (pty *Pty) Write(p []byte) (n int, err error) {
-	if !isWin10 {
-		return pty.tty.(*winpty.WinPTY).StdIn.Write(p)
-	} else {
-		return pty.tty.(*conpty.ConPty).Write(p)
+		tty, err := conpty.Start(shellPath, conpty.ConPtyWorkDir(path))
+		return &conPty{tty: tty}, err
 	}
 }
 
-func (pty *Pty) Read(p []byte) (n int, err error) {
-	if !isWin10 {
-		return pty.tty.(*winpty.WinPTY).StdOut.Read(p)
-	} else {
-		return pty.tty.(*conpty.ConPty).Read(p)
-	}
+func (w *winPTY) Write(p []byte) (n int, err error) {
+	return w.tty.StdIn.Write(p)
 }
 
-func (pty *Pty) Setsize(cols, rows uint32) error {
-	if !isWin10 {
-		pty.tty.(*winpty.WinPTY).SetSize(cols, rows)
-		return nil
-	} else {
-		return pty.tty.(*conpty.ConPty).Resize(int(cols), int(rows))
-	}
+func (w *winPTY) Read(p []byte) (n int, err error) {
+	return w.tty.StdOut.Read(p)
 }
 
-func (pty *Pty) Close() error {
-	if !isWin10 {
-		pty.tty.(*winpty.WinPTY).Close()
-		return nil
-	} else {
-		if err := pty.tty.(*conpty.ConPty).Close(); err != nil {
-			return err
-		}
-		return nil
-	}
+func (w *winPTY) Setsize(cols, rows uint32) error {
+	w.tty.SetSize(cols, rows)
+	return nil
+}
+
+func (w *winPTY) Close() error {
+	w.tty.Close()
+	return nil
+}
+
+func (c *conPty) Write(p []byte) (n int, err error) {
+	return c.tty.Write(p)
+}
+
+func (c *conPty) Read(p []byte) (n int, err error) {
+	return c.tty.Read(p)
+}
+
+func (c *conPty) Setsize(cols, rows uint32) error {
+	c.tty.Resize(int(cols), int(rows))
+	return nil
+}
+
+func (c *conPty) Close() error {
+	if err := c.tty.Close(); err != nil {
+		return err
+	}
+	return nil
 }