worker/internal/workercluster/fsm.go

package workercluster

import (
	"encoding/json"
	"fmt"
	"io"
	"sync"

	"github.com/hashicorp/raft"
)

// FSM is the in-memory replicated state machine for the worker Raft
// cluster. It holds:
//
//   - ConfigVersion + the adopted critical-check list.
//   - ObserverSet with a version stamp.
//   - One IncidentState per CheckID.
//   - Outbox metadata (capped at outboxCap).
//   - Membership cache mirrored from raft's own configuration so readers
//     can iterate without going through raft.GetConfiguration.
//   - One WorkerDiagnostics per WorkerID.
//   - The last PartitionState and CentralWitnessReport.
//
// Concurrency: Apply is invoked serially by the Raft library, so the
// FSM does not lock inside Apply. External readers (Stats, snapshot
// Persist) acquire fsm.mu as readers. The mutex is never held while a
// disk write is in progress.
//
// Snapshot layout: the entire FSM state is JSON-encoded into a single
// snapshot record. This is fine for the small, append-mostly state the
// plan describes (section 6.1).
type FSM struct {
	mu sync.RWMutex

	ConfigVersion uint64
	Config        []CriticalCheckConfig
	ObserverSet   ObserverSet
	Incidents     map[int64]IncidentState
	Outbox        []OutboxMeta
	Membership    map[string]Member
	Diagnostics   map[string]WorkerDiagnostics
	Partition     PartitionState
	Witness       CentralWitnessReport

	// outboxSeq is a monotonically increasing counter handed out to
	// each appended OutboxMeta entry.
	outboxSeq uint64

	// lastIndex tracks the highest applied log index, used by the
	// metadata fields on outbox/diagnostics entries.
	lastIndex uint64
}

const outboxCap = 1024

// NewFSM returns an empty FSM ready to receive Apply calls or a Restore
// snapshot.
func NewFSM() *FSM {
	return &FSM{
		Incidents:   make(map[int64]IncidentState),
		Membership:  make(map[string]Member),
		Diagnostics: make(map[string]WorkerDiagnostics),
		ObserverSet: ObserverSet{},
		Partition:   PartitionState{State: "steady"},
	}
}

