[Raft] Use separate consensus instances per store

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,19 @@ import (
 	"context"
 
 	"github.com/interuss/dss/pkg/aux_/repos"
+	auxraftparams "github.com/interuss/dss/pkg/aux_/store/raftstore/params"
 	"github.com/interuss/dss/pkg/raftstore"
+	"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, func() repos.Repository { return &repo{} })
 }

pkg/raftstore/consensus/consensus.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"sync"
 	"time"
 
 	"github.com/interuss/dss/pkg/logging"
@@ -29,18 +30,27 @@ type Consensus struct {
 	server    *http.Server
 
 	storage *storage
+	commitC chan<- EntryCommit
+
+	once            sync.Once
+	shutdownTimeout time.Duration
 
 	confState     raftpb.ConfState
 	snapshotIndex uint64
 	appliedIndex  uint64
 }
 
-func NewConsensus(ctx context.Context, logger *zap.Logger, peers map[uint64]*url.URL, connectParams params.ConnectParameters) (*Consensus, error) {
-	storage, old, err := newStorage(ctx, logger.With(zap.String("component", "storage")), connectParams.DataDir, connectParams.NodeID, connectParams.SnapshotCatchupEntries)
+func NewConsensus(ctx context.Context, logger *zap.Logger, 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")
 	}
 
+	peers, err := connectParams.PeerMap()
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to parse peer map")
+	}
+
 	nodeUrl, ok := peers[connectParams.NodeID]
 	if !ok {
 		return nil, stacktrace.NewError("node ID %d not found in peers map", connectParams.NodeID)
@@ -63,6 +73,9 @@ func NewConsensus(ctx context.Context, logger *zap.Logger, peers map[uint64]*url
 		node:   node,
 
 		storage: storage,
+		commitC: commitC,
+
+		shutdownTimeout: 2 * connectParams.ElectionInterval(),
 	}
 
 	err = consensus.initTransport(ctx, connectParams.NodeID, connectParams.ClusterID, peers)
@@ -85,21 +98,27 @@ func NewConsensus(ctx context.Context, logger *zap.Logger, peers map[uint64]*url
 			consensus.logger.Error("handleReady exited with error, shutting down consensus", zap.Error(err))
 		}
 
-		shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*connectParams.ElectionInterval())
+		consensus.Stop(context.Background())
+	}()
+
+	return consensus, nil
+}
+
+func (c *Consensus) Stop(ctx context.Context) {
+	c.once.Do(func() {
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), c.shutdownTimeout)
 		defer cancel()
-		if shutdownErr := consensus.server.Shutdown(shutdownCtx); shutdownErr != nil {
-			consensus.logger.Error("failed to shutdown http server", zap.Error(shutdownErr))
+		if shutdownErr := c.server.Shutdown(shutdownCtx); shutdownErr != nil {
+			c.logger.Error("failed to shutdown http server", zap.Error(shutdownErr))
 		} else {
-			consensus.logger.Info("http server shutdown complete")
+			c.logger.Info("http server shutdown complete")
 		}
 
-		consensus.transport.Stop()
-		consensus.logger.Info("transport stopped")
-		consensus.node.Stop()
-		consensus.logger.Info("raft node stopped")
-	}()
-
-	return consensus, nil
+		c.transport.Stop()
+		c.logger.Info("transport stopped")
+		c.node.Stop()
+		c.logger.Info("raft node stopped")
+	})
 }
 
 func peersList(peers map[uint64]*url.URL) []raft.Peer {
@@ -185,10 +204,7 @@ func (c *Consensus) handleReady(tickInterval time.Duration) error {
 					return stacktrace.NewError("snapshot index %d shall be greater than current applied index %d", rd.Snapshot.Metadata.Index, c.appliedIndex)
 				}
 
-				err = c.dispatchSnapshot(rd.Snapshot.Data)
-				if err != nil {
-					return stacktrace.Propagate(err, "failed to dispatch snapshot")
-				}
+				c.commitC <- EntryCommit{SnapshotData: rd.Snapshot.Data}
 
 				c.confState = rd.Snapshot.Metadata.ConfState
 				c.snapshotIndex = rd.Snapshot.Metadata.Index
@@ -218,11 +234,6 @@ func (c *Consensus) publishEntries(_ []raftpb.Entry) error {
 	return nil
 }
 
-// TODO implement
-func (c *Consensus) dispatchSnapshot(_ []byte) error {
-	return nil
-}
-
 func (c *Consensus) entriesToApply(entries []raftpb.Entry) ([]raftpb.Entry, error) {
 	if len(entries) == 0 {
 		return entries, nil
@@ -272,8 +283,3 @@ func (c *Consensus) ReportUnreachable(id uint64) {
 func (c *Consensus) ReportSnapshot(id uint64, status raft.SnapshotStatus) {
 	c.node.ReportSnapshot(id, status)
 }
-
-// RegisterStore allows registering a snapshot provider function for a specific store
-func (c *Consensus) RegisterStore(name string, provider snapshotProvider) {
-	c.storage.registerSnapshotProvider(name, provider)
-}

pkg/raftstore/consensus/proposal.go

@@ -0,0 +1,7 @@
+package consensus
+
+type EntryCommit struct {
+	// TODO add fields to represent the committed entry
+
+	SnapshotData []byte
+}

pkg/raftstore/consensus/storage.go

@@ -2,7 +2,6 @@ package consensus
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
 	"os"
@@ -31,16 +30,16 @@ type storage struct {
 
 	wal *wal.WAL
 
-	snapper   *snap.Snapshotter
-	providers map[string]snapshotProvider
+	snapper  *snap.Snapshotter
+	snapshot snapshotProvider
 
 	snapshotCatchUpEntries uint64
 }
 
 // newStorage initializes the storage by loading the latest snapshot and wal entries from the disk
 // and applies them to the Raft memory storage.
 // It returns the initialized storage, a boolean indicating whether the storage was pre-existent or an error.
-func newStorage(ctx context.Context, logger *zap.Logger, dataDir string, nodeID uint64, snapshotCatchUpEntries uint64) (*storage, bool, error) {
+func newStorage(ctx context.Context, logger *zap.Logger, dataDir string, nodeID uint64, provider snapshotProvider, snapshotCatchUpEntries uint64) (*storage, bool, error) {
 	logger = logging.WithValuesFromContext(ctx, logger)
 
 	// load the latest snapshot
@@ -124,8 +123,9 @@ func newStorage(ctx context.Context, logger *zap.Logger, dataDir string, nodeID
 
 		wal: w,
 
-		snapper:                snapper,
-		providers:              make(map[string]snapshotProvider),
+		snapper:  snapper,
+		snapshot: provider,
+
 		snapshotCatchUpEntries: snapshotCatchUpEntries,
 	}, ok, nil
 }
@@ -177,33 +177,9 @@ func (s *storage) save(snapshot raftpb.Snapshot) error {
 	return s.wal.ReleaseLockTo(snapshot.Metadata.Index)
 }
 
-func (s *storage) registerSnapshotProvider(name string, provider func() ([]byte, error)) {
-	s.providers[name] = provider
-}
-
-// getSnapshot calls all registered snapshot providers and combines their data into a single snapshot.
-func (s *storage) getSnapshot() ([]byte, error) {
-	parts := make(map[string][]byte)
-	for name, provider := range s.providers {
-		data, err := provider()
-		if err != nil {
-			return nil, stacktrace.Propagate(err, "failed to get snapshot data from %q", name)
-		}
-
-		parts[name] = data
-	}
-
-	return json.Marshal(parts)
-}
-
-// snapshotter returns the snapshotter used by the storage.
-func (s *storage) snapshotter() *snap.Snapshotter {
-	return s.snapper
-}
-
 func (s *storage) triggerSnapshot(appliedIndex uint64, confState *raftpb.ConfState) error {
 	s.logger.Info("triggering snapshot", zap.Uint64("appliedIndex", appliedIndex))
-	data, err := s.getSnapshot()
+	data, err := s.snapshot()
 	if err != nil {
 		return stacktrace.Propagate(err, "failed to get snapshot data")
 	}

pkg/raftstore/params/params.go

@@ -3,6 +3,7 @@ package params
 import (
 	"flag"
 	"net/url"
+	"path/filepath"
 	"strconv"
 	"strings"
 	"time"
@@ -123,7 +124,6 @@ var (
 func init() {
 	flag.Uint64Var(&connectParameters.NodeID, "raft_node_id", 0, "Raft node ID for this instance (must be non-zero and unique within the cluster).")
 	flag.Uint64Var(&connectParameters.ClusterID, "raft_cluster_id", 1, "ID of the cluster, used to isolate different Raft clusters running in the same network (must be the same for all nodes in the cluster).")
-	flag.StringVar(&connectParameters.Peers, "raft_peers", "", `Comma-separated "nodeID=peerURL" pairs for all cluster members, including the current node, e.g. "1=http://node1:9021,2=http://node2:9021,3=http://node3:9021"`)
 	flag.StringVar(&connectParameters.DataDir, "raft_datadir", defaultDataDir, "Directory for raft data (WAL segments and snapshots), required for restarts. These should not be deleted while the node is running or across restarts unless the node is being permanently shut down.")
 
 	flag.Uint64Var(&connectParameters.SnapshotCatchupEntries, "raft_snapshot_catchup_entries", defaultSnapshotCatchupEntries,
@@ -142,6 +142,11 @@ func init() {
 }
 
 // GetConnectParameters returns a ConnectParameters instance that gets populated from well-known CLI flags.
-func GetConnectParameters() ConnectParameters {
-	return connectParameters
+func GetConnectParameters(subfolder string) (ConnectParameters, error) {
+	if connectParameters.NodeID == 0 {
+		return ConnectParameters{}, stacktrace.NewError("--raft_node_id is required and must be non-zero")
+	}
+	p := connectParameters
+	p.DataDir = filepath.Join(connectParameters.DataDir, subfolder)
+	return p, nil
 }

pkg/raftstore/store.go

@@ -2,52 +2,29 @@ package raftstore
 
 import (
 	"context"
-	"sync"
 
 	"github.com/interuss/dss/pkg/raftstore/consensus"
 	raftparams "github.com/interuss/dss/pkg/raftstore/params"
 	"github.com/interuss/stacktrace"
 	"go.uber.org/zap"
 )
 
-var (
-	sharedConsensus     *consensus.Consensus
-	sharedConsensusOnce sync.Once
-	sharedConsensusErr  error
-)
-
 type Store[R any] struct {
 	newRepo   func() R
 	consensus *consensus.Consensus
 }
 
-func Init[R any](ctx context.Context, logger *zap.Logger, newRepo func() R) (*Store[R], error) {
-	// scd, rid and aux will share the same consensus instance, so we initialize it once.
-	sharedConsensusOnce.Do(func() {
-		params := raftparams.GetConnectParameters()
-		peers, err := params.PeerMap()
-		if err != nil {
-			sharedConsensusErr = stacktrace.Propagate(err, "failed to parse peer map")
-			return
-		}
-
-		sharedConsensus, sharedConsensusErr = consensus.NewConsensus(ctx, logger, peers, params)
-		if sharedConsensusErr != nil {
-			sharedConsensusErr = stacktrace.Propagate(sharedConsensusErr, "failed to initialize consensus")
-		}
-	})
-	if sharedConsensusErr != nil {
-		return nil, sharedConsensusErr
+func Init[R any](ctx context.Context, logger *zap.Logger, params raftparams.ConnectParameters, newRepo func() R) (*Store[R], error) {
+	commitC := make(chan consensus.EntryCommit)
+	consensusInstance, err := consensus.NewConsensus(ctx, logger, params, func() ([]byte, error) { return nil, nil }, commitC)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to initialize consensus")
 	}
-
-	// TODO: implement
-	sharedConsensus.RegisterStore("provider", func() ([]byte, error) {
-		return nil, nil
-	})
+	// TODO: start consumer goroutine reading from commitC
 
 	return &Store[R]{
 		newRepo:   newRepo,
-		consensus: sharedConsensus,
+		consensus: consensusInstance,
 	}, nil
 }
 

pkg/rid/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 = "rid_raft_peers"
+
+var peers string
+
+func init() {
+	flag.StringVar(&peers, peersFlag, "", `Comma-separated "nodeID=peerURL" pairs for the rid store, e.g. "1=http://node1:9011,2=http://node2:9011,3=http://node3:9011"`)
+}
+
+func GetConnectParameters() (raftparams.ConnectParameters, error) {
+	if peers == "" {
+		return raftparams.ConnectParameters{}, stacktrace.NewError("--%s is required", peersFlag)
+	}
+
+	p, err := raftparams.GetConnectParameters("rid")
+	if err != nil {
+		return raftparams.ConnectParameters{}, err
+	}
+	p.Peers = peers
+	return p, nil
+}

pkg/rid/store/raftstore/store.go

@@ -5,12 +5,18 @@ import (
 
 	"github.com/interuss/dss/pkg/raftstore"
 	"github.com/interuss/dss/pkg/rid/repos"
+	ridraftparams "github.com/interuss/dss/pkg/rid/store/raftstore/params"
+	"github.com/interuss/stacktrace"
 	"go.uber.org/zap"
 )
 
 // 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) {
-	return raftstore.Init[repos.Repository](ctx, logger, func() repos.Repository { return &repo{} })
+	params, err := ridraftparams.GetConnectParameters()
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to get rid raft parameters")
+	}
+	return raftstore.Init(ctx, logger, params, func() repos.Repository { return &repo{} })
 }