[raft/rid] Implement ISA repo methods

cmds/core-service/main.go

@@ -24,6 +24,7 @@ import (
 	auxs "github.com/interuss/dss/pkg/aux_/store"
 	"github.com/interuss/dss/pkg/build"
 	dsserr "github.com/interuss/dss/pkg/errors"
+	requestlocality "github.com/interuss/dss/pkg/locality"
 	"github.com/interuss/dss/pkg/logging"
 	"github.com/interuss/dss/pkg/rid/application"
 	rid_v1 "github.com/interuss/dss/pkg/rid/server/v1"
@@ -100,7 +101,7 @@ func createAuxServer(ctx context.Context, locality string, publicEndpoint string
 		return nil, stacktrace.NewError("Public endpoint not set")
 	}
 
-	auxStore, err := auxs.Init(ctx, logger, true)
+	auxStore, err := auxs.Init(ctx, logger, true, locality)
 	if err != nil {
 		return nil, err
 	}
@@ -121,7 +122,7 @@ func createAuxServer(ctx context.Context, locality string, publicEndpoint string
 
 func createRIDServers(ctx context.Context, locality string, logger *zap.Logger) (*rid_v1.Server, *rid_v2.Server, error) {
 
-	ridStore, err := rids.Init(ctx, logger, true)
+	ridStore, err := rids.Init(ctx, logger, true, locality)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -153,9 +154,9 @@ func createRIDServers(ctx context.Context, locality string, logger *zap.Logger)
 		}, nil
 }
 
-func createSCDServer(ctx context.Context, logger *zap.Logger) (*scd.Server, error) {
+func createSCDServer(ctx context.Context, logger *zap.Logger, locality string) (*scd.Server, error) {
 
-	scdStore, err := scds.Init(ctx, logger, true)
+	scdStore, err := scds.Init(ctx, logger, true, locality)
 	if err != nil {
 		return nil, err
 	}
@@ -352,7 +353,7 @@ func RunHTTPServer(ctx context.Context, ctxCanceler func(), address, locality st
 
 	// Initialize strategic conflict detection
 	if *enableSCD {
-		scdV1Server, err = createSCDServer(ctx, logger)
+		scdV1Server, err = createSCDServer(ctx, logger, locality)
 		if err != nil {
 			return stacktrace.Propagate(err, "Failed to create strategic conflict detection server")
 		}
@@ -366,7 +367,8 @@ 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 = 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.

cmds/db-manager/cleanup/evict.go

@@ -55,12 +55,12 @@ func evict(cmd *cobra.Command, _ []string) error {
 
 	logger := logging.WithValuesFromContext(ctx, logging.Logger)
 
-	scdStore, err := scds.Init(ctx, logger, false)
+	scdStore, err := scds.Init(ctx, logger, false, *locality)
 	if err != nil {
 		return err
 	}
 
-	ridStore, err := rids.Init(ctx, logger, false)
+	ridStore, err := rids.Init(ctx, logger, false, *locality)
 	if err != nil {
 		return err
 	}

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,11 +24,10 @@ 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) (*raftstore.Store[repos.Repository], error) {
+func Init(ctx context.Context, logger *zap.Logger, locality string) (*raftstore.Store[repos.Repository], error) {
 	params, err := auxraftparams.GetConnectParameters()
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "failed to get aux raft parameters")
@@ -39,8 +38,8 @@ func Init(ctx context.Context, logger *zap.Logger) (*raftstore.Store[repos.Repos
 		return nil, stacktrace.Propagate(err, "failed to initialize aux memstore")
 	}
 
-	r := &repo{memStore: memStore, memRepo: memStore.GetRepo()}
-	store, err := raftstore.Init(ctx, logger.With(zap.String("service", "aux_")), params, r, nil)
+	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) (*raftstore.Store[repos.Repos
 
 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/aux_/store/store.go

@@ -19,12 +19,12 @@ import (
 type Store = dssstore.Store[repos.Repository]
 
 // Init selects and initializes the aux store backend.
-func Init(ctx context.Context, logger *zap.Logger, withCheckCron bool) (Store, error) {
+func Init(ctx context.Context, logger *zap.Logger, withCheckCron bool, locality string) (Store, error) {
 	switch storeType := params.GetStoreParameters().StoreType; storeType {
 	case params.SQLStoreType:
 		return auxsqlstore.Init(ctx, logger, withCheckCron)
 	case params.RaftStoreType:
-		return auxraftstore.Init(ctx, logger)
+		return auxraftstore.Init(ctx, logger, locality)
 	case params.MemStoreType:
 		return auxmemstore.Init(ctx, logger)
 	default:

pkg/locality/locality.go

@@ -0,0 +1,36 @@
+package locality
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/interuss/stacktrace"
+)
+
+type key struct{}
+
+// MustFromContext returns the request locality from the context and panics if it is not
+// present, which is a programming error.
+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
+}
+
+// NewContext returns a new context with the given locality.
+func NewContext(ctx context.Context, locality string) context.Context {
+	return context.WithValue(ctx, key{}, locality)
+}
+
+// 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 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(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.