Add ability to skip publishing to partial message capable peers

Author: MarcoPolo

gossipsub.go

@@ -1183,6 +1183,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)
@@ -1259,8 +1266,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

@@ -4377,5 +4377,102 @@ outer:
 	if len(missing) != 0 {
 		t.Errorf("Expected no missing parts, got %v", missing)
 	}
+}
+
+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)
+
+	partialExt := make([]*partialmessages.PartialMessageExtension, 2)
+	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
+			},
+			EagerIWantLimitPerHeartbeat: 0,
+			IWantLimitPerHeartbeat:      1,
+			NewPartialMessage: func(topic string, groupID []byte) (partialmessages.PartialMessage, error) {
+				return &minimalTestPartialMessage{
+					Group: groupID,
+					onExtended: func(m *minimalTestPartialMessage) {
+						t.Logf("Received new part and extended partial message")
+					},
+				}, nil
+			},
+		}
+	}
+
+	psubs := make([]*PubSub, 0, len(hosts)-1)
+	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

@@ -1365,7 +1365,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)
 			}
@@ -1504,7 +1504,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,
+	}
+
 	t.p.eval <- func() {
 		t.p.rt.Preprocess(t.p.host.ID(), []*Message{msg})
 	}