[raft/rid] Implement ISA repo methods

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/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/models/geo.go

@@ -1,6 +1,7 @@
 package models
 
 import (
+	"encoding/json"
 	"time"
 
 	"github.com/golang/geo/s2"
@@ -46,6 +47,83 @@ type Volume3D struct {
 	Footprint Geometry
 }
 
+type Volume3DJSON struct {
+	AltitudeHi *float32      `json:"altitude_hi,omitempty"`
+	AltitudeLo *float32      `json:"altitude_lo,omitempty"`
+	Footprint  *geometryJSON `json:"footprint,omitempty"`
+}
+
+type geometryType string
+
+const (
+	circle  geometryType = "circle"
+	polygon geometryType = "polygon"
+	cells   geometryType = "cells"
+)
+
+// geometryJSON is a helper struct for marshaling and unmarshaling Geometry types to/from JSON.
+type geometryJSON struct {
+	Type    geometryType `json:"type"`
+	Polygon *GeoPolygon  `json:"polygon,omitempty"`
+	Circle  *GeoCircle   `json:"circle,omitempty"`
+	Cells   []s2.CellID  `json:"cells,omitempty"`
+}
+
+func (v Volume3D) MarshalJSON() ([]byte, error) {
+	w := Volume3DJSON{AltitudeHi: v.AltitudeHi, AltitudeLo: v.AltitudeLo}
+	if v.Footprint != nil {
+		switch f := v.Footprint.(type) {
+		case *GeoPolygon:
+			w.Footprint = &geometryJSON{Type: polygon, Polygon: f}
+
+		case *GeoCircle:
+			w.Footprint = &geometryJSON{Type: circle, Circle: f}
+
+		case precomputedCellGeometry:
+			cellsResult := make([]s2.CellID, 0, len(f))
+			for id := range f {
+				cellsResult = append(cellsResult, id)
+			}
+			w.Footprint = &geometryJSON{Type: cells, Cells: cellsResult}
+
+		default:
+			return nil, stacktrace.NewError("Volume3D: unsupported Footprint type %T for JSON marshaling", v.Footprint)
+		}
+	}
+
+	return json.Marshal(w)
+}
+
+func (v *Volume3D) UnmarshalJSON(data []byte) error {
+	var w Volume3DJSON
+	if err := json.Unmarshal(data, &w); err != nil {
+		return err
+	}
+	v.AltitudeHi = w.AltitudeHi
+	v.AltitudeLo = w.AltitudeLo
+	if w.Footprint != nil {
+		switch w.Footprint.Type {
+		case polygon:
+			v.Footprint = w.Footprint.Polygon
+
+		case circle:
+			v.Footprint = w.Footprint.Circle
+
+		case cells:
+			pcg := make(precomputedCellGeometry, len(w.Footprint.Cells))
+			for _, id := range w.Footprint.Cells {
+				pcg[id] = struct{}{}
+			}
+
+			v.Footprint = pcg
+		default:
+			return stacktrace.NewError("Volume3D: unknown geometry type %q", w.Footprint.Type)
+		}
+	}
+
+	return nil
+}
+
 // Geometry models a geometry.
 type Geometry interface {
 	// CalculateCovering returns an s2 cell covering for a geometry.

pkg/models/models.go

@@ -1,6 +1,7 @@
 package models
 
 import (
+	"encoding/json"
 	"strconv"
 	"time"
 
@@ -175,3 +176,26 @@ func (v *Version) ToTimestamp() *time.Time {
 	}
 	return &v.t
 }
+
+func (v *Version) MarshalJSON() ([]byte, error) {
+	return json.Marshal(v.String())
+}
+
+func (v *Version) UnmarshalJSON(data []byte) error {
+	var s string
+	if err := json.Unmarshal(data, &s); err != nil {
+		return err
+	}
+
+	if s == "" {
+		return nil
+	}
+
+	parsed, err := VersionFromString(s)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to unmarshal version")
+	}
+
+	*v = *parsed
+	return nil
+}

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"