add partial messages to gossipsub router

Author: MarcoPolo

extensions.go

@@ -1,12 +1,17 @@
 package pubsub
 
 import (
+	"errors"
+	"iter"
+
+	"github.com/libp2p/go-libp2p-pubsub/partialmessages"
 	pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb"
 	"github.com/libp2p/go-libp2p/core/peer"
 )
 
 type PeerExtensions struct {
-	TestExtension bool
+	TestExtension   bool
+	PartialMessages bool
 }
 
 type TestExtensionConfig struct {
@@ -37,6 +42,7 @@ func peerExtensionsFromRPC(rpc *RPC) PeerExtensions {
 	out := PeerExtensions{}
 	if hasPeerExtensions(rpc) {
 		out.TestExtension = rpc.Control.Extensions.GetTestExtension()
+		out.PartialMessages = rpc.Control.Extensions.GetPartialMessages()
 	}
 	return out
 }
@@ -50,6 +56,14 @@ func (pe *PeerExtensions) ExtendRPC(rpc *RPC) *RPC {
 			TestExtension: &pe.TestExtension,
 		}
 	}
+	if pe.PartialMessages {
+		if rpc.Control == nil {
+			rpc.Control = &pubsub_pb.ControlMessage{}
+		}
+		rpc.Control.Extensions = &pubsub_pb.ControlExtensions{
+			PartialMessages: &pe.PartialMessages,
+		}
+	}
 	return rpc
 }
 
@@ -59,8 +73,9 @@ type extensionsState struct {
 	sentExtensions    map[peer.ID]struct{}
 	reportMisbehavior func(peer.ID)
 	sendRPC           func(p peer.ID, r *RPC, urgent bool)
+	testExtension     *testExtension
 
-	testExtension *testExtension
+	partialMessagesExtension *partialmessages.PartialMessageExtension
 }
 
 func newExtensionsState(myExtensions PeerExtensions, reportMisbehavior func(peer.ID), sendRPC func(peer.ID, *RPC, bool)) *extensionsState {
@@ -132,14 +147,73 @@ func (es *extensionsState) extensionsAddPeer(id peer.ID) {
 	if es.myExtensions.TestExtension && es.peerExtensions[id].TestExtension {
 		es.testExtension.AddPeer(id)
 	}
+
+	if es.myExtensions.PartialMessages && es.peerExtensions[id].PartialMessages {
+		es.partialMessagesExtension.AddPeer(id)
+	}
 }
 
 // extensionsRemovePeer is always called after extensionsAddPeer.
 func (es *extensionsState) extensionsRemovePeer(id peer.ID) {
+	if es.myExtensions.PartialMessages && es.peerExtensions[id].PartialMessages {
+		es.partialMessagesExtension.RemovePeer(id)
+	}
 }
 
 func (es *extensionsState) extensionsHandleRPC(rpc *RPC) {
 	if es.myExtensions.TestExtension && es.peerExtensions[rpc.from].TestExtension {
 		es.testExtension.HandleRPC(rpc.from, rpc.TestExtension)
 	}
+
+	if es.myExtensions.PartialMessages && es.peerExtensions[rpc.from].PartialMessages && rpc.Partial != nil {
+		es.partialMessagesExtension.HandleRPC(rpc.from, rpc.Partial)
+	}
+}
+
+func (es *extensionsState) Heartbeat() {
+	if es.myExtensions.PartialMessages {
+		es.partialMessagesExtension.Heartbeat()
+	}
+}
+
+func WithPartialMessagesExtension(pm *partialmessages.PartialMessageExtension) Option {
+	return func(ps *PubSub) error {
+		gs, ok := ps.rt.(*GossipSubRouter)
+		if !ok {
+			return errors.New("pubsub router is not gossipsub")
+		}
+		pm.Init(routerForPartialMessage{gs})
+
+		gs.extensions.myExtensions.PartialMessages = true
+		gs.extensions.partialMessagesExtension = pm
+		return nil
+	}
+}
+
+type routerForPartialMessage struct {
+	gs *GossipSubRouter
+}
+
+// MeshPeers implements partialmessages.Router.
+func (r routerForPartialMessage) MeshPeers(topic string) iter.Seq[peer.ID] {
+	return func(yield func(peer.ID) bool) {
+		for peer := range r.gs.mesh[topic] {
+			if exts := r.gs.extensions.peerExtensions[peer]; exts.PartialMessages {
+				if !yield(peer) {
+					return
+				}
+			}
+		}
+	}
+}
+
+// SendRPC implements partialmessages.Router.
+func (r routerForPartialMessage) SendRPC(p peer.ID, rpc *pubsub_pb.PartialMessagesExtension, urgent bool) {
+	r.gs.sendRPC(p, &RPC{
+		RPC: pubsub_pb.RPC{
+			Partial: rpc,
+		},
+	}, urgent)
 }
+
+var _ partialmessages.Router = routerForPartialMessage{}

gossipsub.go

@@ -1833,6 +1833,8 @@ func (gs *GossipSubRouter) heartbeat() {
 
 	// advance the message history window
 	gs.mcache.Shift()
+
+	gs.extensions.Heartbeat()
 }
 
 func (gs *GossipSubRouter) clearIHaveCounters() {