danvics/Pulse
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add health/metrics server and proper cleanup to unified agent

Browse source 
- Add /healthz (liveness) and /readyz (readiness) endpoints
- Add /metrics endpoint with Prometheus metrics (pulse_agent_info, pulse_agent_up)
- Properly call dockerAgent.Close() on shutdown
- New config: -health-addr flag and PULSE_HEALTH_ADDR env (default :9191)
- Set to empty string to disable health server
This commit is contained in:
rcourtman 2025-12-02 22:42:05 +00:00
parent 7ab68cf682
commit b6221dcc27
1 changed files with 119 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

138
cmd/pulse-agent/main.go
View file

@ -4,12 +4,17 @@ import (
"context"
"flag"
"fmt"
"net/http"
"os"
"os/signal"
"strings"
"sync/atomic"
"syscall"
"time"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/promauto"
"github.com/prometheus/client_golang/prometheus/promhttp"
"github.com/rcourtman/pulse-go-rewrite/internal/agentupdate"
"github.com/rcourtman/pulse-go-rewrite/internal/dockeragent"
"github.com/rcourtman/pulse-go-rewrite/internal/hostagent"
@ -20,6 +25,17 @@ import (
var (
Version = "dev"
// Prometheus metrics
agentInfo = promauto.NewGaugeVec(prometheus.GaugeOpts{
Name: "pulse_agent_info",
Help: "Information about the Pulse agent",
}, []string{"version", "host_enabled", "docker_enabled"})
agentUp = promauto.NewGauge(prometheus.GaugeOpts{
Name: "pulse_agent_up",
Help: "Whether the Pulse agent is running (1 = up, 0 = down)",
})
)
type multiValue []string
@ -48,11 +64,10 @@ func main() {
logger.Fatal().Err(err).Msg("Windows service failed")
}
if ranAsService {
// Service handled everything, exit normally
return
}
// 4. Setup Context & Signal Handling (for non-service mode)
// 4. Setup Context & Signal Handling
ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer cancel()
@ -66,7 +81,21 @@ func main() {
Bool("auto_update", !cfg.DisableAutoUpdate).
Msg("Starting Pulse Unified Agent")
// 5. Start Auto-Updater
// 5. Set prometheus info metric
agentInfo.WithLabelValues(
Version,
fmt.Sprintf("%t", cfg.EnableHost),
fmt.Sprintf("%t", cfg.EnableDocker),
).Set(1)
agentUp.Set(1)
// 6. Start Health/Metrics Server
var ready atomic.Bool
if cfg.HealthAddr != "" {
startHealthServer(ctx, cfg.HealthAddr, &ready, &logger)
}
// 7. Start Auto-Updater
updater := agentupdate.New(agentupdate.Config{
PulseURL: cfg.PulseURL,
APIToken: cfg.APIToken,
@ -83,25 +112,22 @@ func main() {
return nil
})
// 6. Start Host Agent (if enabled)
// 8. Start Host Agent (if enabled)
if cfg.EnableHost {
hostCfg := hostagent.Config{
PulseURL: cfg.PulseURL,
APIToken: cfg.APIToken,
Interval: cfg.Interval,
HostnameOverride: cfg.HostnameOverride,
AgentID: cfg.AgentID, // Shared ID? Or separate? Usually separate for now.
AgentID: cfg.AgentID,
AgentType: "unified",
AgentVersion: Version, // Pass unified agent version
AgentVersion: Version,
Tags: cfg.Tags,
InsecureSkipVerify: cfg.InsecureSkipVerify,
LogLevel: cfg.LogLevel,
Logger: &logger,
}
// If AgentID is set globally, we might want to suffix it or let the agents derive their own.
// For now, let's pass it through. If it's empty, agents derive their own.
agent, err := hostagent.New(hostCfg)
if err != nil {
logger.Fatal().Err(err).Msg("Failed to initialize host agent")
@ -113,7 +139,8 @@ func main() {
})
}
// 7. Start Docker Agent (if enabled)
// 9. Start Docker Agent (if enabled)
var dockerAgent *dockeragent.Agent
if cfg.EnableDocker {
dockerCfg := dockeragent.Config{
PulseURL: cfg.PulseURL,
@ -122,12 +149,11 @@ func main() {
HostnameOverride: cfg.HostnameOverride,
AgentID: cfg.AgentID,
AgentType: "unified",
AgentVersion: Version, // Pass unified agent version
AgentVersion: Version,
InsecureSkipVerify: cfg.InsecureSkipVerify,
DisableAutoUpdate: true, // Unified agent handles updates
DisableAutoUpdate: true,
LogLevel: cfg.LogLevel,
Logger: &logger,
// Docker specific defaults
SwarmScope: "node",
IncludeContainers: true,
IncludeServices: true,
@ -135,28 +161,92 @@ func main() {
CollectDiskMetrics: true,
}
agent, err := dockeragent.New(dockerCfg)
dockerAgent, err = dockeragent.New(dockerCfg)
if err != nil {
logger.Fatal().Err(err).Msg("Failed to initialize docker agent")
}
// Docker agent has a Close method we should call, but errgroup doesn't make defer easy here.
// Ideally we wrap this. For now, we rely on OS cleanup or context cancellation.
g.Go(func() error {
logger.Info().Msg("Docker agent module started")
return agent.Run(ctx)
return dockerAgent.Run(ctx)
})
}
// 8. Wait for all agents to exit
// Mark as ready after all agents started
ready.Store(true)
// 10. Wait for all agents to exit
if err := g.Wait(); err != nil && err != context.Canceled {
logger.Error().Err(err).Msg("Agent terminated with error")
agentUp.Set(0)
cleanupDockerAgent(dockerAgent, &logger)
os.Exit(1)
}
// 11. Cleanup
agentUp.Set(0)
cleanupDockerAgent(dockerAgent, &logger)
logger.Info().Msg("Pulse Unified Agent stopped")
}
func cleanupDockerAgent(agent *dockeragent.Agent, logger *zerolog.Logger) {
if agent == nil {
return
}
if err := agent.Close(); err != nil {
logger.Warn().Err(err).Msg("Failed to close docker agent")
}
}
func startHealthServer(ctx context.Context, addr string, ready *atomic.Bool, logger *zerolog.Logger) {
mux := http.NewServeMux()
// Liveness probe - always returns 200 if server is running
mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
w.Write([]byte("ok"))
})
// Readiness probe - returns 200 only when agents are initialized
mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) {
if ready.Load() {
w.WriteHeader(http.StatusOK)
w.Write([]byte("ok"))
} else {
w.WriteHeader(http.StatusServiceUnavailable)
w.Write([]byte("not ready"))
}
})
// Prometheus metrics
mux.Handle("/metrics", promhttp.Handler())
srv := &http.Server{
Addr: addr,
Handler: mux,
ReadTimeout: 5 * time.Second,
WriteTimeout: 10 * time.Second,
IdleTimeout: 30 * time.Second,
}
go func() {
<-ctx.Done()
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := srv.Shutdown(shutdownCtx); err != nil && err != http.ErrServerClosed {
logger.Warn().Err(err).Msg("Failed to shut down health server")
}
}()
go func() {
logger.Info().Str("addr", addr).Msg("Health/metrics server listening")
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
logger.Warn().Err(err).Msg("Health server stopped unexpectedly")
}
}()
}
type Config struct {
PulseURL string
APIToken string
@ -174,6 +264,9 @@ type Config struct {
// Auto-update
DisableAutoUpdate bool
// Health/metrics server
HealthAddr string
}
func loadConfig() Config {
@ -189,6 +282,7 @@ func loadConfig() Config {
envEnableHost := utils.GetenvTrim("PULSE_ENABLE_HOST")
envEnableDocker := utils.GetenvTrim("PULSE_ENABLE_DOCKER")
envDisableAutoUpdate := utils.GetenvTrim("PULSE_DISABLE_AUTO_UPDATE")
envHealthAddr := utils.GetenvTrim("PULSE_HEALTH_ADDR")
// Defaults
defaultInterval := 30 * time.Second
@ -208,6 +302,11 @@ func loadConfig() Config {
defaultEnableDocker = utils.ParseBool(envEnableDocker)
}
defaultHealthAddr := envHealthAddr
if defaultHealthAddr == "" {
defaultHealthAddr = ":9191"
}
// Flags
urlFlag := flag.String("url", envURL, "Pulse server URL")
tokenFlag := flag.String("token", envToken, "Pulse API token")
@ -220,6 +319,7 @@ func loadConfig() Config {
enableHostFlag := flag.Bool("enable-host", defaultEnableHost, "Enable Host Agent module")
enableDockerFlag := flag.Bool("enable-docker", defaultEnableDocker, "Enable Docker Agent module")
disableAutoUpdateFlag := flag.Bool("disable-auto-update", utils.ParseBool(envDisableAutoUpdate), "Disable automatic updates")
healthAddrFlag := flag.String("health-addr", defaultHealthAddr, "Health/metrics server address (empty to disable)")
showVersion := flag.Bool("version", false, "Print the agent version and exit")
var tagFlags multiValue
@ -263,10 +363,10 @@ func loadConfig() Config {
EnableHost: *enableHostFlag,
EnableDocker: *enableDockerFlag,
DisableAutoUpdate: *disableAutoUpdateFlag,
HealthAddr: strings.TrimSpace(*healthAddrFlag),
}
}
// Helpers (duplicated from existing agents for now, to be moved to shared pkg later)
func gatherTags(env string, flags []string) []string {
tags := make([]string, 0)
if env != "" {