diff --git a/beszel/internal/agent/agent.go b/beszel/internal/agent/agent.go index c7d976a..fe04135 100644 --- a/beszel/internal/agent/agent.go +++ b/beszel/internal/agent/agent.go @@ -4,42 +4,36 @@ package agent import ( "beszel" "beszel/internal/entities/system" - "context" "log/slog" "os" "strings" "sync" "time" - - "github.com/shirou/gopsutil/v4/common" ) type Agent struct { - sync.Mutex // Used to lock agent while collecting data - debug bool // true if LOG_LEVEL is set to debug - zfs bool // true if system has arcstats - memCalc string // Memory calculation formula - fsNames []string // List of filesystem device names being monitored - fsStats map[string]*system.FsStats // Keeps track of disk stats for each filesystem - netInterfaces map[string]struct{} // Stores all valid network interfaces - netIoStats system.NetIoStats // Keeps track of bandwidth usage - dockerManager *dockerManager // Manages Docker API requests - sensorsContext context.Context // Sensors context to override sys location - sensorsWhitelist map[string]struct{} // List of sensors to monitor - primarySensor string // Value of PRIMARY_SENSOR env var - systemInfo system.Info // Host system info - gpuManager *GPUManager // Manages GPU data - cache *SessionCache // Cache for system stats based on primary session ID + sync.Mutex // Used to lock agent while collecting data + debug bool // true if LOG_LEVEL is set to debug + zfs bool // true if system has arcstats + memCalc string // Memory calculation formula + fsNames []string // List of filesystem device names being monitored + fsStats map[string]*system.FsStats // Keeps track of disk stats for each filesystem + netInterfaces map[string]struct{} // Stores all valid network interfaces + netIoStats system.NetIoStats // Keeps track of bandwidth usage + dockerManager *dockerManager // Manages Docker API requests + sensorConfig *SensorConfig // Sensors config + systemInfo system.Info // Host system info + gpuManager *GPUManager // Manages GPU data + cache *SessionCache // Cache for system stats based on primary session ID } func NewAgent() *Agent { agent := &Agent{ - sensorsContext: context.Background(), - fsStats: make(map[string]*system.FsStats), - cache: NewSessionCache(69 * time.Second), + fsStats: make(map[string]*system.FsStats), + cache: NewSessionCache(69 * time.Second), } agent.memCalc, _ = GetEnv("MEM_CALC") - agent.primarySensor, _ = GetEnv("PRIMARY_SENSOR") + agent.sensorConfig = agent.newSensorConfig() // Set up slog with a log level determined by the LOG_LEVEL env var if logLevelStr, exists := GetEnv("LOG_LEVEL"); exists { switch strings.ToLower(logLevelStr) { @@ -55,24 +49,6 @@ func NewAgent() *Agent { slog.Debug(beszel.Version) - // Set sensors context (allows overriding sys location for sensors) - if sysSensors, exists := GetEnv("SYS_SENSORS"); exists { - slog.Info("SYS_SENSORS", "path", sysSensors) - agent.sensorsContext = context.WithValue(agent.sensorsContext, - common.EnvKey, common.EnvMap{common.HostSysEnvKey: sysSensors}, - ) - } - - // Set sensors whitelist - if sensors, exists := GetEnv("SENSORS"); exists { - agent.sensorsWhitelist = make(map[string]struct{}) - for sensor := range strings.SplitSeq(sensors, ",") { - if sensor != "" { - agent.sensorsWhitelist[sensor] = struct{}{} - } - } - } - // initialize system info / docker manager agent.initializeSystemInfo() agent.initializeDiskInfo() diff --git a/beszel/internal/agent/sensors.go b/beszel/internal/agent/sensors.go new file mode 100644 index 0000000..6d3292e --- /dev/null +++ b/beszel/internal/agent/sensors.go @@ -0,0 +1,142 @@ +package agent + +import ( + "beszel/internal/entities/system" + "context" + "log/slog" + "path" + "strconv" + "strings" + + "github.com/shirou/gopsutil/v4/common" + "github.com/shirou/gopsutil/v4/sensors" +) + +type SensorConfig struct { + context context.Context + sensors map[string]struct{} + primarySensor string + isBlacklist bool + hasWildcards bool +} + +func (a *Agent) newSensorConfig() *SensorConfig { + primarySensor, _ := GetEnv("PRIMARY_SENSOR") + sysSensors, _ := GetEnv("SYS_SENSORS") + sensors, _ := GetEnv("SENSORS") + + return a.newSensorConfigWithEnv(primarySensor, sysSensors, sensors) +} + +// newSensorConfigWithEnv creates a SensorConfig with the provided environment variables +func (a *Agent) newSensorConfigWithEnv(primarySensor, sysSensors, sensors string) *SensorConfig { + config := &SensorConfig{ + context: context.Background(), + primarySensor: primarySensor, + } + + // Set sensors context (allows overriding sys location for sensors) + if sysSensors != "" { + slog.Info("SYS_SENSORS", "path", sysSensors) + config.context = context.WithValue(config.context, + common.EnvKey, common.EnvMap{common.HostSysEnvKey: sysSensors}, + ) + } + + // Set sensors whitelist + if sensors != "" { + // handle blacklist + if strings.HasPrefix(sensors, "-") { + config.isBlacklist = true + sensors = sensors[1:] + } + + config.sensors = make(map[string]struct{}) + for sensor := range strings.SplitSeq(sensors, ",") { + sensor = strings.TrimSpace(sensor) + if sensor != "" { + config.sensors[sensor] = struct{}{} + if strings.Contains(sensor, "*") { + config.hasWildcards = true + } + } + } + } + + return config +} + +// updateTemperatures updates the agent with the latest sensor temperatures +func (a *Agent) updateTemperatures(systemStats *system.Stats) { + // skip if sensors whitelist is set to empty string + if a.sensorConfig.sensors != nil && len(a.sensorConfig.sensors) == 0 { + slog.Debug("Skipping temperature collection") + return + } + + // reset high temp + a.systemInfo.DashboardTemp = 0 + + // get sensor data + temps, _ := sensors.TemperaturesWithContext(a.sensorConfig.context) + slog.Debug("Temperature", "sensors", temps) + + // return if no sensors + if len(temps) == 0 { + return + } + + systemStats.Temperatures = make(map[string]float64, len(temps)) + for i, sensor := range temps { + // skip if temperature is unreasonable + if sensor.Temperature <= 0 || sensor.Temperature >= 200 { + continue + } + sensorName := sensor.SensorKey + if _, ok := systemStats.Temperatures[sensorName]; ok { + // if key already exists, append int to key + sensorName = sensorName + "_" + strconv.Itoa(i) + } + // skip if not in whitelist or blacklist + if !isValidSensor(sensorName, a.sensorConfig) { + continue + } + // set dashboard temperature + if a.sensorConfig.primarySensor == "" { + a.systemInfo.DashboardTemp = max(a.systemInfo.DashboardTemp, sensor.Temperature) + } else if a.sensorConfig.primarySensor == sensorName { + a.systemInfo.DashboardTemp = sensor.Temperature + } + systemStats.Temperatures[sensorName] = twoDecimals(sensor.Temperature) + } +} + +// isValidSensor checks if a sensor is valid based on the sensor name and the sensor config +func isValidSensor(sensorName string, config *SensorConfig) bool { + // If no sensors configuration, everything is valid + if config.sensors == nil { + return true + } + + // Exact match - return true if whitelist, false if blacklist + if _, exactMatch := config.sensors[sensorName]; exactMatch { + return !config.isBlacklist + } + + // If no wildcards, return false if blacklist, true if whitelist + if !config.hasWildcards { + return config.isBlacklist + } + + // Check for wildcard patterns + for pattern := range config.sensors { + if !strings.Contains(pattern, "*") { + continue + } + if match, _ := path.Match(pattern, sensorName); match { + return !config.isBlacklist + } + } + + return config.isBlacklist +} diff --git a/beszel/internal/agent/sensors_test.go b/beszel/internal/agent/sensors_test.go new file mode 100644 index 0000000..20f4020 --- /dev/null +++ b/beszel/internal/agent/sensors_test.go @@ -0,0 +1,338 @@ +//go:build testing +// +build testing + +package agent + +import ( + "context" + "os" + "testing" + + "github.com/shirou/gopsutil/v4/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestIsValidSensor(t *testing.T) { + tests := []struct { + name string + sensorName string + config *SensorConfig + expectedValid bool + }{ + { + name: "Whitelist - sensor in list", + sensorName: "cpu_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"cpu_temp": {}}, + isBlacklist: false, + }, + expectedValid: true, + }, + { + name: "Whitelist - sensor not in list", + sensorName: "gpu_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"cpu_temp": {}}, + isBlacklist: false, + }, + expectedValid: false, + }, + { + name: "Blacklist - sensor in list", + sensorName: "cpu_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"cpu_temp": {}}, + isBlacklist: true, + }, + expectedValid: false, + }, + { + name: "Blacklist - sensor not in list", + sensorName: "gpu_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"cpu_temp": {}}, + isBlacklist: true, + }, + expectedValid: true, + }, + { + name: "Whitelist with wildcard - matching pattern", + sensorName: "core_0_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"core_*_temp": {}}, + isBlacklist: false, + hasWildcards: true, + }, + expectedValid: true, + }, + { + name: "Whitelist with wildcard - non-matching pattern", + sensorName: "gpu_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"core_*_temp": {}}, + isBlacklist: false, + hasWildcards: true, + }, + expectedValid: false, + }, + { + name: "Blacklist with wildcard - matching pattern", + sensorName: "core_0_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"core_*_temp": {}}, + isBlacklist: true, + hasWildcards: true, + }, + expectedValid: false, + }, + { + name: "Blacklist with wildcard - non-matching pattern", + sensorName: "gpu_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"core_*_temp": {}}, + isBlacklist: true, + hasWildcards: true, + }, + expectedValid: true, + }, + { + name: "Nil sensor config", + sensorName: "any_temp", + config: &SensorConfig{ + sensors: nil, + }, + expectedValid: true, + }, + { + name: "Mixed patterns in whitelist - exact match", + sensorName: "cpu_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"cpu_temp": {}, "core_*_temp": {}}, + isBlacklist: false, + hasWildcards: true, + }, + expectedValid: true, + }, + { + name: "Mixed patterns in whitelist - wildcard match", + sensorName: "core_1_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"cpu_temp": {}, "core_*_temp": {}}, + isBlacklist: false, + hasWildcards: true, + }, + expectedValid: true, + }, + { + name: "Mixed patterns in blacklist - exact match", + sensorName: "cpu_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"cpu_temp": {}, "core_*_temp": {}}, + isBlacklist: true, + hasWildcards: true, + }, + expectedValid: false, + }, + { + name: "Mixed patterns in blacklist - wildcard match", + sensorName: "core_1_temp", + config: &SensorConfig{ + sensors: map[string]struct{}{"cpu_temp": {}, "core_*_temp": {}}, + isBlacklist: true, + hasWildcards: true, + }, + expectedValid: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := isValidSensor(tt.sensorName, tt.config) + assert.Equal(t, tt.expectedValid, result, "isValidSensor(%q, config) returned unexpected result", tt.sensorName) + }) + } +} + +func TestNewSensorConfigWithEnv(t *testing.T) { + agent := &Agent{} + + tests := []struct { + name string + primarySensor string + sysSensors string + sensors string + expectedConfig *SensorConfig + }{ + { + name: "Empty configuration", + primarySensor: "", + sysSensors: "", + sensors: "", + expectedConfig: &SensorConfig{ + context: context.Background(), + primarySensor: "", + sensors: nil, + isBlacklist: false, + hasWildcards: false, + }, + }, + { + name: "Primary sensor only", + primarySensor: "cpu_temp", + sysSensors: "", + sensors: "", + expectedConfig: &SensorConfig{ + context: context.Background(), + primarySensor: "cpu_temp", + sensors: nil, + isBlacklist: false, + hasWildcards: false, + }, + }, + { + name: "Whitelist sensors", + primarySensor: "cpu_temp", + sysSensors: "", + sensors: "cpu_temp,gpu_temp", + expectedConfig: &SensorConfig{ + context: context.Background(), + primarySensor: "cpu_temp", + sensors: map[string]struct{}{ + "cpu_temp": {}, + "gpu_temp": {}, + }, + isBlacklist: false, + hasWildcards: false, + }, + }, + { + name: "Blacklist sensors", + primarySensor: "cpu_temp", + sysSensors: "", + sensors: "-cpu_temp,gpu_temp", + expectedConfig: &SensorConfig{ + context: context.Background(), + primarySensor: "cpu_temp", + sensors: map[string]struct{}{ + "cpu_temp": {}, + "gpu_temp": {}, + }, + isBlacklist: true, + hasWildcards: false, + }, + }, + { + name: "Sensors with wildcard", + primarySensor: "cpu_temp", + sysSensors: "", + sensors: "cpu_*,gpu_temp", + expectedConfig: &SensorConfig{ + context: context.Background(), + primarySensor: "cpu_temp", + sensors: map[string]struct{}{ + "cpu_*": {}, + "gpu_temp": {}, + }, + isBlacklist: false, + hasWildcards: true, + }, + }, + { + name: "With SYS_SENSORS path", + primarySensor: "cpu_temp", + sysSensors: "/custom/path", + sensors: "cpu_temp", + expectedConfig: &SensorConfig{ + primarySensor: "cpu_temp", + sensors: map[string]struct{}{ + "cpu_temp": {}, + }, + isBlacklist: false, + hasWildcards: false, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := agent.newSensorConfigWithEnv(tt.primarySensor, tt.sysSensors, tt.sensors) + + // Check primary sensor + assert.Equal(t, tt.expectedConfig.primarySensor, result.primarySensor) + + // Check sensor map + if tt.expectedConfig.sensors == nil { + assert.Nil(t, result.sensors) + } else { + assert.Equal(t, len(tt.expectedConfig.sensors), len(result.sensors)) + for sensor := range tt.expectedConfig.sensors { + _, exists := result.sensors[sensor] + assert.True(t, exists, "Sensor %s should exist in the result", sensor) + } + } + + // Check flags + assert.Equal(t, tt.expectedConfig.isBlacklist, result.isBlacklist) + assert.Equal(t, tt.expectedConfig.hasWildcards, result.hasWildcards) + + // Check context + if tt.sysSensors != "" { + // Verify context contains correct values + envMap, ok := result.context.Value(common.EnvKey).(common.EnvMap) + require.True(t, ok, "Context should contain EnvMap") + sysPath, ok := envMap[common.HostSysEnvKey] + require.True(t, ok, "EnvMap should contain HostSysEnvKey") + assert.Equal(t, tt.sysSensors, sysPath) + } + }) + } +} + +func TestNewSensorConfig(t *testing.T) { + // Save original environment variables + originalPrimary, hasPrimary := os.LookupEnv("BESZEL_AGENT_PRIMARY_SENSOR") + originalSys, hasSys := os.LookupEnv("BESZEL_AGENT_SYS_SENSORS") + originalSensors, hasSensors := os.LookupEnv("BESZEL_AGENT_SENSORS") + + // Restore environment variables after the test + defer func() { + // Clean up test environment variables + os.Unsetenv("BESZEL_AGENT_PRIMARY_SENSOR") + os.Unsetenv("BESZEL_AGENT_SYS_SENSORS") + os.Unsetenv("BESZEL_AGENT_SENSORS") + + // Restore original values if they existed + if hasPrimary { + os.Setenv("BESZEL_AGENT_PRIMARY_SENSOR", originalPrimary) + } + if hasSys { + os.Setenv("BESZEL_AGENT_SYS_SENSORS", originalSys) + } + if hasSensors { + os.Setenv("BESZEL_AGENT_SENSORS", originalSensors) + } + }() + + // Set test environment variables + os.Setenv("BESZEL_AGENT_PRIMARY_SENSOR", "test_primary") + os.Setenv("BESZEL_AGENT_SYS_SENSORS", "/test/path") + os.Setenv("BESZEL_AGENT_SENSORS", "test_sensor1,test_*,test_sensor3") + + agent := &Agent{} + result := agent.newSensorConfig() + + // Verify results + assert.Equal(t, "test_primary", result.primarySensor) + assert.NotNil(t, result.sensors) + assert.Equal(t, 3, len(result.sensors)) + assert.True(t, result.hasWildcards) + assert.False(t, result.isBlacklist) + + // Check that sys sensors path is in context + envMap, ok := result.context.Value(common.EnvKey).(common.EnvMap) + require.True(t, ok, "Context should contain EnvMap") + sysPath, ok := envMap[common.HostSysEnvKey] + require.True(t, ok, "EnvMap should contain HostSysEnvKey") + assert.Equal(t, "/test/path", sysPath) +} diff --git a/beszel/internal/agent/system.go b/beszel/internal/agent/system.go index bb6f041..c325dba 100644 --- a/beszel/internal/agent/system.go +++ b/beszel/internal/agent/system.go @@ -16,7 +16,6 @@ import ( "github.com/shirou/gopsutil/v4/host" "github.com/shirou/gopsutil/v4/mem" psutilNet "github.com/shirou/gopsutil/v4/net" - "github.com/shirou/gopsutil/v4/sensors" ) // Sets initial / non-changing values about the host system @@ -203,7 +202,7 @@ func (a *Agent) getSystemStats() system.Stats { for _, gpu := range gpuData { if gpu.Temperature > 0 { systemStats.Temperatures[gpu.Name] = gpu.Temperature - if a.primarySensor == gpu.Name { + if a.sensorConfig.primarySensor == gpu.Name { a.systemInfo.DashboardTemp = gpu.Temperature } } @@ -224,52 +223,6 @@ func (a *Agent) getSystemStats() system.Stats { return systemStats } -func (a *Agent) updateTemperatures(systemStats *system.Stats) { - // skip if sensors whitelist is set to empty string - if a.sensorsWhitelist != nil && len(a.sensorsWhitelist) == 0 { - slog.Debug("Skipping temperature collection") - return - } - - // reset high temp - a.systemInfo.DashboardTemp = 0 - - // get sensor data - temps, _ := sensors.TemperaturesWithContext(a.sensorsContext) - slog.Debug("Temperature", "sensors", temps) - - // return if no sensors - if len(temps) == 0 { - return - } - - systemStats.Temperatures = make(map[string]float64, len(temps)) - for i, sensor := range temps { - // skip if temperature is unreasonable - if sensor.Temperature <= 0 || sensor.Temperature >= 200 { - continue - } - sensorName := sensor.SensorKey - if _, ok := systemStats.Temperatures[sensorName]; ok { - // if key already exists, append int to key - sensorName = sensorName + "_" + strconv.Itoa(i) - } - // skip if not in whitelist - if a.sensorsWhitelist != nil { - if _, nameInWhitelist := a.sensorsWhitelist[sensorName]; !nameInWhitelist { - continue - } - } - // set dashboard temperature - if a.primarySensor == "" { - a.systemInfo.DashboardTemp = max(a.systemInfo.DashboardTemp, sensor.Temperature) - } else if a.primarySensor == sensorName { - a.systemInfo.DashboardTemp = sensor.Temperature - } - systemStats.Temperatures[sensorName] = twoDecimals(sensor.Temperature) - } -} - // Returns the size of the ZFS ARC memory cache in bytes func getARCSize() (uint64, error) { file, err := os.Open("/proc/spl/kstat/zfs/arcstats")