Vetting/internal/orchestrator/statemachine.go

package orchestrator

import (
	"fmt"

	"vetting/internal/model"
)

// Trigger is an event that drives a state transition.
type Trigger string

const (
	TriggerStartRequested   Trigger = "StartRequested"   // user clicks Start Vetting
	TriggerDispatched       Trigger = "Dispatched"       // dispatcher picked this run (manual-WoL override path; dormant in happy path)
	TriggerRebootCommanded  Trigger = "RebootCommanded"  // dispatcher (or heartbeat race) told the reporter to reboot
	TriggerPXEObserved      Trigger = "PXEObserved"      // iPXE fetched cmdline for MAC
	TriggerAgentClaimed     Trigger = "AgentClaimed"     // agent POSTed /claim with valid token
	TriggerStageFailed      Trigger = "StageFailed"      // a stage reported failure
	TriggerStageCompleted   Trigger = "StageCompleted"   // a stage reported success → advance
	TriggerAllStagesPassed  Trigger = "AllStagesPassed"  // final stage passed
	TriggerOperatorReleased  Trigger = "OperatorReleased"  // user clicked Release on a held run
	TriggerOperatorOverride  Trigger = "OperatorOverride"  // user overrode a held stage; re-enter it
	TriggerOperatorCancelled Trigger = "OperatorCancelled" // user clicked Cancel on an active run
)

// stageStates maps the canonical stage name (from DefaultStageOrder)
// to the matching RunState. Named differently for historical reasons:
// the first stage is "Inventory" (stage row name) but the run state is
// "InventoryCheck". Later stages share a name with their state.
var stageStates = map[string]model.RunState{
	"Inventory":    model.StateInventoryCheck,
	"SpecValidate": model.StateSpecValidate,
	"SMART":        model.StateSMART,
	"CPUStress":    model.StateCPUStress,
	"Storage":      model.StateStorage,
	"Network":      model.StateNetwork,
	"GPU":          model.StateGPU,
	"PSU":          model.StatePSU,
	"Reporting":    model.StateReporting,
}

// stageOrder is the sequence of RunStates the run walks through from
// first stage to Completed. Kept in sync with store.DefaultStageOrder.
var stageOrder = []model.RunState{
	model.StateInventoryCheck,
	model.StateSpecValidate,
	model.StateSMART,
	model.StateCPUStress,
	model.StateStorage,
	model.StateNetwork,
	model.StateGPU,
	model.StatePSU,
	model.StateReporting,
}

type transition struct {
	from []model.RunState
	to   model.RunState
}

var table = map[Trigger]transition{
	TriggerStartRequested:   {from: []model.RunState{model.StateRegistered}, to: model.StateQueued},
	TriggerDispatched:       {from: []model.RunState{model.StateQueued}, to: model.StateWaitingWoL},
	TriggerRebootCommanded:  {from: []model.RunState{model.StateQueued}, to: model.StateWaitingReboot},
	TriggerPXEObserved:      {from: []model.RunState{model.StateWaitingReboot, model.StateWaitingWoL, model.StateBooting}, to: model.StateBooting},
	TriggerAgentClaimed:     {from: []model.RunState{model.StateBooting, model.StateWaitingReboot, model.StateWaitingWoL}, to: model.StateInventoryCheck},
	TriggerStageFailed:       {from: allActiveStates(), to: model.StateFailedHolding},
	TriggerAllStagesPassed:   {from: []model.RunState{model.StateReporting}, to: model.StateCompleted},
	TriggerOperatorReleased:  {from: []model.RunState{model.StateFailedHolding}, to: model.StateReleased},
	TriggerOperatorCancelled: {from: allActiveStates(), to: model.StateCancelled},
}

// Next computes the target state for a trigger against the current state.
// StageCompleted is handled specially: it advances through stageOrder.
func Next(current model.RunState, t Trigger) (model.RunState, error) {
	if t == TriggerStageCompleted {
		return nextStageState(current)
	}
	tr, ok := table[t]
	if !ok {
		return "", fmt.Errorf("unknown trigger %q", t)
	}
	for _, s := range tr.from {
		if s == current {
			return tr.to, nil
		}
	}
	return "", fmt.Errorf("trigger %q not allowed from %q", t, current)
}

// NextForOverride returns the state we should jump to when the operator
// overrides a held stage. It's separate from the generic table because
// the target depends on the failed_stage, not on the current state
// (which is always FailedHolding).
func NextForOverride(current model.RunState, failedStage string) (model.RunState, error) {
	if current != model.StateFailedHolding {
		return "", fmt.Errorf("override not allowed from %q", current)
	}
	s, ok := stageStates[failedStage]
	if !ok {
		return "", fmt.Errorf("override: unknown failed stage %q", failedStage)
	}
	return s, nil
}

// StateForStage returns the RunState that corresponds to a stage name.
// Used by handlers that receive a stage name and want to guard against
// stale/out-of-order agent reports.
func StateForStage(name string) (model.RunState, bool) {
	s, ok := stageStates[name]
	return s, ok
}

func nextStageState(current model.RunState) (model.RunState, error) {
	for i, s := range stageOrder {
		if s == current {
			if i+1 >= len(stageOrder) {
				return model.StateCompleted, nil
			}
			return stageOrder[i+1], nil
		}
	}
	return "", fmt.Errorf("StageCompleted not valid from %q", current)
}

func allActiveStates() []model.RunState {
	return []model.RunState{
		model.StateQueued, model.StateWaitingWoL, model.StateWaitingReboot, model.StateBooting,
		model.StateInventoryCheck, model.StateSpecValidate, model.StateSMART,
		model.StateCPUStress, model.StateStorage, model.StateNetwork,
		model.StateGPU, model.StatePSU, model.StateReporting,
	}
}