refactor: store the node instance-state record in the local state envelope (#921)

spinifex/daemon/jetstream.go

@@ -51,21 +51,28 @@ type KVSyncObserver interface {
 //
 // The instance-state bucket carries two record types, so it is one kvstore
 // bucket with two typed views over it: they share its handle, and a recovery
-// driven through either repairs both.
+// driven through either repairs both. A node's running set is stored as
+// LocalState, the same envelope the local state file uses.
 type JetStreamManager struct {
 	js        jetstream.JetStream
-	stateB    *kvstore.Bucket           // spinifex-instance-state
-	nodeState *kvstore.Store[nodeState] // node.<id> records
-	stopped   *kvstore.Store[vm.VM]     // instance.<id> records
-	term      *kvstore.Store[vm.VM]     // spinifex-terminated-instances
-	clusterKV jetstream.KeyValue        // spinifex-cluster-state
+	stateB    *kvstore.Bucket            // spinifex-instance-state
+	nodeState *kvstore.Store[LocalState] // node.<id> records
+	stopped   *kvstore.Store[vm.VM]      // instance.<id> records
+	term      *kvstore.Store[vm.VM]      // spinifex-terminated-instances
+	clusterKV jetstream.KeyValue         // spinifex-cluster-state
 	replicas  int
 	obs       KVSyncObserver
 }
 
-// nodeState is the envelope a node's running set is stored in.
-type nodeState struct {
-	VMS map[string]*vm.VM `json:"vms"`
+// checkNodeStateVersion reports whether a node record read from KV is one this
+// binary can parse. Version 0 predates record versioning and reads as current;
+// a version above current was written by a newer node and is not guessed at.
+func checkNodeStateVersion(key string, version int) error {
+	if version > LocalStateSchemaVersion {
+		return fmt.Errorf("instance state %s: record schema_version %d is newer than this node understands (%d)",
+			key, version, LocalStateSchemaVersion)
+	}
+	return nil
 }
 
 // SetSyncObserver registers obs to receive best-effort KV sync outcomes. Pass
@@ -138,7 +145,7 @@ func terminatedInstanceConfig(replicas int) kvstore.Config {
 // handle, so a recovery driven through either repairs both.
 func (m *JetStreamManager) setInstanceStateBucket(b *kvstore.Bucket) {
 	m.stateB = b
-	m.nodeState = kvstore.On[nodeState](b)
+	m.nodeState = kvstore.On[LocalState](b)
 	m.stopped = kvstore.On[vm.VM](b)
 }
 
@@ -394,7 +401,8 @@ func (m *JetStreamManager) WriteState(nodeID string, vms map[string]*vm.VM) erro
 	}
 
 	key := InstanceStatePrefix + nodeID
-	if err := m.nodeState.Set(context.Background(), key, &nodeState{VMS: vms}); err != nil {
+	record := LocalState{SchemaVersion: LocalStateSchemaVersion, VMS: vms}
+	if err := m.nodeState.Set(context.Background(), key, &record); err != nil {
 		return err
 	}
 
@@ -442,9 +450,11 @@ func (m *JetStreamManager) WriteStateBytesBestEffort(nodeID string, jsonData []b
 }
 
 // marshalInstanceState produces the JSON wire form of vms, for the best-effort
-// path that marshals under the caller's lock and commits lock-free.
+// path that marshals under the caller's lock and commits lock-free. The KV
+// record and the local state file are the same envelope, so this is the same
+// marshaller: the two cannot drift.
 func marshalInstanceState(vms map[string]*vm.VM) ([]byte, error) {
-	return json.Marshal(nodeState{VMS: vms})
+	return MarshalLocalState(vms)
 }
 
 // LoadState loads the instance state from the KV store for the given node.
@@ -466,6 +476,9 @@ func (m *JetStreamManager) LoadState(nodeID string) (map[string]*vm.VM, bool, er
 		}
 		return nil, false, err
 	}
+	if err := checkNodeStateVersion(key, state.SchemaVersion); err != nil {
+		return nil, false, err
+	}
 	if state.VMS == nil {
 		state.VMS = make(map[string]*vm.VM)
 	}

spinifex/daemon/jetstream_test.go

@@ -1795,3 +1795,105 @@ func TestJetStreamManager_UpdateStoppedInstance_AbsentKeyIsNotFound(t *testing.T
 	})
 	require.ErrorIs(t, err, kvstore.ErrNotFound)
 }
+
+// seedNodeRecord writes raw bytes at a node's key, bypassing the typed view, so
+// a test can present a record shape this binary would never write itself.
+func seedNodeRecord(t *testing.T, jsm *JetStreamManager, nodeID string, raw []byte) {
+	t.Helper()
+	kv, err := jsm.stateB.KV(t.Context())
+	require.NoError(t, err)
+	_, err = kv.Put(t.Context(), InstanceStatePrefix+nodeID, raw)
+	require.NoError(t, err)
+}
+
+// A node record predates schema versioning, so it carries no version at all.
+// Reading it as current is what makes an in-place upgrade need no migration.
+func TestJetStreamManager_LoadState_UnversionedRecordStillReads(t *testing.T) {
+	nc, err := nats.Connect(sharedJSNATSURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc, 1)
+	require.NoError(t, err)
+	require.NoError(t, jsm.InitKVBucket())
+
+	seedNodeRecord(t, jsm, "legacy-node",
+		[]byte(`{"vms":{"i-legacy-001":{"id":"i-legacy-001","status":"running"}}}`))
+
+	loaded, found, err := jsm.LoadState("legacy-node")
+	require.NoError(t, err, "A record written before versioning must still read")
+	require.True(t, found)
+	require.Len(t, loaded, 1)
+	assert.Equal(t, vm.StateRunning, loaded["i-legacy-001"].Status)
+}
+
+// A record from a newer binary is refused rather than parsed on a guess: this
+// is the per-record guard the bucket-level version key cannot provide, since
+// that is checked once at open by whichever node opened it.
+func TestJetStreamManager_LoadState_FutureRecordIsRefused(t *testing.T) {
+	nc, err := nats.Connect(sharedJSNATSURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc, 1)
+	require.NoError(t, err)
+	require.NoError(t, jsm.InitKVBucket())
+
+	seedNodeRecord(t, jsm, "future-node",
+		[]byte(`{"schema_version":99,"vms":{"i-future-001":{"id":"i-future-001"}}}`))
+
+	_, _, err = jsm.LoadState("future-node")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "schema_version 99")
+	assert.Contains(t, err.Error(), "newer than this node understands")
+}
+
+// Both writers stamp the version: the typed WriteState and the best-effort
+// pre-marshalled path, which is the one that would silently skip it.
+func TestJetStreamManager_WriteState_StampsSchemaVersion(t *testing.T) {
+	nc, err := nats.Connect(sharedJSNATSURL)
+	require.NoError(t, err)
+	defer nc.Close()
+
+	jsm, err := NewJetStreamManager(nc, 1)
+	require.NoError(t, err)
+	require.NoError(t, jsm.InitKVBucket())
+
+	vms := map[string]*vm.VM{"i-stamp-001": {ID: "i-stamp-001", Status: vm.StateRunning}}
+	require.NoError(t, jsm.WriteState("stamp-node", vms))
+
+	kv, err := jsm.stateB.KV(t.Context())
+	require.NoError(t, err)
+
+	read := func(nodeID string) LocalState {
+		t.Helper()
+		entry, err := kv.Get(t.Context(), InstanceStatePrefix+nodeID)
+		require.NoError(t, err)
+		var got LocalState
+		require.NoError(t, json.Unmarshal(entry.Value(), &got))
+		return got
+	}
+
+	assert.Equal(t, LocalStateSchemaVersion, read("stamp-node").SchemaVersion)
+
+	data, err := marshalInstanceState(vms)
+	require.NoError(t, err)
+	jsm.WriteStateBytesBestEffort("bytes-node", data, 5*time.Second)
+	assert.Equal(t, LocalStateSchemaVersion, read("bytes-node").SchemaVersion,
+		"the best-effort path marshals its own bytes and must stamp the same version")
+}
+
+// The KV node record and the local state file are one envelope. Asserted on the
+// bytes rather than the type, because a second marshaller is how they drift.
+func TestMarshalInstanceState_MatchesLocalStateFileBytes(t *testing.T) {
+	vms := map[string]*vm.VM{
+		"i-envelope-001": {ID: "i-envelope-001", Status: vm.StateRunning, InstanceType: "t3.micro"},
+	}
+
+	kvBytes, err := marshalInstanceState(vms)
+	require.NoError(t, err)
+	fileBytes, err := MarshalLocalState(vms)
+	require.NoError(t, err)
+
+	assert.JSONEq(t, string(fileBytes), string(kvBytes))
+}

