Merge pull request #230 from aarshkshah1992/feat/configurable-outbound-msg-queue-size

Configurable outbound peer queue sizes

pubsub.go

@@ -3,6 +3,7 @@ package pubsub
 import (
 	"context"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"math/rand"
 	"sync"
@@ -45,6 +46,9 @@ type PubSub struct {
 
 	disc *discover
 
+	// size of the outbound message channel that we maintain for each peer
+	peerOutboundQueueSize int
+
 	// incoming messages from other peers
 	incoming chan *RPC
 
@@ -174,6 +178,7 @@ func NewPubSub(ctx context.Context, h host.Host, rt PubSubRouter, opts ...Option
 		rt:                    rt,
 		val:                   newValidation(),
 		disc:                  &discover{},
+		peerOutboundQueueSize: 32,
 		signID:                h.ID(),
 		signKey:               h.Peerstore().PrivKey(h.ID()),
 		signStrict:            true,
@@ -232,6 +237,18 @@ func NewPubSub(ctx context.Context, h host.Host, rt PubSubRouter, opts ...Option
 	return ps, nil
 }
 
+// WithPeerOutboundQueueSize is an option to set the buffer size for outbound messages to a peer
+// We start dropping messages to a peer if the outbound queue if full
+func WithPeerOutboundQueueSize(size int) Option {
+	return func(p *PubSub) error {
+		if size <= 0 {
+			return errors.New("outbound queue size must always be positive")
+		}
+		p.peerOutboundQueueSize = size
+		return nil
+	}
+}
+
 // WithMessageSigning enables or disables message signing (enabled by default).
 func WithMessageSigning(enabled bool) Option {
 	return func(p *PubSub) error {
@@ -327,7 +344,7 @@ func (p *PubSub) processLoop(ctx context.Context) {
 				continue
 			}
 
-			messages := make(chan *RPC, 32)
+			messages := make(chan *RPC, p.peerOutboundQueueSize)
 			messages <- p.getHelloPacket()
 			go p.handleNewPeer(ctx, pid, messages)
 			p.peers[pid] = messages
@@ -366,7 +383,7 @@ func (p *PubSub) processLoop(ctx context.Context) {
 				// still connected, must be a duplicate connection being closed.
 				// we respawn the writer as we need to ensure there is a stream active
 				log.Warning("peer declared dead but still connected; respawning writer: ", pid)
-				messages := make(chan *RPC, 32)
+				messages := make(chan *RPC, p.peerOutboundQueueSize)
 				messages <- p.getHelloPacket()
 				go p.handleNewPeer(ctx, pid, messages)
 				p.peers[pid] = messages