[Raft] Add generic raftstore

cmds/db-manager/cleanup/evict.go

@@ -14,6 +14,7 @@ import (
 	scdmodels "github.com/interuss/dss/pkg/scd/models"
 	scdrepos "github.com/interuss/dss/pkg/scd/repos"
 	scds "github.com/interuss/dss/pkg/scd/store"
+	dssstore "github.com/interuss/dss/pkg/store"
 	"github.com/interuss/stacktrace"
 	"github.com/spf13/cobra"
 	"github.com/spf13/pflag"
@@ -100,7 +101,7 @@ func evict(cmd *cobra.Command, _ []string) error {
 		}
 		return nil
 	}
-	if err = scdStore.Transact(ctx, scdAction); err != nil {
+	if _, err = scdStore.Transact(ctx, dssstore.NewFuncOperation(scdAction)); err != nil {
 		return fmt.Errorf("failed to execute SCD transaction: %w", err)
 	}
 
@@ -145,7 +146,7 @@ func evict(cmd *cobra.Command, _ []string) error {
 
 		return nil
 	}
-	if err = ridStore.Transact(ctx, ridAction); err != nil {
+	if _, err = ridStore.Transact(ctx, dssstore.NewFuncOperation(ridAction)); err != nil {
 		return fmt.Errorf("failed to execute RID transaction: %w", err)
 	}
 

pkg/aux_/actions/registry.go

@@ -0,0 +1,10 @@
+package actions
+
+import (
+	"github.com/interuss/dss/pkg/aux_/repos"
+	dssstore "github.com/interuss/dss/pkg/store"
+)
+
+// Registry maps operation IDs to their handlers.
+// TODO: implement
+var Registry = map[string]dssstore.OperationHandler[repos.Repository]{}

pkg/aux_/store/raftstore/store.go

@@ -3,9 +3,12 @@ package raftstore
 import (
 	"context"
 
+	"github.com/interuss/dss/pkg/aux_/actions"
 	"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"
 )
@@ -18,5 +21,19 @@ func Init(ctx context.Context, logger *zap.Logger) (*raftstore.Store[repos.Repos
 	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, func() repos.Repository { return &repo{} })
+	return raftstore.Init(ctx, logger.With(zap.String("service", "aux_")), params, &repo{}, actions.Registry)
+}
+
+func (r *repo) GetRepo() repos.Repository { return r }
+
+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/sqlstore/store.go

@@ -3,6 +3,7 @@ package sqlstore
 import (
 	"context"
 
+	"github.com/interuss/dss/pkg/aux_/actions"
 	"github.com/interuss/dss/pkg/aux_/repos"
 	"github.com/interuss/dss/pkg/logging"
 	dssql "github.com/interuss/dss/pkg/sql"
@@ -39,5 +40,6 @@ func Init(ctx context.Context, logger *zap.Logger, withCheckCron bool) (*sqlstor
 				version:   version,
 			}
 		},
+		Registry: actions.Registry,
 	}, withCheckCron)
 }

pkg/memstore/store.go

@@ -12,6 +12,7 @@ import (
 
 	dsserr "github.com/interuss/dss/pkg/errors"
 	"github.com/interuss/dss/pkg/logging"
+	"github.com/interuss/dss/pkg/store"
 	"github.com/interuss/stacktrace"
 	"go.uber.org/zap"
 )
@@ -50,8 +51,8 @@ func Init[R any](ctx context.Context, logger *zap.Logger, name string, r MemRepo
 	return store, nil
 }
 
