feat: Integrate Kademlia DHT protocol to pubsub

Author: infrmtcs

consensus/consensus.go

@@ -1,10 +1,7 @@
 package consensus
 
 import (
-	"context"
 	"errors"
-	"fmt"
-	"strings"
 
 	"github.com/NethermindEth/juno/blockchain"
 	"github.com/NethermindEth/juno/builder"
@@ -26,7 +23,6 @@ import (
 	"github.com/libp2p/go-libp2p/core/crypto"
 	"github.com/libp2p/go-libp2p/core/host"
 	"github.com/libp2p/go-libp2p/core/peer"
-	"github.com/sourcegraph/conc/pool"
 )
 
 type ConsensusServices struct {
@@ -47,6 +43,7 @@ func Init(
 	timeoutFn driver.TimeoutFn,
 	hostAddress string,
 	hostPrivateKey crypto.PrivKey,
+	bootstrapPeersFn func() []peer.AddrInfo,
 ) (ConsensusServices, error) {
 	chainHeight, err := blockchain.Height()
 	if err != nil && !errors.Is(err, db.ErrKeyNotFound) {
@@ -82,7 +79,7 @@ func Init(
 		return ConsensusServices{}, err
 	}
 
-	p2p := p2p.New(host, logger, &builder, &proposalStore, currentHeight, &config.DefaultBufferSizes)
+	p2p := p2p.New(host, logger, &builder, &proposalStore, currentHeight, &config.DefaultBufferSizes, bootstrapPeersFn)
 
 	commitListener := driver.NewCommitListener(logger, &proposalStore, proposer, p2p)
 	driver := driver.New(logger, tendermintDB, stateMachine, commitListener, p2p, timeoutFn)
@@ -96,31 +93,6 @@ func Init(
 	}, nil
 }
 
-func Connect(ctx context.Context, host host.Host, peers string) error {
-	if peers == "" {
-		return nil
-	}
-
-	pool := pool.New().WithErrors().WithFirstError()
-
-	for peerStr := range strings.SplitSeq(peers, ",") {
-		pool.Go(func() error {
-			peerAddr, err := peer.AddrInfoFromString(peerStr)
-			if err != nil {
-				return fmt.Errorf("unable to parse peer address %q: %w", peerStr, err)
-			}
-
-			if err := host.Connect(ctx, *peerAddr); err != nil {
-				return fmt.Errorf("unable to connect to %q: %w", peerStr, err)
-			}
-
-			return nil
-		})
-	}
-
-	return pool.Wait()
-}
-
 func toValue(value *felt.Felt) starknet.Value {
 	return starknet.Value(*value)
 }

consensus/integtest/integ_test.go

@@ -2,8 +2,10 @@ package integtest
 
 import (
 	"fmt"
+	"maps"
+	"slices"
+	gosync "sync"
 	"testing"
-	"time"
 
 	"github.com/NethermindEth/juno/blockchain"
 	"github.com/NethermindEth/juno/consensus"
@@ -16,7 +18,7 @@ import (
 	"github.com/NethermindEth/juno/sync"
 	"github.com/NethermindEth/juno/utils"
 	"github.com/NethermindEth/juno/vm"
-	"github.com/libp2p/go-libp2p/core/host"
+	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/sourcegraph/conc"
 	"github.com/sourcegraph/conc/iter"
 	"github.com/stretchr/testify/assert"
@@ -27,7 +29,6 @@ import (
 const (
 	commitBufferSize = 1024
 	maxLag           = 10
-	gst              = 2 * time.Second // 2 Heartbeats should be enough for the nodes to connect to its peers
 	hostAddress      = "/ip4/0.0.0.0/tcp/0"
 )
 
@@ -78,18 +79,15 @@ func loadGenesis(t *testing.T, log *utils.ZapLogger) (core.StateDiff, map[felt.F
 
 func initNode(
 	t *testing.T,
-	gstChan chan struct{},
 	index int,
 	logger *utils.ZapLogger,
 	commits chan commit,
 	cfg *testConfig,
 	genesisDiff core.StateDiff,
 	genesisClasses map[felt.Felt]core.Class,
-) host.Host {
+	ready chan struct{},
+) networkNodeConfig {
 	t.Helper()
-	defer func() {
-		logger.Debug("finished initNode")
-	}()
 
 	mockServices := consensus.InitMockServices(0, 0, index, cfg.nodeCount)
 
@@ -98,6 +96,12 @@ func initNode(
 	bc := getBlockchain(t, genesisDiff, genesisClasses)
 	vm := vm.New(false, logger)
 
+	adjacentNodes := make(map[int]peer.AddrInfo)
+	bootstrapPeers := gosync.OnceValue(func() []peer.AddrInfo {
+		<-ready
+		return slices.Collect(maps.Values(adjacentNodes))
+	})
+
 	services, err := consensus.Init(
 		logger,
 		consensusDB,
@@ -108,6 +112,7 @@ func initNode(
 		mockServices.TimeoutFn,
 		hostAddress,
 		mockServices.PrivateKey,
+		bootstrapPeers,
 	)
 	require.NoError(t, err)
 
@@ -119,15 +124,17 @@ func initNode(
 		require.NoError(t, services.P2P.Run(t.Context()))
 	})
 	wg.Go(func() {
-		<-gstChan
 		require.NoError(t, services.Driver.Run(t.Context()))
 	})
 	wg.Go(func() {
 		writeBlock(t, index, bc, services.CommitListener, commits)
 	})
 	t.Cleanup(wg.Wait)
 
-	return services.Host
+	return networkNodeConfig{
+		host:          services.Host,
+		adjacentNodes: adjacentNodes,
+	}
 }
 
 func writeBlock(
@@ -233,20 +240,16 @@ func runTest(t *testing.T, cfg testConfig) {
 
 	commits := make(chan commit, commitBufferSize)
 
-	gstChan := make(chan struct{})
+	ready := make(chan struct{})
 
 	hosts := make(networkConfig, honestNodeCount)
 	iterator := iter.Iterator[networkNodeConfig]{MaxGoroutines: honestNodeCount}
 	iterator.ForEachIdx(hosts, func(i int, nodeConfig *networkNodeConfig) {
-		*nodeConfig = networkNodeConfig{
-			host:          initNode(t, gstChan, i, logger, commits, &cfg, genesisDiff, genesisClasses),
-			adjacentNodes: make(map[int]struct{}),
-		}
+		*nodeConfig = initNode(t, i, logger, commits, &cfg, genesisDiff, genesisClasses, ready)
 	})
-	hosts.setup(t, cfg.networkSetup)
 
-	time.Sleep(gst)
-	close(gstChan)
+	cfg.networkSetup(t, hosts)
+	close(ready)
 
 	assertCommits(t, commits, cfg, logger)
 }

consensus/integtest/p2p.go

@@ -1,22 +1,17 @@
 package integtest
 
 import (
-	"fmt"
 	"math/rand"
-	"slices"
-	"strings"
 	"testing"
 
-	"github.com/NethermindEth/juno/consensus"
 	"github.com/libp2p/go-libp2p/core/host"
-	"github.com/multiformats/go-multiaddr"
-	"github.com/sourcegraph/conc/iter"
+	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/stretchr/testify/require"
 )
 
 type networkNodeConfig struct {
 	host          host.Host
-	adjacentNodes map[int]struct{}
+	adjacentNodes map[int]peer.AddrInfo
 }
 
 func (c networkNodeConfig) isConnected(to int) bool {
@@ -26,35 +21,8 @@ func (c networkNodeConfig) isConnected(to int) bool {
 
 type networkConfig []networkNodeConfig
 
-func (c networkConfig) setup(t *testing.T, configFn networkConfigFn) {
-	t.Helper()
-
-	configFn(t, c)
-
-	iter.Iterator[networkNodeConfig]{MaxGoroutines: len(c)}.ForEach(c, func(thisNodeConfig *networkNodeConfig) {
-		peers := make([]string, 0, len(thisNodeConfig.adjacentNodes))
-
-		for otherNode := range thisNodeConfig.adjacentNodes {
-			otherNodeHost := c[otherNode].host
-
-			localAddrIndex := slices.IndexFunc(otherNodeHost.Addrs(), func(addr multiaddr.Multiaddr) bool {
-				return addr[0].Value() == "127.0.0.1"
-			})
-			require.GreaterOrEqual(t, localAddrIndex, 0)
-			localAddr := otherNodeHost.Addrs()[localAddrIndex]
-
-			peerAddr := fmt.Sprintf("%s/p2p/%s", localAddr, otherNodeHost.ID())
-			peers = append(peers, peerAddr)
-		}
-
-		if len(peers) > 0 {
-			require.NoError(t, consensus.Connect(t.Context(), thisNodeConfig.host, strings.Join(peers, ",")))
-		}
-	})
-}
-
 func (c networkConfig) add(from, to int) {
-	c[from].adjacentNodes[to] = struct{}{}
+	c[from].adjacentNodes[to] = peer.AddrInfo{ID: c[to].host.ID(), Addrs: c[to].host.Addrs()}
 }
 
 //nolint:gosec // The whole package is for testing purpose only, so it's safe to use weak random.

consensus/p2p/config/config.go

@@ -2,8 +2,15 @@ package config
 
 import "time"
 
+type TopicBufferSizes struct {
+	// Inbound
+	Subscription int
+	Output       int
+	// Outbound
+	ProtoBroadcaster int
+}
+
 type BufferSizes struct {
-	ProposalSubscription      int
 	VoteSubscription          int
 	ProposalDemux             int
 	ProposalCommitNotifier    int
@@ -13,19 +20,20 @@ type BufferSizes struct {
 	PrecommitOutput           int
 	ProposalProtoBroadcaster  int
 	VoteProtoBroadcaster      int
+	PubSubQueueSize           int
 	RetryInterval             time.Duration
 }
 
 var DefaultBufferSizes = BufferSizes{
-	ProposalSubscription:      1024,
 	VoteSubscription:          1024,
 	ProposalDemux:             1024,
-	ProposalCommitNotifier:    32,
-	ProposalSingleStreamInput: 32,
-	ProposalOutputs:           32,
+	ProposalCommitNotifier:    1024,
+	ProposalSingleStreamInput: 1024,
+	ProposalOutputs:           1024,
 	PrevoteOutput:             1024,
 	PrecommitOutput:           1024,
 	ProposalProtoBroadcaster:  1024,
 	VoteProtoBroadcaster:      1024,
+	PubSubQueueSize:           1024,
 	RetryInterval:             1 * time.Second,
 }

consensus/p2p/p2p.go

@@ -15,17 +15,23 @@ import (
 	"github.com/NethermindEth/juno/core/address"
 	"github.com/NethermindEth/juno/service"
 	"github.com/NethermindEth/juno/utils"
+	dht "github.com/libp2p/go-libp2p-kad-dht"
 	pubsub "github.com/libp2p/go-libp2p-pubsub"
 	"github.com/libp2p/go-libp2p/core/host"
+	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/p2p/discovery/routing"
 	"github.com/sourcegraph/conc"
 )
 
 type topicName string
 
 const (
-	proposalTopicName topicName = "consensus_proposals"
-	voteTopicName     topicName = "consensus_votes"
-	gossipSubHistory            = 60
+	protocolPrefix                = "starknet"
+	chainID                       = "1" // TODO: Make this configurable
+	consensusProtocolID           = "consensus"
+	proposalTopicName   topicName = "consensus_proposals"
+	voteTopicName       topicName = "consensus_votes"
+	gossipSubHistory              = 60
 )
 
 type P2P[V types.Hashable[H], H types.Hash, A types.Addr] interface {
@@ -36,12 +42,14 @@ type P2P[V types.Hashable[H], H types.Hash, A types.Addr] interface {
 }
 
 type p2p[V types.Hashable[H], H types.Hash, A types.Addr] struct {
-	host            host.Host
-	log             utils.Logger
-	commitNotifier  chan types.Height
-	broadcasters    Broadcasters[V, H, A]
-	listeners       Listeners[V, H, A]
-	topicAttachment map[topicName][]attachedToTopic
+	host             host.Host
+	log              utils.Logger
+	commitNotifier   chan types.Height
+	broadcasters     Broadcasters[V, H, A]
+	listeners        Listeners[V, H, A]
+	topicAttachment  map[topicName][]attachedToTopic
+	pubSubQueueSize  int
+	bootstrapPeersFn func() []peer.AddrInfo
 }
 
 type attachedToTopic interface {
@@ -55,6 +63,7 @@ func New(
 	proposalStore *proposal.ProposalStore[starknet.Hash],
 	currentHeight types.Height,
 	bufferSizeConfig *config.BufferSizes,
+	bootstrapPeersFn func() []peer.AddrInfo,
 ) P2P[starknet.Value, starknet.Hash, address.Address] {
 	commitNotifier := make(chan types.Height, bufferSizeConfig.ProposalCommitNotifier)
 
@@ -104,12 +113,14 @@ func New(
 	}
 
 	return &p2p[starknet.Value, starknet.Hash, address.Address]{
-		host:            host,
-		log:             log,
-		commitNotifier:  commitNotifier,
-		broadcasters:    broadcasters,
-		listeners:       listeners,
-		topicAttachment: topicAttachment,
+		host:             host,
+		log:              log,
+		commitNotifier:   commitNotifier,
+		broadcasters:     broadcasters,
+		listeners:        listeners,
+		topicAttachment:  topicAttachment,
+		pubSubQueueSize:  bufferSizeConfig.PubSubQueueSize,
+		bootstrapPeersFn: bootstrapPeersFn,
 	}
 }
 
@@ -122,7 +133,27 @@ func (p *p2p[V, H, A]) getGossipSubOptions() pubsub.Option {
 }
 
 func (p *p2p[V, H, A]) Run(ctx context.Context) error {
-	gossipSub, err := pubsub.NewGossipSub(ctx, p.host, p.getGossipSubOptions())
+	dht, err := dht.New(
+		ctx,
+		p.host,
+		dht.ProtocolPrefix("/"+protocolPrefix),
+		dht.ProtocolExtension("/"+chainID),
+		dht.ProtocolExtension("/"+consensusProtocolID),
+		dht.BootstrapPeersFunc(p.bootstrapPeersFn),
+		dht.Mode(dht.ModeServer),
+	)
+	if err != nil {
+		return fmt.Errorf("unable to create dht with error: %w", err)
+	}
+
+	gossipSub, err := pubsub.NewGossipSub(
+		ctx,
+		p.host,
+		p.getGossipSubOptions(),
+		pubsub.WithPeerOutboundQueueSize(p.pubSubQueueSize),
+		pubsub.WithValidateQueueSize(p.pubSubQueueSize),
+		pubsub.WithDiscovery(routing.NewRoutingDiscovery(dht)),
+	)
 	if err != nil {
 		return fmt.Errorf("unable to create gossipsub with error: %w", err)
 	}