feat: add SPREAD anonymous gossip protocol (opt-in)

extensions.go

@@ -6,13 +6,15 @@ import (
 
 	"github.com/libp2p/go-libp2p-pubsub/partialmessages"
 	pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb"
+	"github.com/libp2p/go-libp2p-pubsub/vivaldi"
 	"github.com/libp2p/go-libp2p/core/peer"
 	"github.com/libp2p/go-libp2p/core/protocol"
 )
 
 type PeerExtensions struct {
 	TestExtension   bool
 	PartialMessages bool
+	Spread          bool
 }
 
 type TestExtensionConfig struct {
@@ -44,6 +46,7 @@ func peerExtensionsFromRPC(rpc *RPC) PeerExtensions {
 	if hasPeerExtensions(rpc) {
 		out.TestExtension = rpc.Control.Extensions.GetTestExtension()
 		out.PartialMessages = rpc.Control.Extensions.GetPartialMessages()
+		out.Spread = rpc.Control.Extensions.GetSpread()
 	}
 	return out
 }
@@ -67,6 +70,15 @@ func (pe *PeerExtensions) ExtendRPC(rpc *RPC) *RPC {
 		}
 		rpc.Control.Extensions.PartialMessages = &pe.PartialMessages
 	}
+	if pe.Spread {
+		if rpc.Control == nil {
+			rpc.Control = &pubsub_pb.ControlMessage{}
+		}
+		if rpc.Control.Extensions == nil {
+			rpc.Control.Extensions = &pubsub_pb.ControlExtensions{}
+		}
+		rpc.Control.Extensions.Spread = &pe.Spread
+	}
 	return rpc
 }
 