spinifex/daemon/local_state.go

@@ -12,7 +12,8 @@ import (
 )
 
 const (
-	// LocalStateSchemaVersion guards on-disk compatibility; bump on breaking changes.
+	// LocalStateSchemaVersion guards compatibility of a node's instance state
+	// record, on disk and in KV; bump on breaking changes.
 	LocalStateSchemaVersion = 1
 
 	// DefaultLocalStateDir is the default directory under DataDir for local state files.
@@ -25,8 +26,9 @@ const (
 	DefaultDataDir = "/var/lib/spinifex"
 )
 
-// LocalState is the on-disk representation of a node's instance state.
-// SchemaVersion gates compatibility; unknown versions are rejected.
+// LocalState is a node's instance state, both on disk and as the node.<id>
+// record in the instance-state KV bucket. SchemaVersion gates compatibility;
+// unknown versions are rejected.
 type LocalState struct {
 	SchemaVersion int               `json:"schema_version"`
 	VMS           map[string]*vm.VM `json:"vms"`
@@ -110,6 +112,9 @@ func ReadLocalState(path string) (*LocalState, error) {
 		return nil, fmt.Errorf("parse local state %s: %w", path, err)
 	}
 
+	// Stricter than the KV read, deliberately: the file has carried a version
+	// since it was introduced, so a 0 here is corruption rather than an old
+	// record, and reading it as current would be a guess.
 	if state.SchemaVersion != LocalStateSchemaVersion {
 		return nil, fmt.Errorf("local state %s: unknown schema_version %d (expected %d)",
 			path, state.SchemaVersion, LocalStateSchemaVersion)