DDI/agent/core/agent.go

// SPDX-License-Identifier: Apache-2.0

package core

import (
	"context"
	"crypto/ed25519"
	"errors"
	"fmt"
	"log/slog"
	"reflect"
	"slices"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"git.beugsiud.de/beug.k/DDI/internal/bundle"
)

// Timing defaults of the agent that are not configurable.
const (
	healthRetryInterval  = time.Second
	reportRetryInterval  = 30 * time.Second
	reportTimeout        = 30 * time.Second
	statusTimeout        = 10 * time.Second
	serviceVersionMaxAge = 10 * time.Minute
	// maxMessageLen bounds the message of a deployment report.
	maxMessageLen = 4096
)

// Agent drives a role: it applies the revisions announced by a Source and
// reports outcomes, health and leases.
type Agent struct {
	kind      Kind
	cfg       *Config
	role      Role
	store     *Store
	agentID   string
	bundleKey ed25519.PublicKey
	source    Source
	log       *slog.Logger

	// Intervals, shortened by tests.
	healthRetry time.Duration
	reportRetry time.Duration

	// roleMu serializes calls into the role: only one apply at a time, and
	// health checks for heartbeats are skipped while an apply runs.
	roleMu sync.Mutex

	mu        sync.Mutex
	st        *State
	current   *Revision // applied revision, nil if none or not loadable
	health    Health
	healthAt  time.Time
	version   string
	versionAt time.Time

	qmu   sync.Mutex
	queue []applyRequest
	wake  chan struct{}

	flushCh chan struct{}
	leaseCh chan struct{}
}

// applyRequest asks for an apply of a revision.
type applyRequest struct {
	revision     uint64
	trigger      Trigger
	deploymentID string
}

// NewAgent prepares an agent. The identity provides the agent ID and the
// bundle signing key.
func NewAgent(kind Kind, cfg *Config, role Role, store *Store, id *Identity, src Source, log *slog.Logger) (*Agent, error) {
	if role.Kind() != kind {
		return nil, fmt.Errorf("role %q does not match agent role %q", role.Kind(), kind)
	}
	st, err := store.LoadState()
	if err != nil {
		return nil, err
	}
	a := &Agent{
		kind: kind, cfg: cfg, role: role, store: store, agentID: id.AgentID,
		bundleKey: id.BundleKey, source: src, log: log, st: st,
		healthRetry: healthRetryInterval, reportRetry: reportRetryInterval,
		wake: make(chan struct{}, 1), flushCh: make(chan struct{}, 1), leaseCh: make(chan struct{}, 1),
	}
	if st.Applied != 0 {
		rev, err := store.LoadRevision(st.Applied)
		if err != nil {
			log.Warn("applied revision is not available locally", "revision", st.Applied, "error", err)
		} else {
			a.current = &rev
		}
	}
	return a, nil
}

// Run works until ctx is canceled. The managed service keeps running with
// its current configuration afterwards.
func (a *Agent) Run(ctx context.Context) error {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	a.prune()
	var wg sync.WaitGroup
	wg.Go(func() { a.applyLoop(ctx) })
	wg.Go(func() { a.reportLoop(ctx) })
	if ls, ok := a.role.(LeaseSource); ok {
		wg.Go(func() { a.leaseLoop(ctx, ls) })
	}
	err := a.source.Run(ctx, a)
	stopped := ctx.Err() != nil
	cancel()
	wg.Wait()
	if stopped {
		return nil
	}
	return err
}

// state returns a copy of the state (without the report queue and the
// handled deployments).
func (a *Agent) state() State {
	a.mu.Lock()
	defer a.mu.Unlock()
	st := *a.st
	st.PendingReports = nil
	st.HandledDeployments = nil
	return st
}

// Connected implements SourceHandler.
func (a *Agent) Connected() {
	notify(a.flushCh)
	notify(a.leaseCh)
}

// Desired implements SourceHandler.
func (a *Agent) Desired(revision uint64, autoApply bool) {
	if revision == 0 {
		return
	}
	if !autoApply {
		a.log.Debug("desired revision not applied in manual mode", "revision", revision)
		return
	}
	a.enqueue(applyRequest{revision: revision, trigger: TriggerPull})
}

// Push implements SourceHandler.
func (a *Agent) Push(revision uint64, deploymentID string) {
	a.enqueue(applyRequest{revision: revision, trigger: TriggerPush, deploymentID: deploymentID})
}

