[raft/aux] Add aux Raftstore

cmds/core-service/main.go

@@ -92,6 +92,14 @@ func createKeyResolver() (auth.KeyResolver, error) {
 }
 
 func createAuxServer(ctx context.Context, locality string, publicEndpoint string, opts params.Options, logger *zap.Logger) (*aux.Server, error) {
+	if locality == "" {
+		return nil, stacktrace.NewError("Locality not set")
+	}
+
+	if publicEndpoint == "" {
+		return nil, stacktrace.NewError("Public endpoint not set")
+	}
+
 	auxStore, err := auxs.Init(ctx, logger, true)
 	if err != nil {
 		return nil, err

pkg/aux_/pool_participants.go

@@ -8,6 +8,7 @@ import (
 	restapi "github.com/interuss/dss/pkg/api/auxv1"
 	"github.com/interuss/dss/pkg/aux_/models"
 	dsserr "github.com/interuss/dss/pkg/errors"
+	"github.com/interuss/dss/pkg/timestamp"
 	"github.com/interuss/stacktrace"
 )
 
@@ -76,7 +77,7 @@ func (a *Server) PutDSSInstancesHeartbeat(ctx context.Context, req *restapi.PutD
 		return resp
 	}
 
-	if req.Source == nil {
+	if req.Source == nil || *req.Source == "" {
 		resp.Response400 = &restapi.ErrorResponse{Message: dsserr.Handle(ctx, stacktrace.Propagate(err, "Source not set"))}
 		return resp
 	}
@@ -94,6 +95,9 @@ func (a *Server) PutDSSInstancesHeartbeat(ctx context.Context, req *restapi.PutD
 			return resp
 		}
 		heartbeat.Timestamp = &ts
+	} else {
+		now := timestamp.MustGetRequestTimestamp(ctx)
+		heartbeat.Timestamp = &now
 	}
 
 	if req.NextHeartbeatExpectedBefore != nil {
@@ -103,6 +107,11 @@ func (a *Server) PutDSSInstancesHeartbeat(ctx context.Context, req *restapi.PutD
 			return resp
 		}
 		heartbeat.NextHeartbeatExpectedBefore = &ts
+
+		if heartbeat.NextHeartbeatExpectedBefore.Before(*heartbeat.Timestamp) {
+			resp.Response400 = &restapi.ErrorResponse{Message: dsserr.Handle(ctx, stacktrace.NewErrorWithCode(dsserr.BadRequest, "Cannot expect the timestamp of the next heartbeat before the timestamp of the new heartbeat"))}
+			return resp
+		}
 	}
 
 	err = repo.RecordHeartbeat(ctx, heartbeat)

pkg/aux_/repos/repos.go

@@ -14,12 +14,19 @@ type Misc interface {
 }
 
 // aux_.repos.DSSMetadata abstracts pool-information interactions with the DSS metadata repository.
+//
+// Implementations do not validate their arguments: callers are responsible for ensuring their correctness.
 type DSSMetadata interface {
-	// SaveOwnMetadata store our metadata into the pool participants
+	// SaveOwnMetadata stores our metadata into the pool participants.
+	// locality and publicEndpoint must both be non-empty.
 	SaveOwnMetadata(ctx context.Context, locality string, publicEndpoint string) error
 	// GetDSSMetadata returns all DSS metadata of pool participants
 	GetDSSMetadata(ctx context.Context) ([]*auxmodels.DSSMetadata, error)
-	// Record a new Timestamp
+	// RecordHeartbeat records a new heartbeat.
+	// hearthbeat.Locality and hearthbeat.Source must both be non-empty
+	// hearthbeat.Timestamp must be set
+	// if hearthbeat.NextHeartbeatExpectedBefore is set, it must not be before
+	// hearthbeat.Timestamp.
 	RecordHeartbeat(ctx context.Context, hearthbeat auxmodels.Heartbeat) error
 }
 

pkg/aux_/store/memstore/dss.go

@@ -11,13 +11,6 @@ import (
 )
 
 func (r *repo) SaveOwnMetadata(ctx context.Context, loc string, publicEndpoint string) error {
-	if loc == "" {
-		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Locality not set")
-	}
-	if publicEndpoint == "" {
-		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Public endpoint not set")
-	}
-
 	now := timestamp.MustGetRequestTimestamp(ctx)
 
 	r.state.Participants[locality(loc)] = &participant{
@@ -62,23 +55,7 @@ func (r *repo) GetDSSMetadata(_ context.Context) ([]*auxmodels.DSSMetadata, erro
 	return metadata, nil
 }
 
-func (r *repo) RecordHeartbeat(ctx context.Context, hb auxmodels.Heartbeat) error {
-	if hb.Locality == "" {
-		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Locality not set")
-	}
-	if hb.Source == "" {
-		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Source not set")
-	}
-
-	if hb.Timestamp == nil {
-		now := timestamp.MustGetRequestTimestamp(ctx).UTC()
-		hb.Timestamp = &now
-	}
-
-	if hb.NextHeartbeatExpectedBefore != nil && hb.NextHeartbeatExpectedBefore.Before(*hb.Timestamp) {
-		return stacktrace.NewErrorWithCode(dsserr.BadRequest, "Cannot expect the timestamp of the next heartbeat before the timestamp of the new heartbeat")
-	}
-
+func (r *repo) RecordHeartbeat(_ context.Context, hb auxmodels.Heartbeat) error {
 	r.state.Heartbeats[heartbeatKey{Locality: locality(hb.Locality), Source: hb.Source}] = &heartbeat{
 		Timestamp:                   hb.Timestamp,
 		NextHeartbeatExpectedBefore: hb.NextHeartbeatExpectedBefore,

pkg/aux_/store/memstore/dss_test.go

@@ -6,24 +6,13 @@ import (
 	"time"
 
 	auxmodels "github.com/interuss/dss/pkg/aux_/models"
-	dsserr "github.com/interuss/dss/pkg/errors"
 	"github.com/interuss/dss/pkg/timestamp"
-	"github.com/interuss/stacktrace"
 	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/require"
 )
 
 var fakeClock = clockwork.NewFakeClock()
 
-func TestSaveOwnMetadataValidation(t *testing.T) {
-	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
-	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()
 	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
@@ -59,42 +48,6 @@ func TestSaveOwnMetadataUpsert(t *testing.T) {
 	require.Equal(t, "https://new.example.com", md[0].PublicEndpoint)
 }
 
-func TestRecordHeartbeatValidation(t *testing.T) {
-	ctx := context.Background()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
-	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()
-	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())
-	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()
 	ctx = timestamp.WithRequestTimestamp(ctx, fakeClock.Now())

pkg/aux_/store/raftstore/dss.go

@@ -2,25 +2,62 @@ package raftstore
 
 import (
 	"context"
+	"encoding/json"
+	"strconv"
 
 	auxmodels "github.com/interuss/dss/pkg/aux_/models"
-	dsserr "github.com/interuss/dss/pkg/errors"
+	raftparams "github.com/interuss/dss/pkg/raftstore/params"
 	"github.com/interuss/stacktrace"
 )
 
-// SaveOwnMetadata returns nil instead of dsserr.NotImplemented because it is needed to allow the server to startup.
-func (r *repo) SaveOwnMetadata(_ context.Context, locality string, publicEndpoint string) error {
-	return nil
+type saveOwnMetadataPayload struct {
+	Locality       string `json:"locality"`
+	PublicEndpoint string `json:"public_endpoint"`
 }
 
-func (r *repo) GetDSSMetadata(_ context.Context) ([]*auxmodels.DSSMetadata, error) {
-	return nil, stacktrace.NewErrorWithCode(dsserr.NotImplemented, "GetDSSMetadata not implemented for raftstore")
+func (r *repo) SaveOwnMetadata(ctx context.Context, locality string, publicEndpoint string) error {
+	payload := saveOwnMetadataPayload{
+		Locality:       locality,
+		PublicEndpoint: publicEndpoint,
+	}
+
+	buf, err := json.Marshal(payload)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to marshal payload")
+	}
+
+	_, err = r.consensus.HandleClientRequest(ctx, saveOwnMetadata, buf, false)
+	return err
+}
+
+func (r *repo) GetDSSMetadata(ctx context.Context) ([]*auxmodels.DSSMetadata, error) {
+	result, err := r.consensus.HandleClientRequest(ctx, getDSSMetadata, nil, true)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to propose %s", getDSSMetadata)
+	}
+
+	if res, ok := result.([]*auxmodels.DSSMetadata); ok {
+		return res, nil
+	}
+
+	return nil, stacktrace.NewError("unexpected result type: %T", result)
 }
 
-func (r *repo) RecordHeartbeat(_ context.Context, heartbeat auxmodels.Heartbeat) error {
-	return stacktrace.NewErrorWithCode(dsserr.NotImplemented, "RecordHeartbeat not implemented for raftstore")
+func (r *repo) RecordHeartbeat(ctx context.Context, heartbeat auxmodels.Heartbeat) error {
+	buf, err := json.Marshal(heartbeat)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to marshal heartbeat")
+	}
+
+	_, err = r.consensus.HandleClientRequest(ctx, recordHeartbeat, buf, false)
+	return err
 }
 
 func (r *repo) GetDSSAirspaceRepresentationID(_ context.Context) (string, error) {
-	return "", stacktrace.NewErrorWithCode(dsserr.NotImplemented, "GetDSSAirspaceRepresentationID not implemented for raftstore")
+	connectParameters, err := raftparams.GetConnectParameters("aux")
+	if err != nil {
+		return "", stacktrace.Propagate(err, "failed to get aux raft parameters")
+	}
+
+	return strconv.Itoa(int(connectParameters.ClusterID)), nil
 }

pkg/aux_/store/raftstore/store.go

@@ -2,38 +2,86 @@ package raftstore
 
 import (
 	"context"
+	"encoding/json"
 
-	"github.com/interuss/dss/pkg/aux_/actions"
+	auxmodels "github.com/interuss/dss/pkg/aux_/models"
 	"github.com/interuss/dss/pkg/aux_/repos"
+	auxmemstore "github.com/interuss/dss/pkg/aux_/store/memstore"
 	auxraftparams "github.com/interuss/dss/pkg/aux_/store/raftstore/params"
-	dsserr "github.com/interuss/dss/pkg/errors"
+	"github.com/interuss/dss/pkg/memstore"
 	"github.com/interuss/dss/pkg/raftstore"
 	"github.com/interuss/dss/pkg/raftstore/consensus"
 	"github.com/interuss/stacktrace"
 	"go.uber.org/zap"
 )
 
+const (
+	saveOwnMetadata consensus.RequestType = "saveOwnMetadata"
+	getDSSMetadata  consensus.RequestType = "getDSSMetadata"
+	recordHeartbeat consensus.RequestType = "recordHeartbeat"
+)
+
 // repo is a full implementation of aux_.repos.Repository for Raft-based storage.
-type repo struct{}
+type repo struct {
+	consensus *consensus.Consensus
+	memStore  *memstore.Store[repos.Repository]
+	memRepo   repos.Repository
+}
 
 func Init(ctx context.Context, logger *zap.Logger) (*raftstore.Store[repos.Repository], error) {
 	params, err := auxraftparams.GetConnectParameters()
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "failed to get aux raft parameters")
 	}
-	return raftstore.Init(ctx, logger.With(zap.String("service", "aux_")), params, &repo{}, actions.Registry)
+
+	memStore, err := auxmemstore.Init(ctx, logger)
+	if err != nil {
+		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)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to initialize aux raftstore")
+	}
+
+	r.consensus = store.Consensus
+
+	return store, nil
 }
 
 func (r *repo) GetRepo() repos.Repository { return r }
 
 func (r *repo) GetSnapshot() ([]byte, error) {
-	return nil, stacktrace.NewErrorWithCode(dsserr.NotImplemented, "not implemented yet")
+	return r.memStore.GetSnapshot()
 }
 
-func (r *repo) RestoreFromSnapshot([]byte) error {
-	return stacktrace.NewErrorWithCode(dsserr.NotImplemented, "not implemented yet")
+func (r *repo) RestoreFromSnapshot(data []byte) error {
+	return r.memStore.RestoreFromSnapshot(data)
 }
 
-func (r *repo) Apply(_ context.Context, _ consensus.Proposal) (any, error) {
-	return nil, stacktrace.NewErrorWithCode(dsserr.NotImplemented, "not implemented yet")
+func (r *repo) Apply(ctx context.Context, proposal consensus.Proposal) (any, error) {
+	switch proposal.RequestType {
+	case saveOwnMetadata:
+		var payload saveOwnMetadataPayload
+		if err := json.Unmarshal(proposal.Value, &payload); err != nil {
+			return nil, stacktrace.Propagate(err, "failed to unmarshal %s payload", saveOwnMetadata)
+		}
+
+		return nil, r.memRepo.SaveOwnMetadata(ctx, payload.Locality, payload.PublicEndpoint)
+
+	case getDSSMetadata:
+		return r.memRepo.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)
+
+	default:
+		return nil, stacktrace.NewError("unknown request type: %q", proposal.RequestType)
+	}
 }