From 64d64010b1d89c1e5a2b8fb33538975fd605aa7a Mon Sep 17 00:00:00 2001 From: Maximilien Cuony Date: Mon, 15 Jun 2026 15:07:26 +0200 Subject: [PATCH 1/2] [raft/memstore] Add aux memstore --- pkg/aux_/store/memstore/dss.go | 69 ++++++++++++++- pkg/aux_/store/memstore/dss_test.go | 125 ++++++++++++++++++++++++++++ pkg/aux_/store/memstore/store.go | 28 ++++++- 3 files changed, 216 insertions(+), 6 deletions(-) create mode 100644 pkg/aux_/store/memstore/dss_test.go diff --git a/pkg/aux_/store/memstore/dss.go b/pkg/aux_/store/memstore/dss.go index 38fa4d5bf..7198b3936 100644 --- a/pkg/aux_/store/memstore/dss.go +++ b/pkg/aux_/store/memstore/dss.go @@ -2,6 +2,8 @@ package memstore import ( "context" + "database/sql" + "time" auxmodels "github.com/interuss/dss/pkg/aux_/models" dsserr "github.com/interuss/dss/pkg/errors" @@ -9,17 +11,76 @@ import ( ) func (r *repo) SaveOwnMetadata(_ context.Context, locality string, publicEndpoint string) error { - return stacktrace.NewErrorWithCode(dsserr.NotImplemented, "SaveOwnMetadata not implemented for memstore") + if locality == "" { + return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Locality not set") + } + if publicEndpoint == "" { + return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Public endpoint not set") + } + + r.participants[locality] = &participant{ + publicEndpoint: publicEndpoint, + updatedAt: time.Now(), + } + return nil } func (r *repo) GetDSSMetadata(_ context.Context) ([]*auxmodels.DSSMetadata, error) { - return nil, stacktrace.NewErrorWithCode(dsserr.NotImplemented, "GetDSSMetadata not implemented for memstore") + metadata := make([]*auxmodels.DSSMetadata, 0, len(r.participants)) + for locality, p := range r.participants { + updatedAt := p.updatedAt + m := &auxmodels.DSSMetadata{ + Locality: locality, + 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 != locality { + continue + } + if !found || hb.Timestamp.After(*latest.Timestamp) { + latest = hb + found = true + } + } + + if found { + m.LatestTimestamp.Source = sql.NullString{String: latest.Source, Valid: true} + m.LatestTimestamp.Timestamp = latest.Timestamp + m.LatestTimestamp.NextHeartbeatExpectedBefore = latest.NextHeartbeatExpectedBefore + m.LatestTimestamp.Reporter = sql.NullString{String: latest.Reporter, Valid: true} + } + + metadata = append(metadata, m) + } + return metadata, nil } func (r *repo) RecordHeartbeat(_ context.Context, heartbeat auxmodels.Heartbeat) error { - return stacktrace.NewErrorWithCode(dsserr.NotImplemented, "RecordHeartbeat not implemented for memstore") + if heartbeat.Locality == "" { + return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Locality not set") + } + if heartbeat.Source == "" { + return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Source not set") + } + + if heartbeat.Timestamp == nil { + now := time.Now() + 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: heartbeat.Locality, source: heartbeat.Source}] = heartbeat + return nil } func (r *repo) GetDSSAirspaceRepresentationID(_ context.Context) (string, error) { - return "", stacktrace.NewErrorWithCode(dsserr.NotImplemented, "GetDSSAirspaceRepresentationID not implemented for memstore") + return "", stacktrace.NewErrorWithCode(dsserr.NotImplemented, "GetDSSAirspaceRepresentationID not implementable for memstore") } diff --git a/pkg/aux_/store/memstore/dss_test.go b/pkg/aux_/store/memstore/dss_test.go new file mode 100644 index 000000000..43563d27b --- /dev/null +++ b/pkg/aux_/store/memstore/dss_test.go @@ -0,0 +1,125 @@ +package memstore + +import ( + "context" + "testing" + "time" + + auxmodels "github.com/interuss/dss/pkg/aux_/models" + dsserr "github.com/interuss/dss/pkg/errors" + "github.com/interuss/stacktrace" + "github.com/stretchr/testify/require" +) + +func TestSaveOwnMetadataValidation(t *testing.T) { + ctx := context.Background() + r := newRepo() + + require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.SaveOwnMetadata(ctx, "", "https://example.com"))) + require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.SaveOwnMetadata(ctx, "dss-1", ""))) +} + +func TestSaveOwnMetadataRoundTrip(t *testing.T) { + ctx := context.Background() + r := newRepo() + + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com")) + + md, err := r.GetDSSMetadata(ctx) + require.NoError(t, err) + + require.Len(t, md, 1) + require.Equal(t, "dss-1", md[0].Locality) + require.Equal(t, "https://example.com", md[0].PublicEndpoint) + require.NotNil(t, md[0].UpdatedAt) + + // No heartbeat recorded yet. + require.False(t, md[0].LatestTimestamp.Source.Valid) + require.Nil(t, md[0].LatestTimestamp.Timestamp) +} + +func TestSaveOwnMetadataUpsert(t *testing.T) { + ctx := context.Background() + r := newRepo() + + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://old.example.com")) + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://new.example.com")) + + md, err := r.GetDSSMetadata(ctx) + require.NoError(t, err) + + require.Len(t, md, 1) + require.Equal(t, "https://new.example.com", md[0].PublicEndpoint) +} + +func TestRecordHeartbeatValidation(t *testing.T) { + ctx := context.Background() + r := newRepo() + + require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Source: "source1"}))) + require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1"}))) + + ts := time.Now() + before := ts.Add(-time.Minute) + err := r.RecordHeartbeat(ctx, auxmodels.Heartbeat{ + Locality: "dss-1", + Source: "source1", + Timestamp: &ts, + NextHeartbeatExpectedBefore: &before, + }) + + require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(err)) +} + +func TestRecordHeartbeatDefaultsTimestamp(t *testing.T) { + ctx := context.Background() + r := newRepo() + + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com")) + require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1"})) + + md, err := r.GetDSSMetadata(ctx) + require.NoError(t, err) + + require.Len(t, md, 1) + require.True(t, md[0].LatestTimestamp.Source.Valid) + require.NotNil(t, md[0].LatestTimestamp.Timestamp) +} + +func TestGetDSSMetadataPicksLatestHeartbeat(t *testing.T) { + ctx := context.Background() + r := newRepo() + + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com")) + + older := time.Now().Add(-time.Hour) + newer := time.Now() + require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1", Timestamp: &older, Reporter: "uss1"})) + require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source2", Timestamp: &newer, Reporter: "uss2"})) + + md, err := r.GetDSSMetadata(ctx) + require.NoError(t, err) + + require.Len(t, md, 1) + require.True(t, md[0].LatestTimestamp.Timestamp.Equal(newer)) + require.Equal(t, "source2", md[0].LatestTimestamp.Source.String) + require.Equal(t, "uss2", md[0].LatestTimestamp.Reporter.String) +} + +func TestGetDSSMetadataUpdatesHeartbeatPerSource(t *testing.T) { + ctx := context.Background() + r := newRepo() + + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com")) + + first := time.Now().Add(-time.Hour) + second := time.Now() + require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1", Timestamp: &first})) + require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1", Timestamp: &second})) + + md, err := r.GetDSSMetadata(ctx) + require.NoError(t, err) + + require.Len(t, md, 1) + require.True(t, md[0].LatestTimestamp.Timestamp.Equal(second)) +} diff --git a/pkg/aux_/store/memstore/store.go b/pkg/aux_/store/memstore/store.go index c2b875f9b..f94a1ffc2 100644 --- a/pkg/aux_/store/memstore/store.go +++ b/pkg/aux_/store/memstore/store.go @@ -2,17 +2,41 @@ package memstore import ( "context" + "time" + auxmodels "github.com/interuss/dss/pkg/aux_/models" "github.com/interuss/dss/pkg/aux_/repos" "github.com/interuss/dss/pkg/memstore" "go.uber.org/zap" ) // repo is a full implementation of aux_.repos.Repository for memory-based storage. -type repo struct{} +type repo struct { + // participants holds pool participants metadata, keyed by locality. + participants map[string]*participant + // heartbeats holds the latest heartbeat per (locality, source). + heartbeats map[heartbeatKey]auxmodels.Heartbeat +} + +type participant struct { + publicEndpoint string + updatedAt time.Time +} + +type heartbeatKey struct { + locality string + source string +} + +func newRepo() *repo { + return &repo{ + participants: map[string]*participant{}, + heartbeats: map[heartbeatKey]auxmodels.Heartbeat{}, + } +} func Init(ctx context.Context, logger *zap.Logger) (*memstore.Store[repos.Repository], error) { - return memstore.Init(ctx, logger, "aux_", &repo{}) + return memstore.Init(ctx, logger, "aux_", newRepo()) } func (r *repo) GetRepo() repos.Repository { return r } From 49f28517b751b39a1e400fa22577e5735f72193c Mon Sep 17 00:00:00 2001 From: Maximilien Cuony Date: Tue, 16 Jun 2026 11:05:14 +0200 Subject: [PATCH 2/2] [raft/memstore] Add snapshop capability to aux memstore --- pkg/aux_/store/memstore/dss.go | 22 ++++----- pkg/aux_/store/memstore/snapshot.go | 35 ++++++++++++++ pkg/aux_/store/memstore/snapshot_test.go | 59 ++++++++++++++++++++++++ pkg/aux_/store/memstore/store.go | 24 +++++++--- pkg/memstore/store.go | 2 + pkg/rid/store/memstore/snapshot.go | 13 ++++++ pkg/scd/store/memstore/snapshot.go | 13 ++++++ 7 files changed, 150 insertions(+), 18 deletions(-) create mode 100644 pkg/aux_/store/memstore/snapshot.go create mode 100644 pkg/aux_/store/memstore/snapshot_test.go create mode 100644 pkg/rid/store/memstore/snapshot.go create mode 100644 pkg/scd/store/memstore/snapshot.go diff --git a/pkg/aux_/store/memstore/dss.go b/pkg/aux_/store/memstore/dss.go index 7198b3936..79aeeee3a 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, locality string, publicEndpoin return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Public endpoint not set") } - r.participants[locality] = &participant{ - publicEndpoint: publicEndpoint, - updatedAt: time.Now(), + r.state.Participants[locality] = &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 locality, p := range r.participants { - updatedAt := p.updatedAt + metadata := make([]*auxmodels.DSSMetadata, 0, len(r.state.Participants)) + for locality, p := range r.state.Participants { + updatedAt := p.UpdatedAt m := &auxmodels.DSSMetadata{ Locality: locality, - 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 != locality { + for key, hb := range r.state.Heartbeats { + if key.Locality != locality { 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: heartbeat.Locality, source: heartbeat.Source}] = heartbeat + r.state.Heartbeats[heartbeatKey{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 f94a1ffc2..52ed59d12 100644 --- a/pkg/aux_/store/memstore/store.go +++ b/pkg/aux_/store/memstore/store.go @@ -12,6 +12,15 @@ import ( // repo is a full implementation of aux_.repos.Repository for memory-based storage. type repo struct { + state state +} + +// state is the serializable in-memory state. +type state struct { + // Participants holds pool participants metadata, keyed by locality. + Participants map[string]*participant + // Heartbeats holds the latest heartbeat per (locality, source). + Heartbeats map[heartbeatKey]auxmodels.Heartbeat // participants holds pool participants metadata, keyed by locality. participants map[string]*participant // heartbeats holds the latest heartbeat per (locality, source). @@ -19,20 +28,21 @@ type repo struct { } type participant struct { - publicEndpoint string - updatedAt time.Time + PublicEndpoint string + UpdatedAt time.Time } type heartbeatKey struct { - locality string - source string + Locality string + Source string } func newRepo() *repo { return &repo{ - participants: map[string]*participant{}, - heartbeats: map[heartbeatKey]auxmodels.Heartbeat{}, - } + state: state{ + Participants: map[string]*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 9c850f2c3..2fe7a12e6 100644 --- a/pkg/memstore/store.go +++ b/pkg/memstore/store.go @@ -18,6 +18,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") +}