@@ -91,6 +103,7 @@ type extensionsState struct {
 	testExtension     *testExtension
 
 	partialMessagesExtension partialMessageInterface
+	spreadState              *SpreadState
 }
 
 func newExtensionsState(myExtensions PeerExtensions, reportMisbehavior func(peer.ID), sendRPC func(peer.ID, *RPC, bool)) *extensionsState {
@@ -101,6 +114,7 @@ func newExtensionsState(myExtensions PeerExtensions, reportMisbehavior func(peer
 		reportMisbehavior: reportMisbehavior,
 		sendRPC:           sendRPC,
 		testExtension:     nil,
+		spreadState:       NewSpreadState(),
 	}
 }
 
@@ -168,13 +182,20 @@ func (es *extensionsState) extensionsOnNewOutboundStream(id peer.ID) {
 	if es.myExtensions.TestExtension && es.peerExtensions[id].TestExtension {
 		es.testExtension.OnNewOutboundStream(id)
 	}
+
+	if es.myExtensions.Spread && es.peerExtensions[id].Spread {
+		es.spreadState.AddPeer(id)
+	}
 }
 
 // extensionsOnClosedOutboundStream is always called after extensionsOnNewOutboundStream.
 func (es *extensionsState) extensionsOnClosedOutboundStream(id peer.ID) {
 	if es.myExtensions.PartialMessages && es.peerExtensions[id].PartialMessages {
 		es.partialMessagesExtension.OnClosedOutboundStream(id)
 	}
+	if es.myExtensions.Spread && es.peerExtensions[id].Spread {
+		es.spreadState.RemovePeer(id)
+	}
 }
 
 func (es *extensionsState) extensionsHandleRPC(rpc *RPC) error {
@@ -189,6 +210,18 @@ func (es *extensionsState) extensionsHandleRPC(rpc *RPC) error {
 		}
 	}
 
+	// Update SPREAD topic membership when peers announce subscriptions
+	if es.myExtensions.Spread && es.peerExtensions[rpc.from].Spread && rpc.Subscriptions != nil {
+		for _, sub := range rpc.GetSubscriptions() {
+			topic := sub.GetTopicid()
+			if sub.GetSubscribe() {
+				es.spreadState.AddPeerTopic(topic, rpc.from)
+			} else {
+				es.spreadState.RemovePeerTopic(topic, rpc.from)
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -254,6 +287,41 @@ func PublishPartial[PeerState any](ps *PubSub, topic string, groupID []byte, pub
 	}
 }
 
+// WithVivaldi wires a Vivaldi service into the spread extension and starts
+// the runner. Passing a nil service disables Vivaldi.
+func WithVivaldi(vsvc *vivaldi.Service, cfg *VivaldiConfig) Option {
+	return func(ps *PubSub) error {
+		gs, ok := ps.rt.(*GossipSubRouter)
+		if !ok {
+			return errors.New("pubsub router is not gossipsub")
+		}
+		if gs.extensions == nil || gs.extensions.spreadState == nil {
+			return errors.New("spread extension state not initialized")
+		}
+		gs.extensions.spreadState.ConfigureVivaldi(vsvc, cfg)
+		// Start runner if service is provided
+		if vsvc != nil {
+			gs.extensions.spreadState.StartVivaldiRunner()
+		}
+		return nil
+	}
+}
+
+// WithSpreadClusteringConfig configures spread clustering behaviour.
+func WithSpreadClusteringConfig(cfg *SpreadClusteringConfig) Option {
+	return func(ps *PubSub) error {
+		gs, ok := ps.rt.(*GossipSubRouter)
+		if !ok {
+			return errors.New("pubsub router is not gossipsub")
+		}
+		if gs.extensions == nil || gs.extensions.spreadState == nil {
+			return errors.New("spread extension state not initialized")
+		}
+		gs.extensions.spreadState.ConfigureClustering(cfg)
+		return nil
+	}
+}
+
 type partialMessageRouter struct {
 	gs *GossipSubRouter
 }

gossip_protocols.go

@@ -0,0 +1,13 @@
+package pubsub
+
+// Gossip protocol choices selectable via WithProtocolChoice. GOSSIPSUB is the
+// default behaviour; SPREAD enables the SPREAD anonymous-gossip selection for
+// messages marked with msg.Spread.
+const (
+	GOSSIPSUB = 1
+	SPREAD    = 2
+)
+
+// GossipProtocolChoice selects which forwarding strategy the GossipSub router
+// uses. See the GOSSIPSUB and SPREAD constants.
+type GossipProtocolChoice int

gossipsub.go

@@ -294,26 +294,28 @@ func NewGossipSubWithRouter(ctx context.Context, h host.Host, rt PubSubRouter, o
 func DefaultGossipSubRouter(h host.Host) *GossipSubRouter {
 	params := DefaultGossipSubParams()
 	rt := &GossipSubRouter{
-		peers:           make(map[peer.ID]protocol.ID),
-		mesh:            make(map[string]map[peer.ID]struct{}),
-		fanout:          make(map[string]map[peer.ID]struct{}),
-		lastpub:         make(map[string]int64),
-		gossip:          make(map[peer.ID][]*pb.ControlIHave),
-		control:         make(map[peer.ID]*pb.ControlMessage),
-		backoff:         make(map[string]map[peer.ID]time.Time),
-		peerhave:        make(map[peer.ID]int),
-		peerdontwant:    make(map[peer.ID]int),
-		unwanted:        make(map[peer.ID]map[checksum]int),
-		iasked:          make(map[peer.ID]int),
-		outbound:        make(map[peer.ID]bool),
-		connect:         make(chan connectInfo, params.MaxPendingConnections),
-		cab:             pstoremem.NewAddrBook(),
-		mcache:          NewMessageCache(params.HistoryGossip, params.HistoryLength),
-		protos:          GossipSubDefaultProtocols,
-		feature:         GossipSubDefaultFeatures,
-		tagTracer:       newTagTracer(h.ConnManager()),
-		params:          params,
-		reducePXRecords: defaultPXRecordReducer,
+		peers:                       make(map[peer.ID]protocol.ID),
+		mesh:                        make(map[string]map[peer.ID]struct{}),
+		fanout:                      make(map[string]map[peer.ID]struct{}),
+		lastpub:                     make(map[string]int64),
+		gossip:                      make(map[peer.ID][]*pb.ControlIHave),
+		control:                     make(map[peer.ID]*pb.ControlMessage),
+		backoff:                     make(map[string]map[peer.ID]time.Time),
+		peerhave:                    make(map[peer.ID]int),
+		peerdontwant:                make(map[peer.ID]int),
+		unwanted:                    make(map[peer.ID]map[checksum]int),
+		iasked:                      make(map[peer.ID]int),
+		outbound:                    make(map[peer.ID]bool),
+		connect:                     make(chan connectInfo, params.MaxPendingConnections),
+		cab:                         pstoremem.NewAddrBook(),
+		mcache:                      NewMessageCache(params.HistoryGossip, params.HistoryLength),
+		protos:                      GossipSubDefaultProtocols,
+		feature:                     GossipSubDefaultFeatures,
+		tagTracer:                   newTagTracer(h.ConnManager()),
+		params:                      params,
+		reducePXRecords:             defaultPXRecordReducer,
+		spreadPropagation:           NewSpreadPropagation(nil),
+		spreadRelayDuplicateCounter: make(map[spreadDuplicateKey]int),
 	}
 
 	rt.extensions = newExtensionsState(PeerExtensions{}, func(p peer.ID) {
@@ -364,6 +366,35 @@ func DefaultGossipSubParams() GossipSubParams {
 	}
 }
 
+// WithProtocolChoice is a gossipsub router option that selects the forwarding
+// strategy. The default is GOSSIPSUB; pass SPREAD to enable SPREAD selection
+// for messages marked with msg.Spread.
+func WithProtocolChoice(choice GossipProtocolChoice) Option {
+	return func(p *PubSub) error {
+		// If p can be casted to GossipSub, set the protocol choice on the router.
+		gs, ok := p.rt.(*GossipSubRouter)
+		if !ok {
+			return fmt.Errorf("cannot set protocol choice since pubsub router is not gossipsub")
+		}
+		gs.gossipProtocolChoice = choice
+		return nil
+	}
+}
+
+// WithSpreadPropagationConfig is a gossipsub router option that sets the SPREAD
+// peer-selection parameters (intra/inter-cluster fanout and probabilities).
+// Unset or invalid fields fall back to the package defaults.
+func WithSpreadPropagationConfig(cfg *SpreadConfig) Option {
+	return func(ps *PubSub) error {
+		gs, ok := ps.rt.(*GossipSubRouter)
+		if !ok {
+			return fmt.Errorf("pubsub router is not gossipsub")
+		}
+		gs.spreadPropagation = NewSpreadPropagation(cfg)
+		return nil
+	}
+}
+
 // WithPeerScore is a gossipsub router option that enables peer scoring.
 func WithPeerScore(params *PeerScoreParams, thresholds *PeerScoreThresholds) Option {
 	return func(ps *PubSub) error {
@@ -650,6 +681,14 @@ type GossipSubRouter struct {
 	// number of heartbeats since the beginning of time; this allows us to amortize some resource
 	// clean up -- eg backoff clean up.
 	heartbeatTicks uint64
+
+	// GossipProtocolChoice
+	gossipProtocolChoice GossipProtocolChoice
+
+	// SPREAD propagation selector.
+	spreadPropagation *SpreadPropagation
+	// Per-(topic, message ID) duplicate SPREAD re-propagation counters.
+	spreadRelayDuplicateCounter map[spreadDuplicateKey]int
 }
 
 var _ BatchPublisher = &GossipSubRouter{}
@@ -659,6 +698,11 @@ type connectInfo struct {
 	spr *record.Envelope
 }
 
+type spreadDuplicateKey struct {
+	topic string
+	msgID string
+}
+
 func (gs *GossipSubRouter) Protocols() []protocol.ID {
 	return gs.protos
 }
@@ -1349,6 +1393,54 @@ func (gs *GossipSubRouter) Publish(msg *Message) {
 	}
 }
 
+func (gs *GossipSubRouter) canRelaySeenSpreadDuplicate(msg *Message) bool {
+	if msg == nil || !msg.Spread {
+		return false
+	}
+	if gs.spreadPropagation == nil || gs.extensions == nil || !gs.extensions.myExtensions.Spread {
+		return false
+	}
+	return gs.spreadPropagation.config.DuplicateRepropagation > 0
+}
+
+func (gs *GossipSubRouter) hasSpreadDuplicateRelayBudget(topic, msgID string) bool {
+	maxRelays := gs.spreadPropagation.config.DuplicateRepropagation
+	if maxRelays <= 0 {
+		return false
+	}
+
+	key := spreadDuplicateKey{topic: topic, msgID: msgID}
+	return gs.spreadRelayDuplicateCounter[key] < maxRelays
+}
+
+func (gs *GossipSubRouter) consumeSpreadDuplicateRelayBudget(topic, msgID string) bool {
+	maxRelays := gs.spreadPropagation.config.DuplicateRepropagation
+	if maxRelays <= 0 {
+		return false
+	}
+
+	key := spreadDuplicateKey{topic: topic, msgID: msgID}
+	used := gs.spreadRelayDuplicateCounter[key]
+	if used >= maxRelays {
+		return false
+	}
+
+	gs.spreadRelayDuplicateCounter[key] = used + 1
+	return true
+}
+
+func (gs *GossipSubRouter) cleanupSpreadDuplicateRelayBudget() {
+	if len(gs.spreadRelayDuplicateCounter) == 0 {
+		return
+	}
+	for key := range gs.spreadRelayDuplicateCounter {
+		if gs.p.seenMessage(key.msgID) {
+			continue
+		}
+		delete(gs.spreadRelayDuplicateCounter, key)
+	}
+}
+
 func (gs *GossipSubRouter) rpcs(msg *Message) iter.Seq2[peer.ID, *RPC] {
 	return func(yield func(peer.ID, *RPC) bool) {
 		gs.mcache.Put(msg)
@@ -1405,7 +1497,39 @@ func (gs *GossipSubRouter) rpcs(msg *Message) iter.Seq2[peer.ID, *RPC] {
 			}
 		}
 
+		// If this message was marked as SPREAD and SPREAD mode is enabled on this router,
+		// override the default selection with the SPREAD-selected peers.
+		if gs.gossipProtocolChoice == SPREAD && msg.Spread {
+			useAngular := false
+			if gs.spreadPropagation != nil && gs.spreadPropagation.config != nil {
+				useAngular = gs.spreadPropagation.config.UseAngularInterPeers
+			}
+			clusterPeers, interClusterPeers := gs.extensions.spreadState.GetPropagationPeers(topic, gs.p.host.ID(), useAngular)
+			if gs.spreadPropagation == nil {
+				gs.spreadPropagation = NewSpreadPropagation(nil)
+			}
+			selected := gs.spreadPropagation.GetPeersForPropagation(from, clusterPeers, interClusterPeers)
+
+			spreadTosend := make(map[peer.ID]struct{}, len(selected))
+			for _, p := range selected {
+				if p == from || p == peer.ID(msg.GetFrom()) {
+					continue
+				}
+				spreadTosend[p] = struct{}{}
+			}
+
+			if len(spreadTosend) > 0 {
+				tosend = spreadTosend
+			}
+		}
+
 		out := rpcWithMessages(msg.Message)
+		// Set RPC-level spread extension if both the router is configured to use
+		// SPREAD by default and the message explicitly requests spread.
+		if gs.gossipProtocolChoice == SPREAD && msg.Spread {
+			v := true
+			out.Spread = &pb.SpreadExtension{SourceIsSpreadNode: &v}
+		}
 		for pid := range tosend {
 			if pid == from || pid == peer.ID(msg.GetFrom()) {
 				continue
@@ -1909,6 +2033,7 @@ func (gs *GossipSubRouter) heartbeat() {
 
 	// advance the message history window
 	gs.mcache.Shift()
+	gs.cleanupSpreadDuplicateRelayBudget()
 
 	gs.extensions.Heartbeat()
 }