refactor(store): stage 3 of the InstanceInfo split - versioned session records with transparent legacy upgrade

pkg/sablier/session_record.go

@@ -0,0 +1,47 @@
+package sablier
+
+import "encoding/json"
+
+// SessionRecordVersion is the current schema version of persisted session
+// records. Bump it when the persisted shape changes, and handle the previous
+// versions in SessionRecord.UnmarshalJSON.
+const SessionRecordVersion = 1
+
+// SessionRecord is the persisted form of one session entry, used by every
+// Store implementation. It exists so the storage schema is explicit and
+// versioned instead of being whatever the domain struct happens to marshal
+// as: Version gates future schema evolution, and UnmarshalJSON transparently
+// upgrades records written before versioning existed (a bare InstanceInfo
+// document), so live state survives the upgrade.
+type SessionRecord struct {
+	Version  int          `json:"v"`
+	Instance InstanceInfo `json:"instance"`
+}
+
+// NewSessionRecord wraps an instance snapshot in the current schema version.
+func NewSessionRecord(instance InstanceInfo) SessionRecord {
+	return SessionRecord{Version: SessionRecordVersion, Instance: instance}
+}
+
+// UnmarshalJSON decodes a session record, upgrading legacy payloads: records
+// written before versioning were a bare InstanceInfo document (no "v" key).
+// Anything that is valid JSON but neither shape decodes to a zero record,
+// which callers detect through the empty Instance.Name (mirroring the
+// pre-versioning tolerance for foreign keys sharing the keyspace).
+func (r *SessionRecord) UnmarshalJSON(b []byte) error {
+	type plain SessionRecord
+	var p plain
+	if err := json.Unmarshal(b, &p); err != nil {
+		return err
+	}
+	if p.Version == 0 {
+		var legacy InstanceInfo
+		if err := json.Unmarshal(b, &legacy); err != nil {
+			return err
+		}
+		*r = SessionRecord{Version: SessionRecordVersion, Instance: legacy}
+		return nil
+	}
+	*r = SessionRecord(p)
+	return nil
+}

pkg/sablier/session_record_test.go

@@ -0,0 +1,65 @@
+package sablier
+
+import (
+	"encoding/json"
+	"testing"
+
+	"gotest.tools/v3/assert"
+)
+
+func TestSessionRecordRoundTrip(t *testing.T) {
+	rec := NewSessionRecord(InstanceInfo{Name: "web", Status: InstanceStatusReady, CurrentReplicas: 1, DesiredReplicas: 1})
+	assert.Equal(t, SessionRecordVersion, rec.Version)
+
+	b, err := json.Marshal(rec)
+	assert.NilError(t, err)
+	assert.Assert(t, containsSub(string(b), `"v":1`), "persisted record must carry its schema version: %s", b)
+
+	var got SessionRecord
+	assert.NilError(t, json.Unmarshal(b, &got))
+	assert.DeepEqual(t, rec, got)
+}
+
+// TestSessionRecordLegacyUpgrade pins the migration path: records persisted
+// before versioning were a bare InstanceInfo document; they must decode into
+// a current-version record with every field intact, so live sessions survive
+// the upgrade (both in valkey and in the state.json snapshot).
+func TestSessionRecordLegacyUpgrade(t *testing.T) {
+	legacy := `{"name":"web","currentReplicas":1,"desiredReplicas":2,"status":"ready","enabled":"true","groups":["default"],"runningHours":"09:00-18:00"}`
+
+	var got SessionRecord
+	assert.NilError(t, json.Unmarshal([]byte(legacy), &got))
+
+	assert.Equal(t, SessionRecordVersion, got.Version)
+	assert.Equal(t, "web", got.Instance.Name)
+	assert.Equal(t, int32(1), got.Instance.CurrentReplicas)
+	assert.Equal(t, int32(2), got.Instance.DesiredReplicas)
+	assert.Equal(t, InstanceStatusReady, got.Instance.Status)
+	assert.Equal(t, "true", got.Instance.Enabled)
+	assert.DeepEqual(t, []string{"default"}, got.Instance.Groups)
+	assert.Equal(t, "09:00-18:00", got.Instance.RunningHours)
+}
+
+func TestSessionRecordForeignPayloads(t *testing.T) {
+	t.Run("invalid JSON errors", func(t *testing.T) {
+		var got SessionRecord
+		assert.Assert(t, json.Unmarshal([]byte(`not json`), &got) != nil)
+	})
+	t.Run("foreign object decodes to a zero record", func(t *testing.T) {
+		// Stores sharing their keyspace skip entries whose Instance.Name does
+		// not match the key; a foreign JSON object must therefore decode to a
+		// record with an empty name rather than erroring the enumeration.
+		var got SessionRecord
+		assert.NilError(t, json.Unmarshal([]byte(`{"foo":"bar"}`), &got))
+		assert.Equal(t, "", got.Instance.Name)
+	})
+}
+
+func containsSub(s, sub string) bool {
+	for i := 0; i+len(sub) <= len(s); i++ {
+		if s[i:i+len(sub)] == sub {
+			return true
+		}
+	}
+	return false
+}

