[raft/memstore] Add aux memstore

cmds/core-service/main.go

@@ -32,6 +32,7 @@ import (
 	scds "github.com/interuss/dss/pkg/scd/store"
 	"github.com/interuss/dss/pkg/store"
 	"github.com/interuss/dss/pkg/store/params"
+	"github.com/interuss/dss/pkg/timestamp"
 	"github.com/interuss/dss/pkg/version"
 	"github.com/interuss/dss/pkg/versioning"
 	"github.com/interuss/stacktrace"
@@ -340,6 +341,7 @@ 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.Middleware(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

@@ -100,7 +100,7 @@ func evict(cmd *cobra.Command, _ []string) error {
 		}
 		return nil
 	}
-	if err = scdStore.Transact(ctx, scdAction); err != nil {
+	if _, err = scdStore.Transact(ctx, "", nil, scdAction); err != nil {
 		return fmt.Errorf("failed to execute SCD transaction: %w", err)
 	}
 
@@ -145,7 +145,7 @@ func evict(cmd *cobra.Command, _ []string) error {
 
 		return nil
 	}
-	if err = ridStore.Transact(ctx, ridAction); err != nil {
+	if _, err = ridStore.Transact(ctx, "", nil, ridAction); err != nil {
 		return fmt.Errorf("failed to execute RID transaction: %w", err)
 	}
 

pkg/aux_/store/memstore/doc.go

@@ -0,0 +1,3 @@
+// Package aux_.store.memstore provides a full implementation of store.Store[aux_.repos.Repository]
+// storing data in memory. It is meant to be used by raftstore.
+package memstore

pkg/aux_/store/memstore/dss.go

@@ -0,0 +1,86 @@
+package memstore
+
+import (
+	"context"
+	"database/sql"
+	"time"
+
+	auxmodels "github.com/interuss/dss/pkg/aux_/models"
+	dsserr "github.com/interuss/dss/pkg/errors"
+	"github.com/interuss/stacktrace"
+)
+
+func (r *repo) SaveOwnMetadata(_ context.Context, locality string, publicEndpoint string) error {
+	if locality == "" {
+		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Locality not set")
+	}
+	if publicEndpoint == "" {
+		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Public endpoint not set")
+	}
+
+	r.participants[locality] = &participant{
+		publicEndpoint: publicEndpoint,
+		updatedAt:      time.Now(),
+	}
+	return nil
+}
+
+func (r *repo) GetDSSMetadata(_ context.Context) ([]*auxmodels.DSSMetadata, error) {
+	metadata := make([]*auxmodels.DSSMetadata, 0, len(r.participants))
+	for locality, p := range r.participants {
+		updatedAt := p.updatedAt
+		m := &auxmodels.DSSMetadata{
+			Locality:       locality,
+			PublicEndpoint: p.publicEndpoint,
+			UpdatedAt:      &updatedAt,
+		}
+
+		// Find the latest heartbeat across all sources for this locality.
+		var latest auxmodels.Heartbeat
+		found := false
+		for key, hb := range r.heartbeats {
+			if key.locality != locality {
+				continue
+			}
+			if !found || hb.Timestamp.After(*latest.Timestamp) {
+				latest = hb
+				found = true
+			}
+		}
+
+		if found {
+			m.LatestTimestamp.Source = sql.NullString{String: latest.Source, Valid: true}
+			m.LatestTimestamp.Timestamp = latest.Timestamp
+			m.LatestTimestamp.NextHeartbeatExpectedBefore = latest.NextHeartbeatExpectedBefore
+			m.LatestTimestamp.Reporter = sql.NullString{String: latest.Reporter, Valid: true}
+		}
+
+		metadata = append(metadata, m)
+	}
+	return metadata, nil
+}
+
+func (r *repo) RecordHeartbeat(_ context.Context, heartbeat auxmodels.Heartbeat) error {
+	if heartbeat.Locality == "" {
+		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Locality not set")
+	}
+	if heartbeat.Source == "" {
+		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Source not set")
+	}
+
+	if heartbeat.Timestamp == nil {
+		now := time.Now()
+		heartbeat.Timestamp = &now
+	}
+
+	if heartbeat.NextHeartbeatExpectedBefore != nil && heartbeat.NextHeartbeatExpectedBefore.Before(*heartbeat.Timestamp) {
+		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Cannot expect the timestamp of the next heartbeat before the timestamp of the new heartbeat")
+	}
+
+	r.heartbeats[heartbeatKey{locality: heartbeat.Locality, source: heartbeat.Source}] = heartbeat
+	return nil
+}
+
+func (r *repo) GetDSSAirspaceRepresentationID(_ context.Context) (string, error) {
+	return "", stacktrace.NewErrorWithCode(dsserr.NotImplemented, "GetDSSAirspaceRepresentationID not implementable for memstore")
+}

pkg/aux_/store/memstore/dss_test.go

@@ -0,0 +1,125 @@
+package memstore
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	auxmodels "github.com/interuss/dss/pkg/aux_/models"
+	dsserr "github.com/interuss/dss/pkg/errors"
+	"github.com/interuss/stacktrace"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSaveOwnMetadataValidation(t *testing.T) {
+	ctx := context.Background()
+	r := newRepo()
+
+	require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.SaveOwnMetadata(ctx, "", "https://example.com")))
+	require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.SaveOwnMetadata(ctx, "dss-1", "")))
+}
+
+func TestSaveOwnMetadataRoundTrip(t *testing.T) {
+	ctx := context.Background()
+	r := newRepo()
+
+	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
+
+	md, err := r.GetDSSMetadata(ctx)
+	require.NoError(t, err)
+
+	require.Len(t, md, 1)
+	require.Equal(t, "dss-1", md[0].Locality)
+	require.Equal(t, "https://example.com", md[0].PublicEndpoint)
+	require.NotNil(t, md[0].UpdatedAt)
+
+	// No heartbeat recorded yet.
+	require.False(t, md[0].LatestTimestamp.Source.Valid)
+	require.Nil(t, md[0].LatestTimestamp.Timestamp)
+}
+
+func TestSaveOwnMetadataUpsert(t *testing.T) {
+	ctx := context.Background()
+	r := newRepo()
+
+	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://old.example.com"))
+	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://new.example.com"))
+
+	md, err := r.GetDSSMetadata(ctx)
+	require.NoError(t, err)
+
+	require.Len(t, md, 1)
+	require.Equal(t, "https://new.example.com", md[0].PublicEndpoint)
+}
+
+func TestRecordHeartbeatValidation(t *testing.T) {
+	ctx := context.Background()
+	r := newRepo()
+
+	require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Source: "source1"})))
+	require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1"})))
+
+	ts := time.Now()
+	before := ts.Add(-time.Minute)
+	err := r.RecordHeartbeat(ctx, auxmodels.Heartbeat{
+		Locality:                    "dss-1",
+		Source:                      "source1",
+		Timestamp:                   &ts,
+		NextHeartbeatExpectedBefore: &before,
+	})
+
+	require.Equal(t, dsserr.BadRequest, stacktrace.GetCode(err))
+}
+
+func TestRecordHeartbeatDefaultsTimestamp(t *testing.T) {
+	ctx := context.Background()
+	r := newRepo()
+
+	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
+	require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1"}))
+
+	md, err := r.GetDSSMetadata(ctx)
+	require.NoError(t, err)
+
+	require.Len(t, md, 1)
+	require.True(t, md[0].LatestTimestamp.Source.Valid)
+	require.NotNil(t, md[0].LatestTimestamp.Timestamp)
+}
+
+func TestGetDSSMetadataPicksLatestHeartbeat(t *testing.T) {
+	ctx := context.Background()
+	r := newRepo()
+
+	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
+
+	older := time.Now().Add(-time.Hour)
+	newer := time.Now()
+	require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1", Timestamp: &older, Reporter: "uss1"}))
+	require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source2", Timestamp: &newer, Reporter: "uss2"}))
+
+	md, err := r.GetDSSMetadata(ctx)
+	require.NoError(t, err)
+
+	require.Len(t, md, 1)
+	require.True(t, md[0].LatestTimestamp.Timestamp.Equal(newer))
+	require.Equal(t, "source2", md[0].LatestTimestamp.Source.String)
+	require.Equal(t, "uss2", md[0].LatestTimestamp.Reporter.String)
+}
+
+func TestGetDSSMetadataUpdatesHeartbeatPerSource(t *testing.T) {
+	ctx := context.Background()
+	r := newRepo()
+
+	require.NoError(t, r.SaveOwnMetadata(ctx, "dss-1", "https://example.com"))
+
+	first := time.Now().Add(-time.Hour)
+	second := time.Now()
+	require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1", Timestamp: &first}))
+	require.NoError(t, r.RecordHeartbeat(ctx, auxmodels.Heartbeat{Locality: "dss-1", Source: "source1", Timestamp: &second}))
+
+	md, err := r.GetDSSMetadata(ctx)
+	require.NoError(t, err)
+
+	require.Len(t, md, 1)
+	require.True(t, md[0].LatestTimestamp.Timestamp.Equal(second))
+}

