ITWorx-Pulse-Public/cmd/agent/agent.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"time"

	"github.com/itworx/pulse/internal/agentprotocol"
	"github.com/itworx/pulse/internal/agentstore"
	"github.com/itworx/pulse/internal/array"
	"github.com/itworx/pulse/internal/container"
	"github.com/itworx/pulse/internal/disk"
	"github.com/itworx/pulse/internal/host"
	"github.com/itworx/pulse/internal/pool"
	"github.com/itworx/pulse/internal/process"
	"github.com/itworx/pulse/internal/runtimeconfig"
	"github.com/itworx/pulse/internal/share"
)

const (
	// maxLoopInterval bounds one scheduling iteration so the heartbeat is refreshed at
	// least every 10 seconds even when collection is configured to run far less often.
	// See docs/operations/WORKER_AGENT_HEALTHCHECK_CONTRACT.md point 3.
	maxLoopInterval = 5 * time.Second
	// operationTimeout bounds every blocking call inside one iteration — collection and
	// the database write. It sits well under the 10 second heartbeat window so a hung
	// database cannot stall the loop into a false "unhealthy" restart, and equally
	// cannot hide a real hang: the call is abandoned and reported.
	operationTimeout = 4 * time.Second
	// heartbeatFileMode keeps the liveness file readable only by the agent's own user;
	// the healthcheck script runs as the same user.
	heartbeatFileMode = 0o600
)

// snapshotSource is the narrow view of the collector the runtime needs. It keeps the
// loop testable without a procfs tree.
type snapshotSource interface {
	Host(context.Context) (host.RawSnapshot, error)
	Processes(context.Context) (process.RawSnapshot, error)
}

type containerSnapshotSource interface {
	Containers(context.Context) (container.RawSnapshot, error)
}
type arraySnapshotSource interface {
	Array(context.Context) (array.RawSnapshot, error)
}
type diskSnapshotSource interface {
	Disks(context.Context) (disk.RawSnapshot, error)
}
type poolSnapshotSource interface {
	Pools(context.Context) (pool.RawSnapshot, error)
}
type shareSnapshotSource interface {
	Shares(context.Context) (share.RawSnapshot, error)
}

// capability binds one telemetry surface to the collection that produces it. Every
// capability the agent announces is read-only; there is no write path in this binary.
type capability struct {
	id      agentstore.Capability
	version string
	collect func(context.Context) (json.RawMessage, time.Time, error)
}

// tickerFactory produces the loop's tick channel. Tests replace it with a channel they
// drive by hand so loop behaviour is asserted without sleeping.
type tickerFactory func(time.Duration) (<-chan time.Time, func())

type agent struct {
	agentID          string
	writer           agentstore.Writer
	logger           *slog.Logger
	now              func() time.Time
	collectInterval  time.Duration
	loopInterval     time.Duration
	operationTimeout time.Duration
	heartbeatPath    string
	capabilities     []capability
	newTicker        tickerFactory

	// nextCollect is the earliest time the next collection pass may run. It is only
	// touched from the loop goroutine.
	nextCollect time.Time
}

func realTicker(interval time.Duration) (<-chan time.Time, func()) {
	ticker := time.NewTicker(interval)
	return ticker.C, ticker.Stop
}

func newAgent(config runtimeconfig.AgentConfig, source snapshotSource, writer agentstore.Writer, logger *slog.Logger) *agent {
	loopInterval := config.CollectInterval
	if loopInterval > maxLoopInterval {
		loopInterval = maxLoopInterval
	}
	return &agent{
		agentID:          config.AgentID,
		writer:           writer,
		logger:           logger,
		now:              time.Now,
		collectInterval:  config.CollectInterval,
		loopInterval:     loopInterval,
		operationTimeout: operationTimeout,
		heartbeatPath:    config.Service.HeartbeatFile,
		capabilities:     capabilities(source),
		newTicker:        realTicker,
	}
}

