KJZH001/nezhahq-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(pty): switch to using interface (#34)

Browse Source 
* refactor(pty): switch to using interface

for better flexibility (although may not be useful)

* chore: reduce error messages

* delete redundant messages
This commit is contained in:
UUBulb 2024-07-07 10:18:59 +08:00 committed by GitHub
parent fb7b45527b
commit 09c4ef1764
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 134 additions and 94 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
pkg/gpu/gpu.go
View File

@ -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
}

8
pkg/gpu/gpu_darwin.go
View File

@ -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
}

120
pkg/monitor/monitor.go
View File

@ -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 {

89
pkg/pty/pty_windows.go
View File

@ -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
}