refactor: improve GPU data parsing

- Use byte-based regex matching instead of string-based matching
- Increase buffer size for GPU data
- Switch to `bufio.Scanner`
This commit is contained in:
henrygd
2025-03-04 00:15:10 -05:00
parent 9e56f4611f
commit 31431fd211

View File

@@ -3,6 +3,7 @@ package agent
import ( import (
"beszel/internal/entities/system" "beszel/internal/entities/system"
"bufio" "bufio"
"bytes"
"encoding/json" "encoding/json"
"fmt" "fmt"
"os/exec" "os/exec"
@@ -75,7 +76,7 @@ func (c *gpuCollector) collect() error {
scanner := bufio.NewScanner(stdout) scanner := bufio.NewScanner(stdout)
if c.buf == nil { if c.buf == nil {
c.buf = make([]byte, 0, 4*1024) c.buf = make([]byte, 0, 10*1024)
} }
scanner.Buffer(c.buf, bufio.MaxScanTokenSize) scanner.Buffer(c.buf, bufio.MaxScanTokenSize)
@@ -110,27 +111,26 @@ func (gm *GPUManager) getJetsonParser() func(output []byte) bool {
if !ok { if !ok {
return true return true
} }
data := string(output)
// Parse RAM usage // Parse RAM usage
ramMatches := ramPattern.FindStringSubmatch(data) ramMatches := ramPattern.FindSubmatch(output)
if ramMatches != nil { if ramMatches != nil {
gpuData.MemoryUsed, _ = strconv.ParseFloat(ramMatches[1], 64) gpuData.MemoryUsed, _ = strconv.ParseFloat(string(ramMatches[1]), 64)
gpuData.MemoryTotal, _ = strconv.ParseFloat(ramMatches[2], 64) gpuData.MemoryTotal, _ = strconv.ParseFloat(string(ramMatches[2]), 64)
} }
// Parse GR3D (GPU) usage // Parse GR3D (GPU) usage
gr3dMatches := gr3dPattern.FindStringSubmatch(data) gr3dMatches := gr3dPattern.FindSubmatch(output)
if gr3dMatches != nil { if gr3dMatches != nil {
gpuData.Usage, _ = strconv.ParseFloat(gr3dMatches[1], 64) gpuData.Usage, _ = strconv.ParseFloat(string(gr3dMatches[1]), 64)
} }
// Parse temperature // Parse temperature
tempMatches := tempPattern.FindStringSubmatch(data) tempMatches := tempPattern.FindSubmatch(output)
if tempMatches != nil { if tempMatches != nil {
gpuData.Temperature, _ = strconv.ParseFloat(tempMatches[1], 64) gpuData.Temperature, _ = strconv.ParseFloat(string(tempMatches[1]), 64)
} }
// Parse power usage // Parse power usage
powerMatches := powerPattern.FindStringSubmatch(data) powerMatches := powerPattern.FindSubmatch(output)
if powerMatches != nil { if powerMatches != nil {
power, _ := strconv.ParseFloat(powerMatches[2], 64) power, _ := strconv.ParseFloat(string(powerMatches[2]), 64)
gpuData.Power = power / 1000 gpuData.Power = power / 1000
} }
gpuData.Count++ gpuData.Count++
@@ -142,8 +142,10 @@ func (gm *GPUManager) getJetsonParser() func(output []byte) bool {
func (gm *GPUManager) parseNvidiaData(output []byte) bool { func (gm *GPUManager) parseNvidiaData(output []byte) bool {
gm.Lock() gm.Lock()
defer gm.Unlock() defer gm.Unlock()
scanner := bufio.NewScanner(bytes.NewReader(output))
var valid bool var valid bool
for line := range strings.Lines(string(output)) { for scanner.Scan() {
line := scanner.Text() // Or use scanner.Bytes() for []byte
fields := strings.Split(strings.TrimSpace(line), ", ") fields := strings.Split(strings.TrimSpace(line), ", ")
if len(fields) < 7 { if len(fields) < 7 {
continue continue