[raft/memstore] Add snapshop capability to aux memstore

pkg/aux_/store/memstore/dss.go

@@ -18,28 +18,28 @@ func (r *repo) SaveOwnMetadata(_ context.Context, loc string, publicEndpoint str
 		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Public endpoint not set")
 	}
 
-	r.participants[locality(loc)] = &participant{
-		publicEndpoint: publicEndpoint,
-		updatedAt:      time.Now(),
+	r.state.Participants[locality(loc)] = &participant{
+		PublicEndpoint: publicEndpoint,
+		UpdatedAt:      time.Now().UTC(),
 	}
 	return nil
 }
 
 func (r *repo) GetDSSMetadata(_ context.Context) ([]*auxmodels.DSSMetadata, error) {
-	metadata := make([]*auxmodels.DSSMetadata, 0, len(r.participants))
-	for loc, p := range r.participants {
-		updatedAt := p.updatedAt
+	metadata := make([]*auxmodels.DSSMetadata, 0, len(r.state.Participants))
+	for loc, p := range r.state.Participants {
+		updatedAt := p.UpdatedAt
 		m := &auxmodels.DSSMetadata{
 			Locality:       string(loc),
-			PublicEndpoint: p.publicEndpoint,
+			PublicEndpoint: p.PublicEndpoint,
 			UpdatedAt:      &updatedAt,
 		}
 
 		// Find the latest heartbeat across all sources for this locality.
 		var latest auxmodels.Heartbeat
 		found := false
-		for key, hb := range r.heartbeats {
-			if key.locality != loc {
+		for key, hb := range r.state.Heartbeats {
+			if key.Locality != loc {
 				continue
 			}
 			if !found || hb.Timestamp.After(*latest.Timestamp) {
@@ -69,15 +69,15 @@ func (r *repo) RecordHeartbeat(_ context.Context, heartbeat auxmodels.Heartbeat)
 	}
 
 	if heartbeat.Timestamp == nil {
-		now := time.Now()
+		now := time.Now().UTC()
 		heartbeat.Timestamp = &now
 	}
 
 	if heartbeat.NextHeartbeatExpectedBefore != nil && heartbeat.NextHeartbeatExpectedBefore.Before(*heartbeat.Timestamp) {
 		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Cannot expect the timestamp of the next heartbeat before the timestamp of the new heartbeat")
 	}
 
-	r.heartbeats[heartbeatKey{locality: locality(heartbeat.Locality), source: heartbeat.Source}] = heartbeat
+	r.state.Heartbeats[heartbeatKey{Locality: locality(heartbeat.Locality), Source: heartbeat.Source}] = heartbeat
 	return nil
 }
 

pkg/aux_/store/memstore/snapshot.go

@@ -0,0 +1,35 @@
+package memstore
+
+import (
+	"bytes"
+	"encoding/gob"
+
+	"github.com/interuss/stacktrace"
+)
+
+const snapshotVersion = 1
+
+type snapshotEnvelope struct {
+	Version int
+	State   state
+}
+
+func (r *repo) GetSnapshot() ([]byte, error) {
+	var buf bytes.Buffer
+	if err := gob.NewEncoder(&buf).Encode(snapshotEnvelope{Version: snapshotVersion, State: r.state}); err != nil {
+		return nil, stacktrace.Propagate(err, "Failed to encode memstore snapshot")
+	}
+	return buf.Bytes(), nil
+}
+
+func (r *repo) RestoreFromSnapshot(data []byte) error {
+	var env snapshotEnvelope
+	if err := gob.NewDecoder(bytes.NewReader(data)).Decode(&env); err != nil {
+		return stacktrace.Propagate(err, "Failed to decode memstore snapshot")
+	}
+	if env.Version != snapshotVersion {
+		return stacktrace.NewError("Unsupported memstore snapshot version %d, expected %d", env.Version, snapshotVersion)
+	}
+	r.state = env.State
+	return nil
+}

pkg/aux_/store/memstore/snapshot_test.go

@@ -0,0 +1,59 @@
+package memstore
+
+import (
+	"bytes"
+	"context"
+	"encoding/gob"
+	"testing"
+	"time"
+
+	auxmodels "github.com/interuss/dss/pkg/aux_/models"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSnapshotRoundTrip(t *testing.T) {
+	ctx := context.Background()
+	src := newRepo()
+	require.NoError(t, src.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
+	ts := time.Now().UTC()
+	require.NoError(t, src.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source-1", Timestamp: &ts, Reporter: "uss-1"}))
+
+	data, err := src.GetSnapshot()
+	require.NoError(t, err)
+
+	dst := newRepo()
+	require.NoError(t, dst.RestoreFromSnapshot(data))
+
+	want, err := src.GetDSSMetadata(ctx)
+	require.NoError(t, err)
+	got, err := dst.GetDSSMetadata(ctx)
+	require.NoError(t, err)
+	require.Equal(t, want, got)
+}
+
+func TestRestoreFromSnapshotReplacesState(t *testing.T) {
+	ctx := context.Background()
+	src := newRepo()
+	require.NoError(t, src.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
+	data, err := src.GetSnapshot()
+	require.NoError(t, err)
+
+	dst := newRepo()
+	require.NoError(t, dst.SaveOwnMetadata(ctx, "dss-2", "https://other.example.com"))
+	require.NoError(t, dst.RestoreFromSnapshot(data))
+
+	md, err := dst.GetDSSMetadata(ctx)
+	require.NoError(t, err)
+	require.Len(t, md, 1)
+	require.Equal(t, "dss-1", md[0].Locality)
+}
+
+func TestRestoreFromSnapshotInvalidData(t *testing.T) {
+	require.Error(t, newRepo().RestoreFromSnapshot([]byte("random value that is definitely not valid")))
+}
+
+func TestRestoreFromSnapshotVersionMismatch(t *testing.T) {
+	var buf bytes.Buffer
+	require.NoError(t, gob.NewEncoder(&buf).Encode(snapshotEnvelope{Version: snapshotVersion + 1}))
+	require.Error(t, newRepo().RestoreFromSnapshot(buf.Bytes()))
+}

pkg/aux_/store/memstore/store.go

@@ -14,27 +14,33 @@ type locality string
 
 // repo is a full implementation of aux_.repos.Repository for memory-based storage.
 type repo struct {
-	// participants holds pool participants metadata, keyed by locality.
-	participants map[locality]*participant
-	// heartbeats holds the latest heartbeat per (locality, source).
-	heartbeats map[heartbeatKey]auxmodels.Heartbeat
+	state state
+}
+
+// state is the serializable in-memory state.
+type state struct {
+	// Participants holds pool participants metadata, keyed by locality.
+	Participants map[locality]*participant
+	// Heartbeats holds the latest heartbeat per (locality, source).
+	Heartbeats map[heartbeatKey]auxmodels.Heartbeat
 }
 
 type participant struct {
-	publicEndpoint string
-	updatedAt      time.Time
+	PublicEndpoint string
+	UpdatedAt      time.Time
 }
 
 type heartbeatKey struct {
-	locality locality
-	source   string
+	Locality locality
+	Source   string
 }
 
 func newRepo() *repo {
 	return &repo{
-		participants: map[locality]*participant{},
-		heartbeats:   map[heartbeatKey]auxmodels.Heartbeat{},
-	}
+		state: state{
+			Participants: map[locality]*participant{},
+			Heartbeats:   map[heartbeatKey]auxmodels.Heartbeat{},
+		}}
 }
 
 func Init(ctx context.Context, logger *zap.Logger) (*memstore.Store[repos.Repository], error) {

pkg/memstore/store.go

@@ -19,6 +19,8 @@ import (
 
 type MemRepo[R any] interface {
 	GetRepo() R
+	GetSnapshot() ([]byte, error)
+	RestoreFromSnapshot([]byte) error
 }
 
 type Store[R any] struct {

pkg/rid/store/memstore/snapshot.go

@@ -0,0 +1,13 @@
+package memstore
+
+import (
+	"github.com/interuss/stacktrace"
+)
+
+func (r *repo) GetSnapshot() ([]byte, error) {
+	return nil, stacktrace.NewError("GetSnapshot not yet implemented for rid")
+}
+
+func (r *repo) RestoreFromSnapshot(data []byte) error {
+	return stacktrace.NewError("RestoreFromSnapshot not yet implemented for rid")
+}

pkg/scd/store/memstore/snapshot.go

@@ -0,0 +1,13 @@
+package memstore
+
+import (
+	"github.com/interuss/stacktrace"
+)
+
+func (r *repo) GetSnapshot() ([]byte, error) {
+	return nil, stacktrace.NewError("GetSnapshot not yet implemented for rid")
+}
+
+func (r *repo) RestoreFromSnapshot(data []byte) error {
+	return stacktrace.NewError("RestoreFromSnapshot not yet implemented for rid")
+}