feat: Rewrite WAL and avoid missing entries

Author: infrmtcs

consensus/consensus.go

@@ -54,7 +54,7 @@ func Init(
 
 	proposalStore := proposal.ProposalStore[starknet.Hash]{}
 	proposer := proposer.New(logger, &builder, &proposalStore, *nodeAddress, toValue)
-	stateMachine := tendermint.New(tendermintDB, logger, *nodeAddress, proposer, validators, currentHeight)
+	stateMachine := tendermint.New(logger, *nodeAddress, proposer, validators, currentHeight)
 
 	p2p := p2p.New(host, logger, &builder, &proposalStore, currentHeight, &config.DefaultBufferSizes, bootstrapPeersFn)
 

consensus/db/db.go

@@ -33,8 +33,6 @@ type walMsgCount uint32
 // 2. Right before we broadcast a message
 //
 // We call Delete when we start a new height and commit a block
-//
-//go:generate mockgen -destination=../mocks/mock_db.go -package=mocks github.com/NethermindEth/juno/consensus/db TendermintDB
 type TendermintDB[V types.Hashable[H], H types.Hash, A types.Addr] interface {
 	// Flush writes the accumulated batch operations to the underlying database.
 	Flush() error

consensus/driver/driver.go

@@ -2,6 +2,7 @@ package driver
 
 import (
 	"context"
+	"fmt"
 	"time"
 
 	"github.com/NethermindEth/juno/consensus/db"
@@ -62,16 +63,40 @@ func (d *Driver[V, H, A]) Run(ctx context.Context) error {
 		}
 	}()
 
-	d.stateMachine.ReplayWAL()
+	if err := d.replay(ctx); err != nil {
+		return err
+	}
+
+	return d.listen(ctx)
+}
+
+func (d *Driver[V, H, A]) replay(ctx context.Context) error {
+	for walEntry, err := range d.db.LoadAllEntries() {
+		if err != nil {
+			return fmt.Errorf("failed to load WAL entries: %w", err)
+		}
 
+		if _, err := d.execute(ctx, true, d.stateMachine.ProcessWAL(walEntry)); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (d *Driver[V, H, A]) listen(ctx context.Context) error {
 	for {
 		select {
 		case <-ctx.Done():
 			return nil
 		default:
 		}
+
 		actions := d.stateMachine.ProcessStart(0)
-		isCommitted := d.execute(ctx, actions)
+		isCommitted, err := d.execute(ctx, false, actions)
+		if err != nil {
+			return err
+		}
 
 		// Todo: check message signature everytime a message is received.
 		// For the time being it can be assumed the signature is correct.
@@ -100,46 +125,69 @@ func (d *Driver[V, H, A]) Run(ctx context.Context) error {
 				actions = d.stateMachine.ProcessPrecommit(p)
 			}
 
-			isCommitted = d.execute(ctx, actions)
+			isCommitted, err = d.execute(ctx, false, actions)
+			if err != nil {
+				return err
+			}
 		}
 	}
 }
 
 // This function executes the actions returned by the stateMachine.
 // It returns true if a commit action was executed. This is to notify the caller to start a new height with round 0.
+// Note: `WriteWAL` actions are generated as part of processing the event itself, so there's no
+// need to write them to the WAL again here. `isReplaying` is used to disable the writing of WAL.
 func (d *Driver[V, H, A]) execute(
 	ctx context.Context,
-	executingActions []actions.Action[V, H, A],
-) (isCommitted bool) {
-	for _, action := range executingActions {
+	isReplaying bool,
+	resultActions []actions.Action[V, H, A],
+) (isCommitted bool, err error) {
+	for _, action := range resultActions {
+		if !isReplaying && action.RequiresWALFlush() {
+			if err := d.db.Flush(); err != nil {
+				return false, fmt.Errorf("failed to flush WAL: %w", err)
+			}
+		}
+
 		switch action := action.(type) {
+		case *actions.WriteWAL[V, H, A]:
+			if !isReplaying {
+				if err := d.db.SetWALEntry(action.Entry); err != nil {
+					return false, fmt.Errorf("failed to write WAL: %w", err)
+				}
+			}
+
 		case *actions.BroadcastProposal[V, H, A]:
 			d.broadcasters.ProposalBroadcaster.Broadcast(ctx, (*types.Proposal[V, H, A])(action))
+
 		case *actions.BroadcastPrevote[H, A]:
 			d.broadcasters.PrevoteBroadcaster.Broadcast(ctx, (*types.Prevote[H, A])(action))
+
 		case *actions.BroadcastPrecommit[H, A]:
 			d.broadcasters.PrecommitBroadcaster.Broadcast(ctx, (*types.Precommit[H, A])(action))
+
 		case *actions.ScheduleTimeout:
-			d.scheduledTms[types.Timeout(*action)] = time.AfterFunc(d.getTimeout(action.Step, action.Round), func() {
-				select {
-				case <-ctx.Done():
-				case d.timeoutsCh <- types.Timeout(*action):
-				}
-			})
-		case *actions.Commit[V, H, A]:
-			if err := d.db.Flush(); err != nil {
-				d.log.Fatalf("failed to flush WAL during commit", "height", action.Height, "round", action.Round, "err", err)
-			}
+			d.setTimeout(ctx, types.Timeout(*action))
 
+		case *actions.Commit[V, H, A]:
 			d.log.Debugw("Committing", "height", action.Height, "round", action.Round)
 			d.commitListener.OnCommit(ctx, action.Height, *action.Value)
 
 			if err := d.db.DeleteWALEntries(action.Height); err != nil {
-				d.log.Errorw("failed to delete WAL messages during commit", "height", action.Height, "round", action.Round, "err", err)
+				return true, fmt.Errorf("failed to delete WAL messages during commit: %w", err)
 			}
 
-			return true
+			return true, nil
 		}
 	}
-	return false
+	return false, nil
+}
+
+func (d *Driver[V, H, A]) setTimeout(ctx context.Context, timeout types.Timeout) {
+	d.scheduledTms[timeout] = time.AfterFunc(d.getTimeout(timeout.Step, timeout.Round), func() {
+		select {
+		case <-ctx.Done():
+		case d.timeoutsCh <- timeout:
+		}
+	})
 }

consensus/driver/driver_test.go

@@ -157,7 +157,6 @@ func TestDriver(t *testing.T) {
 	stateMachine := mocks.NewMockStateMachine[starknet.Value, starknet.Hash, starknet.Address](
 		ctrl,
 	)
-	stateMachine.EXPECT().ReplayWAL().AnyTimes().Return() // ignore WAL replay logic here
 
 	commitAction := starknet.Commit(getRandProposal(random))
 

consensus/mocks/mock_db.go

@@ -25,14 +25,15 @@ type (
 	MessageHeader = types.MessageHeader[Address]
 
 	Action             = actions.Action[Value, Hash, Address]
+	WriteWAL           = actions.WriteWAL[Value, Hash, Address]
 	BroadcastProposal  = actions.BroadcastProposal[Value, Hash, Address]
 	BroadcastPrevote   = actions.BroadcastPrevote[Hash, Address]
 	BroadcastPrecommit = actions.BroadcastPrecommit[Hash, Address]
 	Commit             = actions.Commit[Value, Hash, Address]
 
 	WALEntry     = wal.Entry[Value, Hash, Address]
-	WALProposal  = wal.WALProposal[Value, Hash, Address]
-	WALPrevote   = wal.WALPrevote[Hash, Address]
-	WALPrecommit = wal.WALPrecommit[Hash, Address]
-	WALTimeout   = wal.WALTimeout
+	WALProposal  = wal.Proposal[Value, Hash, Address]
+	WALPrevote   = wal.Prevote[Hash, Address]
+	WALPrecommit = wal.Precommit[Hash, Address]
+	WALTimeout   = wal.Timeout
 )

consensus/mocks/mock_state_machine.go

@@ -65,7 +65,7 @@ func TestSync(t *testing.T) {
 			cancel()
 		})
 
-	stateMachine := tendermint.New(tmDB, logger, nodeAddr, mockApp, allNodes, types.Height(0))
+	stateMachine := tendermint.New(logger, nodeAddr, mockApp, allNodes, types.Height(0))
 
 	proposalCh := make(chan *starknet.Proposal)
 	prevoteCh := make(chan *starknet.Prevote)