fix: route every running-state write through the local file (#912)

spinifex/daemon/daemon.go

@@ -2261,7 +2261,7 @@ func (d *Daemon) initJetStream() error {
 		retryDelay = min(retryDelay*2, 10*time.Second)
 	}
 
-	d.stateStore = newStateStoreAdapter(d.jsManager)
+	d.stateStore = newStateStoreAdapter(d.jsManager, d.persistState)
 
 	return nil
 }
@@ -2584,22 +2584,28 @@ func (d *Daemon) localStatePath() string {
 // WriteState persists instance state. Local file is the source of truth; KV is
 // best-effort. Both forms are marshalled inside vmMgr.View to avoid data races.
 func (d *Daemon) WriteState() error {
+	var err error
+	d.vmMgr.View(func(vms map[string]*vm.VM) {
+		err = d.persistState(d.node, vms)
+	})
+	return err
+}
+
+// persistState is the single path both stores move through: every caller that
+// changes the running set reaches it, so the file cannot fall behind the map.
+// Callers hold the manager lock, so vms cannot change mid-encode; the two locks
+// are always taken manager-first to keep one order across every writer.
+func (d *Daemon) persistState(nodeID string, vms map[string]*vm.VM) error {
 	d.stateWriteMu.Lock()
 	defer d.stateWriteMu.Unlock()
 
-	var (
-		localData, kvData []byte
-		marshalErr        error
-	)
-	d.vmMgr.View(func(vms map[string]*vm.VM) {
-		localData, marshalErr = MarshalLocalState(vms)
-		if marshalErr != nil {
-			return
-		}
-		kvData, marshalErr = marshalInstanceState(vms)
-	})
-	if marshalErr != nil {
-		return fmt.Errorf("marshal state: %w", marshalErr)
+	localData, err := MarshalLocalState(vms)
+	if err != nil {
+		return fmt.Errorf("marshal state: %w", err)
+	}
+	kvData, err := marshalInstanceState(vms)
+	if err != nil {
+		return fmt.Errorf("marshal state: %w", err)
 	}
 
 	// The KV write is independent of local disk health (JetStream, not this
@@ -2615,7 +2621,7 @@ func (d *Daemon) WriteState() error {
 	}
 
 	if d.jsManager != nil {
-		d.jsManager.WriteStateBytesBestEffort(d.node, kvData, kvSyncTimeout)
+		d.jsManager.WriteStateBytesBestEffort(nodeID, kvData, kvSyncTimeout)
 	}
 
 	if localErr != nil {

spinifex/daemon/daemon_handlers_test.go

@@ -118,7 +118,7 @@ func createFullTestDaemonWithJetStream(t *testing.T, natsURL string) *Daemon {
 	require.NoError(t, err)
 	err = daemon.jsManager.InitTerminatedInstanceBucket()
 	require.NoError(t, err)
-	daemon.stateStore = newStateStoreAdapter(daemon.jsManager)
+	daemon.stateStore = newStateStoreAdapter(daemon.jsManager, daemon.persistState)
 
 	// Re-bind the instance service so describe-stopped/terminated handlers see
 	// the KV that was just initialised.

spinifex/daemon/daemon_test.go

@@ -1695,7 +1695,7 @@ func TestTerminatedTeardownReaper_SelfHealsFailedVolumeTeardown(t *testing.T) {
 	jsManager, err := NewJetStreamManager(daemon.natsConn, 1)
 	require.NoError(t, err)
 	require.NoError(t, jsManager.InitTerminatedInstanceBucket())
-	daemon.stateStore = newStateStoreAdapter(jsManager)
+	daemon.stateStore = newStateStoreAdapter(jsManager, daemon.persistState)
 
 	// DeleteVolume's snapshot-reference guard requires a non-nil snapshotKV.
 	js, err := jetstream.New(daemon.natsConn)

spinifex/daemon/state_test.go

@@ -10,6 +10,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/testutil"
@@ -45,7 +46,7 @@ func createDaemonWithJetStream(t *testing.T) *Daemon {
 	require.NoError(t, daemon.jsManager.InitKVBucket())
 	require.NoError(t, daemon.jsManager.InitClusterStateBucket())
 	require.NoError(t, daemon.jsManager.InitTerminatedInstanceBucket())
-	daemon.stateStore = newStateStoreAdapter(daemon.jsManager)
+	daemon.stateStore = newStateStoreAdapter(daemon.jsManager, daemon.persistState)
 
 	// Wire just enough vm.Deps for manager-driven state operations to work
 	// (migrate, MarkFailed, Restore classification). Full wiring (network
@@ -592,6 +593,97 @@ func simulateCleanRestore(t *testing.T, daemon *Daemon) {
 	daemon.vmMgr.Restore()
 }
 
+// TestUpdateAndPersist_MutationsReachTheLocalFile covers the API mutations of a
+// running instance at their shared choke point. Each case mirrors the mutation a
+// real handler applies; all of them used to persist to KV alone, leaving the
+// file to whatever a later, unrelated writer happened to flush.
+func TestUpdateAndPersist_MutationsReachTheLocalFile(t *testing.T) {
+	tests := []struct {
+		name   string
+		mutate func(*vm.VM)
+		verify func(*testing.T, *vm.VM)
+	}{
+		{
+			name:   "associate IAM instance profile",
+			mutate: func(v *vm.VM) { v.IamInstanceProfileArn = "arn:aws:iam::000000000000:instance-profile/app" },
+			verify: func(t *testing.T, v *vm.VM) {
+				assert.Equal(t, "arn:aws:iam::000000000000:instance-profile/app", v.IamInstanceProfileArn)
+			},
+		},
+		{
+			name: "modify instance attribute",
+			mutate: func(v *vm.VM) {
+				v.RunInstancesInput = &ec2.RunInstancesInput{DisableApiTermination: aws.Bool(true)}
+			},
+			verify: func(t *testing.T, v *vm.VM) {
+				require.NotNil(t, v.RunInstancesInput)
+				assert.Equal(t, aws.Bool(true), v.RunInstancesInput.DisableApiTermination)
+			},
+		},
+		{
+			name: "set instance monitoring",
+			mutate: func(v *vm.VM) {
+				v.RunInstancesInput = &ec2.RunInstancesInput{
+					Monitoring: &ec2.RunInstancesMonitoringEnabled{Enabled: aws.Bool(true)},
+				}
+			},
+			verify: func(t *testing.T, v *vm.VM) {
+				require.NotNil(t, v.RunInstancesInput)
+				require.NotNil(t, v.RunInstancesInput.Monitoring)
+				assert.Equal(t, aws.Bool(true), v.RunInstancesInput.Monitoring.Enabled)
+			},
+		},
+		{
+			name: "set instance tags",
+			mutate: func(v *vm.VM) {
+				v.Instance = &ec2.Instance{
+					Tags: []*ec2.Tag{{Key: aws.String("Name"), Value: aws.String("web")}},
+				}
+			},
+			verify: func(t *testing.T, v *vm.VM) {
+				require.NotNil(t, v.Instance)
+				require.Len(t, v.Instance.Tags, 1)
+				assert.Equal(t, "web", aws.StringValue(v.Instance.Tags[0].Value))
+			},
+		},
+		{
+			name: "set spot lineage",
+			mutate: func(v *vm.VM) {
+				v.InstanceLifecycle = ec2.InstanceLifecycleTypeSpot
+				v.SpotInstanceRequestId = "sir-abc"
+			},
+			verify: func(t *testing.T, v *vm.VM) {
+				assert.Equal(t, ec2.InstanceLifecycleTypeSpot, v.InstanceLifecycle)
+				assert.Equal(t, "sir-abc", v.SpotInstanceRequestId)
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			daemon := createDaemonWithJetStream(t)
+			daemon.vmMgr.Insert(&vm.VM{ID: "i-mutate", Status: vm.StateRunning})
+			require.NoError(t, daemon.WriteState())
+
+			found, err := daemon.vmMgr.UpdateAndPersist("i-mutate", func(v *vm.VM) bool {
+				tt.mutate(v)
+				return true
+			})
+			require.NoError(t, err)
+			require.True(t, found)
+
+			// Read the file directly: no second writer runs in between, so
+			// anything missing here is a write path that skipped it.
+			state, err := ReadLocalState(daemon.localStatePath())
+			require.NoError(t, err)
+			require.NotNil(t, state)
+			persisted, ok := state.VMS["i-mutate"]
+			require.True(t, ok, "the mutated instance must be in the local file")
+			tt.verify(t, persisted)
+		})
+	}
+}
+
 // TestRestoreInstances_LocalRecordSurvivesMissingKVKey pins the boot-path
 // precedence. startLocal loads the local file, then startCluster's restore
 // pulls from KV. An absent KV key means "this node has no record here", which

spinifex/daemon/vm_adapters.go

@@ -28,19 +28,22 @@ import (
 )
 
 // stateStoreAdapter satisfies vm.StateStore by delegating to JetStreamManager.
-// Tolerates nil JetStreamManager during early boot.
+// Tolerates nil JetStreamManager during early boot. The running set is the one
+// exception: it goes through the daemon's persist path so that the local file,
+// not just KV, sees every change the vm package makes.
 type stateStoreAdapter struct {
-	js *JetStreamManager
+	js      *JetStreamManager
+	persist func(nodeID string, vms map[string]*vm.VM) error
 }
 
 var _ vm.StateStore = (*stateStoreAdapter)(nil)
 
-func newStateStoreAdapter(js *JetStreamManager) *stateStoreAdapter {
-	return &stateStoreAdapter{js: js}
+func newStateStoreAdapter(js *JetStreamManager, persist func(string, map[string]*vm.VM) error) *stateStoreAdapter {
+	return &stateStoreAdapter{js: js, persist: persist}
 }
 
 func (a *stateStoreAdapter) SaveRunningState(nodeID string, snapshot map[string]*vm.VM) error {
-	return a.js.WriteState(nodeID, snapshot)
+	return a.persist(nodeID, snapshot)
 }
 
 func (a *stateStoreAdapter) LoadRunningState(nodeID string) (map[string]*vm.VM, bool, error) {