[raft/memstore] Add Checkpoint/Restore

Author: the-glu

pkg/aux_/store/memstore/store.go

@@ -7,6 +7,7 @@ import (
 	auxmodels "github.com/interuss/dss/pkg/aux_/models"
 	"github.com/interuss/dss/pkg/aux_/repos"
 	"github.com/interuss/dss/pkg/memstore"
+	"github.com/interuss/stacktrace"
 	"go.uber.org/zap"
 )
 
@@ -50,3 +51,33 @@ func Init(ctx context.Context, logger *zap.Logger) (*memstore.Store[repos.Reposi
 }
 
 func (r *repo) GetRepo() repos.Repository { return r }
+
+// clone returns a copy of s with independent maps and participant records.
+func (s state) clone() state {
+	ps := make(map[string]*participant, len(s.Participants))
+	for k, v := range s.Participants {
+		cp := *v
+		ps[k] = &cp
+	}
+	hb := make(map[heartbeatKey]auxmodels.Heartbeat, len(s.Heartbeats))
+	for k, v := range s.Heartbeats {
+		hb[k] = v
+	}
+	return state{Participants: ps, Heartbeats: hb}
+}
+
+// Checkpoint returns a fast, restorable in-memory copy of the current state.
+func (r *repo) Checkpoint() any {
+	return r.state.clone()
+}
+
+// Restore replaces the current state with a checkpoint previously returned by
+// Checkpoint. The checkpoint is copied, so it stays reusable.
+func (r *repo) Restore(cp any) error {
+	s, ok := cp.(state)
+	if !ok {
+		return stacktrace.NewError("Invalid checkpoint type %T", cp)
+	}
+	r.state = s.clone()
+	return nil
+}

pkg/aux_/store/memstore/store_test.go

@@ -0,0 +1,51 @@
+package memstore
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCheckpointRestore(t *testing.T) {
+	ctx := context.Background()
+	r := newRepo()
+
+	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
+
+	cp := 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.
+	require.NoError(t, r.Restore(cp))
+	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()
+	r := newRepo()
+
+	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://old.example.com"))
+
+	cp := r.Checkpoint()
+
+	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://new.example.com"))
+
+	require.NoError(t, r.Restore(cp))
+	md, err := r.GetDSSMetadata(ctx)
+	require.NoError(t, err)
+	require.Len(t, md, 1)
+	require.Equal(t, "https://old.example.com", md[0].PublicEndpoint)
+}
+
+func TestRestoreInvalidType(t *testing.T) {
+	require.Error(t, newRepo().Restore("not a checkpoint"))
+}

pkg/memstore/store.go

@@ -22,6 +22,8 @@ type MemRepo[R any] interface {
 	GetRepo() R
 	GetSnapshot() ([]byte, error)
 	RestoreFromSnapshot([]byte) error
+	Checkpoint() any
+	Restore(any) error
 }
 
 type Store[R any] struct {
@@ -62,6 +64,16 @@ func (s *Store[R]) Interact(_ context.Context) (R, error) {
 	return s.memRepo.GetRepo(), nil
 }
 
+// Checkpoint returns a fast, restorable in-memory copy of the current state.
+func (s *Store[R]) Checkpoint() any {
+	return s.memRepo.Checkpoint()
+}
+
+// Restore replaces the current state with a checkpoint returned by Checkpoint.
+func (s *Store[R]) Restore(cp any) error {
+	return s.memRepo.Restore(cp)
+}
+
 func (s *Store[R]) Close() error {
 	return nil
 }

pkg/rid/store/memstore/store.go

@@ -134,3 +134,37 @@ func cloneFloat32(f *float32) *float32 {
 	v := *f
 	return &v
 }
+
+// clone returns a copy of s with independent maps and records. Cell slices and
+// time pointers are shared, as they are never mutated in place.
+func (s state) clone() state {
+	isas := make(map[dssmodels.ID]*isaRecord, len(s.ISAs))
+	for id, rec := range s.ISAs {
+		cp := *rec
+		isas[id] = &cp
+	}
+	subs := make(map[dssmodels.ID]*subscriptionRecord, len(s.Subscriptions))
+	for id, rec := range s.Subscriptions {
+		cp := *rec
+		subs[id] = &cp
+	}
+	return state{ISAs: isas, Subscriptions: subs}
+}
+
+// Checkpoint returns a fast, restorable in-memory copy of the current state.
+// Unlike GetSnapshot it does not serialize, so it is cheap but only valid
+// in-process.
+func (r *repo) Checkpoint() any {
+	return r.state.clone()
+}
+
+// Restore replaces the current state with a checkpoint previously returned by
+// Checkpoint. The checkpoint is copied, so it stays reusable.
+func (r *repo) Restore(cp any) error {
+	s, ok := cp.(state)
+	if !ok {
+		return stacktrace.NewError("Invalid checkpoint type %T", cp)
+	}
+	r.state = s.clone()
+	return nil
+}

pkg/rid/store/memstore/store_test.go

@@ -45,3 +45,53 @@ func TestDatabaseEnsuresBeginsBeforeExpires(t *testing.T) {
 	})
 	require.Error(t, err)
 }
+
+func TestCheckpointRestoreISA(t *testing.T) {
+	ctx := context.Background()
+	repo := setUpStore(t)
+
+	_, err := repo.InsertISA(ctx, serviceArea)
+	require.NoError(t, err)
+
+	cp := 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.
+	require.NoError(t, repo.Restore(cp))
+	back, err := repo.GetISA(ctx, serviceArea.ID, false)
+	require.NoError(t, err)
+	require.NotNil(t, back)
+}
+
+func TestCheckpointIsolatesNotificationIndex(t *testing.T) {
+	ctx := context.Background()
+	repo := setUpStore(t)
+
+	sub, err := repo.InsertSubscription(ctx, subscriptionsPool[0].input)
+	require.NoError(t, err)
+
+	cp := 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)
+
+	require.NoError(t, repo.Restore(cp))
+	restored, err := repo.GetSubscription(ctx, sub.ID)
+	require.NoError(t, err)
+	require.Equal(t, sub.NotificationIndex, restored.NotificationIndex)
+}
+
+func TestRestoreInvalidType(t *testing.T) {
+	require.Error(t, setUpStore(t).Restore("not a checkpoint"))
+}

pkg/scd/store/memstore/store.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/interuss/dss/pkg/memstore"
 	"github.com/interuss/dss/pkg/scd/repos"
+	"github.com/interuss/stacktrace"
 	"go.uber.org/zap"
 )
 
@@ -16,3 +17,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() any {
+	return nil
+}
+
+func (r *repo) Restore(any) error {
+	return stacktrace.NewError("Restore not yet implemented for scd")
+}