[raft/memstore] Add Checkpoint/Restore

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
 }
 

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,15 +13,16 @@ 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.
 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,16 +35,71 @@ 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) {
 	return memstore.Init(ctx, logger, "aux_", newRepo())
 }
 
 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()
+}

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)
+}

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
 }

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()
+}

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)
+}