monlet/agent/internal/app/app.go

package app

import (
	"context"
	"errors"
	"log/slog"
	"sync"
	"time"

	"github.com/monlet/agent/internal/client"
	"github.com/monlet/agent/internal/config"
	"github.com/monlet/agent/internal/metrics"
	"github.com/monlet/agent/internal/scheduler"
	"github.com/monlet/agent/internal/spool"
)

type App struct {
	Cfg     *config.Config
	AgentID string
	Version string
	Client  *client.Client
	Spool   *spool.Spool
	Metrics *metrics.Metrics
	Log     *slog.Logger
}

func (a *App) Run(ctx context.Context) error {
	events := make(chan scheduler.Event, 256)

	a.Metrics.ChecksConfigured.Set(float64(len(a.Cfg.Checks)))
	a.Metrics.BuildInfo.WithLabelValues(a.Version, a.Cfg.Mode).Set(1)
	for _, c := range a.Cfg.Checks {
		a.Metrics.CheckInterval.WithLabelValues(c.ID).Set(c.Interval.Duration.Seconds())
	}

	var wg sync.WaitGroup

	if a.Cfg.ExposesMetrics() {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := a.Metrics.Serve(ctx, a.Cfg.Metrics.Listen); err != nil {
				a.Log.Error("metrics server", "err", err)
			}
		}()
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		scheduler.Run(ctx, a.AgentID, a.Cfg.Checks, events, scheduler.Hooks{
			OnSkipped: func(id string) { a.Metrics.CheckSkipped.WithLabelValues(id).Inc() },
			OnResult: func(ev scheduler.Event) {
				a.Metrics.CheckRuns.WithLabelValues(ev.CheckID, ev.Status).Inc()
				a.Metrics.CheckDuration.WithLabelValues(ev.CheckID).Observe(float64(ev.DurationMs) / 1000.0)
				a.Metrics.CheckStatus.WithLabelValues(ev.CheckID).Set(metrics.StatusCode(ev.Status))
				a.Metrics.CheckExitCode.WithLabelValues(ev.CheckID).Set(float64(ev.ExitCode))
				ts := float64(ev.ObservedAt.Unix())
				a.Metrics.CheckLastRun.WithLabelValues(ev.CheckID).Set(ts)
				if ev.Status == "ok" {
					a.Metrics.CheckLastSuccess.WithLabelValues(ev.CheckID).Set(ts)
				}
			},
		})
	}()

	if a.Cfg.PushesToServer() {
		wg.Add(1)
		go func() {
			defer wg.Done()
			a.spoolWriter(ctx, events)
		}()
		wg.Add(1)
		go func() {
			defer wg.Done()
			a.heartbeatLoop(ctx)
		}()
		wg.Add(1)
		go func() {
			defer wg.Done()
			a.sendLoop(ctx)
		}()
	} else {
		// drain scheduler events to avoid blocking it
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				select {
				case <-ctx.Done():
					return
				case <-events:
				}
			}
		}()
	}

	wg.Wait()
	return nil
}

func (a *App) spoolWriter(ctx context.Context, events <-chan scheduler.Event) {
	for {
		select {
		case <-ctx.Done():
			return
		case ev := <-events:
			if err := a.Spool.Enqueue(ev); err != nil {
				a.Log.Error("spool enqueue", "err", err)
			}
			a.Metrics.SpoolDepth.Set(float64(a.Spool.Len()))
			a.Metrics.SpoolBytes.Set(float64(a.Spool.Bytes()))
		}
	}
}

func (a *App) heartbeatLoop(ctx context.Context) {
	interval := a.Cfg.Server.HeartbeatInterval.Duration
	attempt := 0
	for {
		hb := client.HeartbeatRequest{
			AgentID:    a.AgentID,
			ObservedAt: time.Now().UTC(),
			Hostname:   a.Cfg.Hostname,
			Version:    a.Version,
			Mode:       a.Cfg.Mode,
		}
		a.Metrics.SendAttempts.WithLabelValues("heartbeat").Inc()
		err := a.Client.SendHeartbeat(ctx, hb)
		if err != nil {
			a.Metrics.SendFailures.WithLabelValues("heartbeat").Inc()
			if errors.Is(err, context.Canceled) {
				return
			}
			a.Log.Warn("heartbeat send failed", "err", err)
			if client.Retryable(err) {
				attempt++
				if sleep(ctx, client.Delay(attempt)) {
					return
				}
				continue
			}
		} else {
			a.Metrics.LastHeartbeat.Set(float64(time.Now().Unix()))
		}
		attempt = 0
		if sleep(ctx, interval) {
			return
		}
	}
}

func (a *App) sendLoop(ctx context.Context) {
	interval := a.Cfg.Server.BatchInterval.Duration
	attempt := 0
	for {
		items, err := a.Spool.Peek(client.MaxBatchEvents)
		if err != nil {
			a.Log.Error("spool peek", "err", err)
		}
		if len(items) == 0 {
			if sleep(ctx, interval) {
				return
			}
			continue
		}
		evs := make([]scheduler.Event, len(items))
		seqs := make([]uint64, len(items))
		for i, it := range items {
			evs[i] = it.Event
			seqs[i] = it.Seq
		}
		a.Metrics.SendAttempts.WithLabelValues("events").Inc()
		resp, err := a.Client.SendEvents(ctx, evs)
		if err != nil {
			a.Metrics.SendFailures.WithLabelValues("events").Inc()
			if errors.Is(err, context.Canceled) {
				return
			}
			a.Log.Warn("events send failed", "err", err)
			if client.Retryable(err) {
				attempt++
				if sleep(ctx, client.Delay(attempt)) {
					return
				}
				continue
			}
			// non-retryable: drop to avoid hot loop
			a.Spool.Commit(seqs)
			attempt = 0
			continue
		}
		attempt = 0
		a.Spool.Commit(seqs)
		a.Metrics.EventsAccepted.Add(float64(resp.Accepted))
		a.Metrics.EventsDedup.Add(float64(resp.Deduplicated))
		a.Metrics.SpoolDepth.Set(float64(a.Spool.Len()))
		a.Metrics.SpoolBytes.Set(float64(a.Spool.Bytes()))
	}
}

// sleep returns true if context was cancelled while waiting.
func sleep(ctx context.Context, d time.Duration) bool {
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-ctx.Done():
		return true
	case <-t.C:
		return false
	}
}