// Status implements SourceHandler. It runs a health check unless an apply
// is in progress, in which case the last result is returned.
func (a *Agent) Status(ctx context.Context) AgentStatus {
	if a.roleMu.TryLock() {
		a.refreshStatus(ctx)
		a.roleMu.Unlock()
	}
	a.mu.Lock()
	defer a.mu.Unlock()
	return AgentStatus{
		Applied: a.st.Applied, LastGood: a.st.LastGood,
		Health: a.health, CheckedAt: a.healthAt, ServiceVersion: a.version,
	}
}

// refreshStatus runs a health check and, if stale, reads the service
// version. The caller holds roleMu.
func (a *Agent) refreshStatus(ctx context.Context) {
	ctx, cancel := context.WithTimeout(ctx, statusTimeout)
	defer cancel()
	a.mu.Lock()
	cur := a.current
	needVersion := a.version == "" || time.Since(a.versionAt) > serviceVersionMaxAge
	a.mu.Unlock()
	h := a.role.Health(ctx, cur)
	var version string
	if needVersion {
		v, err := a.role.ServiceVersion(ctx)
		if err != nil {
			a.log.Warn("cannot determine the service version", "error", err)
		}
		version = v
	}
	a.mu.Lock()
	a.health, a.healthAt = h, time.Now()
	if needVersion && version != "" {
		a.version, a.versionAt = version, time.Now()
	}
	a.mu.Unlock()
}

// enqueue adds a request. A pull is dropped if an equal or newer revision
// is already queued, a push if the same deployment is already queued; a
// push replaces queued pulls up to its revision.
func (a *Agent) enqueue(r applyRequest) {
	a.qmu.Lock()
	if r.trigger == TriggerPull {
		for _, q := range a.queue {
			if q.revision >= r.revision {
				a.qmu.Unlock()
				return
			}
		}
	} else {
		if r.deploymentID != "" && slices.ContainsFunc(a.queue, func(q applyRequest) bool {
			return q.trigger == TriggerPush && q.deploymentID == r.deploymentID
		}) {
			a.qmu.Unlock()
			return
		}
		a.queue = slices.DeleteFunc(a.queue, func(q applyRequest) bool {
			return q.trigger == TriggerPull && q.revision <= r.revision
		})
	}
	a.queue = append(a.queue, r)
	a.qmu.Unlock()
	notify(a.wake)
}

func (a *Agent) dequeue() (applyRequest, bool) {
	a.qmu.Lock()
	defer a.qmu.Unlock()
	if len(a.queue) == 0 {
		return applyRequest{}, false
	}
	r := a.queue[0]
	a.queue = a.queue[1:]
	return r, true
}

// applyLoop processes requests one at a time.
func (a *Agent) applyLoop(ctx context.Context) {
	for {
		r, ok := a.dequeue()
		if !ok {
			select {
			case <-ctx.Done():
				return
			case <-a.wake:
				continue
			}
		}
		if ctx.Err() != nil {
			return
		}
		a.handle(ctx, r)
	}
}

// handle decides whether a request leads to an apply.
func (a *Agent) handle(ctx context.Context, r applyRequest) {
	log := a.log.With("revision", r.revision, "trigger", r.trigger, "deployment_id", r.deploymentID)
	if r.trigger == TriggerPush && a.skipHandled(log, r) {
		return
	}
	st := a.state()
	switch {
	case r.trigger == TriggerPull && r.revision <= st.Applied:
		log.Debug("revision already applied or older", "applied", st.Applied)
		return
	case r.trigger == TriggerPull && r.revision <= st.LastUnsuccessful:
		log.Debug("revision failed before; waiting for a newer revision or an explicit push",
			"last_unsuccessful", st.LastUnsuccessful)
		return
	case r.revision == 0 || r.revision < st.Applied:
		now := time.Now()
		log.Warn("refusing to apply a revision older than the applied one", "applied", st.Applied)
		a.finish(log, DeploymentReport{
			DeploymentID: r.deploymentID, Revision: r.revision, Trigger: r.trigger, Outcome: OutcomeRejected,
			Message:        fmt.Sprintf("revision %d is older than the applied revision %d", r.revision, st.Applied),
			ActiveRevision: st.Applied, StartedAt: now, FinishedAt: now,
		}, nil)
		return
	}
	a.apply(ctx, r, st, log)
}

