make crawler protocol messenger configurable (#1128)

Author: guillaumemichel

crawler/crawler.go

@@ -52,7 +52,13 @@ func NewDefaultCrawler(host host.Host, opts ...Option) (*DefaultCrawler, error)
 		}
 	}
 
-	pm, err := pb.NewProtocolMessenger(&messageSender{h: host, protocols: o.protocols, timeout: o.perMsgTimeout})
+	var err error
+	var pm *pb.ProtocolMessenger
+	if o.msgSenderBuilder != nil {
+		pm, err = pb.NewProtocolMessenger(o.msgSenderBuilder(host, o.protocols))
+	} else {
+		pm, err = pb.NewProtocolMessenger(&messageSender{h: host, protocols: o.protocols, timeout: o.perMsgTimeout})
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -192,7 +198,7 @@ func (c *DefaultCrawler) Run(ctx context.Context, startingPeers []*peer.AddrInfo
 	// Start worker goroutines
 	var wg sync.WaitGroup
 	wg.Add(c.parallelism)
-	for i := 0; i < c.parallelism; i++ {
+	for range c.parallelism {
 		go func() {
 			defer wg.Done()
 			for p := range jobs {

crawler/options.go

@@ -1,20 +1,24 @@
 package crawler
 
 import (
+	"slices"
 	"time"
 
 	"github.com/libp2p/go-libp2p-kad-dht/amino"
+	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
+	"github.com/libp2p/go-libp2p/core/host"
 	"github.com/libp2p/go-libp2p/core/protocol"
 )
 
 // Option DHT Crawler option type.
 type Option func(*options) error
 
 type options struct {
-	protocols      []protocol.ID
-	parallelism    int
-	connectTimeout time.Duration
-	perMsgTimeout  time.Duration
+	protocols        []protocol.ID
+	parallelism      int
+	connectTimeout   time.Duration
+	perMsgTimeout    time.Duration
+	msgSenderBuilder func(h host.Host, protos []protocol.ID) pb.MessageSenderWithDisconnect
 }
 
 // defaults are the default crawler options. This option will be automatically
@@ -31,7 +35,7 @@ var defaults = func(o *options) error {
 // WithProtocols defines the ordered set of protocols the crawler will use to talk to other nodes
 func WithProtocols(protocols []protocol.ID) Option {
 	return func(o *options) error {
-		o.protocols = append([]protocol.ID{}, protocols...)
+		o.protocols = slices.Clone(protocols)
 		return nil
 	}
 }
@@ -59,3 +63,12 @@ func WithConnectTimeout(timeout time.Duration) Option {
 		return nil
 	}
 }
+
+// WithCustomMessageSender configures the pb.MessageSender of the IpfsDHT to use the
+// custom implementation of the pb.MessageSender
+func WithCustomMessageSender(messageSenderBuilder func(h host.Host, protos []protocol.ID) pb.MessageSenderWithDisconnect) Option {
+	return func(o *options) error {
+		o.msgSenderBuilder = messageSenderBuilder
+		return nil
+	}
+}