add Topic option to skip publishing to partial message-capable peers

Author: MarcoPolo

gossipsub.go

@@ -1295,6 +1295,13 @@ func (gs *GossipSubRouter) PublishBatch(messages []*Message, opts *BatchPublishO
 	}
 }
 
+func (gs *GossipSubRouter) skipPartialMessageCapablePeers(topicID string) bool {
+	if t, ok := gs.p.myTopics[topicID]; ok {
+		return t.skipPublishingToPartialMessageCapablePeers
+	}
+	return false
+}
+
 func (gs *GossipSubRouter) Publish(msg *Message) {
 	for p, rpc := range gs.rpcs(msg) {
 		gs.sendRPC(p, rpc, false)
@@ -1371,8 +1378,9 @@ func (gs *GossipSubRouter) rpcs(msg *Message) iter.Seq2[peer.ID, *RPC] {
 		}
 
 		out := rpcWithMessages(msg.Message)
+		skipPMPeers := gs.skipPartialMessageCapablePeers(msg.GetTopic())
 		for pid := range tosend {
-			if pid == from || pid == peer.ID(msg.GetFrom()) {
+			if pid == from || pid == peer.ID(msg.GetFrom()) || (skipPMPeers && gs.extensions.peerExtensions[pid].PartialMessages) {
 				continue
 			}
 

gossipsub_test.go

@@ -62,15 +62,6 @@ func getGossipsubs(ctx context.Context, hs []host.Host, opts ...Option) []*PubSu
 	return psubs
 }
 
-func getGossipsubsOptFn(ctx context.Context, hs []host.Host, optFn func(int, host.Host) []Option) []*PubSub {
-	var psubs []*PubSub
-	for i, h := range hs {
-		opts := optFn(i, h)
-		psubs = append(psubs, getGossipsub(ctx, h, opts...))
-	}
-	return psubs
-}
-
 func TestGossipSubParamsValidate(t *testing.T) {
 	params := DefaultGossipSubParams()
 	params.Dhi = 1
@@ -92,6 +83,15 @@ func TestGossipSubBootstrapParamsValidate(t *testing.T) {
 	}
 }
 
+func getGossipsubsOptFn(ctx context.Context, hs []host.Host, optFn func(int, host.Host) []Option) []*PubSub {
+	var psubs []*PubSub
+	for i, h := range hs {
+		opts := optFn(i, h)
+		psubs = append(psubs, getGossipsub(ctx, h, opts...))
+	}
+	return psubs
+}
+
 func TestSparseGossipsub(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -4685,3 +4685,112 @@ func TestPartialMessages(t *testing.T) {
 		}
 	}
 }
+
+func TestSkipPublishingToPeersWithPartialMessageSupport(t *testing.T) {
+	topicName := "test-topic"
+
+	// 3 hosts.
+	// hosts[0]: Publisher. Supports partial messages
+	// hosts[1]: Subscriber. Supports partial messages
+	// hosts[2]: Alternate publisher. Does not support partial messages. Only
+	//   connected to hosts[0]
+	hosts := getDefaultHosts(t, 3)
+
+	const hostsWithPartialMessageSupport = 2
+	partialExt := make([]*partialmessages.PartialMessageExtension, hostsWithPartialMessageSupport)
+	// A list of maps from topic+groupID to partialMessage. One map per peer
+	partialMessageStore := make([]map[string]*minimalTestPartialMessage, hostsWithPartialMessageSupport)
+	for i := range hostsWithPartialMessageSupport {
+		partialMessageStore[i] = make(map[string]*minimalTestPartialMessage)
+	}
+
+	// Only hosts with partial message support
+	psubs := make([]*PubSub, 0, len(hosts)-1)
+
+	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug}))
+
+	for i := range partialExt {
+		partialExt[i] = &partialmessages.PartialMessageExtension{
+			Logger: logger,
+			ValidateRPC: func(from peer.ID, rpc *pb.PartialMessagesExtension) error {
+				return nil
+			},
+			OnIncomingRPC: func(from peer.ID, topic string, groupID, iwant, ihave, partialMessageBytes []byte) {
+				pm, ok := partialMessageStore[i][topic+string(groupID)]
+				if !ok {
+					pm = &minimalTestPartialMessage{
+						Group: groupID,
+					}
+					partialMessageStore[i][topic+string(groupID)] = pm
+				}
+				if publishOpts := pm.onIncomingRPC(from, topic, groupID, iwant, ihave, partialMessageBytes); publishOpts != nil {
+					go psubs[i].PublishPartialMessage(topic, pm, *publishOpts)
+				}
+			},
+		}
+	}
+
+	for i, h := range hosts[:2] {
+		psub := getGossipsub(context.Background(), h, WithPartialMessagesExtension(partialExt[i]))
+		psubs = append(psubs, psub)
+	}
+
+	nonPartialPubsub := getGossipsub(context.Background(), hosts[2])
+
+	denseConnect(t, hosts[:2])
+	time.Sleep(2 * time.Second)
+
+	// Connect nonPartialPubsub to the publisher
+	connect(t, hosts[0], hosts[2])
+
+	var topics []*Topic
+	var subs []*Subscription
+	for _, psub := range psubs {
+		topic, err := psub.Join(topicName, WithSkipPublishingToPartialMessageCapablePeers())
+		if err != nil {
+			t.Fatal(err)
+		}
+		topics = append(topics, topic)
+		s, err := topic.Subscribe()
+		if err != nil {
+			t.Fatal(err)
+		}
+		subs = append(subs, s)
+	}
+
+	topicForNonPartial, err := nonPartialPubsub.Join(topicName)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Wait for subscriptions to propagate
+	time.Sleep(time.Second)
+
+	topics[0].Publish(context.Background(), []byte("Hello"))
+
+	// Publish from another peer, the publisher (psub[0]) should not forward this to psub[1].
+	topicForNonPartial.Publish(context.Background(), []byte("from non-partial"))
+
+	recvdMessage := make(chan struct{}, 1)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	go func() {
+		msg, err := subs[1].Next(ctx)
+		if err == context.Canceled {
+			return
+		}
+		if err != nil {
+			t.Log(err)
+			t.Fail()
+			return
+		}
+		t.Log("Received msg", string(msg.Data))
+		recvdMessage <- struct{}{}
+	}()
+
+	select {
+	case <-recvdMessage:
+		t.Fatal("Received message")
+	case <-time.After(2 * time.Second):
+	}
+}