// skipHandled reports whether a push belongs to a deployment that was
// handled before, which the server resends after a reconnect or announces
// while polling as long as it has no report. The revision is not applied
// again: if the report is still queued, the report loop delivers it;
// otherwise the remembered report is queued again so that the server can
// close the deployment.
func (a *Agent) skipHandled(log *slog.Logger, r applyRequest) bool {
	a.mu.Lock()
	prev, ok := a.st.handledDeployment(r.deploymentID)
	if !ok {
		a.mu.Unlock()
		return false
	}
	queued := slices.ContainsFunc(a.st.PendingReports, func(p DeploymentReport) bool {
		return p.DeploymentID == r.deploymentID
	})
	var err error
	if !queued {
		a.st.PendingReports = append(a.st.PendingReports, prev)
		a.st.UpdatedAt = time.Now().UTC()
		err = a.store.SaveState(a.st)
	}
	a.mu.Unlock()
	if err != nil {
		log.Error("cannot save state", "error", err)
	}
	log.Info("deployment already handled; not applying it again",
		"outcome", prev.Outcome, "handled_revision", prev.Revision, "report_queued", queued)
	notify(a.flushCh)
	return true
}

// apply runs the pipeline of docs/adr/0021: fetch, verify, store, Validate,
// Activate, Health, rollback. Shutdown aborts it until Activate starts; from
// then on it is completed (bounded by the configured timeouts) so that the
// service is never left between two revisions.
func (a *Agent) apply(ctx context.Context, r applyRequest, st State, log *slog.Logger) {
	started := time.Now()
	report := DeploymentReport{
		DeploymentID: r.deploymentID, Revision: r.revision, Trigger: r.trigger,
		ActiveRevision: st.Applied, StartedAt: started,
	}
	log.Info("applying revision", "applied", st.Applied, "last_good", st.LastGood)

	reject := func(msg string, err error) {
		report.Outcome = OutcomeRejected
		report.Message = fmt.Sprintf("%s: %v", msg, err)
		log.Error("revision rejected", "error", report.Message)
		a.finish(log, report, func(s *State) { s.LastUnsuccessful = max(s.LastUnsuccessful, r.revision) })
	}

	fctx, cancel := context.WithTimeout(ctx, a.cfg.ApplyTimeout)
	b, err := a.source.FetchBundle(fctx, r.revision)
	cancel()
	if err != nil {
		if ctx.Err() != nil {
			return
		}
		log.Warn("cannot fetch bundle", "error", err)
		if r.trigger == TriggerPush {
			// Transport problems do not block later pulls of this revision.
			report.Outcome = OutcomeRejected
			report.Message = "fetch bundle: " + err.Error()
			a.finish(log, report, nil)
		}
		return
	}
	minRev := st.Applied + 1
	if r.trigger == TriggerPush {
		minRev = st.Applied // an explicit push may re-apply the current revision
	}
	m, contents, err := a.verify(b, r.revision, minRev)
	if err != nil {
		reject("bundle verification failed", err)
		return
	}
	rev, err := a.store.StoreRevision(m, b.Manifest, b.Signature, contents)
	if err != nil {
		reject("store revision", err)
		return
	}
	if ctx.Err() != nil {
		log.Info("shutdown before validation; revision not applied")
		return
	}

	a.roleMu.Lock()
	defer a.roleMu.Unlock()
	rctx := context.WithoutCancel(ctx)

	vctx, cancel := context.WithTimeout(rctx, a.cfg.ApplyTimeout)
	err = a.role.Validate(vctx, rev)
	cancel()
	if err != nil {
		reject("validation failed", err)
		return
	}
	if ctx.Err() != nil {
		log.Info("shutdown before activation; revision not applied")
		return
	}

	a.mu.Lock()
	prev := a.current
	a.mu.Unlock()
	actx, cancel := context.WithTimeout(rctx, a.cfg.ApplyTimeout)
	err = a.role.Activate(actx, rev, prev)
	cancel()
	activated := err == nil
	var reason string
	if activated {
		h := a.waitHealthy(rctx, &rev)
		if h.State == HealthOK {
			report.Outcome = OutcomeSucceeded
			report.ActiveRevision = rev.Number
			log.Info("revision applied")
			a.finish(log, report, func(s *State) {
				s.Applied, s.LastGood = rev.Number, rev.Number
				a.current = &rev
				a.health, a.healthAt = h, time.Now()
			})
			return
		}
		reason = "health check failed: " + describeHealth(h)
	} else {
		reason = "activation failed: " + err.Error()
	}
	log.Error("revision unhealthy, rolling back", "error", reason, "last_good", st.LastGood)
	a.rollback(rctx, log, report, rev, activated, reason, st)
}

