[Snapshot-Agent Memory-Regions 2/N] Adding slots field to state-machine (#165)

* snapshot-agent: slot-aware snapshot/restore transitions in the state machine

StartSnapshotSlot/StartRestoreSlot extend StartSnapshot/StartRestore
with a named snapshot slot, recorded as the job's loaded slot on
success, for backends where several snapshots of one process coexist:

- restore while RUNNING only short-circuits to "already-running" when
  the requested slot is the loaded one; a different slot proceeds (live
  slot swap);
- a FAULTED job may be snapshotted or restored again, since faults are
  typically transient (dead workload PID, timed-out checkpoint client)
  and should not require an agent redeploy.

With slot == "" both entry points behave exactly like the existing
ones, and StartSnapshot/StartRestore now delegate to them, so current
backends (CUDA/app) are unchanged — pinned by slot_test.go alongside
the new slot-swap and fault-recovery coverage (gotest.tools joins
go.mod for it). Nothing calls the slot variants yet.

Signed-off-by: Edwinhr716 <edandres249@gmail.com>

* snapshot-agent: fold slot tests into the state-machine test tables

Review feedback on the slot slice: slot_test.go invented its own
conventions — one function per scenario, sequence-style drivers, and
gotest.tools asserts — next to a state-manager_test.go that already
does table-driven transitions with injected initial state. Since every
slot scenario reduces to a single transition once State and Slot are
seeded through the test exports, the coverage now lives in the
TestStartSnapshotSlot/TestStartRestoreSlot tables beside
TestStartSnapshot/TestStartRestore (same runner shape, plus slot
columns and an InternalJobSlot read-side export). Net new pin: a
failed worker leaves the previously loaded slot recorded.
gotest.tools leaves go.mod again.

Signed-off-by: Edwinhr716 <edandres249@gmail.com>

---------

Signed-off-by: Edwinhr716 <edandres249@gmail.com>

Author: Edwinhr716

pkg/snapshot-agent/state-machine/export_test.go

@@ -25,3 +25,17 @@ func (sm *StateManager) InternalMu() *sync.RWMutex {
 func (sm *StateManager) InternalGetOrCreateJob(jobID, group string) *Job {
 	return sm.getOrCreateJob(jobID, group)
 }
+
+// InternalJobSlot reads the job's loaded slot under the locks the async
+// operation goroutines write it under.
+func (sm *StateManager) InternalJobSlot(jobID string) string {
+	sm.mu.RLock()
+	job := sm.jobs[jobID]
+	sm.mu.RUnlock()
+	if job == nil {
+		return ""
+	}
+	job.mu.Lock()
+	defer job.mu.Unlock()
+	return job.Slot
+}

pkg/snapshot-agent/state-machine/state-manager.go

@@ -1,6 +1,7 @@
 package statemachine
 
 import (
+	"log/slog"
 	"sync"
 	"time"
 
@@ -24,7 +25,11 @@ type Job struct {
 	Group string
 	State pb.JobState
 	PIDs  []int
-	mu    sync.Mutex
+	// Slot is the snapshot slot currently loaded on the device, for backends
+	// with named snapshot slots (memory-regions). Empty for backends without
+	// slot semantics.
+	Slot string
+	mu   sync.Mutex
 }
 
 // Operation represents a long-running snapshot or restore task.
@@ -82,6 +87,17 @@ func (sm *StateManager) RegisterJob(jobID, group string) {
 
 // StartSnapshot initiates a snapshot operation if the job state allows it.
 func (sm *StateManager) StartSnapshot(jobID, group string, worker func() error) (string, error) {
+	return sm.StartSnapshotSlot(jobID, group, "", worker)
+}
+
+// StartSnapshotSlot is StartSnapshot for backends with named snapshot slots
+// (memory-regions): slot names the snapshot being taken and is recorded as
+// the job's loaded slot on success. With slot == "" the behavior is exactly
+// StartSnapshot's. With a non-empty slot, a FAULTED job may also be
+// snapshotted: faults are typically transient (dead workload PID, timed-out
+// cr_client) and a fresh attempt should reset the job rather than requiring
+// an agent redeploy.
+func (sm *StateManager) StartSnapshotSlot(jobID, group, slot string, worker func() error) (string, error) {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 	job := sm.getOrCreateJob(jobID, group)
@@ -94,8 +110,11 @@ func (sm *StateManager) StartSnapshot(jobID, group string, worker func() error)
 		return "", status.Errorf(codes.Aborted, "job %s is already transitioning", jobID)
 	}
 
-	// 2. State Validation: Only allow snapshotting of RUNNING jobs
-	if job.State != pb.JobState_JOB_STATE_RUNNING {
+	// 2. Fault Recovery (slot-aware backends only)
+	if slot != "" && job.State == pb.JobState_JOB_STATE_FAULTED {
+		slog.Warn("Job is FAULTED; allowing new snapshot to reset it", "jobID", jobID, "slot", slot)
+	} else if job.State != pb.JobState_JOB_STATE_RUNNING {
+		// 3. State Validation: Only allow snapshotting of RUNNING jobs
 		return "", status.Errorf(codes.FailedPrecondition, "cannot snapshot job %s in state %s (must be RUNNING)", jobID, job.State)
 	}
 
@@ -132,6 +151,7 @@ func (sm *StateManager) StartSnapshot(jobID, group string, worker func() error)
 			op.Status = pb.OperationStatus_OPERATION_STATUS_COMPLETE
 			op.StorageBytes = 1024
 			job.State = pb.JobState_JOB_STATE_SAVED
+			job.Slot = slot
 		}
 	}()
 
@@ -140,15 +160,26 @@ func (sm *StateManager) StartSnapshot(jobID, group string, worker func() error)
 
 // StartRestore initiates a restore operation if the job state allows it.
 func (sm *StateManager) StartRestore(jobID, group string, worker func() error) (string, error) {
+	return sm.StartRestoreSlot(jobID, group, "", worker)
+}
+
+// StartRestoreSlot is StartRestore for backends with named snapshot slots
+// (memory-regions). With slot == "" the behavior is exactly StartRestore's.
+// With a non-empty slot:
+//   - a RUNNING job only short-circuits to "already-running" when the
+//     requested slot is already the loaded one; restoring a different slot
+//     proceeds (live slot swap, the core memory-regions use case);
+//   - a FAULTED job may be restored (fault recovery; see StartSnapshotSlot).
+func (sm *StateManager) StartRestoreSlot(jobID, group, slot string, worker func() error) (string, error) {
 	sm.mu.Lock()
 	defer sm.mu.Unlock()
 	job := sm.getOrCreateJob(jobID, group)
 
 	job.mu.Lock()
 	defer job.mu.Unlock()
 
-	// 1. Redundancy Optimization
-	if job.State == pb.JobState_JOB_STATE_RUNNING {
+	// 1. Redundancy Optimization: the requested state is already live.
+	if job.State == pb.JobState_JOB_STATE_RUNNING && job.Slot == slot {
 		return "already-running", nil
 	}
 
@@ -157,8 +188,12 @@ func (sm *StateManager) StartRestore(jobID, group string, worker func() error) (
 		return "", status.Errorf(codes.Aborted, "job %s is already transitioning", jobID)
 	}
 
-	// 3. State Validation: Only allow restoring of SAVED jobs
-	if job.State != pb.JobState_JOB_STATE_SAVED {
+	// 3. Fault Recovery (slot-aware backends only)
+	if slot != "" && job.State == pb.JobState_JOB_STATE_FAULTED {
+		slog.Warn("Job is FAULTED; allowing new restore to reset it", "jobID", jobID, "slot", slot)
+	} else if job.State != pb.JobState_JOB_STATE_SAVED && (slot == "" || job.State != pb.JobState_JOB_STATE_RUNNING) {
+		// 4. State Validation: restores need a SAVED job — or, for
+		// slot-aware backends, a RUNNING job swapping to a different slot.
 		return "", status.Errorf(codes.FailedPrecondition, "cannot restore job %s in state %s (must be SAVED)", jobID, job.State)
 	}
 
@@ -194,6 +229,7 @@ func (sm *StateManager) StartRestore(jobID, group string, worker func() error) (
 		} else {
 			op.Status = pb.OperationStatus_OPERATION_STATUS_COMPLETE
 			job.State = pb.JobState_JOB_STATE_RUNNING
+			job.Slot = slot
 			op.SnapshotDeviceBytes = 1024
 		}
 	}()