pkg/aux_/store/memstore/store.go

@@ -0,0 +1,42 @@
+package memstore
+
+import (
+	"context"
+	"time"
+
+	auxmodels "github.com/interuss/dss/pkg/aux_/models"
+	"github.com/interuss/dss/pkg/aux_/repos"
+	"github.com/interuss/dss/pkg/memstore"
+	"go.uber.org/zap"
+)
+
+// repo is a full implementation of aux_.repos.Repository for memory-based storage.
+type repo struct {
+	// participants holds pool participants metadata, keyed by locality.
+	participants map[string]*participant
+	// heartbeats holds the latest heartbeat per (locality, source).
+	heartbeats map[heartbeatKey]auxmodels.Heartbeat
+}
+
+type participant struct {
+	publicEndpoint string
+	updatedAt      time.Time
+}
+
+type heartbeatKey struct {
+	locality string
+	source   string
+}
+
+func newRepo() *repo {
+	return &repo{
+		participants: map[string]*participant{},
+		heartbeats:   map[heartbeatKey]auxmodels.Heartbeat{},
+	}
+}
+
+func Init(ctx context.Context, logger *zap.Logger) (*memstore.Store[repos.Repository], error) {
+	return memstore.Init(ctx, logger, "aux_", newRepo())
+}
+
+func (r *repo) GetRepo() repos.Repository { return r }

