Add support for supportsPartial field in SubOpts

This lets us distinguish who can support a partial message request
versus who requests partial message on a per topic basis.

Author: MarcoPolo

comm.go

@@ -37,9 +37,9 @@ func (p *PubSub) getHelloPacket() *RPC {
 			requestPartial = ts.requestPartialMessages
 		}
 		as := &pb.RPC_SubOpts{
-			Topicid:   proto.String(t),
-			Subscribe: proto.Bool(true),
-			Partial:   &requestPartial,
+			Topicid:         proto.String(t),
+			Subscribe:       proto.Bool(true),
+			RequestsPartial: &requestPartial,
 		}
 		rpc.Subscriptions = append(rpc.Subscriptions, as)
 	}

extensions.go

@@ -199,6 +199,11 @@ type partialMessageRouter struct {
 	gs *GossipSubRouter
 }
 
+// PeerRequestsPartial implements partialmessages.Router.
+func (r partialMessageRouter) PeerRequestsPartial(peer peer.ID, topic string) bool {
+	return r.gs.peerRequestsPartial(peer, topic)
+}
+
 // MeshPeers implements partialmessages.Router.
 func (r partialMessageRouter) MeshPeers(topic string) iter.Seq[peer.ID] {
 	return func(yield func(peer.ID) bool) {
@@ -212,9 +217,9 @@ func (r partialMessageRouter) MeshPeers(topic string) iter.Seq[peer.ID] {
 		}
 
 		for peer := range peerSet {
-			if exts := r.gs.extensions.peerExtensions[peer]; exts.PartialMessages {
-				if peerStates, ok := r.gs.p.topics[topic]; ok && peerStates[peer].requestsPartial {
-					// Check that the peer wanted partial messages
+			if r.gs.extensions.peerExtensions[peer].PartialMessages {
+				if r.gs.peerSupportsPartial(peer, topic) {
+					// Peer supports partial messages
 					if !yield(peer) {
 						return
 					}

gossipsub.go

@@ -864,7 +864,7 @@ func (gs *GossipSubRouter) Preprocess(from peer.ID, msgs []*Message) {
 				// We don't send IDONTWANT to the peer that sent us the messages
 				continue
 			}
-			if gs.peerWantsPartial(p, topic) {
+			if gs.peerRequestsPartial(p, topic) {
 				// Don't send IDONTWANT to peers that are using partial messages
 				// for this topic
 				continue
@@ -1381,7 +1381,7 @@ func (gs *GossipSubRouter) rpcs(msg *Message) iter.Seq2[peer.ID, *RPC] {
 			if pid == from || pid == peer.ID(msg.GetFrom()) {
 				continue
 			}
-			if gs.peerWantsPartial(pid, topic) {
+			if gs.peerRequestsPartial(pid, topic) {
 				// The peer requested partial messages. We'll skip sending them full messages
 				continue
 			}
@@ -1393,7 +1393,12 @@ func (gs *GossipSubRouter) rpcs(msg *Message) iter.Seq2[peer.ID, *RPC] {
 	}
 }
 
-func (gs *GossipSubRouter) peerWantsPartial(p peer.ID, topic string) bool {
+func (gs *GossipSubRouter) peerSupportsPartial(p peer.ID, topic string) bool {
+	peerStates, ok := gs.p.topics[topic]
+	return ok && gs.extensions.myExtensions.PartialMessages && peerStates[p].requestsPartial
+}
+
+func (gs *GossipSubRouter) peerRequestsPartial(p peer.ID, topic string) bool {
 	peerStates, ok := gs.p.topics[topic]
 	return ok && gs.extensions.myExtensions.PartialMessages && peerStates[p].requestsPartial
 }

partialmessages/partialmsgs.go

@@ -133,6 +133,7 @@ type PublishOptions struct {
 type Router interface {
 	SendRPC(p peer.ID, r *pb.PartialMessagesExtension, urgent bool)
 	MeshPeers(topic string) iter.Seq[peer.ID]
+	PeerRequestsPartial(peer peer.ID, topic string) bool
 }
 
 func (e *PartialMessageExtension) groupState(topic string, groupID []byte, peerInitiated bool) (*partialMessageStatePerTopicGroup, error) {
@@ -200,6 +201,7 @@ func (e *PartialMessageExtension) PublishPartial(topic string, partial Message,
 	}
 	for p := range peers {
 		log := e.Logger.With("peer", p)
+		requestedPartial := e.router.PeerRequestsPartial(p, topic)
 
 		var rpc pb.PartialMessagesExtension
 		var sendRPC bool
@@ -212,7 +214,7 @@ func (e *PartialMessageExtension) PublishPartial(topic string, partial Message,
 		}
 
 		// Try to fulfill any wants from the peer
-		if pState.partsMetadata != nil {
+		if requestedPartial && pState.partsMetadata != nil {
 			// This peer has previously asked for a certain part. We'll give
 			// them what we can.
 			pm, err := partial.PartialMessageBytes(pState.partsMetadata)
@@ -234,7 +236,7 @@ func (e *PartialMessageExtension) PublishPartial(topic string, partial Message,
 		// Only send the eager push to the peer if:
 		//   - we didn't reply to an explicit request
 		//   - we have something to eager push
-		if !inResponseToIWant && len(opts.EagerPush) > 0 {
+		if requestedPartial && !inResponseToIWant && len(opts.EagerPush) > 0 {
 			log.Debug("Eager pushing")
 			sendRPC = true
 			rpc.PartialMessage = opts.EagerPush

partialmessages/partialmsgs_test.go

@@ -24,6 +24,11 @@ type testRouter struct {
 	meshPeers func(topic string) iter.Seq[peer.ID]
 }
 
+// PeerRequestsPartial implements Router.
+func (r *testRouter) PeerRequestsPartial(peer peer.ID, topic string) bool {
+	return true
+}
+
 func (r *testRouter) SendRPC(p peer.ID, rpc *pubsub_pb.PartialMessagesExtension, urgent bool) {
 	r.sendRPC(p, rpc, urgent)
 }

pubsub.go

@@ -47,6 +47,7 @@ type ProtocolMatchFn = func(protocol.ID) func(protocol.ID) bool
 
 type peerTopicState struct {
 	requestsPartial bool
+	supportsPartial bool
 }
 
 type peerOutgoingStream struct {
@@ -1197,15 +1198,18 @@ func (p *PubSub) handleRemoveRelay(topic string) {
 // Only called from processLoop.
 func (p *PubSub) announce(topic string, sub bool) {
 	var requestPartialMessages bool
+	var supportsPartialMessages bool
 	if sub {
 		if t, ok := p.myTopics[topic]; ok {
 			requestPartialMessages = t.requestPartialMessages
+			supportsPartialMessages = t.supportsPartialMessages
 		}
 	}
 	subopt := &pb.RPC_SubOpts{
-		Topicid:   &topic,
-		Subscribe: &sub,
-		Partial:   &requestPartialMessages,
+		Topicid:         &topic,
+		Subscribe:       &sub,
+		RequestsPartial: &requestPartialMessages,
+		SupportsPartial: &supportsPartialMessages,
 	}
 
 	out := rpcWithSubs(subopt)
@@ -1248,15 +1252,18 @@ func (p *PubSub) doAnnounceRetry(pid peer.ID, topic string, sub bool) {
 	}
 
 	var requestPartialMessages bool
+	var supportsPartialMessages bool
 	if sub {
 		if t, ok := p.myTopics[topic]; ok {
 			requestPartialMessages = t.requestPartialMessages
+			supportsPartialMessages = t.supportsPartialMessages
 		}
 	}
 	subopt := &pb.RPC_SubOpts{
-		Topicid:   &topic,
-		Subscribe: &sub,
-		Partial:   &requestPartialMessages,
+		Topicid:         &topic,
+		Subscribe:       &sub,
+		RequestsPartial: &requestPartialMessages,
+		SupportsPartial: &supportsPartialMessages,
 	}
 
 	out := rpcWithSubs(subopt)
@@ -1360,16 +1367,19 @@ func (p *PubSub) handleIncomingRPC(rpc *RPC) {
 				p.topics[t] = tmap
 			}
 
-			if _, ok = tmap[rpc.from]; !ok {
-				tmap[rpc.from] = peerTopicState{requestsPartial: subopt.GetPartial()}
+			pts := peerTopicState{
+				requestsPartial: subopt.GetRequestsPartial(),
+				// If they peer requested partial, they support it by default
+				supportsPartial: subopt.GetRequestsPartial() || subopt.GetSupportsPartial(),
+			}
+			_, seenBefore := tmap[rpc.from]
+			tmap[rpc.from] = pts
+			if !seenBefore {
+				tmap[rpc.from] = pts
 				if topic, ok := p.myTopics[t]; ok {
 					peer := rpc.from
 					topic.sendNotification(PeerEvent{PeerJoin, peer})
 				}
-			} else {
-				s := tmap[rpc.from]
-				s.requestsPartial = subopt.GetPartial()
-				tmap[rpc.from] = s
 			}
 		} else {
 			tmap, ok := p.topics[t]
@@ -1565,6 +1575,22 @@ func RequestPartialMessages() TopicOpt {
 			return errors.New("partial messages are not enabled")
 		}
 		t.requestPartialMessages = true
+		t.supportsPartialMessages = true
+		return nil
+	}
+}
+
+func SupportsPartialMessages() TopicOpt {
+	return func(t *Topic) error {
+		gs, ok := t.p.rt.(*GossipSubRouter)
+		if !ok {
+			return errors.New("partial messages only supported by gossipsub")
+		}
+
+		if !gs.extensions.myExtensions.PartialMessages {
+			return errors.New("partial messages are not enabled")
+		}
+		t.supportsPartialMessages = true
 		return nil
 	}
 }

topic.go

@@ -33,7 +33,8 @@ type Topic struct {
 	mux    sync.RWMutex
 	closed bool
 
-	requestPartialMessages bool
+	requestPartialMessages  bool
+	supportsPartialMessages bool
 }
 
 // String returns the topic associated with t