[raftstore] Embed memstore and use checkpoint

cmds/core-service/main.go

@@ -368,9 +368,9 @@ func RunHTTPServer(ctx context.Context, ctxCanceler func(), address, locality st
 	handler = authorizer.TokenMiddleware(handler)
 	handler = http.TimeoutHandler(handler, *timeout, "request timeout")
 	handler = logging.HTTPMiddleware(logger, *dumpRequests, handler)
-	handler = timestamp.RequestTimestampMiddleware(handler)
+	handler = timestamp.Middleware(handler)
 	handler = random.Middleware(handler)
-	handler = requestlocality.LocalityMiddleware(locality)(handler)
+	handler = requestlocality.Middleware(locality)(handler)
 
 	if *enableMetrics || *enableTracing {
 		// We use the default settings; the APIRouter handler will override the span value accordingly, as it has more information.

pkg/aux_/pool_participants.go

@@ -96,7 +96,7 @@ func (a *Server) PutDSSInstancesHeartbeat(ctx context.Context, req *restapi.PutD
 		}
 		heartbeat.Timestamp = &ts
 	} else {
-		now := timestamp.MustGetRequestTimestamp(ctx)
+		now := timestamp.MustFromContext(ctx)
 		heartbeat.Timestamp = &now
 	}
 

pkg/aux_/store/memstore/dss.go

@@ -11,7 +11,7 @@ import (
 )
 
 func (r *repo) SaveOwnMetadata(ctx context.Context, loc string, publicEndpoint string) error {
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	r.state.Participants[locality(loc)] = &participant{
 		PublicEndpoint: publicEndpoint,

pkg/aux_/store/memstore/dss_test.go

@@ -15,7 +15,7 @@ var fakeClock = clockwork.NewFakeClock()
 
 func TestSaveOwnMetadataRoundTrip(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
@@ -35,7 +35,7 @@ func TestSaveOwnMetadataRoundTrip(t *testing.T) {
 
 func TestSaveOwnMetadataUpsert(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://old.example.com"))
@@ -50,7 +50,7 @@ func TestSaveOwnMetadataUpsert(t *testing.T) {
 
 func TestGetDSSMetadataPicksLatestHeartbeat(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
@@ -71,7 +71,7 @@ func TestGetDSSMetadataPicksLatestHeartbeat(t *testing.T) {
 
 func TestGetDSSMetadataUpdatesHeartbeatPerSource(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))

pkg/aux_/store/memstore/snapshot_test.go

@@ -17,7 +17,7 @@ import (
 
 func TestSnapshotRoundTrip(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	src := newRepo()
 	require.NoError(t, src.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
 	ts := time.Now().UTC()
@@ -40,7 +40,7 @@ func TestSnapshotRoundTrip(t *testing.T) {
 
 func TestRestoreFromSnapshotReplacesState(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	src := newRepo()
 	require.NoError(t, src.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
 	data, err := src.GetSnapshot()

pkg/aux_/store/memstore/store_test.go

@@ -10,7 +10,7 @@ import (
 
 func TestCheckpointRestore(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 
 	r := newRepo()
 
@@ -34,7 +34,7 @@ func TestCheckpointRestore(t *testing.T) {
 
 func TestCheckpointIsolatesUpsert(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	r := newRepo()
 
 	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://old.example.com"))

pkg/aux_/store/raftstore/store.go

@@ -24,8 +24,7 @@ const (
 // repo is a full implementation of aux_.repos.Repository for Raft-based storage.
 type repo struct {
 	consensus *consensus.Consensus
-	memStore  *memstore.Store[repos.Repository]
-	memRepo   repos.Repository
+	*memstore.Store[repos.Repository]
 }
 
 func Init(ctx context.Context, logger *zap.Logger, locality string) (*raftstore.Store[repos.Repository], error) {
@@ -39,7 +38,7 @@ func Init(ctx context.Context, logger *zap.Logger, locality string) (*raftstore.
 		return nil, stacktrace.Propagate(err, "failed to initialize aux memstore")
 	}
 
-	r := &repo{memStore: memStore, memRepo: memStore.GetRepo()}
+	r := &repo{Store: memStore}
 	store, err := raftstore.Init(ctx, logger.With(zap.String("service", "aux_")), locality, params, r, nil)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "failed to initialize aux raftstore")
@@ -52,14 +51,6 @@ func Init(ctx context.Context, logger *zap.Logger, locality string) (*raftstore.
 
 func (r *repo) GetRepo() repos.Repository { return r }
 
-func (r *repo) GetSnapshot() ([]byte, error) {
-	return r.memStore.GetSnapshot()
-}
-
-func (r *repo) RestoreFromSnapshot(data []byte) error {
-	return r.memStore.RestoreFromSnapshot(data)
-}
-
 func (r *repo) Apply(ctx context.Context, proposal consensus.Proposal) (any, error) {
 	switch proposal.RequestType {
 	case saveOwnMetadata:
@@ -68,18 +59,18 @@ func (r *repo) Apply(ctx context.Context, proposal consensus.Proposal) (any, err
 			return nil, stacktrace.Propagate(err, "failed to unmarshal %s payload", saveOwnMetadata)
 		}
 
-		return nil, r.memRepo.SaveOwnMetadata(ctx, payload.Locality, payload.PublicEndpoint)
+		return nil, r.Store.GetRepo().SaveOwnMetadata(ctx, payload.Locality, payload.PublicEndpoint)
 
 	case getDSSMetadata:
-		return r.memRepo.GetDSSMetadata(ctx)
+		return r.Store.GetRepo().GetDSSMetadata(ctx)
 
 	case recordHeartbeat:
 		var heartbeat auxmodels.Heartbeat
 		if err := json.Unmarshal(proposal.Value, &heartbeat); err != nil {
 			return nil, stacktrace.Propagate(err, "failed to unmarshal %s payload", recordHeartbeat)
 		}
 
-		return nil, r.memRepo.RecordHeartbeat(ctx, heartbeat)
+		return nil, r.Store.GetRepo().RecordHeartbeat(ctx, heartbeat)
 
 	default:
 		return nil, stacktrace.NewError("unknown request type: %q", proposal.RequestType)

pkg/errors/errors.go

@@ -16,10 +16,6 @@ const (
 	// larger than the max area allowed. See geo/s2.go.
 	AreaTooLarge = stacktrace.ErrorCode(iota)
 
-	// MissingOVNs is the error to signal that an AirspaceConflictResponse should
-	// be returned rather than the standard error response.
-	MissingOVNs
-
 	// AlreadyExists is used when attempting to create a resource that already
 	// exists.
 	AlreadyExists

pkg/locality/locality.go

@@ -7,30 +7,30 @@ import (
 	"github.com/interuss/stacktrace"
 )
 
-type localityKey struct{}
+type key struct{}
 
-// MustGetRequestLocality returns the request locality from the context and panics if it is not
+// MustFromContext returns the request locality from the context and panics if it is not
 // present, which is a programming error.
-func MustGetRequestLocality(ctx context.Context) string {
-	locality, ok := ctx.Value(localityKey{}).(string)
+func MustFromContext(ctx context.Context) string {
+	locality, ok := ctx.Value(key{}).(string)
 	if !ok {
 		panic(stacktrace.NewError("request locality not present in context"))
 	}
 
 	return locality
 }
 
-// WithRequestLocality returns a new context with the given locality.
-func WithRequestLocality(ctx context.Context, locality string) context.Context {
-	return context.WithValue(ctx, localityKey{}, locality)
+// NewContext returns a new context with the given locality.
+func NewContext(ctx context.Context, locality string) context.Context {
+	return context.WithValue(ctx, key{}, locality)
 }
 
-// LocalityMiddleware is an HTTP middleware that stamps each incoming request with this
+// Middleware is an HTTP middleware that stamps each incoming request with this
 // DSS instance's locality so that locality-dependent operations execute deterministically across nodes.
-func LocalityMiddleware(locality string) func(http.Handler) http.Handler {
+func Middleware(locality string) func(http.Handler) http.Handler {
 	return func(next http.Handler) http.Handler {
 		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			next.ServeHTTP(w, r.WithContext(WithRequestLocality(r.Context(), locality)))
+			next.ServeHTTP(w, r.WithContext(NewContext(r.Context(), locality)))
 		})
 	}
 }

pkg/raftstore/consensus/proposal.go

@@ -35,7 +35,7 @@ type Proposal struct {
 }
 
 func (c *Consensus) newProposal(ctx context.Context, requestType RequestType, value []byte, readOnly bool) Proposal {
-	timestamp := timestamp.MustGetRequestTimestamp(ctx)
+	timestamp := timestamp.MustFromContext(ctx)
 	seed := random.MustFromContext(ctx)
 
 	return Proposal{

pkg/raftstore/store.go

@@ -5,6 +5,7 @@ import (
 
 	"github.com/interuss/dss/pkg/locality"
 	"github.com/interuss/dss/pkg/logging"
+	"github.com/interuss/dss/pkg/memstore"
 	"github.com/interuss/dss/pkg/raftstore/consensus"
 	raftparams "github.com/interuss/dss/pkg/raftstore/params"
 	"github.com/interuss/dss/pkg/random"
@@ -15,19 +16,12 @@ import (
 )
 
 type RaftRepo[R any] interface {
-	GetRepo() R
+	memstore.MemRepo[R]
+
 	// Apply is called on every committed entry. The proposal must be applied atomically.
 	// The any return mirrors store.OperationHandler.Execute: different requests yield different
 	// concrete result types. Callers recover the type via store.TransactWithResult.
 	Apply(ctx context.Context, proposal consensus.Proposal) (any, error)
-
-	// GetSnapshot returns a serialized view of current state, suitable
-	// for restoring via RestoreFromSnapshot.
-	GetSnapshot() ([]byte, error)
-
-	// RestoreFromSnapshot replaces all state with the snapshot in data.
-	// data is always the output of a prior GetSnapshot.
-	RestoreFromSnapshot(data []byte) error
 }
 
 type Store[R any] struct {
@@ -116,10 +110,15 @@ func (s *Store[R]) processCommits(ctx context.Context, commitCh <-chan consensus
 				continue
 			}
 
-			proposalCtx := timestamp.WithRequestTimestamp(ctx, commit.Prop.Timestamp)
-			proposalCtx = locality.WithRequestLocality(proposalCtx, commit.Prop.Locality)
+			proposalCtx := timestamp.NewContext(ctx, commit.Prop.Timestamp)
+			proposalCtx = locality.NewContext(proposalCtx, commit.Prop.Locality)
 			proposalCtx = random.NewContext(proposalCtx, commit.Prop.Seed)
+			s.raftRepo.Checkpoint()
 			result, err := s.raftRepo.Apply(proposalCtx, commit.Prop)
+			if err != nil {
+				s.logger.Warn("failed to apply proposal, rolling back", zap.String("proposal_id", commit.Prop.ID), zap.String("proposal_type", string(commit.Prop.RequestType)), zap.Error(err))
+				s.raftRepo.Restore()
+			}
 			commit.Done <- consensus.ProposalResult{Result: result, Error: err}
 		}
 	}

pkg/rid/actions/registry.go → pkg/rid/operations/registry.go

@@ -1,4 +1,4 @@
-package actions
+package operations
 
 import (
 	"github.com/interuss/dss/pkg/rid/repos"

pkg/rid/actions/subscription.go → pkg/rid/operations/subscription.go

@@ -1,4 +1,4 @@
-package actions
+package operations
 
 import (
 	"context"
@@ -16,16 +16,16 @@ func init() {
 	Registry[ridv1.DeleteSubscriptionOperationID] = dssstore.OperationHandler[repos.Repository]{
 		Encode:  dssstore.EncodeJSON,
 		Decode:  dssstore.DecodeJSON[*ridv1.DeleteSubscriptionRequest],
-		Execute: ExecuteDeleteSubscription,
+		Execute: executeDeleteSubscription,
 	}
 	Registry[ridv2.DeleteSubscriptionOperationID] = dssstore.OperationHandler[repos.Repository]{
 		Encode:  dssstore.EncodeJSON,
 		Decode:  dssstore.DecodeJSON[*ridv2.DeleteSubscriptionRequest],
-		Execute: ExecuteDeleteSubscription,
+		Execute: executeDeleteSubscription,
 	}
 }
 
-func ExecuteDeleteSubscription(ctx context.Context, repo repos.Repository, request dssstore.OperationRequest) (any, error) {
+func executeDeleteSubscription(ctx context.Context, repo repos.Repository, request dssstore.OperationRequest) (any, error) {
 	var (
 		rawID      string
 		rawVersion string

pkg/rid/store/memstore/identification_service_area.go

@@ -61,7 +61,7 @@ func (r *repo) InsertISA(ctx context.Context, isa *ridmodels.IdentificationServi
 		return nil, stacktrace.NewError("ISA with id %s already exists", isa.ID)
 	}
 
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	rec := isaRecordFromModel(isa, now)
 	r.state.ISAs[isa.ID] = rec
@@ -77,7 +77,7 @@ func (r *repo) UpdateISA(ctx context.Context, isa *ridmodels.IdentificationServi
 		return nil, nil
 	}
 
-	now := timestamp.MustGetRequestTimestamp(ctx)
+	now := timestamp.MustFromContext(ctx)
 
 	rec := isaRecordFromModel(isa, now)
 	rec.Owner = prev.Owner // It's not possible to update the owner of an ISA, this ensure it's to changed to a new value.

pkg/rid/store/memstore/identification_service_area_test.go

@@ -34,7 +34,7 @@ var (
 
 func TestStoreSearchISAs(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	cells := s2.CellUnion{
 		s2.CellID(17106221850767130624),
 		s2.CellID(17106221885126868992),
@@ -137,7 +137,7 @@ func TestStoreSearchISAs(t *testing.T) {
 
 func TestBadVersion(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	saOut1, err := repo.InsertISA(ctx, serviceArea)
@@ -159,7 +159,7 @@ func TestBadVersion(t *testing.T) {
 
 func TestStoreExpiredISA(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	saOut, err := repo.InsertISA(ctx, serviceArea)
@@ -194,7 +194,7 @@ func TestStoreExpiredISA(t *testing.T) {
 
 func TestStoreDeleteISAs(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	// Insert the ISA.
@@ -215,7 +215,7 @@ func TestStoreDeleteISAs(t *testing.T) {
 
 func TestStoreISAWithNoGeoData(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	endTime := fakeClock.Now().Add(24 * time.Hour)
@@ -230,7 +230,7 @@ func TestStoreISAWithNoGeoData(t *testing.T) {
 
 func TestListExpiredISAs(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	// Insert ISA with endtime 1 day from now
@@ -261,7 +261,7 @@ func TestListExpiredISAs(t *testing.T) {
 
 func TestListExpiredISAsWithEmptyWriter(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	// Insert ISA with endtime 1 day from now
@@ -294,7 +294,7 @@ func TestListExpiredISAsWithEmptyWriter(t *testing.T) {
 
 func TestStoreCountISAs(t *testing.T) {
 	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
+	ctx = timestamp.NewContext(ctx, fakeClock.Now())
 	repo := setUpStore(t)
 
 	// Insert the ISA.