// rollback restores the last good revision after a failed activation.
func (a *Agent) rollback(ctx context.Context, log *slog.Logger, report DeploymentReport, rev Revision, activated bool, reason string, st State) {
	// What is installed if the rollback cannot restore anything: the new
	// revision if it was activated, otherwise (the role undid its own
	// changes) the previous one.
	installed := st.Applied
	if activated {
		installed = rev.Number
	}
	fail := func(msg string) {
		report.Outcome = OutcomeFailed
		report.Message = reason + "; " + msg
		report.ActiveRevision = installed
		log.Error("deployment failed", "error", report.Message)
		a.finish(log, report, func(s *State) {
			s.LastUnsuccessful = max(s.LastUnsuccessful, rev.Number)
			if installed != s.Applied {
				s.Applied = installed
				a.current = &rev
			}
		})
	}
	if st.LastGood == 0 || st.LastGood == rev.Number {
		fail("no last good revision to roll back to")
		return
	}
	good, err := a.store.LoadRevision(st.LastGood)
	if err != nil {
		fail(fmt.Sprintf("rollback to revision %d impossible: %v", st.LastGood, err))
		return
	}
	actx, cancel := context.WithTimeout(ctx, a.cfg.ApplyTimeout)
	err = a.role.Activate(actx, good, &rev)
	cancel()
	if err != nil {
		fail(fmt.Sprintf("rollback to revision %d failed: %v", good.Number, err))
		return
	}
	installed = good.Number
	h := a.waitHealthy(ctx, &good)
	if h.State != HealthOK {
		report.Outcome = OutcomeFailed
		report.Message = fmt.Sprintf("%s; revision %d restored but unhealthy: %s", reason, good.Number, describeHealth(h))
	} else {
		report.Outcome = OutcomeRolledBack
		report.Message = fmt.Sprintf("%s; rolled back to revision %d", reason, good.Number)
	}
	report.ActiveRevision = good.Number
	log.Warn("deployment rolled back", "outcome", report.Outcome, "active_revision", good.Number)
	a.finish(log, report, func(s *State) {
		s.LastUnsuccessful = max(s.LastUnsuccessful, rev.Number)
		s.Applied = good.Number
		a.current = &good
		a.health, a.healthAt = h, time.Now()
	})
}

// verify checks a fetched bundle against the bundle signing key.
func (a *Agent) verify(b *Bundle, revision, minRev uint64) (*bundle.Manifest, map[string][]byte, error) {
	contents := make(map[string][]byte, len(b.Artifacts))
	for _, art := range b.Artifacts {
		if _, dup := contents[art.Path]; dup {
			return nil, nil, fmt.Errorf("duplicate artifact %s", art.Path)
		}
		contents[art.Path] = art.Content
	}
	m, err := bundle.Verify(a.bundleKey, b.Manifest, b.Signature, contents, bundle.Expect{
		AgentID: a.agentID, Role: string(a.kind), MinRevision: minRev,
	})
	if err != nil {
		return nil, nil, err
	}
	if m.Revision != revision {
		return nil, nil, fmt.Errorf("bundle holds revision %d instead of %d", m.Revision, revision)
	}
	for _, f := range m.Files {
		if reservedArtifact(f.Path) {
			return nil, nil, fmt.Errorf("artifact path %s is reserved", f.Path)
		}
	}
	return m, contents, nil
}

// waitHealthy calls Health every healthRetry until it reports OK or the
// health timeout has passed, and returns the last result.
func (a *Agent) waitHealthy(ctx context.Context, rev *Revision) Health {
	ctx, cancel := context.WithTimeout(ctx, a.cfg.HealthTimeout)
	defer cancel()
	t := time.NewTimer(0)
	defer t.Stop()
	for {
		h := a.role.Health(ctx, rev)
		if h.State == HealthOK {
			return h
		}
		t.Reset(a.healthRetry)
		select {
		case <-ctx.Done():
			return h
		case <-t.C:
		}
	}
}