-func (s *Store[R]) Transact(ctx context.Context, _ func(context.Context, R) error) error {
-	return stacktrace.NewErrorWithCode(dsserr.NotImplemented, "Transact not implemented for memstore")
+func (s *Store[R]) Transact(ctx context.Context, _ store.OperationRequest) (any, error) {
+	return nil, stacktrace.NewErrorWithCode(dsserr.NotImplemented, "Transact not implemented for memstore")
 }
 
 func (s *Store[R]) Interact(_ context.Context) (R, error) {

pkg/raftstore/consensus/consensus.go

@@ -131,7 +131,7 @@ func (c *Consensus) Stop(ctx context.Context) {
 }
 
 // HandleClientRequest blocks until the proposal is committed and applied / dropped or until ctx is cancelled.
-func (c *Consensus) HandleClientRequest(ctx context.Context, requestType string, value any, readOnly bool) (any, error) {
+func (c *Consensus) HandleClientRequest(ctx context.Context, requestType string, value []byte, readOnly bool) (any, error) {
 	proposal, err := c.newProposal(ctx, requestType, value, readOnly)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "failed to create proposal")

pkg/raftstore/consensus/proposal.go

@@ -2,11 +2,11 @@ package consensus
 
 import (
 	"context"
-	"encoding/json"
 	"sync"
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/interuss/dss/pkg/timestamp"
 	"github.com/interuss/stacktrace"
 )
 
@@ -30,17 +30,16 @@ type Proposal struct {
 	ReadOnly bool `json:"read_only"`
 }
 
-func (c *Consensus) newProposal(ctx context.Context, requestType string, payload any, readOnly bool) (Proposal, error) {
-	// TODO - Fetch timestamp from context
-	value, err := json.Marshal(payload)
-	if err != nil {
-		return Proposal{}, stacktrace.Propagate(err, "failed to serialize proposal payload")
+func (c *Consensus) newProposal(ctx context.Context, requestType string, value []byte, readOnly bool) (Proposal, error) {
+	timestamp, err := timestamp.RequestTimestampFromContext(ctx)
+	if err != nil || timestamp.IsZero() {
+		return Proposal{}, stacktrace.Propagate(err, "failed to get timestamp from context")
 	}
 
 	return Proposal{
 		ID:          uuid.NewString(),
 		NodeID:      c.nodeID,
-		Timestamp:   time.Now().UTC(),
+		Timestamp:   timestamp,
 		RequestType: requestType,
 		Value:       value,
 		ReadOnly:    readOnly,

pkg/raftstore/store.go

@@ -2,45 +2,109 @@ package raftstore
 
 import (
 	"context"
+	"sync"
 
+	"github.com/interuss/dss/pkg/logging"
 	"github.com/interuss/dss/pkg/raftstore/consensus"
 	raftparams "github.com/interuss/dss/pkg/raftstore/params"
+	"github.com/interuss/dss/pkg/store"
+	"github.com/interuss/dss/pkg/timestamp"
 	"github.com/interuss/stacktrace"
 	"go.uber.org/zap"
 )
 
+type RaftRepo[R any] interface {
+	GetRepo() R
+	// Apply is called on every committed entry. The proposal must be applied atomically.
+	Apply(ctx context.Context, proposal consensus.Proposal) (any, error)
+	GetSnapshot() ([]byte, error)
+	RestoreFromSnapshot(data []byte) error
+}
+
 type Store[R any] struct {
-	newRepo   func() R
-	consensus *consensus.Consensus
+	logger *zap.Logger
+
+	raftRepo RaftRepo[R]
+	cancel   context.CancelFunc
+	registry map[string]store.OperationHandler[R]
+
+	Consensus *consensus.Consensus
+
+	wg sync.WaitGroup
 }
 
-func Init[R any](ctx context.Context, logger *zap.Logger, params raftparams.ConnectParameters, newRepo func() R) (*Store[R], error) {
+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) {
+	ctx, cancel := context.WithCancel(ctx)
+
+	store := &Store[R]{
+		raftRepo: r,
+		logger:   logging.WithValuesFromContext(ctx, logger),
+		cancel:   cancel,
+		registry: registry,
+	}
 	commitC := make(chan consensus.EntryCommit)
-	consensusInstance, err := consensus.NewConsensus(ctx, logger, params, func() ([]byte, error) { return nil, nil }, commitC)
+	store.wg.Go(func() { store.processCommits(ctx, commitC) })
+
+	consensusInstance, err := consensus.NewConsensus(ctx, logger, params, r.GetSnapshot, commitC)
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "failed to initialize consensus")
 	}
-	// TODO: start consumer goroutine reading from commitC
 
-	return &Store[R]{
-		newRepo:   newRepo,
-		consensus: consensusInstance,
-	}, nil
+	store.Consensus = consensusInstance
+
+	return store, nil
 }
 
 // Transact proposes the entry to Raft and blocks until it is committed and applied.
-func (s *Store[R]) Transact(ctx context.Context, f func(context.Context, R) error) error {
-	// TODO: implement
-	return nil
+func (s *Store[R]) Transact(ctx context.Context, request store.OperationRequest) (any, error) {
+	handler, ok := s.registry[request.OperationID()]
+	if !ok {
+		return nil, stacktrace.NewError("no handler registered for operation %q", request.OperationID())
+	}
+	payload, err := handler.Encode(request)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to encode op %q", request.OperationID())
+	}
+	return s.Consensus.HandleClientRequest(ctx, request.OperationID(), payload, handler.IsReadOnly)
 }
 
-// Interact returns a repository that can be used to query the store without proposing a Raft entry.
+// Interact returns the underlying Raft repo which, for every operation, will propose it to Raft and return the results.
 func (s *Store[R]) Interact(_ context.Context) (R, error) {
-	return s.newRepo(), nil
+	return s.raftRepo.GetRepo(), nil
 }
 
-// Close shuts down the consensus instance.
+// Close shuts down the consensus instance and processCommits loop.
 func (s *Store[R]) Close() error {
-	// TODO: implement
+	s.Consensus.Stop(context.Background())
+	s.cancel()
+	s.logger.Info("waiting for commit processing goroutine to exit")
+	s.wg.Wait()
 	return nil
 }
+
+// processCommits reads committed entries from the consensus layer and applies them via Apply.
+func (s *Store[R]) processCommits(ctx context.Context, commitCh <-chan consensus.EntryCommit) {
+	for {
+		select {
+		case <-ctx.Done():
+			s.logger.Info("stopping commit processing loop")
+			return
+		case commit, ok := <-commitCh:
+			if !ok {
+				s.logger.Info("commit channel closed, stopping commit processing loop")
+				return
+			}
+
+			if commit.SnapshotData != nil {
+				if err := s.raftRepo.RestoreFromSnapshot(commit.SnapshotData); err != nil {
+					s.logger.Fatal("failed to restore from snapshot", zap.Error(err))
+				}
+				continue
+			}
+
+			proposalCtx := timestamp.WithRequestTimestamp(ctx, commit.Prop.Timestamp)
+			result, err := s.raftRepo.Apply(proposalCtx, commit.Prop)
+			commit.Done <- consensus.ProposalResult{Result: result, Error: err}
+		}
+	}
+}

pkg/rid/actions/registry.go

@@ -0,0 +1,10 @@
+package actions
+
+import (
+	"github.com/interuss/dss/pkg/rid/repos"
+	dssstore "github.com/interuss/dss/pkg/store"
+)
+
+// Registry maps operation IDs to their handlers.
+// TODO: implement
+var Registry = map[string]dssstore.OperationHandler[repos.Repository]{}

pkg/rid/application/application_test.go

@@ -13,6 +13,8 @@ import (
 	dssql "github.com/interuss/dss/pkg/sql"
 	"github.com/interuss/dss/pkg/sqlstore"
 	"github.com/interuss/dss/pkg/sqlstore/params"
+	dssstore "github.com/interuss/dss/pkg/store"
+	"github.com/interuss/stacktrace"
 	"github.com/jonboulle/clockwork"
 	"github.com/stretchr/testify/require"
 
@@ -35,8 +37,11 @@ func (s *mockRepo) Interact(ctx context.Context) (repos.Repository, error) {
 	return s, nil
 }
 
-func (s *mockRepo) Transact(ctx context.Context, f func(ctx context.Context, repo repos.Repository) error) error {
-	return f(ctx, s)
+func (s *mockRepo) Transact(ctx context.Context, request dssstore.OperationRequest) (any, error) {
+	if fn, ok := request.(*dssstore.FuncOperation[repos.Repository]); ok {
+		return nil, fn.Execute(ctx, s)
+	}
+	return nil, stacktrace.NewError("unsupported operation type %T", request)
 }
 
 func (s *mockRepo) Close() error {

pkg/rid/application/isa.go

@@ -10,6 +10,7 @@ import (
 	dssmodels "github.com/interuss/dss/pkg/models"
 	ridmodels "github.com/interuss/dss/pkg/rid/models"
 	"github.com/interuss/dss/pkg/rid/repos"
+	"github.com/interuss/dss/pkg/store"
 	"github.com/interuss/stacktrace"
 )
 
@@ -63,7 +64,7 @@ func (a *app) DeleteISA(ctx context.Context, id dssmodels.ID, owner dssmodels.Ow
 		subs []*ridmodels.Subscription
 	)
 	// The following will automatically retry TXN retry errors.
-	err := a.store.Transact(ctx, func(ctx context.Context, repo repos.Repository) error {
+	_, err := a.store.Transact(ctx, store.NewFuncOperation(func(ctx context.Context, repo repos.Repository) error {
 		old, err := repo.GetISA(ctx, id, true)
 		switch {
 		case err != nil:
@@ -88,7 +89,7 @@ func (a *app) DeleteISA(ctx context.Context, id dssmodels.ID, owner dssmodels.Ow
 			return stacktrace.Propagate(err, "Error updating notification indices")
 		}
 		return nil
-	})
+	}))
 	return ret, subs, err // No need to Propagate this error as this stack layer does not add useful information
 }
 
@@ -104,7 +105,7 @@ func (a *app) InsertISA(ctx context.Context, isa *ridmodels.IdentificationServic
 		subs []*ridmodels.Subscription
 	)
 	// The following will automatically retry TXN retry errors.
-	err := a.store.Transact(ctx, func(ctx context.Context, repo repos.Repository) error {
+	_, err := a.store.Transact(ctx, store.NewFuncOperation(func(ctx context.Context, repo repos.Repository) error {
 		// ensure it doesn't exist yet
 		old, err := repo.GetISA(ctx, isa.ID, false)
 		if err != nil {
@@ -126,7 +127,7 @@ func (a *app) InsertISA(ctx context.Context, isa *ridmodels.IdentificationServic
 			return stacktrace.Propagate(err, "Error inserting ISA")
 		}
 		return nil
-	})
+	}))
 	return ret, subs, err // No need to Propagate this error as this stack layer does not add useful information
 }
 
@@ -138,7 +139,7 @@ func (a *app) UpdateISA(ctx context.Context, isa *ridmodels.IdentificationServic
 		subs []*ridmodels.Subscription
 	)
 	// The following will automatically retry TXN retry errors.
-	err := a.store.Transact(ctx, func(ctx context.Context, repo repos.Repository) error {
+	_, err := a.store.Transact(ctx, store.NewFuncOperation(func(ctx context.Context, repo repos.Repository) error {
 		var err error
 
 		old, err := repo.GetISA(ctx, isa.ID, true)
@@ -176,7 +177,7 @@ func (a *app) UpdateISA(ctx context.Context, isa *ridmodels.IdentificationServic
 			return stacktrace.Propagate(err, "Error updating notification indices")
 		}
 		return nil
-	})
+	}))
 
 	return ret, subs, err // No need to Propagate this error as this stack layer does not add useful information
 }