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

	lastDropLog time.Time
}

const (
	dropReasonOverflowEvents = "spool_overflow_events"
	dropReasonOverflowBytes  = "spool_overflow_bytes"
	dropReasonSendRejected   = "send_non_retryable"

	dropLogThrottle = 30 * time.Second
)

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

	// Account drops accumulated during spool Open (e.g. shrunk limits on restart).
	a.recordSpoolDrops(a.Spool.Drops())
	a.Metrics.ChecksConfigured.Set(float64(len(a.Cfg.Checks)))
	a.Metrics.BuildInfo.WithLabelValues(a.Version, boolLabel(a.Cfg.PushesToServer()), boolLabel(a.Cfg.ExposesMetrics())).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.recordSpoolDrops(a.Spool.Drops())
			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,
			Features: client.AgentFeatures{
				Push:    a.Cfg.PushesToServer(),
				Metrics: a.Cfg.ExposesMetrics(),
			},
			Labels: a.Cfg.HeartbeatLabels(a.Version),
		}
		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.Metrics.EventsDropped.WithLabelValues(dropReasonSendRejected).Add(float64(len(seqs)))
			a.Log.Warn("events dropped: server rejected batch", "count", len(seqs), "reason", dropReasonSendRejected)
			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
	}
}

func (a *App) recordSpoolDrops(d spool.DropStats) {
	if d.OverflowEvents == 0 && d.OverflowBytes == 0 {
		return
	}
	if d.OverflowEvents > 0 {
		a.Metrics.EventsDropped.WithLabelValues(dropReasonOverflowEvents).Add(float64(d.OverflowEvents))
	}
	if d.OverflowBytes > 0 {
		a.Metrics.EventsDropped.WithLabelValues(dropReasonOverflowBytes).Add(float64(d.OverflowBytes))
	}
	now := time.Now()
	if now.Sub(a.lastDropLog) >= dropLogThrottle {
		a.Log.Warn("spool drops",
			"events_overflow", d.OverflowEvents,
			"bytes_overflow", d.OverflowBytes,
		)
		a.lastDropLog = now
	}
}

func boolLabel(v bool) string {
	if v {
		return "true"
	}
	return "false"
}