feat(agent): NETWORK env var and support for multiple keys

- merges agent.Run with agent.NewAgent
- separates StartServer method
- bumps go version to 1.24
- add tests
This commit is contained in:
henrygd
2025-02-19 00:32:27 -05:00
parent 1a6a2a64f2
commit d170e7a00d
6 changed files with 771 additions and 100 deletions

View File

@@ -28,29 +28,16 @@ type Agent struct {
}
func NewAgent() *Agent {
newAgent := &Agent{
sensorsContext: context.Background(),
fsStats: make(map[string]*system.FsStats),
agent := &Agent{
fsStats: make(map[string]*system.FsStats),
}
newAgent.memCalc, _ = GetEnv("MEM_CALC")
return newAgent
}
agent.memCalc, _ = GetEnv("MEM_CALC")
// GetEnv retrieves an environment variable with a "BESZEL_AGENT_" prefix, or falls back to the unprefixed key.
func GetEnv(key string) (value string, exists bool) {
if value, exists = os.LookupEnv("BESZEL_AGENT_" + key); exists {
return value, exists
}
// Fallback to the old unprefixed key
return os.LookupEnv(key)
}
func (a *Agent) Run(pubKey []byte, addr string) {
// 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) {
case "debug":
a.debug = true
agent.debug = true
slog.SetLogLoggerLevel(slog.LevelDebug)
case "warn":
slog.SetLogLoggerLevel(slog.LevelWarn)
@@ -64,40 +51,51 @@ func (a *Agent) Run(pubKey []byte, addr string) {
// Set sensors context (allows overriding sys location for sensors)
if sysSensors, exists := GetEnv("SYS_SENSORS"); exists {
slog.Info("SYS_SENSORS", "path", sysSensors)
a.sensorsContext = context.WithValue(a.sensorsContext,
agent.sensorsContext = context.WithValue(agent.sensorsContext,
common.EnvKey, common.EnvMap{common.HostSysEnvKey: sysSensors},
)
} else {
agent.sensorsContext = context.Background()
}
// Set sensors whitelist
if sensors, exists := GetEnv("SENSORS"); exists {
a.sensorsWhitelist = make(map[string]struct{})
agent.sensorsWhitelist = make(map[string]struct{})
for _, sensor := range strings.Split(sensors, ",") {
if sensor != "" {
a.sensorsWhitelist[sensor] = struct{}{}
agent.sensorsWhitelist[sensor] = struct{}{}
}
}
}
// initialize system info / docker manager
a.initializeSystemInfo()
a.initializeDiskInfo()
a.initializeNetIoStats()
a.dockerManager = newDockerManager(a)
agent.initializeSystemInfo()
agent.initializeDiskInfo()
agent.initializeNetIoStats()
agent.dockerManager = newDockerManager(agent)
// initialize GPU manager
if gm, err := NewGPUManager(); err != nil {
slog.Debug("GPU", "err", err)
} else {
a.gpuManager = gm
agent.gpuManager = gm
}
// if debugging, print stats
if a.debug {
slog.Debug("Stats", "data", a.gatherStats())
if agent.debug {
slog.Debug("Stats", "data", agent.gatherStats())
}
a.startServer(pubKey, addr)
return agent
}
// GetEnv retrieves an environment variable with a "BESZEL_AGENT_" prefix, or falls back to the unprefixed key.
func GetEnv(key string) (value string, exists bool) {
if value, exists = os.LookupEnv("BESZEL_AGENT_" + key); exists {
return value, exists
}
// Fallback to the old unprefixed key
return os.LookupEnv(key)
}
func (a *Agent) gatherStats() system.CombinedData {