[Raft] Add generic raftstore

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/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,46 +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(_ context.Context, _ store.OperationRequest) (any, error) {
-	// TODO: implement
-	return nil, 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/store/raftstore/store.go

@@ -3,7 +3,10 @@ package raftstore
 import (
 	"context"
 
+	dsserr "github.com/interuss/dss/pkg/errors"
 	"github.com/interuss/dss/pkg/raftstore"
+	"github.com/interuss/dss/pkg/raftstore/consensus"
+	"github.com/interuss/dss/pkg/rid/actions"
 	"github.com/interuss/dss/pkg/rid/repos"
 	ridraftparams "github.com/interuss/dss/pkg/rid/store/raftstore/params"
 	"github.com/interuss/stacktrace"
@@ -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 rid raft parameters")
 	}
-	return raftstore.Init(ctx, logger.With(zap.String("service", "rid")), params, func() repos.Repository { return &repo{} })
+	return raftstore.Init(ctx, logger.With(zap.String("service", "rid")), 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/scd/store/raftstore/store.go

@@ -3,7 +3,10 @@ package raftstore
 import (
 	"context"
 
+	dsserr "github.com/interuss/dss/pkg/errors"
 	"github.com/interuss/dss/pkg/raftstore"
+	"github.com/interuss/dss/pkg/raftstore/consensus"
+	"github.com/interuss/dss/pkg/scd/actions"
 	"github.com/interuss/dss/pkg/scd/repos"
 	scdraftparams "github.com/interuss/dss/pkg/scd/store/raftstore/params"
 	"github.com/interuss/stacktrace"
@@ -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 scd raft parameters")
 	}
-	return raftstore.Init(ctx, logger.With(zap.String("service", "scd")), params, func() repos.Repository { return &repo{} })
+	return raftstore.Init(ctx, logger.With(zap.String("service", "scd")), 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/store/store.go

@@ -25,9 +25,10 @@ type OperationRequest interface {
 
 // OperationHandler holds the logic for the encoding and execution of a registered operation
 type OperationHandler[R any] struct {
-	Encode  func(req OperationRequest) ([]byte, error)
-	Decode  func(buf []byte) (OperationRequest, error)
-	Execute func(ctx context.Context, repo R, request OperationRequest) (any, error)
+	Encode     func(req OperationRequest) ([]byte, error)
+	Decode     func(buf []byte) (OperationRequest, error)
+	Execute    func(ctx context.Context, repo R, request OperationRequest) (any, error)
+	IsReadOnly bool
 }
 
 // TransactWithResult wraps Store.Transact and casts the result to ResultType, avoiding a cast at every call site.