feat: Implement P2P components

Author: infrmtcs

consensus/driver/driver.go

@@ -1,11 +1,12 @@
 package driver
 
 import (
-	"sync"
+	"context"
 	"time"
 
 	"github.com/NethermindEth/juno/consensus/db"
 	"github.com/NethermindEth/juno/consensus/p2p"
+	"github.com/NethermindEth/juno/consensus/proposer"
 	"github.com/NethermindEth/juno/consensus/tendermint"
 	"github.com/NethermindEth/juno/consensus/types"
 	"github.com/NethermindEth/juno/utils"
@@ -23,39 +24,34 @@ type Driver[V types.Hashable[H], H types.Hash, A types.Addr] struct {
 	db           db.TendermintDB[V, H, A]
 	stateMachine tendermint.StateMachine[V, H, A]
 	blockchain   Blockchain[V, H]
+	p2p          p2p.P2P[V, H, A]
+	proposer     proposer.Proposer[V, H]
 
 	getTimeout timeoutFn
 
-	listeners    p2p.Listeners[V, H, A]
-	broadcasters p2p.Broadcasters[V, H, A]
-
 	scheduledTms map[types.Timeout]*time.Timer
 	timeoutsCh   chan types.Timeout
-
-	wg   sync.WaitGroup
-	quit chan struct{}
 }
 
 func New[V types.Hashable[H], H types.Hash, A types.Addr](
 	log utils.Logger,
 	db db.TendermintDB[V, H, A],
 	stateMachine tendermint.StateMachine[V, H, A],
 	blockchain Blockchain[V, H],
-	listeners p2p.Listeners[V, H, A],
-	broadcasters p2p.Broadcasters[V, H, A],
+	p2p p2p.P2P[V, H, A],
+	proposer proposer.Proposer[V, H],
 	getTimeout timeoutFn,
-) *Driver[V, H, A] {
-	return &Driver[V, H, A]{
+) Driver[V, H, A] {
+	return Driver[V, H, A]{
 		log:          log,
 		db:           db,
 		stateMachine: stateMachine,
 		blockchain:   blockchain,
+		p2p:          p2p,
+		proposer:     proposer,
 		getTimeout:   getTimeout,
-		listeners:    listeners,
-		broadcasters: broadcasters,
 		scheduledTms: make(map[types.Timeout]*time.Timer),
 		timeoutsCh:   make(chan types.Timeout),
-		quit:         make(chan struct{}),
 	}
 }
 
@@ -64,60 +60,52 @@ func New[V types.Hashable[H], H types.Hash, A types.Addr](
 // these messages and returns a set of actions to be executed by the Driver.
 // The Driver executes these actions (namely broadcasting messages
 // and triggering scheduled timeouts).
-func (d *Driver[V, H, A]) Start() {
+func (d *Driver[V, H, A]) Run(ctx context.Context) error {
 	d.stateMachine.ReplayWAL()
 
-	d.wg.Add(1)
-	go func() {
-		defer d.wg.Done()
-
-		actions := d.stateMachine.ProcessStart(0)
-		d.execute(actions)
+	listeners := d.p2p.Listeners()
+	broadcasters := d.p2p.Broadcasters()
 
-		// Todo: check message signature everytime a message is received.
-		// For the time being it can be assumed the signature is correct.
+	actions := d.stateMachine.ProcessStart(0)
+	d.execute(ctx, broadcasters, actions)
 
-		for {
-			select {
-			case <-d.quit:
-				return
-			case tm := <-d.timeoutsCh:
-				// Handling of timeouts is priorities over messages
-				delete(d.scheduledTms, tm)
-				actions = d.stateMachine.ProcessTimeout(tm)
-			case p := <-d.listeners.ProposalListener.Listen():
-				actions = d.stateMachine.ProcessProposal(p)
-			case p := <-d.listeners.PrevoteListener.Listen():
-				actions = d.stateMachine.ProcessPrevote(p)
-			case p := <-d.listeners.PrecommitListener.Listen():
-				actions = d.stateMachine.ProcessPrecommit(p)
+	// Todo: check message signature everytime a message is received.
+	// For the time being it can be assumed the signature is correct.
+	for {
+		select {
+		case <-ctx.Done():
+			for _, tm := range d.scheduledTms {
+				tm.Stop()
 			}
-			d.execute(actions)
+			return nil
+		case tm := <-d.timeoutsCh:
+			// Handling of timeouts is priorities over messages
+			delete(d.scheduledTms, tm)
+			actions = d.stateMachine.ProcessTimeout(tm)
+		case p := <-listeners.ProposalListener.Listen():
+			actions = d.stateMachine.ProcessProposal(p)
+		case p := <-listeners.PrevoteListener.Listen():
+			actions = d.stateMachine.ProcessPrevote(p)
+		case p := <-listeners.PrecommitListener.Listen():
+			actions = d.stateMachine.ProcessPrecommit(p)
 		}
-	}()
-}
-
-func (d *Driver[V, H, A]) Stop() {
-	close(d.quit)
-	d.wg.Wait()
-	for _, tm := range d.scheduledTms {
-		tm.Stop()
+		d.execute(ctx, broadcasters, actions)
 	}
 }
 
-func (d *Driver[V, H, A]) execute(actions []types.Action[V, H, A]) {
+func (d *Driver[V, H, A]) execute(ctx context.Context, broadcasters p2p.Broadcasters[V, H, A], actions []types.Action[V, H, A]) {
 	for _, action := range actions {
 		switch action := action.(type) {
 		case *types.BroadcastProposal[V, H, A]:
-			d.broadcasters.ProposalBroadcaster.Broadcast(types.Proposal[V, H, A](*action))
+			broadcasters.ProposalBroadcaster.Broadcast(types.Proposal[V, H, A](*action))
 		case *types.BroadcastPrevote[H, A]:
-			d.broadcasters.PrevoteBroadcaster.Broadcast(types.Prevote[H, A](*action))
+			broadcasters.PrevoteBroadcaster.Broadcast(types.Prevote[H, A](*action))
 		case *types.BroadcastPrecommit[H, A]:
-			d.broadcasters.PrecommitBroadcaster.Broadcast(types.Precommit[H, A](*action))
+			broadcasters.PrecommitBroadcaster.Broadcast(types.Precommit[H, A](*action))
 		case *types.ScheduleTimeout:
 			d.scheduledTms[types.Timeout(*action)] = time.AfterFunc(d.getTimeout(action.Step, action.Round), func() {
 				select {
-				case <-d.quit:
+				case <-ctx.Done():
 				case d.timeoutsCh <- types.Timeout(*action):
 				}
 			})
@@ -126,7 +114,10 @@ func (d *Driver[V, H, A]) execute(actions []types.Action[V, H, A]) {
 				d.log.Fatalf("failed to flush WAL during commit", "height", action.Height, "round", action.Round, "err", err)
 			}
 
+			d.log.Debugw("Committing", "height", action.Height, "round", action.Round)
 			d.blockchain.Commit(action.Height, *action.Value)
+			d.proposer.OnCommit(ctx, action.Height, *action.Value)
+			d.p2p.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)

consensus/driver/driver_test.go

@@ -1,8 +1,8 @@
 package driver_test
 
 import (
+	"context"
 	"math/rand"
-	"sync"
 	"testing"
 	"time"
 
@@ -16,6 +16,7 @@ import (
 	"github.com/NethermindEth/juno/core/hash"
 	"github.com/NethermindEth/juno/db/pebble"
 	"github.com/NethermindEth/juno/utils"
+	"github.com/sourcegraph/conc"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/mock/gomock"
@@ -40,70 +41,6 @@ type expectedBroadcast struct {
 	precommits []starknet.Precommit
 }
 
-type mockListener[M starknet.Message] struct {
-	ch chan M
-}
-
-func newMockListener[M starknet.Message](ch chan M) *mockListener[M] {
-	return &mockListener[M]{
-		ch: ch,
-	}
-}
-
-func (m *mockListener[M]) Listen() <-chan M {
-	return m.ch
-}
-
-type mockBroadcaster[M starknet.Message] struct {
-	wg                  sync.WaitGroup
-	mu                  sync.Mutex
-	broadcastedMessages []M
-}
-
-func (m *mockBroadcaster[M]) Broadcast(msg M) {
-	m.mu.Lock()
-	defer m.mu.Unlock()
-	m.broadcastedMessages = append(m.broadcastedMessages, msg)
-	m.wg.Done()
-}
-
-type mockBlockchain struct {
-	t              *testing.T
-	expectedCommit *starknet.Commit
-}
-
-func (m *mockBlockchain) Commit(height types.Height, value starknet.Value) {
-	require.Equal(m.t, m.expectedCommit.Value, &value)
-	require.Equal(m.t, m.expectedCommit.Height, height)
-}
-
-func newMockBlockchain(t *testing.T, expectedCommit *starknet.Commit) blockchain {
-	return &mockBlockchain{
-		t:              t,
-		expectedCommit: expectedCommit,
-	}
-}
-
-func mockListeners(
-	proposalCh chan starknet.Proposal,
-	prevoteCh chan starknet.Prevote,
-	precommitCh chan starknet.Precommit,
-) listeners {
-	return listeners{
-		ProposalListener:  newMockListener(proposalCh),
-		PrevoteListener:   newMockListener(prevoteCh),
-		PrecommitListener: newMockListener(precommitCh),
-	}
-}
-
-func mockBroadcasters() broadcasters {
-	return broadcasters{
-		ProposalBroadcaster:  &mockBroadcaster[starknet.Proposal]{},
-		PrevoteBroadcaster:   &mockBroadcaster[starknet.Prevote]{},
-		PrecommitBroadcaster: &mockBroadcaster[starknet.Precommit]{},
-	}
-}
-
 func mockTimeoutFn(step types.Step, round types.Round) time.Duration {
 	return 1 * time.Millisecond
 }
@@ -217,20 +154,20 @@ func TestDriver(t *testing.T) {
 	proposalCh := make(chan starknet.Proposal)
 	prevoteCh := make(chan starknet.Prevote)
 	precommitCh := make(chan starknet.Precommit)
-	broadcasters := mockBroadcasters()
 
 	stateMachine := mocks.NewMockStateMachine[starknet.Value, hash.Hash, starknet.Address](ctrl)
 	stateMachine.EXPECT().ReplayWAL().AnyTimes().Return() // ignore WAL replay logic here
 
 	commitAction := starknet.Commit(getRandProposal(random))
+	p2p := newMockP2P(proposalCh, prevoteCh, precommitCh)
 
 	driver := driver.New(
 		utils.NewNopZapLogger(),
 		newTendermintDB(t),
 		stateMachine,
 		newMockBlockchain(t, &commitAction),
-		mockListeners(proposalCh, prevoteCh, precommitCh),
-		broadcasters,
+		p2p,
+		newMockProposer(),
 		mockTimeoutFn,
 	)
 
@@ -261,25 +198,30 @@ func TestDriver(t *testing.T) {
 	stateMachine.EXPECT().ProcessTimeout(inputTimeoutPrevote).Return(generateAndRegisterRandomActions(random, expectedBroadcast))
 	stateMachine.EXPECT().ProcessTimeout(inputTimeoutPrecommit).Return(generateAndRegisterRandomActions(random, expectedBroadcast))
 
-	increaseBroadcasterWaitGroup(expectedBroadcast.proposals, broadcasters.ProposalBroadcaster)
-	increaseBroadcasterWaitGroup(expectedBroadcast.prevotes, broadcasters.PrevoteBroadcaster)
-	increaseBroadcasterWaitGroup(expectedBroadcast.precommits, broadcasters.PrecommitBroadcaster)
+	increaseBroadcasterWaitGroup(expectedBroadcast.proposals, p2p.Broadcasters().ProposalBroadcaster)
+	increaseBroadcasterWaitGroup(expectedBroadcast.prevotes, p2p.Broadcasters().PrevoteBroadcaster)
+	increaseBroadcasterWaitGroup(expectedBroadcast.precommits, p2p.Broadcasters().PrecommitBroadcaster)
 
-	driver.Start()
+	ctx, cancel := context.WithCancel(t.Context())
 
-	go func() {
+	wg := conc.NewWaitGroup()
+	wg.Go(func() {
+		require.NoError(t, driver.Run(ctx))
+	})
+	wg.Go(func() {
 		proposalCh <- inputProposalMsg
-	}()
-	go func() {
+	})
+	wg.Go(func() {
 		prevoteCh <- inputPrevoteMsg
-	}()
-	go func() {
+	})
+	wg.Go(func() {
 		precommitCh <- inputPrecommitMsg
-	}()
+	})
+	t.Cleanup(wg.Wait)
 
-	waitAndAssertBroadcaster(t, expectedBroadcast.proposals, broadcasters.ProposalBroadcaster)
-	waitAndAssertBroadcaster(t, expectedBroadcast.prevotes, broadcasters.PrevoteBroadcaster)
-	waitAndAssertBroadcaster(t, expectedBroadcast.precommits, broadcasters.PrecommitBroadcaster)
+	waitAndAssertBroadcaster(t, expectedBroadcast.proposals, p2p.Broadcasters().ProposalBroadcaster)
+	waitAndAssertBroadcaster(t, expectedBroadcast.prevotes, p2p.Broadcasters().PrevoteBroadcaster)
+	waitAndAssertBroadcaster(t, expectedBroadcast.precommits, p2p.Broadcasters().PrecommitBroadcaster)
 
-	driver.Stop()
+	cancel()
 }

consensus/driver/mock_blockchain_test.go

@@ -0,0 +1,26 @@
+package driver_test
+
+import (
+	"testing"
+
+	"github.com/NethermindEth/juno/consensus/starknet"
+	"github.com/NethermindEth/juno/consensus/types"
+	"github.com/stretchr/testify/require"
+)
+
+type mockBlockchain struct {
+	t              *testing.T
+	expectedCommit *starknet.Commit
+}
+
+func (m *mockBlockchain) Commit(height types.Height, value starknet.Value) {
+	require.Equal(m.t, m.expectedCommit.Value, &value)
+	require.Equal(m.t, m.expectedCommit.Height, height)
+}
+
+func newMockBlockchain(t *testing.T, expectedCommit *starknet.Commit) blockchain {
+	return &mockBlockchain{
+		t:              t,
+		expectedCommit: expectedCommit,
+	}
+}