[Raft] Consensus entries publishing and proposal

pkg/aux_/store/raftstore/store.go

@@ -18,5 +18,5 @@ 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, params, func() repos.Repository { return &repo{} })
+	return raftstore.Init(ctx, logger.With(zap.String("service", "aux_")), params, func() repos.Repository { return &repo{} })
 }

pkg/raftstore/consensus/consensus.go

@@ -2,6 +2,7 @@ package consensus
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
@@ -34,6 +35,10 @@ type Consensus struct {
 
 	once            sync.Once
 	shutdownTimeout time.Duration
+	serverErrC      chan error    // http server errors
+	stopC           chan struct{} // when closed, signals shutdown to the raft updates consumer goroutine
+
+	tracker *proposalsTracker
 
 	confState     raftpb.ConfState
 	snapshotIndex uint64
@@ -76,6 +81,10 @@ func NewConsensus(ctx context.Context, logger *zap.Logger, connectParams params.
 		commitC: commitC,
 
 		shutdownTimeout: 2 * connectParams.ElectionInterval(),
+		serverErrC:      make(chan error, 1),
+		stopC:           make(chan struct{}),
+
+		tracker: newProposalsTracker(),
 	}
 
 	err = consensus.initTransport(ctx, connectParams.NodeID, connectParams.ClusterID, peers)
@@ -92,20 +101,20 @@ func NewConsensus(ctx context.Context, logger *zap.Logger, connectParams params.
 	consensus.snapshotIndex = snap.Metadata.Index
 	consensus.appliedIndex = snap.Metadata.Index
 
-	go func() {
-		err := consensus.handleReady(connectParams.TickInterval)
-		if err != nil {
-			consensus.logger.Error("handleReady exited with error, shutting down consensus", zap.Error(err))
-		}
+	if len(snap.Data) > 0 {
+		consensus.commitC <- EntryCommit{SnapshotData: snap.Data}
+	}
 
-		consensus.Stop(context.Background())
-	}()
+	consensus.startRaftUpdatesConsumer(connectParams.TickInterval, connectParams.SnapshotIntervalEntries)
 
 	return consensus, nil
 }
 
 func (c *Consensus) Stop(ctx context.Context) {
 	c.once.Do(func() {
+		c.logger.Info("stopping consensus")
+		close(c.stopC)
+
 		shutdownCtx, cancel := context.WithTimeout(context.Background(), c.shutdownTimeout)
 		defer cancel()
 		if shutdownErr := c.server.Shutdown(shutdownCtx); shutdownErr != nil {
@@ -121,6 +130,36 @@ 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) {
+	proposal, err := c.newProposal(ctx, requestType, value, readOnly)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to create proposal")
+	}
+
+	buf, err := json.Marshal(proposal)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "failed to marshal proposal")
+	}
+
+	applied := c.tracker.track(proposal.ID)
+
+	err = c.node.Propose(ctx, buf)
+	if err != nil {
+		c.tracker.untrack(proposal.ID, ProposalResult{Error: err})
+		return nil, stacktrace.Propagate(err, "failed to propose value to Raft")
+	}
+
+	select {
+	case res := <-applied:
+		return res.Result, res.Error
+
+	case <-ctx.Done():
+		c.tracker.untrack(proposal.ID, ProposalResult{Error: ctx.Err()})
+		return nil, ctx.Err()
+	}
+}
+
 func peersList(peers map[uint64]*url.URL) []raft.Peer {
 	result := make([]raft.Peer, 0, len(peers))
 	for id := range peers {
@@ -139,7 +178,9 @@ func (c *Consensus) initTransport(ctx context.Context, nodeID uint64, clusterID
 		Raft:        c,
 		ServerStats: v2stats.NewServerStats(nodeIDStr, nodeIDStr),
 		LeaderStats: v2stats.NewLeaderStats(c.logger, nodeIDStr),
-		ErrorC:      make(chan error),
+		// ErrorC is buffered to make sure rafthttp can send errors even if the raft updates consumer is not listening yet.
+		// The buffer size of 1 is sufficient because the transport will only send one error before stopping.
+		ErrorC: make(chan error, 1),
 	}
 
 	err := transport.Start()
@@ -161,6 +202,8 @@ func (c *Consensus) initTransport(ctx context.Context, nodeID uint64, clusterID
 		return stacktrace.NewError("node ID %d not found in peers map", nodeID)
 	}
 
+	c.transport = transport
+
 	c.server = &http.Server{
 		Addr:    listeningAddr,
 		Handler: transport.Handler(),
@@ -170,70 +213,179 @@ func (c *Consensus) initTransport(ctx context.Context, nodeID uint64, clusterID
 		err := c.server.ListenAndServe()
 		if err != nil && !errors.Is(err, http.ErrServerClosed) {
 			c.logger.Error("http server error", zap.Error(err))
-			c.transport.ErrorC <- err
+			c.serverErrC <- err
 		}
 	}()
 
-	c.transport = transport
 	return nil
 }
 
-// handleReady processes the Ready channel of the Raft node and applies committed entries to the state machine
-func (c *Consensus) handleReady(tickInterval time.Duration) error {
-	ticker := time.NewTicker(tickInterval)
-	defer ticker.Stop()
-
-	for {
-		select {
-		case <-ticker.C:
-			c.node.Tick()
-		case err := <-c.transport.ErrorC:
-			return stacktrace.Propagate(err, "transport error")
-		case rd, ok := <-c.node.Ready():
-			if !ok {
-				return stacktrace.NewError("could not read from Ready(), shutting down handler")
-			}
+// startRaftUpdatesConsumer starts a goroutine that processes the Ready channel of the Raft node and applies committed entries to the state machine
+func (c *Consensus) startRaftUpdatesConsumer(tickInterval time.Duration, snapshotInterval uint64) {
+	go func() {
+		defer c.Stop(context.Background())
+
+		ticker := time.NewTicker(tickInterval)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-c.stopC:
+				return
+			case <-ticker.C:
+				c.node.Tick()
+			case raftUpdate, ok := <-c.node.Ready():
+				if !ok {
+					c.logger.Error("could not read from Ready()")
+					return
+				}
 
-			err := c.storage.handleReceivedState(rd.Snapshot, rd.HardState, rd.Entries)
-			if err != nil {
-				return stacktrace.Propagate(err, "failed to handle received snapshot")
-			}
+				err := c.storage.handleReceivedState(raftUpdate.Snapshot, raftUpdate.HardState, raftUpdate.Entries)
+				if err != nil {
+					c.logger.Error("failed to handle received snapshot", zap.Error(err))
+					return
+				}
+
+				if !raft.IsEmptySnap(raftUpdate.Snapshot) {
+					if raftUpdate.Snapshot.Metadata.Index <= c.appliedIndex {
+						c.logger.Error("snapshot index shall be greater than current applied index",
+							zap.Uint64("snapshotIndex", raftUpdate.Snapshot.Metadata.Index), zap.Uint64("appliedIndex", c.appliedIndex))
+						return
+					}
+
+					c.commitC <- EntryCommit{SnapshotData: raftUpdate.Snapshot.Data}
 
-			if !raft.IsEmptySnap(rd.Snapshot) {
-				if rd.Snapshot.Metadata.Index <= c.appliedIndex {
-					return stacktrace.NewError("snapshot index %d shall be greater than current applied index %d", rd.Snapshot.Metadata.Index, c.appliedIndex)
+					c.confState = raftUpdate.Snapshot.Metadata.ConfState
+					c.snapshotIndex = raftUpdate.Snapshot.Metadata.Index
+					c.appliedIndex = raftUpdate.Snapshot.Metadata.Index
 				}
 
-				c.commitC <- EntryCommit{SnapshotData: rd.Snapshot.Data}
+				c.updateSnapshotConfState(raftUpdate.Messages)
+				c.transport.Send(raftUpdate.Messages)
+
+				entries, err := c.entriesToApply(raftUpdate.CommittedEntries)
+				if err != nil {
+					c.logger.Error("failed to get entries to apply", zap.Error(err))
+					return
+				}
+
+				err = c.publishEntries(entries, snapshotInterval)
+				if err != nil {
+					c.logger.Error("failed to publish entries", zap.Error(err))
+					return
+				}
 
-				c.confState = rd.Snapshot.Metadata.ConfState
-				c.snapshotIndex = rd.Snapshot.Metadata.Index
-				c.appliedIndex = rd.Snapshot.Metadata.Index
+				c.node.Advance()
+			case err := <-c.transport.ErrorC:
+				c.logger.Error("transport error", zap.Error(err))
+				return
+			case err := <-c.serverErrC:
+				c.logger.Error("http server error", zap.Error(err))
+				return
 			}
+		}
+	}()
+}
+
+func (c *Consensus) publishEntries(entries []raftpb.Entry, snapshotInterval uint64) error {
+	if len(entries) == 0 {
+		return nil
+	}
 
-			c.updateSnapshotConfState(rd.Messages)
-			c.transport.Send(rd.Messages)
+	c.logger.Info("publishing entries", zap.Int("numEntries", len(entries)), zap.Uint64("firstIndex", entries[0].Index), zap.Uint64("lastIndex", entries[len(entries)-1].Index))
 
-			entries, err := c.entriesToApply(rd.CommittedEntries)
+	var triggerSnapshot bool
+	var err error
+	var wg sync.WaitGroup
+	for _, entry := range entries {
+		switch entry.Type {
+		case raftpb.EntryNormal:
+			err := c.submitNormalEntryToStorage(entry.Data, &wg)
 			if err != nil {
-				return stacktrace.Propagate(err, "failed to get entries to apply")
+				return stacktrace.Propagate(err, "failed to process normal entry")
 			}
-
-			err = c.publishEntries(entries)
+		case raftpb.EntryConfChange:
+			err := c.applyConfigChange(entry.Data)
 			if err != nil {
-				return stacktrace.Propagate(err, "failed to publish entries")
+				return stacktrace.Propagate(err, "failed to process config change entry")
 			}
+		case raftpb.EntryConfChangeV2:
+			triggerSnapshot, err = c.applyConfigChangeV2(entry.Data)
+			if err != nil {
+				return stacktrace.Propagate(err, "failed to process config change v2 entry")
+			}
+		}
+	}
+
+	// wait for all entries to be applied before updating the applied index and potentially triggering a snapshot
+	wg.Wait()
+	c.appliedIndex = entries[len(entries)-1].Index
 
-			c.node.Advance()
+	if triggerSnapshot || c.appliedIndex-c.snapshotIndex >= snapshotInterval {
+		err := c.storage.triggerSnapshot(c.appliedIndex, &c.confState)
+		if err != nil {
+			return stacktrace.Propagate(err, "failed to trigger snapshot")
 		}
+
+		c.snapshotIndex = c.appliedIndex
 	}
+
+	return nil
 }
 
-// TODO implement
-func (c *Consensus) publishEntries(_ []raftpb.Entry) error {
+// submitNormalEntryToStorage passes the proposal to the store and waits for the result to be returned before untracking it.
+func (c *Consensus) submitNormalEntryToStorage(data []byte, wg *sync.WaitGroup) error {
+	if len(data) <= 0 {
+		return nil
+	}
+
+	var proposal Proposal
+	err := json.Unmarshal(data, &proposal)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to unmarshal committed proposal")
+	}
+
+	//if readOnly proposal and we did not initiate it, skip it (noop)
+	if proposal.ReadOnly && !c.tracker.isPending(proposal.ID) {
+		return nil
+	}
+
+	applyDoneC := make(chan ProposalResult, 1)
+	wg.Go(func() {
+		// Ensure that the tracker is cleaned up
+		c.tracker.untrack(proposal.ID, <-applyDoneC)
+	})
+
+	c.commitC <- EntryCommit{Prop: proposal, Done: applyDoneC}
+	return nil
+}
+
+// raftpb.ConfChange is still used internally by Raft, we just need to apply the change to the node.
+// Changes requested by clients are processed by processConfigChangeV2Entry.
+func (c *Consensus) applyConfigChange(data []byte) error {
+	var cc raftpb.ConfChange
+	err := cc.Unmarshal(data)
+	if err != nil {
+		return stacktrace.Propagate(err, "failed to unmarshal config change data")
+	}
+
+	c.confState = *c.node.ApplyConfChange(cc)
 	return nil
 }
 
+func (c *Consensus) applyConfigChangeV2(data []byte) (bool, error) {
+	var cc raftpb.ConfChangeV2
+	err := cc.Unmarshal(data)
+	if err != nil {
+		return false, stacktrace.Propagate(err, "failed to unmarshal config change data")
+	}
+
+	c.confState = *c.node.ApplyConfChange(cc)
+
+	// TODO - implement config changes when triggered by a client
+	return false, nil
+}
+
 func (c *Consensus) entriesToApply(entries []raftpb.Entry) ([]raftpb.Entry, error) {
 	if len(entries) == 0 {
 		return entries, nil