pkg/store/inmemory/inmemory.go

@@ -15,12 +15,14 @@ var _ json.Unmarshaler = (*InMemory)(nil)
 
 func NewInMemory() sablier.Store {
 	return &InMemory{
-		kv: tinykv.New[sablier.InstanceInfo](1*time.Second, nil),
+		// Values are versioned session records; SessionRecord's unmarshaler
+		// transparently upgrades snapshot files written before versioning.
+		kv: tinykv.New[sablier.SessionRecord](1*time.Second, nil),
 	}
 }
 
 type InMemory struct {
-	kv tinykv.KV[sablier.InstanceInfo]
+	kv tinykv.KV[sablier.SessionRecord]
 }
 
 func (i InMemory) UnmarshalJSON(bytes []byte) error {
@@ -36,11 +38,11 @@ func (i InMemory) Get(_ context.Context, s string) (sablier.InstanceInfo, error)
 	if !ok {
 		return sablier.InstanceInfo{}, store.ErrKeyNotFound
 	}
-	return val, nil
+	return val.Instance, nil
 }
 
 func (i InMemory) Put(_ context.Context, state sablier.InstanceInfo, duration time.Duration) error {
-	return i.kv.Put(state.Name, state, duration)
+	return i.kv.Put(state.Name, sablier.NewSessionRecord(state), duration)
 }
 
 func (i InMemory) Delete(_ context.Context, s string) error {
@@ -49,14 +51,14 @@ func (i InMemory) Delete(_ context.Context, s string) error {
 }
 
 func (i InMemory) Range(_ context.Context, f func(sablier.InstanceInfo, time.Time)) error {
-	i.kv.Range(func(_ string, value sablier.InstanceInfo, expiresAt time.Time) {
-		f(value, expiresAt)
+	i.kv.Range(func(_ string, value sablier.SessionRecord, expiresAt time.Time) {
+		f(value.Instance, expiresAt)
 	})
 	return nil
 }
 
 func (i InMemory) OnExpire(_ context.Context, f func(string)) error {
-	i.kv.SetOnExpire(func(k string, _ sablier.InstanceInfo) {
+	i.kv.SetOnExpire(func(k string, _ sablier.SessionRecord) {
 		f(k)
 	})
 	return nil

pkg/store/inmemory/inmemory_record_test.go

@@ -0,0 +1,52 @@
+package inmemory
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/sablierapp/sablier/pkg/sablier"
+	"gotest.tools/v3/assert"
+)
+
+// TestInMemorySnapshotIsVersioned pins the persisted snapshot schema: entries
+// written to state.json carry the session-record version, and a snapshot in
+// the legacy shape (bare InstanceInfo values, as written by older versions)
+// still loads. The load path is additionally covered end to end by the
+// sabliercmd storage tests, whose fixtures use the legacy shape.
+func TestInMemorySnapshotIsVersioned(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	t.Run("snapshot carries the schema version", func(t *testing.T) {
+		t.Parallel()
+		s := NewInMemory()
+		assert.NilError(t, s.Put(ctx, sablier.InstanceInfo{Name: "web", Status: sablier.InstanceStatusReady}, time.Minute))
+
+		raw, err := s.(*InMemory).MarshalJSON()
+		assert.NilError(t, err)
+		assert.Assert(t, strings.Contains(string(raw), `"v":1`), "snapshot entries must be versioned: %s", raw)
+	})
+
+	t.Run("legacy snapshot loads and upgrades", func(t *testing.T) {
+		t.Parallel()
+		legacy := map[string]any{
+			"web": map[string]any{
+				"value":     map[string]any{"name": "web", "currentReplicas": 1, "desiredReplicas": 1, "status": "ready"},
+				"expiresAt": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
+			},
+		}
+		raw, err := json.Marshal(legacy)
+		assert.NilError(t, err)
+
+		s := NewInMemory()
+		assert.NilError(t, s.(*InMemory).UnmarshalJSON(raw))
+
+		info, err := s.Get(ctx, "web")
+		assert.NilError(t, err)
+		assert.Equal(t, "web", info.Name)
+		assert.Equal(t, sablier.InstanceStatusReady, info.Status)
+	})
+}

pkg/store/valkey/valkey.go

@@ -42,17 +42,19 @@ func (v *ValKey) Get(ctx context.Context, s string) (sablier.InstanceInfo, error
 		return sablier.InstanceInfo{}, err
 	}
 
-	var i sablier.InstanceInfo
-	err = json.Unmarshal(b, &i)
+	// Session entries are stored as versioned records; SessionRecord's
+	// unmarshaler transparently upgrades pre-versioning payloads.
+	var r sablier.SessionRecord
+	err = json.Unmarshal(b, &r)
 	if err != nil {
 		return sablier.InstanceInfo{}, err
 	}
 
-	return i, nil
+	return r.Instance, nil
 }
 
 func (v *ValKey) Put(ctx context.Context, state sablier.InstanceInfo, duration time.Duration) error {
-	value, err := json.Marshal(state)
+	value, err := json.Marshal(sablier.NewSessionRecord(state))
 	if err != nil {
 		return err
 	}
@@ -108,18 +110,18 @@ func (v *ValKey) Range(ctx context.Context, f func(sablier.InstanceInfo, time.Ti
 			}
 
 			// The store may share its keyspace with keys that are not Sablier
-			// sessions. Skip anything that is not a valid InstanceInfo, or whose
-			// payload does not belong to this key, instead of aborting the whole
-			// enumeration on a single foreign or corrupt key.
-			var i sablier.InstanceInfo
-			if err = json.Unmarshal(b, &i); err != nil {
+			// sessions. Skip anything that is not a valid session record, or
+			// whose payload does not belong to this key, instead of aborting the
+			// whole enumeration on a single foreign or corrupt key.
+			var r sablier.SessionRecord
+			if err = json.Unmarshal(b, &r); err != nil {
 				continue
 			}
-			if i.Name != key {
+			if r.Instance.Name != key {
 				continue
 			}
 
-			f(i, expiresAt)
+			f(r.Instance, expiresAt)
 		}
 
 		cursor = entry.Cursor

pkg/store/valkey/valkey_record_test.go

@@ -0,0 +1,60 @@
+package valkey
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/sablierapp/sablier/pkg/sablier"
+	"gotest.tools/v3/assert"
+)
+
+// TestValKeySessionRecordMigration pins the persistence schema: new entries
+// are written as versioned session records, and entries written by older
+// versions (bare InstanceInfo payloads) are transparently upgraded on read,
+// so live sessions survive a Sablier upgrade without being dropped.
+func TestValKeySessionRecordMigration(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping test in short mode.")
+	}
+
+	ctx := t.Context()
+	vk := setupValKey(t)
+
+	t.Parallel()
+	t.Run("legacy payload is upgraded on read", func(t *testing.T) {
+		legacy := `{"name":"legacy-instance","currentReplicas":1,"desiredReplicas":1,"status":"ready"}`
+		err := vk.Client.Do(ctx, vk.Client.B().Set().Key("legacy-instance").Value(legacy).Ex(time.Minute).Build()).Error()
+		assert.NilError(t, err)
+
+		info, err := vk.Get(ctx, "legacy-instance")
+		assert.NilError(t, err)
+		assert.Equal(t, "legacy-instance", info.Name)
+		assert.Equal(t, sablier.InstanceStatusReady, info.Status)
+
+		// The legacy entry must also be visible to Range (the metrics path).
+		seen := false
+		err = vk.Range(ctx, func(i sablier.InstanceInfo, _ time.Time) {
+			if i.Name == "legacy-instance" {
+				seen = true
+			}
+		})
+		assert.NilError(t, err)
+		assert.Assert(t, seen, "legacy entry must be enumerated by Range")
+	})
+
+	t.Run("new entries carry the schema version", func(t *testing.T) {
+		err := vk.Put(ctx, sablier.InstanceInfo{Name: "versioned-instance", Status: sablier.InstanceStatusReady}, time.Minute)
+		assert.NilError(t, err)
+
+		raw, err := vk.Client.Do(ctx, vk.Client.B().Get().Key("versioned-instance").Build()).AsBytes()
+		assert.NilError(t, err)
+		var m map[string]any
+		assert.NilError(t, json.Unmarshal(raw, &m))
+		assert.Equal(t, float64(sablier.SessionRecordVersion), m["v"])
+
+		info, err := vk.Get(ctx, "versioned-instance")
+		assert.NilError(t, err)
+		assert.Equal(t, "versioned-instance", info.Name)
+	})
+}