diff --git a/pkg/aux_/store/memstore/dss.go b/pkg/aux_/store/memstore/dss.go index 673ce6748..ad8a9e4ef 100644 --- a/pkg/aux_/store/memstore/dss.go +++ b/pkg/aux_/store/memstore/dss.go @@ -38,20 +38,20 @@ func (r *repo) GetDSSMetadata(_ context.Context) ([]*auxmodels.DSSMetadata, erro } // Find the latest heartbeat across all sources for this locality. - var latest auxmodels.Heartbeat - found := false + var latest *heartbeat + var latestSource string for key, hb := range r.state.Heartbeats { if key.Locality != loc { continue } - if !found || hb.Timestamp.After(*latest.Timestamp) { + if latest == nil || hb.Timestamp.After(*latest.Timestamp) { latest = hb - found = true + latestSource = key.Source } } - if found { - m.LatestTimestamp.Source = sql.NullString{String: latest.Source, Valid: true} + if latest != nil { + m.LatestTimestamp.Source = sql.NullString{String: latestSource, Valid: true} m.LatestTimestamp.Timestamp = latest.Timestamp m.LatestTimestamp.NextHeartbeatExpectedBefore = latest.NextHeartbeatExpectedBefore m.LatestTimestamp.Reporter = sql.NullString{String: latest.Reporter, Valid: true} @@ -62,24 +62,28 @@ func (r *repo) GetDSSMetadata(_ context.Context) ([]*auxmodels.DSSMetadata, erro return metadata, nil } -func (r *repo) RecordHeartbeat(ctx context.Context, heartbeat auxmodels.Heartbeat) error { - if heartbeat.Locality == "" { +func (r *repo) RecordHeartbeat(ctx context.Context, hb auxmodels.Heartbeat) error { + if hb.Locality == "" { return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Locality not set") } - if heartbeat.Source == "" { + if hb.Source == "" { return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Source not set") } - if heartbeat.Timestamp == nil { + if hb.Timestamp == nil { now := timestamp.MustGetRequestTimestamp(ctx).UTC() - heartbeat.Timestamp = &now + hb.Timestamp = &now } - if heartbeat.NextHeartbeatExpectedBefore != nil && heartbeat.NextHeartbeatExpectedBefore.Before(*heartbeat.Timestamp) { + if hb.NextHeartbeatExpectedBefore != nil && hb.NextHeartbeatExpectedBefore.Before(*hb.Timestamp) { return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Cannot expect the timestamp of the next heartbeat before the timestamp of the new heartbeat") } - r.state.Heartbeats[heartbeatKey{Locality: locality(heartbeat.Locality), Source: heartbeat.Source}] = heartbeat + r.state.Heartbeats[heartbeatKey{Locality: locality(hb.Locality), Source: hb.Source}] = &heartbeat{ + Timestamp: hb.Timestamp, + NextHeartbeatExpectedBefore: hb.NextHeartbeatExpectedBefore, + Reporter: hb.Reporter, + } return nil } diff --git a/pkg/aux_/store/memstore/store.go b/pkg/aux_/store/memstore/store.go index 1562b81a6..ba1686162 100644 --- a/pkg/aux_/store/memstore/store.go +++ b/pkg/aux_/store/memstore/store.go @@ -4,7 +4,6 @@ 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" @@ -14,7 +13,8 @@ type locality string // repo is a full implementation of aux_.repos.Repository for memory-based storage. type repo struct { - state state + state state + checkpoint state } // state is the serializable in-memory state. @@ -22,7 +22,7 @@ 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 + Heartbeats map[heartbeatKey]*heartbeat } type participant struct { @@ -35,12 +35,23 @@ type heartbeatKey struct { Source string } +type heartbeat struct { + Timestamp *time.Time + NextHeartbeatExpectedBefore *time.Time + Reporter string +} + func newRepo() *repo { + + state := state{ + Participants: map[locality]*participant{}, + Heartbeats: map[heartbeatKey]*heartbeat{}, + } + return &repo{ - state: state{ - Participants: map[locality]*participant{}, - Heartbeats: map[heartbeatKey]auxmodels.Heartbeat{}, - }} + state: state, + checkpoint: state.clone(), + } } func Init(ctx context.Context, logger *zap.Logger) (*memstore.Store[repos.Repository], error) { @@ -48,3 +59,47 @@ func Init(ctx context.Context, logger *zap.Logger) (*memstore.Store[repos.Reposi } func (r *repo) GetRepo() repos.Repository { return r } + +func clonePtr[T any](v *T) *T { + if v == nil { + return nil + } + return new(*v) +} + +func (h *heartbeat) clone() *heartbeat { + cp := *h + cp.Timestamp = clonePtr(h.Timestamp) + cp.NextHeartbeatExpectedBefore = clonePtr(h.NextHeartbeatExpectedBefore) + return &cp +} + +func (p *participant) clone() *participant { + cp := *p + return &cp +} + +// clone returns a copy of s with independent maps and participant records. +func (s state) clone() state { + ps := make(map[locality]*participant, len(s.Participants)) + for k, v := range s.Participants { + ps[k] = v.clone() + } + hb := make(map[heartbeatKey]*heartbeat, len(s.Heartbeats)) + for k, v := range s.Heartbeats { + hb[k] = v.clone() + } + return state{Participants: ps, Heartbeats: hb} +} + +// Checkpoint ask the repo to store a quick, internal checkpoint with its current state. +// There is at most one check point, any existing checkpoint is overwritten +func (r *repo) Checkpoint() { + r.checkpoint = r.state.clone() +} + +// Restore replaces the current state with the latest checkpoint. May be called multiple time +// to restore the same checkpoint. +func (r *repo) Restore() { + r.state = r.checkpoint.clone() +} diff --git a/pkg/aux_/store/memstore/store_test.go b/pkg/aux_/store/memstore/store_test.go new file mode 100644 index 000000000..1526efb04 --- /dev/null +++ b/pkg/aux_/store/memstore/store_test.go @@ -0,0 +1,51 @@ +package memstore + +import ( + "context" + "testing" + + "github.com/interuss/dss/pkg/timestamp" + "github.com/stretchr/testify/require" +) + +func TestCheckpointRestore(t *testing.T) { + ctx := context.Background() + ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now()) + + r := newRepo() + + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com")) + + r.Checkpoint() + + // Mutate after the checkpoint. + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-2", "https://other.example.com")) + md, err := r.GetDSSMetadata(ctx) + require.NoError(t, err) + require.Len(t, md, 2) + + // Restore drops dss-2. + r.Restore() + md, err = r.GetDSSMetadata(ctx) + require.NoError(t, err) + require.Len(t, md, 1) + require.Equal(t, "dss-1", md[0].Locality) +} + +func TestCheckpointIsolatesUpsert(t *testing.T) { + ctx := context.Background() + ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now()) + r := newRepo() + + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://old.example.com")) + + r.Checkpoint() + + require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://new.example.com")) + + r.Restore() + md, err := r.GetDSSMetadata(ctx) + require.NoError(t, err) + require.Len(t, md, 1) + require.Equal(t, "https://old.example.com", md[0].PublicEndpoint) +} diff --git a/pkg/memstore/store.go b/pkg/memstore/store.go index 106941b04..073ea3335 100644 --- a/pkg/memstore/store.go +++ b/pkg/memstore/store.go @@ -15,6 +15,14 @@ type MemRepo[R any] interface { GetRepo() R GetSnapshot() ([]byte, error) RestoreFromSnapshot([]byte) error + + // Checkpoint ask the repo to store a quick, internal checkpoint with its current state. + // There is at most one check point, any existing checkpoint is overwritten + Checkpoint() + + // Restore replaces the current state with the latest checkpoint. May be called multiple time + // to restore the same checkpoint. + Restore() } // Memstore is a special kind of store: @@ -60,6 +68,14 @@ func (s *Store[R]) Interact(_ context.Context) (R, error) { return s.memRepo.GetRepo(), nil } +func (s *Store[R]) Checkpoint() { + s.memRepo.Checkpoint() +} + +func (s *Store[R]) Restore() { + s.memRepo.Restore() +} + func (s *Store[R]) Close() error { return nil } diff --git a/pkg/rid/store/memstore/store.go b/pkg/rid/store/memstore/store.go index edb7fbba0..531d92c3f 100644 --- a/pkg/rid/store/memstore/store.go +++ b/pkg/rid/store/memstore/store.go @@ -2,6 +2,7 @@ package memstore import ( "context" + "slices" "time" "github.com/golang/geo/s2" @@ -15,7 +16,8 @@ import ( // repo is a full implementation of rid.repos.Repository for memory-based storage. type repo struct { - state state + state state + checkpoint state } // state is the serializable in-memory state. @@ -64,10 +66,13 @@ func newRepo() *repo { } func (r *repo) resetState() { + r.state = state{ ISAs: map[dssmodels.ID]*isaRecord{}, Subscriptions: map[dssmodels.ID]*subscriptionRecord{}, } + + r.Checkpoint() } func Init(ctx context.Context, logger *zap.Logger) (*memstore.Store[repos.Repository], error) { @@ -173,3 +178,49 @@ func listExpired[M any, R expiringRecord[M]](store map[dssmodels.ID]R, writer st } return out } + +func (rec *isaRecord) clone() *isaRecord { + cp := *rec + cp.Cells = slices.Clone(rec.Cells) + cp.StartTime = clonePtr(rec.StartTime) + cp.EndTime = clonePtr(rec.EndTime) + cp.AltitudeHi = clonePtr(rec.AltitudeHi) + cp.AltitudeLo = clonePtr(rec.AltitudeLo) + return &cp +} + +func (rec *subscriptionRecord) clone() *subscriptionRecord { + cp := *rec + cp.Cells = slices.Clone(rec.Cells) + cp.StartTime = clonePtr(rec.StartTime) + cp.EndTime = clonePtr(rec.EndTime) + cp.AltitudeHi = clonePtr(rec.AltitudeHi) + cp.AltitudeLo = clonePtr(rec.AltitudeLo) + return &cp +} + +// clone returns a deep copy of s. May be optimzed in speed by not cloning everything, as long +// rest of the package don't mutate fields, iff speed of this function is important. +func (s state) clone() state { + isas := make(map[dssmodels.ID]*isaRecord, len(s.ISAs)) + for id, rec := range s.ISAs { + isas[id] = rec.clone() + } + subs := make(map[dssmodels.ID]*subscriptionRecord, len(s.Subscriptions)) + for id, rec := range s.Subscriptions { + subs[id] = rec.clone() + } + return state{ISAs: isas, Subscriptions: subs} +} + +// Checkpoint ask the repo to store a quick, internal checkpoint with its current state. +// There is at most one check point, any existing checkpoint is overwritten +func (r *repo) Checkpoint() { + r.checkpoint = r.state.clone() +} + +// Restore replaces the current state with the latest checkpoint. May be called multiple time +// to restore the same checkpoint. +func (r *repo) Restore() { + r.state = r.checkpoint.clone() +} diff --git a/pkg/rid/store/memstore/store_test.go b/pkg/rid/store/memstore/store_test.go index 3bf484463..62752a537 100644 --- a/pkg/rid/store/memstore/store_test.go +++ b/pkg/rid/store/memstore/store_test.go @@ -8,6 +8,7 @@ import ( "github.com/google/uuid" dssmodels "github.com/interuss/dss/pkg/models" ridmodels "github.com/interuss/dss/pkg/rid/models" + "github.com/interuss/dss/pkg/timestamp" "github.com/jonboulle/clockwork" "github.com/stretchr/testify/require" ) @@ -29,6 +30,7 @@ func setUpStore(t *testing.T) *repo { func TestDatabaseEnsuresBeginsBeforeExpires(t *testing.T) { ctx := context.Background() + ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now()) repo := setUpStore(t) var ( @@ -45,3 +47,51 @@ func TestDatabaseEnsuresBeginsBeforeExpires(t *testing.T) { }) require.Error(t, err) } + +func TestCheckpointRestoreISA(t *testing.T) { + ctx := context.Background() + ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now()) + repo := setUpStore(t) + + _, err := repo.InsertISA(ctx, serviceArea) + require.NoError(t, err) + + repo.Checkpoint() + + // Mutate after the checkpoint. + isa, err := repo.GetISA(ctx, serviceArea.ID, false) + require.NoError(t, err) + _, err = repo.DeleteISA(ctx, isa) + require.NoError(t, err) + gone, err := repo.GetISA(ctx, serviceArea.ID, false) + require.NoError(t, err) + require.Nil(t, gone) + + // Restore brings it back. + repo.Restore() + back, err := repo.GetISA(ctx, serviceArea.ID, false) + require.NoError(t, err) + require.NotNil(t, back) +} + +func TestCheckpointIsolatesNotificationIndex(t *testing.T) { + ctx := context.Background() + ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now()) + repo := setUpStore(t) + + sub, err := repo.InsertSubscription(ctx, subscriptionsPool[0].input) + require.NoError(t, err) + + repo.Checkpoint() + + // In-place notification-index bump must not leak into the checkpoint. + updated, err := repo.UpdateNotificationIdxsInCells(ctx, sub.Cells) + require.NoError(t, err) + require.Len(t, updated, 1) + require.Equal(t, sub.NotificationIndex+1, updated[0].NotificationIndex) + + repo.Restore() + restored, err := repo.GetSubscription(ctx, sub.ID) + require.NoError(t, err) + require.Equal(t, sub.NotificationIndex, restored.NotificationIndex) +} diff --git a/pkg/scd/store/memstore/store.go b/pkg/scd/store/memstore/store.go index 45365614a..06d187ee3 100644 --- a/pkg/scd/store/memstore/store.go +++ b/pkg/scd/store/memstore/store.go @@ -16,3 +16,11 @@ func Init(ctx context.Context, logger *zap.Logger) (*memstore.Store[repos.Reposi } func (r *repo) GetRepo() repos.Repository { return r } + +func (r *repo) Checkpoint() { + panic("Checkpoint not yet implemented for scd") +} + +func (r *repo) Restore() { + panic("Restore not yet implemented for scd") +}