support blacklisting and wildcard matching in SENSORS env var (#650)

- Moved sensor related code to sensors.go
- Added SensorConfig struct
- Added newSensorConfig
- Added tests
This commit is contained in:
henrygd
2025-04-17 21:08:05 -04:00
parent a2e9056a00
commit 0526c88ce0
4 changed files with 497 additions and 88 deletions

View File

@@ -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()

View File

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

View File

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

View File

@@ -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")