// capabilities lists what this agent reports. Adding a capability here is the only way
// to widen what the agent reads, which keeps the surface auditable.
func capabilities(source snapshotSource) []capability {
	result := []capability{
		{
			id:      agentstore.CapabilityHost,
			version: host.ContractVersion,
			collect: func(ctx context.Context) (json.RawMessage, time.Time, error) {
				snapshot, err := source.Host(ctx)
				if err != nil {
					return nil, time.Time{}, err
				}
				payload, err := json.Marshal(snapshot)
				if err != nil {
					return nil, time.Time{}, err
				}
				return payload, snapshot.ObservedAt, nil
			},
		},
		{
			id:      agentstore.CapabilityProcesses,
			version: process.ContractVersion,
			collect: func(ctx context.Context) (json.RawMessage, time.Time, error) {
				snapshot, err := source.Processes(ctx)
				if err != nil {
					return nil, time.Time{}, err
				}
				payload, err := json.Marshal(snapshot)
				if err != nil {
					return nil, time.Time{}, err
				}
				return payload, snapshot.ObservedAt, nil
			},
		},
	}
	if containers, ok := source.(containerSnapshotSource); ok {
		result = append(result, capability{
			id: agentstore.CapabilityContainers, version: container.ContractVersion,
			collect: func(ctx context.Context) (json.RawMessage, time.Time, error) {
				snapshot, err := containers.Containers(ctx)
				if err != nil {
					return nil, time.Time{}, err
				}
				payload, err := json.Marshal(snapshot)
				if err != nil {
					return nil, time.Time{}, err
				}
				return payload, snapshot.ObservedAt, nil
			},
		})
	}
	if source, ok := source.(arraySnapshotSource); ok {
		result = append(result, rawCapability(agentstore.CapabilityArray, array.ContractVersion, source.Array))
	}
	if source, ok := source.(diskSnapshotSource); ok {
		result = append(result, rawCapability(agentstore.CapabilityDisks, disk.ContractVersion, source.Disks))
	}
	if source, ok := source.(poolSnapshotSource); ok {
		result = append(result, rawCapability(agentstore.CapabilityPools, pool.ContractVersion, source.Pools))
	}
	if source, ok := source.(shareSnapshotSource); ok {
		result = append(result, rawCapability(agentstore.CapabilityShares, share.ContractVersion, source.Shares))
	}
	return result
}

func rawCapability[T any](id agentstore.Capability, version string, collect func(context.Context) (T, error)) capability {
	return capability{id: id, version: version, collect: func(ctx context.Context) (json.RawMessage, time.Time, error) {
		snapshot, err := collect(ctx)
		if err != nil {
			return nil, time.Time{}, err
		}
		payload, err := json.Marshal(snapshot)
		if err != nil {
			return nil, time.Time{}, err
		}
		observed := observedAt(snapshot)
		return payload, observed, nil
	}}
}

// All raw snapshot contracts carry ObservedAt. Keep this tiny type assertion local to
// the agent rather than introducing a repository-wide generic telemetry abstraction.
func observedAt(snapshot any) time.Time {
	switch value := snapshot.(type) {
	case array.RawSnapshot:
		return value.ObservedAt
	case disk.RawSnapshot:
		return value.ObservedAt
	case pool.RawSnapshot:
		return value.ObservedAt
	case share.RawSnapshot:
		return value.ObservedAt
	default:
		return time.Time{}
	}
}

// hello is the capability announcement. Every entry is read-only, which is not a
// decoration: agentprotocol.Hello.Validate rejects a hello that claims anything else,
// so the agent's own start-up self-check fails loudly if a mutating capability is ever
// added here by mistake.
func (a *agent) hello() agentprotocol.Hello {
	announced := make([]agentprotocol.Capability, 0, len(a.capabilities))
	for _, item := range a.capabilities {
		announced = append(announced, agentprotocol.Capability{
			ID:       string(item.id),
			Version:  item.version,
			ReadOnly: true,
		})
	}
	return agentprotocol.Hello{
		Protocol:     agentprotocol.Version,
		AgentID:      a.agentID,
		ObservedAt:   a.now().UTC(),
		Capabilities: announced,
	}
}