// Apply runs one committed log entry through the FSM. The raft
// library serializes Apply calls, but external readers (Stats,
// Snapshot) can run concurrently, so we still acquire the write lock.
func (f *FSM) Apply(log *raft.Log) interface{} {
	f.mu.Lock()
	defer f.mu.Unlock()

	if log.Type != raft.LogCommand {
		// Configuration changes are managed by the raft library and
		// do not flow through our entry types.
		return nil
	}

	entry, err := DecodeEntry(log.Data)
	if err != nil {
		// An undecodable entry is a bug; raft expects Apply to be
		// deterministic and the FSM to never fail. We log nothing
		// here (the library forwards the response back to the caller
		// via ApplyFuture) and just leave the FSM untouched.
		return fmt.Errorf("fsm: decode entry: %w", err)
	}
	f.lastIndex = log.Index

	switch entry.Type {
	case EntryConfigAdopt:
		var p ConfigAdoptPayload
		if err := json.Unmarshal(entry.Adopted, &p); err != nil {
			return fmt.Errorf("fsm: decode config.adopt: %w", err)
		}
		f.ConfigVersion = p.Version
		f.Config = p.Checks
		return nil

	case EntryObserverSetUpdate:
		var p ObserverSetPayload
		if err := json.Unmarshal(entry.Set, &p); err != nil {
			return fmt.Errorf("fsm: decode observer_set.update: %w", err)
		}
		f.ObserverSet = p.Set
		f.ObserverSet.AdoptedAtIx = log.Index
		return nil

	case EntryMembershipProposeAdd:
		var p MemberPayload
		if err := json.Unmarshal(entry.Member, &p); err != nil {
			return fmt.Errorf("fsm: decode membership.propose_add: %w", err)
		}
		p.Member.JoinedAtIx = log.Index
		f.Membership[p.Member.WorkerID] = p.Member
		return nil

	case EntryMembershipDemote:
		var p MemberPayload
		if err := json.Unmarshal(entry.Member, &p); err != nil {
			return fmt.Errorf("fsm: decode membership.demote: %w", err)
		}
		if cur, ok := f.Membership[p.Member.WorkerID]; ok {
			prev := cur
			p.Prev = &prev
			cur.Role = "observer"
			cur.LastSeenIx = log.Index
			f.Membership[p.Member.WorkerID] = cur
		}
		return nil

	case EntryMembershipRemove:
		var p MemberPayload
		if err := json.Unmarshal(entry.Member, &p); err != nil {
			return fmt.Errorf("fsm: decode membership.remove: %w", err)
		}
		if cur, ok := f.Membership[p.Member.WorkerID]; ok {
			prev := cur
			p.Prev = &prev
			f.Membership[p.Member.WorkerID] = cur
			delete(f.Membership, p.Member.WorkerID)
		}
		return nil

	case EntryIncidentObserve:
		// TODO(phase-1): feed observations into the incident FSM
		// rule. For now we just record the latest observation
		// against the check's IncidentState so callers can read it.
		var p IncidentObservePayload
		if err := json.Unmarshal(entry.Observe, &p); err != nil {
			return fmt.Errorf("fsm: decode incident.observe: %w", err)
		}
		inc := f.Incidents[p.CheckID]
		inc.CheckID = p.CheckID
		if p.Label == "ok" {
			inc.LastOKAtIx = log.Index
		}
		inc.LastConfirmAtIx = log.Index
		f.Incidents[p.CheckID] = inc
		return nil

	case EntryIncidentTransition:
		var p IncidentTransitionPayload
		if err := json.Unmarshal(entry.State, &p); err != nil {
			return fmt.Errorf("fsm: decode incident.transition: %w", err)
		}
		f.Incidents[p.CheckID] = p.State
		return nil

	case EntryOutboxEnqueue:
		var p OutboxPayload
		if err := json.Unmarshal(entry.Outbox, &p); err != nil {
			return fmt.Errorf("fsm: decode outbox.enqueue: %w", err)
		}
		f.outboxSeq++
		p.Entry.Seq = f.outboxSeq
		if p.Entry.State == "" {
			p.Entry.State = "pending"
		}
		f.Outbox = append(f.Outbox, p.Entry)
		if len(f.Outbox) > outboxCap {
			// Drop oldest. The cap is small and metadata-only.
			f.Outbox = f.Outbox[len(f.Outbox)-outboxCap:]
		}
		return nil

	case EntryOutboxDelivered, EntryOutboxAck:
		var p OutboxUpdatePayload
		raw := entry.Outbox
		if len(raw) == 0 {
			raw = entry.State
		}
		if err := json.Unmarshal(raw, &p); err != nil {
			return fmt.Errorf("fsm: decode outbox.update: %w", err)
		}
		for i := range f.Outbox {
			if f.Outbox[i].Seq == p.Seq {
				f.Outbox[i].State = p.State
				f.Outbox[i].Attempts = p.Attempts
				f.Outbox[i].LastError = p.LastError
				break
			}
		}
		return nil

	case EntryPartitionReport:
		var p PartitionReportPayload
		if err := json.Unmarshal(entry.Report, &p); err != nil {
			return fmt.Errorf("fsm: decode partition.report: %w", err)
		}
		p.State.UpdatedIx = log.Index
		f.Partition = p.State
		return nil

	case EntryDiagnosticsUpdate:
		var p DiagnosticsPayload
		if err := json.Unmarshal(entry.Diag, &p); err != nil {
			return fmt.Errorf("fsm: decode diagnostics.update: %w", err)
		}
		p.Diag.LastBeatIx = log.Index
		f.Diagnostics[p.Diag.WorkerID] = p.Diag
		return nil

	default:
		return fmt.Errorf("fsm: unknown entry type %q", entry.Type)
	}
}

// Snapshot returns a snapshot of the current FSM state. The library
// expects Snapshot() to return quickly; the actual encoding happens in
// Persist.
//
// We deep-copy the maps and slices so subsequent Apply calls do not
// mutate the captured state.
func (f *FSM) Snapshot() (raft.FSMSnapshot, error) {
	f.mu.RLock()
	defer f.mu.RUnlock()

	state := fsmState{
		ConfigVersion: f.ConfigVersion,
		Config:        append([]CriticalCheckConfig(nil), f.Config...),
		ObserverSet:   f.ObserverSet,
		Incidents:     cloneIncidents(f.Incidents),
		Outbox:        append([]OutboxMeta(nil), f.Outbox...),
		Membership:    cloneMembers(f.Membership),
		Diagnostics:   cloneDiag(f.Diagnostics),
		Partition:     f.Partition,
		Witness:       f.Witness,
		OutboxSeq:     f.outboxSeq,
		LastIndex:     f.lastIndex,
	}
	return &fsmSnapshot{state: state}, nil
}