// finish records the outcome: it updates the state (update runs under the
// state lock), queues the report, persists both and triggers the report
// delivery.
func (a *Agent) finish(log *slog.Logger, r DeploymentReport, update func(*State)) {
	r.Message = truncate(r.Message, maxMessageLen)
	if r.FinishedAt.IsZero() {
		r.FinishedAt = time.Now()
	}
	a.mu.Lock()
	if update != nil {
		update(a.st)
	}
	last := r
	a.st.LastResult = &last
	a.st.PendingReports = append(a.st.PendingReports, r)
	if r.Trigger == TriggerPush {
		a.st.rememberDeployment(r)
	}
	a.st.UpdatedAt = time.Now().UTC()
	err := a.store.SaveState(a.st)
	a.mu.Unlock()
	if err != nil {
		log.Error("cannot save state", "error", err)
	}
	a.prune()
	notify(a.flushCh)
}

// prune removes old revisions, keeping the applied and last good one.
func (a *Agent) prune() {
	st := a.state()
	if err := a.store.Prune(KeepRevisions, st.Applied, st.LastGood); err != nil {
		a.log.Warn("cannot prune old revisions", "error", err)
	}
}

// reportLoop delivers queued deployment reports.
func (a *Agent) reportLoop(ctx context.Context) {
	t := time.NewTicker(a.reportRetry)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-a.flushCh:
		case <-t.C:
		}
		a.flushReports(ctx)
	}
}

// flushReports sends queued reports in order until one fails.
func (a *Agent) flushReports(ctx context.Context) {
	for ctx.Err() == nil {
		a.mu.Lock()
		if len(a.st.PendingReports) == 0 {
			a.mu.Unlock()
			return
		}
		r := a.st.PendingReports[0]
		a.mu.Unlock()

		rctx, cancel := context.WithTimeout(ctx, reportTimeout)
		err := a.source.ReportDeployment(rctx, r)
		cancel()
		switch {
		case errors.Is(err, ErrReportRefused):
			a.log.Warn("deployment report refused; dropping it", "revision", r.Revision,
				"deployment_id", r.DeploymentID, "error", err)
		case err != nil:
			a.log.Debug("cannot deliver deployment report; will retry", "revision", r.Revision,
				"deployment_id", r.DeploymentID, "error", err)
			return
		default:
			a.log.Info("deployment reported", "revision", r.Revision, "deployment_id", r.DeploymentID, "outcome", r.Outcome)
		}
		a.mu.Lock()
		if i := slices.IndexFunc(a.st.PendingReports, func(p DeploymentReport) bool { return reflect.DeepEqual(p, r) }); i >= 0 {
			a.st.PendingReports = slices.Delete(a.st.PendingReports, i, i+1)
		}
		err = a.store.SaveState(a.st)
		a.mu.Unlock()
		if err != nil {
			a.log.Error("cannot save state", "error", err)
		}
	}
}

// describeHealth renders a non-OK health result for reports.
func describeHealth(h Health) string {
	var b strings.Builder
	b.WriteString(h.State.String())
	if h.Message != "" {
		b.WriteString(": " + h.Message)
	}
	for _, c := range h.Checks {
		if c.State == HealthOK {
			continue
		}
		fmt.Fprintf(&b, "; %s: %s", c.Name, c.State)
		if c.Message != "" {
			b.WriteString(" (" + c.Message + ")")
		}
	}
	return b.String()
}

// String returns the lowercase name of the state.
func (s HealthState) String() string {
	switch s {
	case HealthOK:
		return "ok"
	case HealthDegraded:
		return "degraded"
	case HealthFailed:
		return "failed"
	default:
		return "unknown"
	}
}

// truncate shortens s to at most n bytes without splitting a UTF-8
// sequence.
func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	const marker = " [truncated]"
	cut := n - len(marker)
	for cut > 0 && !utf8.RuneStart(s[cut]) {
		cut--
	}
	return s[:cut] + marker
}

// notify performs a non-blocking send on a 1-buffered channel.
func notify(ch chan struct{}) {
	select {
	case ch <- struct{}{}:
	default:
	}
}

// errRoleNotAvailable is returned by the placeholder role factory.
var errRoleNotAvailable = errors.New("role not available in this build")