// run drives the collection loop until the context is cancelled.
//
// The order is fixed by the healthcheck contract: validate, log, write one heartbeat
// before the first blocking call, then loop. Each iteration does its bounded work and
// ends by refreshing the heartbeat, so a stuck iteration stops the heartbeat instead of
// a background ticker papering over the stall.
func (a *agent) run(ctx context.Context) error {
	if a.writer == nil {
		return errNoWriter
	}
	hello := a.hello()
	if err := hello.Validate(a.now().UTC()); err != nil {
		return fmt.Errorf("agent hello failed its own read-only validation: %w", err)
	}
	announced := make([]string, 0, len(hello.Capabilities))
	for _, item := range hello.Capabilities {
		announced = append(announced, item.ID)
	}
	a.logger.Info("pulse agent started",
		"agent_id", a.agentID,
		"protocol", hello.Protocol,
		"capabilities", announced,
		"read_only", true,
		"collect_interval", a.collectInterval.String(),
		"loop_interval", a.loopInterval.String(),
		"operation_timeout", a.operationTimeout.String(),
		"heartbeat_file", a.heartbeatPath,
	)

	// Contract point 4: a slow-but-healthy cold start must not look like a hang.
	a.heartbeat()

	ticks, stop := a.newTicker(a.loopInterval)
	defer stop()

loop:
	for ctx.Err() == nil {
		a.iterate(ctx)
		select {
		case <-ctx.Done():
			break loop
		case <-ticks:
		}
	}
	a.logger.Info("pulse agent stopped", "agent_id", a.agentID)
	return nil
}

// iterate is one unit of work: collect when due, then heartbeat. An empty iteration —
// nothing due yet — is still a completed iteration and still heartbeats.
func (a *agent) iterate(ctx context.Context) {
	now := a.now()
	if !now.Before(a.nextCollect) {
		a.collectOnce(ctx)
		a.nextCollect = a.now().Add(a.collectInterval)
	}
	a.heartbeat()
}

// collectOnce publishes every capability independently. One failing capability is
// logged and skipped; it neither aborts the pass nor causes a stale or empty snapshot
// to be written in its place. A missing snapshot is exactly what the reader turns into
// Unknown (ADR-0008), which is the honest outcome.
func (a *agent) collectOnce(ctx context.Context) {
	for _, item := range a.capabilities {
		if ctx.Err() != nil {
			return
		}
		started := a.now()
		bytesWritten, err := a.publish(ctx, item)
		if err != nil {
			a.logger.Error("agent capability collection failed",
				"agent_id", a.agentID,
				"capability", string(item.id),
				"duration_ms", a.now().Sub(started).Milliseconds(),
				"error", err.Error(),
			)
			continue
		}
		a.logger.Info("agent snapshot published",
			"agent_id", a.agentID,
			"capability", string(item.id),
			"payload_bytes", bytesWritten,
			"duration_ms", a.now().Sub(started).Milliseconds(),
		)
	}
}

// publish collects and writes one capability under a bounded context.
func (a *agent) publish(ctx context.Context, item capability) (int, error) {
	operation, cancel := context.WithTimeout(ctx, a.operationTimeout)
	defer cancel()

	payload, observedAt, err := item.collect(operation)
	if err != nil {
		return 0, fmt.Errorf("collect %s: %w", item.id, err)
	}
	if len(payload) == 0 {
		return 0, fmt.Errorf("collect %s: empty payload", item.id)
	}
	if len(payload) > agentstore.MaxPayloadBytes {
		// Writing a truncated snapshot would be worse than writing none: the reader
		// cannot tell a truncated payload from a complete one.
		return 0, fmt.Errorf("collect %s: payload of %d bytes exceeds the %d byte limit", item.id, len(payload), agentstore.MaxPayloadBytes)
	}
	if observedAt.IsZero() {
		observedAt = a.now().UTC()
	}
	if err := a.writer.Put(operation, agentstore.Snapshot{
		AgentID:    a.agentID,
		Capability: item.id,
		ObservedAt: observedAt.UTC(),
		Payload:    payload,
	}); err != nil {
		return 0, fmt.Errorf("publish %s: %w", item.id, err)
	}
	return len(payload), nil
}

// heartbeat refreshes the liveness file the compose healthcheck watches. Only the mtime
// matters to the script; the RFC 3339 body exists so `docker exec … cat /tmp/healthy`
// tells an operator something. A write failure is logged and the loop continues:
// crashing on a tmpfs hiccup would turn a cosmetic problem into an outage, and a
// sustained failure surfaces on its own as staleness.
func (a *agent) heartbeat() {
	if a.heartbeatPath == "" {
		return
	}
	content := a.now().UTC().Format(time.RFC3339) + "\n"
	if err := os.WriteFile(a.heartbeatPath, []byte(content), heartbeatFileMode); err != nil {
		a.logger.Error("agent heartbeat write failed",
			"agent_id", a.agentID,
			"heartbeat_file", a.heartbeatPath,
			"error", err.Error(),
		)
	}
}

var errNoWriter = errors.New("agent snapshot writer is required")