pubsub.go

@@ -1376,7 +1376,7 @@ func (p *PubSub) handleIncomingRPC(rpc *RPC) {
 				continue
 			}
 
-			msg := &Message{pmsg, "", rpc.from, nil, false}
+			msg := &Message{Message: pmsg, ID: "", ReceivedFrom: rpc.from, ValidatorData: nil, Local: false}
 			if p.shouldPush(msg) {
 				toPush = append(toPush, msg)
 			}
@@ -1515,7 +1515,16 @@ type rmTopicReq struct {
 	resp  chan error
 }
 
-type TopicOptions struct{}
+type TopicOptions struct {
+	SkipPublishingToPartialMessageCapablePeers bool
+}
+
+func WithSkipPublishingToPartialMessageCapablePeers() TopicOpt {
+	return func(t *Topic) error {
+		t.skipPublishingToPartialMessageCapablePeers = true
+		return nil
+	}
+}
 
 type TopicOpt func(t *Topic) error
 

topic.go

@@ -32,6 +32,8 @@ type Topic struct {
 
 	mux    sync.RWMutex
 	closed bool
+
+	skipPublishingToPartialMessageCapablePeers bool
 }
 
 // String returns the topic associated with t
@@ -348,7 +350,14 @@ func (t *Topic) validate(ctx context.Context, data []byte, opts ...PubOpt) (*Mes
 		}
 	}
 
-	msg := &Message{m, "", t.p.host.ID(), pub.validatorData, pub.local}
+	msg := &Message{
+		Message:       m,
+		ID:            "",
+		ReceivedFrom:  t.p.host.ID(),
+		ValidatorData: pub.validatorData,
+		Local:         pub.local,
+	}
+
 	select {
 	case t.p.eval <- func() {
 		t.p.rt.Preprocess(t.p.host.ID(), []*Message{msg})