diff --git a/pkg/aux_/store/memstore/dss.go b/pkg/aux_/store/memstore/dss.go index 68aca6cac..adc7739d5 100644 --- a/pkg/aux_/store/memstore/dss.go +++ b/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,7 +69,7 @@ func (r *repo) RecordHeartbeat(_ context.Context, heartbeat auxmodels.Heartbeat) } if heartbeat.Timestamp == nil { - now := time.Now() + now := time.Now().UTC() heartbeat.Timestamp = &now } @@ -77,7 +77,7 @@ func (r *repo) RecordHeartbeat(_ context.Context, heartbeat auxmodels.Heartbeat) 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 } diff --git a/pkg/aux_/store/memstore/snapshot.go b/pkg/aux_/store/memstore/snapshot.go new file mode 100644 index 000000000..15deb274b --- /dev/null +++ b/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 +} diff --git a/pkg/aux_/store/memstore/snapshot_test.go b/pkg/aux_/store/memstore/snapshot_test.go new file mode 100644 index 000000000..2085e2488 --- /dev/null +++ b/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())) +} diff --git a/pkg/aux_/store/memstore/store.go b/pkg/aux_/store/memstore/store.go index b5d4f5f22..1562b81a6 100644 --- a/pkg/aux_/store/memstore/store.go +++ b/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) { diff --git a/pkg/memstore/store.go b/pkg/memstore/store.go index 026e3c2c5..ce136488d 100644 --- a/pkg/memstore/store.go +++ b/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 { diff --git a/pkg/rid/store/memstore/snapshot.go b/pkg/rid/store/memstore/snapshot.go new file mode 100644 index 000000000..dea64e6a4 --- /dev/null +++ b/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") +} diff --git a/pkg/scd/store/memstore/snapshot.go b/pkg/scd/store/memstore/snapshot.go new file mode 100644 index 000000000..dea64e6a4 --- /dev/null +++ b/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") +}