[raft/scd] Extract update and create constraint

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
 	}
@@ -151,9 +152,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
 	}
@@ -350,7 +351,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")
 		}
@@ -365,6 +366,7 @@ func RunHTTPServer(ctx context.Context, ctxCanceler func(), address, locality st
 	handler = http.TimeoutHandler(handler, *timeout, "request timeout")
 	handler = logging.HTTPMiddleware(logger, *dumpRequests, handler)
 	handler = timestamp.RequestTimestampMiddleware(handler)
+	handler = requestlocality.LocalityMiddleware(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_/store/raftstore/store.go

@@ -28,7 +28,7 @@ type repo struct {
 	memRepo   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")
@@ -40,7 +40,7 @@ func Init(ctx context.Context, logger *zap.Logger) (*raftstore.Store[repos.Repos
 	}
 
 	r := &repo{memStore: memStore, memRepo: memStore.GetRepo()}
-	store, err := raftstore.Init(ctx, logger.With(zap.String("service", "aux_")), params, r, nil)
+	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")
 	}

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 localityKey struct{}
+
+// MustGetRequestLocality 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)
+	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)
+}
+
+// LocalityMiddleware 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 {
+	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)))
+		})
+	}
+}

pkg/raftstore/consensus/consensus.go

@@ -24,8 +24,9 @@ import (
 type Consensus struct {
 	logger *zap.Logger
 
-	nodeID uint64
-	node   raft.Node
+	locality string
+	nodeID   uint64
+	node     raft.Node
 
 	transport *rafthttp.Transport
 	server    *http.Server
@@ -45,7 +46,7 @@ type Consensus struct {
 	appliedIndex  uint64
 }
 
-func NewConsensus(ctx context.Context, logger *zap.Logger, connectParams params.ConnectParameters, provider snapshotProvider, commitC chan<- EntryCommit) (*Consensus, error) {
+func NewConsensus(ctx context.Context, logger *zap.Logger, locality string, connectParams params.ConnectParameters, provider snapshotProvider, commitC chan<- EntryCommit) (*Consensus, error) {
 	storage, old, err := newStorage(ctx, logger.With(zap.String("component", "storage")), connectParams.DataDir, connectParams.NodeID, provider, connectParams.SnapshotCatchupEntries)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "failed to initialize storage")
@@ -74,11 +75,11 @@ func NewConsensus(ctx context.Context, logger *zap.Logger, connectParams params.
 	consensus := &Consensus{
 		logger: logging.WithValuesFromContext(ctx, logger),
 
-		nodeID: connectParams.NodeID,
-		node:   node,
-
-		storage: storage,
-		commitC: commitC,
+		nodeID:   connectParams.NodeID,
+		node:     node,
+		locality: locality,
+		storage:  storage,
+		commitC:  commitC,
 
 		shutdownTimeout: 2 * connectParams.ElectionInterval(),
 		serverErrC:      make(chan error, 1),

pkg/raftstore/consensus/proposal.go

@@ -18,6 +18,7 @@ type EntryCommit struct {
 
 type Proposal struct {
 	ID          string      `json:"id"`
+	Locality    string      `json:"locality"`
 	NodeID      uint64      `json:"node_id"`
 	Timestamp   time.Time   `json:"timestamp"`
 	RequestType RequestType `json:"request_type"`
@@ -34,6 +35,7 @@ func (c *Consensus) newProposal(ctx context.Context, requestType RequestType, va
 
 	return Proposal{
 		ID:          uuid.NewString(),
+		Locality:    c.locality,
 		NodeID:      c.nodeID,
 		Timestamp:   timestamp.UTC(),
 		RequestType: requestType,

pkg/raftstore/store.go

@@ -3,6 +3,7 @@ package raftstore
 import (
 	"context"
 
+	"github.com/interuss/dss/pkg/locality"
 	"github.com/interuss/dss/pkg/logging"
 	"github.com/interuss/dss/pkg/raftstore/consensus"
 	raftparams "github.com/interuss/dss/pkg/raftstore/params"
@@ -40,7 +41,7 @@ type Store[R any] struct {
 	done chan struct{}
 }
 
-func Init[R any](ctx context.Context, logger *zap.Logger, params raftparams.ConnectParameters, r RaftRepo[R], registry map[string]store.OperationHandler[R]) (*Store[R], error) {
+func Init[R any](ctx context.Context, logger *zap.Logger, locality string, params raftparams.ConnectParameters, r RaftRepo[R], registry map[string]store.OperationHandler[R]) (*Store[R], error) {
 	ctx, cancel := context.WithCancel(ctx)
 
 	store := &Store[R]{
@@ -56,7 +57,7 @@ func Init[R any](ctx context.Context, logger *zap.Logger, params raftparams.Conn
 		store.processCommits(ctx, commitC)
 	}()
 
-	consensusInstance, err := consensus.NewConsensus(ctx, logger, params, r.GetSnapshot, commitC)
+	consensusInstance, err := consensus.NewConsensus(ctx, logger, locality, params, r.GetSnapshot, commitC)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "failed to initialize consensus")
 	}
@@ -115,6 +116,7 @@ func (s *Store[R]) processCommits(ctx context.Context, commitCh <-chan consensus
 			}
 
 			proposalCtx := timestamp.WithRequestTimestamp(ctx, commit.Prop.Timestamp)
+			proposalCtx = locality.WithRequestLocality(proposalCtx, commit.Prop.Locality)
 			result, err := s.raftRepo.Apply(proposalCtx, commit.Prop)
 			commit.Done <- consensus.ProposalResult{Result: result, Error: err}
 		}

pkg/rid/store/raftstore/store.go

@@ -16,12 +16,13 @@ import (
 // repo is a full implementation of rid.repos.Repository for Raft-based storage.
 type repo struct{}
 
-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 := ridraftparams.GetConnectParameters()
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "failed to get rid raft parameters")
 	}
-	return raftstore.Init(ctx, logger.With(zap.String("service", "rid")), params, &repo{}, actions.Registry)
+
+	return raftstore.Init(ctx, logger.With(zap.String("service", "rid")), locality, params, &repo{}, actions.Registry)
 }
 
 func (r *repo) GetRepo() repos.Repository { return r }

pkg/rid/store/store.go

@@ -18,13 +18,13 @@ import (
 type Store = dssstore.Store[repos.Repository]
 
 // Init selects and initializes the rid 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) {
 	storeType := params.GetStoreParameters().StoreType
 	switch storeType {
 	case params.SQLStoreType:
 		return ridsqlstore.Init(ctx, logger, withCheckCron)
 	case params.RaftStoreType:
-		return ridraftstore.Init(ctx, logger)
+		return ridraftstore.Init(ctx, logger, locality)
 	case params.MemStoreType:
 		return ridmemstore.Init(ctx, logger)
 	default:

pkg/scd/actions/availability.go

@@ -0,0 +1,109 @@
+package actions
+
+import (
+	"context"
+
+	restapi "github.com/interuss/dss/pkg/api/scdv1"
+	dsserr "github.com/interuss/dss/pkg/errors"
+	dssmodels "github.com/interuss/dss/pkg/models"
+	scdmodels "github.com/interuss/dss/pkg/scd/models"
+	"github.com/interuss/dss/pkg/scd/repos"
+	dssstore "github.com/interuss/dss/pkg/store"
+	"github.com/interuss/stacktrace"
+	"github.com/jackc/pgx/v5"
+)
+
+func init() {
+	Registry[restapi.GetUssAvailabilityOperationID] = dssstore.OperationHandler[repos.Repository]{
+		Encode:     dssstore.EncodeJSON,
+		Decode:     dssstore.DecodeJSON[*restapi.GetUssAvailabilityRequest],
+		Execute:    ExecuteGetUssAvailability,
+		IsReadOnly: true,
+	}
+	Registry[restapi.SetUssAvailabilityOperationID] = dssstore.OperationHandler[repos.Repository]{
+		Encode:  dssstore.EncodeJSON,
+		Decode:  dssstore.DecodeJSON[*restapi.SetUssAvailabilityRequest],
+		Execute: ExecuteSetUssAvailability,
+	}
+}
+
+func GetDefaultAvailabilityResponse(id dssmodels.Manager) *restapi.UssAvailabilityStatusResponse {
+	return &restapi.UssAvailabilityStatusResponse{
+		Status: restapi.UssAvailabilityStatus{
+			Availability: restapi.UssAvailabilityState_Unknown,
+			Uss:          id.String()},
+		Version: "",
+	}
+}
+
+func ExecuteGetUssAvailability(ctx context.Context, repo repos.Repository, request dssstore.OperationRequest) (any, error) {
+	req, ok := request.(*restapi.GetUssAvailabilityRequest)
+	if !ok {
+		return nil, stacktrace.NewError("unexpected request type %T for operation %q", request, restapi.GetUssAvailabilityOperationID)
+	}
+
+	id := dssmodels.ManagerFromString(req.UssId)
+
+	// Get USS availability from Store
+	ussa, err := repo.GetUssAvailability(ctx, id)
+	if err != nil && err != pgx.ErrNoRows {
+		return nil, stacktrace.Propagate(err, "Could not get USS availability from repo")
+	}
+	if ussa == nil {
+		// Return default availability status "Unknown"
+		return GetDefaultAvailabilityResponse(id), nil
+	}
+
+	return &restapi.UssAvailabilityStatusResponse{
+		Status:  *ussa.ToRest(),
+		Version: ussa.Version.String(),
+	}, nil
+}
+
+func ExecuteSetUssAvailability(ctx context.Context, repo repos.Repository, request dssstore.OperationRequest) (any, error) {
+	req, ok := request.(*restapi.SetUssAvailabilityRequest)
+	if !ok {
+		return nil, stacktrace.NewError("unexpected request type %T for operation %q", request, restapi.SetUssAvailabilityOperationID)
+	}
+
+	// Retrieve USS availability status from request params
+	availability, err := scdmodels.UssAvailabilityStateFromRest(req.Body.Availability)
+	if err != nil {
+		return nil, stacktrace.PropagateWithCode(err, dsserr.BadRequest, "Invalid availability state")
+	}
+	id := dssmodels.ManagerFromString(req.UssId)
+	version := scdmodels.OVN(req.Body.OldVersion)
+	ussareq := &scdmodels.UssAvailabilityStatus{
+		Uss:          id,
+		Availability: availability,
+	}
+
+	old, err := repo.GetUssAvailability(ctx, id)
+	if err != nil && err != pgx.ErrNoRows {
+		return nil, stacktrace.Propagate(err, "Could not get USS availability from repo")
+	}
+	switch {
+	case old == nil && !version.Empty():
+		// The user wants set a new availability status but it already exists.
+		return nil, stacktrace.NewErrorWithCode(dsserr.AlreadyExists, "availability for USS %s already exists", id.String())
+	case old != nil && old.Version != version:
+		// The user wants to update an availability status but the version doesn't match.
+		return nil, stacktrace.Propagate(
+			stacktrace.NewErrorWithCode(dsserr.VersionMismatch, "USS availability version %s is not current", version),
+			"Current version is %s but client specified version %s", old.Version, version)
+	}
+
+	// Upsert the USS availability
+	ussa, err := repo.UpsertUssAvailability(ctx, ussareq)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "Could not upsert USS Availability into repo")
+	}
+	if ussa == nil {
+		return nil, stacktrace.NewError("UpsertUssAvailability returned no USS availability for ID: %s", id)
+	}
+
+	return &restapi.UssAvailabilityStatusResponse{
+		Status:  *ussa.ToRest(),
+		Version: ussa.Version.String(),
+	}, nil
+}