pkg/aux_/store/raftstore/params/params.go

@@ -0,0 +1,29 @@
+package params
+
+import (
+	"flag"
+
+	raftparams "github.com/interuss/dss/pkg/raftstore/params"
+	"github.com/interuss/stacktrace"
+)
+
+const peersFlag = "aux_raft_peers"
+
+var peers string
+
+func init() {
+	flag.StringVar(&peers, peersFlag, "", `Comma-separated "nodeID=peerURL" pairs for the aux store, e.g. "1=http://node1:9031,2=http://node2:9031,3=http://node3:9031"`)
+}
+
+func GetConnectParameters() (raftparams.ConnectParameters, error) {
+	if peers == "" {
+		return raftparams.ConnectParameters{}, stacktrace.NewError("--%s is required", peersFlag)
+	}
+
+	p, err := raftparams.GetConnectParameters("aux")
+	if err != nil {
+		return raftparams.ConnectParameters{}, err
+	}
+	p.Peers = peers
+	return p, nil
+}

pkg/aux_/store/raftstore/store.go

@@ -4,13 +4,37 @@ import (
 	"context"
 
 	"github.com/interuss/dss/pkg/aux_/repos"
+	auxraftparams "github.com/interuss/dss/pkg/aux_/store/raftstore/params"
+	dsserr "github.com/interuss/dss/pkg/errors"
 	"github.com/interuss/dss/pkg/raftstore"
+	"github.com/interuss/dss/pkg/raftstore/consensus"
+	"github.com/interuss/stacktrace"
 	"go.uber.org/zap"
 )
 
 // repo is a full implementation of aux_.repos.Repository for Raft-based storage.
 type repo struct{}
 
 func Init(ctx context.Context, logger *zap.Logger) (*raftstore.Store[repos.Repository], error) {
-	return raftstore.Init[repos.Repository](ctx, logger, func() repos.Repository { return &repo{} })
+	params, err := auxraftparams.GetConnectParameters()
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to get aux raft parameters")
+	}
+	return raftstore.Init(ctx, logger, params, "aux", &repo{})
+}
+
+func (r *repo) GetRepo() repos.Repository { return r }
+
+func (r *repo) IsReadOnly(_ raftstore.RequestType) bool { return false }
+
+func (r *repo) GetSnapshot() ([]byte, error) {
+	return nil, stacktrace.NewErrorWithCode(dsserr.NotImplemented, "not implemented yet")
+}
+
+func (r *repo) RestoreFromSnapshot([]byte) error {
+	return stacktrace.NewErrorWithCode(dsserr.NotImplemented, "not implemented yet")
+}
+
+func (r *repo) Apply(_ context.Context, _ consensus.Proposal) (any, error) {
+	return nil, stacktrace.NewErrorWithCode(dsserr.NotImplemented, "not implemented yet")
 }

pkg/aux_/store/store.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/interuss/dss/pkg/aux_/repos"
+	auxmemstore "github.com/interuss/dss/pkg/aux_/store/memstore"
 	auxraftstore "github.com/interuss/dss/pkg/aux_/store/raftstore"
 	auxsqlstore "github.com/interuss/dss/pkg/aux_/store/sqlstore"
 	dssstore "github.com/interuss/dss/pkg/store"
@@ -24,6 +25,8 @@ func Init(ctx context.Context, logger *zap.Logger, withCheckCron bool) (Store, e
 		return auxsqlstore.Init(ctx, logger, withCheckCron)
 	case params.RaftStoreType:
 		return auxraftstore.Init(ctx, logger)
+	case params.MemStoreType:
+		return auxmemstore.Init(ctx, logger)
 	default:
 		return nil, stacktrace.NewError("Unsupported store type %q for aux", storeType)
 	}