// Restore replaces the FSM state with the snapshot read from sink.
// Called on startup if a snapshot is present, before any Apply.
func (f *FSM) Restore(sink io.ReadCloser) error {
	defer sink.Close() //nolint:errcheck // library contract

	var state fsmState
	if err := json.NewDecoder(sink).Decode(&state); err != nil {
		return fmt.Errorf("fsm: restore decode: %w", err)
	}

	f.mu.Lock()
	defer f.mu.Unlock()

	f.ConfigVersion = state.ConfigVersion
	f.Config = state.Config
	f.ObserverSet = state.ObserverSet
	f.Incidents = state.Incidents
	f.Outbox = state.Outbox
	f.Membership = state.Membership
	f.Diagnostics = state.Diagnostics
	f.Partition = state.Partition
	f.Witness = state.Witness
	f.outboxSeq = state.OutboxSeq
	f.lastIndex = state.LastIndex

	if f.Incidents == nil {
		f.Incidents = make(map[int64]IncidentState)
	}
	if f.Membership == nil {
		f.Membership = make(map[string]Member)
	}
	if f.Diagnostics == nil {
		f.Diagnostics = make(map[string]WorkerDiagnostics)
	}
	return nil
}

// fsmState is the on-the-wire snapshot representation. Versioned so we
// can evolve the FSM shape without breaking older snapshots.
type fsmState struct {
	SchemaVersion int                          `json:"schema_version"`
	ConfigVersion uint64                       `json:"config_version"`
	Config        []CriticalCheckConfig        `json:"config"`
	ObserverSet   ObserverSet                  `json:"observer_set"`
	Incidents     map[int64]IncidentState      `json:"incidents"`
	Outbox        []OutboxMeta                 `json:"outbox"`
	Membership    map[string]Member            `json:"membership"`
	Diagnostics   map[string]WorkerDiagnostics `json:"diagnostics"`
	Partition     PartitionState               `json:"partition"`
	Witness       CentralWitnessReport         `json:"witness"`
	OutboxSeq     uint64                       `json:"outbox_seq"`
	LastIndex     uint64                       `json:"last_index"`
}

// fsmSnapshot wraps fsmState so Persist can stream it to a sink while
// the library holds the FSM lock.
type fsmSnapshot struct {
	state fsmState
}

const fsmSnapshotSchemaVersion = 1

func (s *fsmSnapshot) Persist(sink raft.SnapshotSink) error {
	s.state.SchemaVersion = fsmSnapshotSchemaVersion
	if err := json.NewEncoder(sink).Encode(s.state); err != nil {
		_ = sink.Cancel()
		return fmt.Errorf("fsm: persist encode: %w", err)
	}
	return sink.Close()
}

// Release is a no-op; the snapshot holds no external resources.
func (s *fsmSnapshot) Release() {}

// LastIndex returns the highest applied log index the FSM has seen.
// Exposed for callers that need to compute trigger conditions.
func (f *FSM) LastIndex() uint64 {
	f.mu.RLock()
	defer f.mu.RUnlock()
	return f.lastIndex
}

// SnapshotStats is a small read-only view used by tests and metrics.
type SnapshotStats struct {
	ConfigVersion  uint64                       `json:"config_version"`
	ConfigCount    int                          `json:"config_count"`
	Incidents      int                          `json:"incidents"`
	OutboxLen      int                          `json:"outbox_len"`
	Members        int                          `json:"members"`
	Diagnostics    int                          `json:"diagnostics"`
	Partition      string                       `json:"partition"`
	ObserverSet    ObserverSet                  `json:"observer_set"`
	Membership     map[string]Member            `json:"membership"`
	DiagnosticsMap map[string]WorkerDiagnostics `json:"diagnostics_map"`
}

// Stats returns a point-in-time view of the FSM state. Safe for
// concurrent callers; takes the read lock.
func (f *FSM) Stats() SnapshotStats {
	f.mu.RLock()
	defer f.mu.RUnlock()

	return SnapshotStats{
		ConfigVersion:  f.ConfigVersion,
		ConfigCount:    len(f.Config),
		Incidents:      len(f.Incidents),
		OutboxLen:      len(f.Outbox),
		Members:        len(f.Membership),
		Diagnostics:    len(f.Diagnostics),
		Partition:      f.Partition.State,
		ObserverSet:    f.ObserverSet,
		Membership:     cloneMembers(f.Membership),
		DiagnosticsMap: cloneDiag(f.Diagnostics),
	}
}

func cloneMembers(in map[string]Member) map[string]Member {
	out := make(map[string]Member, len(in))
	for k, v := range in {
		out[k] = v
	}
	return out
}

func cloneDiag(in map[string]WorkerDiagnostics) map[string]WorkerDiagnostics {
	out := make(map[string]WorkerDiagnostics, len(in))
	for k, v := range in {
		out[k] = v
	}
	return out
}

func cloneIncidents(in map[int64]IncidentState) map[int64]IncidentState {
	out := make(map[int64]IncidentState, len(in))
	for k, v := range in {
		out[k] = v
	}
	return out
}