add test for SupportsPartial subscribe option

Creates a mixed network where only one peer request partial, and the
rest either support or don't

Author: MarcoPolo

floodsub_test.go

@@ -57,6 +57,12 @@ func denseConnect(t *testing.T, hosts []host.Host) {
 	connectSome(t, hosts, 10)
 }
 
+func ringConnect(t *testing.T, hosts []host.Host) {
+	for i := range hosts {
+		connect(t, hosts[i], hosts[(i+1)%len(hosts)])
+	}
+}
+
 func connectSome(t *testing.T, hosts []host.Host, d int) {
 	for i, a := range hosts {
 		for j := 0; j < d; j++ {

gossipsub_test.go

@@ -4730,6 +4730,176 @@ func TestPartialMessages(t *testing.T) {
 	}
 }
 
+func TestPeerSupportsPartialMessages(t *testing.T) {
+	// First peer requests partial messages.
+	// other peers supports it.
+	// all should end with the full message.
+
+	topic := "test-topic"
+	const hostCount = 5
+	hosts := getDefaultHosts(t, hostCount)
+	psubs := make([]*PubSub, 0, len(hosts))
+	topics := make([]*Topic, 0, len(hosts))
+	subs := make([]*Subscription, 0, len(hosts))
+
+	gossipsubCtx, closeGossipsub := context.WithCancel(context.Background())
+	go func() {
+		<-gossipsubCtx.Done()
+		for _, h := range hosts {
+			h.Close()
+		}
+	}()
+
+	partialExt := make([]*partialmessages.PartialMessageExtension, hostCount)
+	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelDebug}))
+
+	// A list of maps from topic+groupID to partialMessage. One map per peer
+	// var partialMessageStoreMu sync.Mutex
+	partialMessageStore := make([]map[string]*minimalTestPartialMessage, hostCount)
+	for i := range hostCount {
+		partialMessageStore[i] = make(map[string]*minimalTestPartialMessage)
+	}
+
+	for i := range partialExt {
+		partialExt[i] = &partialmessages.PartialMessageExtension{
+			Logger: logger.With("id", i),
+			ValidateRPC: func(from peer.ID, rpc *pb.PartialMessagesExtension) error {
+				// No validation. Only for this test. In production you should
+				// have some basic fast rules here.
+				return nil
+			},
+			MergePartsMetadata: func(_ string, left, right partialmessages.PartsMetadata) partialmessages.PartsMetadata {
+				return partialmessages.MergeBitmap(left, right)
+			},
+			OnIncomingRPC: func(from peer.ID, rpc *pb.PartialMessagesExtension) error {
+				groupID := rpc.GroupID
+				pm, ok := partialMessageStore[i][topic+string(groupID)]
+				if !ok {
+					pm = &minimalTestPartialMessage{
+						Group: groupID,
+					}
+					partialMessageStore[i][topic+string(groupID)] = pm
+				}
+				if publishOpts := pm.onIncomingRPC(from, rpc); publishOpts != nil {
+					go psubs[i].PublishPartialMessage(topic, pm, *publishOpts)
+					if pm.complete() {
+						encoded, _ := pm.PartialMessageBytes(partialmessages.PartsMetadata([]byte{0}))
+						go func() {
+							err := psubs[i].Publish(topic, encoded)
+							if err != nil {
+								panic(err)
+							}
+						}()
+					}
+
+				}
+				return nil
+			},
+		}
+	}
+
+	for i, h := range hosts {
+		var topicOpts []TopicOpt
+		if i == 0 {
+			topicOpts = append(topicOpts, RequestPartialMessages())
+		} else if i == 1 {
+			// The right neighbor doesn't support partial messages
+		} else {
+			topicOpts = append(topicOpts, SupportsPartialMessages())
+		}
+
+		psub := getGossipsub(gossipsubCtx, h, WithPartialMessagesExtension(partialExt[i]))
+		topic, err := psub.Join(topic, topicOpts...)
+		if err != nil {
+			t.Fatal(err)
+		}
+		sub, err := topic.Subscribe()
+		if err != nil {
+			t.Fatal(err)
+		}
+		psubs = append(psubs, psub)
+		topics = append(topics, topic)
+		subs = append(subs, sub)
+	}
+
+	ringConnect(t, hosts)
+	time.Sleep(2 * time.Second)
+
+	group := []byte("test-group")
+	emptyMsg := &minimalTestPartialMessage{
+		Group: group,
+	}
+	fullMsg := &minimalTestPartialMessage{
+		Group: group,
+		Parts: [2][]byte{
+			[]byte("Hello"),
+			[]byte("World"),
+		},
+	}
+
+	for i := range hosts {
+		if i == 0 {
+			partialMessageStore[i][topic+string(group)] = emptyMsg
+			// first host has no data
+			err := psubs[i].PublishPartialMessage(topic, emptyMsg, partialmessages.PublishOptions{})
+			if err != nil {
+				t.Fatal(err)
+			}
+		} else {
+			if i != 1 {
+				copiedMsg := *fullMsg
+				partialMessageStore[i][topic+string(group)] = &copiedMsg
+				err := psubs[i].PublishPartialMessage(topic, &copiedMsg, partialmessages.PublishOptions{})
+				if err != nil {
+					t.Fatal(err)
+				}
+			}
+			encoded, err := fullMsg.PartialMessageBytes(partialmessages.PartsMetadata([]byte{0}))
+			if err != nil {
+				t.Fatal(err)
+			}
+			err = topics[i].Publish(context.Background(), encoded)
+			if err != nil {
+				t.Fatal(err)
+			}
+		}
+	}
+
+	time.Sleep(2 * time.Second)
+
+	// Close gossipsub before we inspect the state to avoid race conditions
+	closeGossipsub()
+	time.Sleep(1 * time.Second)
+
+	if len(partialMessageStore) != hostCount {
+		t.Errorf("One host is missing the partial message")
+	}
+
+	for i, msgStore := range partialMessageStore {
+		if i == 1 {
+			// Host 1 doesn't support partial messages
+			continue
+		}
+		if len(msgStore) == 0 {
+			t.Errorf("Host %d is missing the partial message", i)
+		}
+		for _, partialMessage := range msgStore {
+			if !partialMessage.complete() {
+				t.Errorf("expected complete message, but %v is incomplete at host %d", partialMessage, i)
+			}
+		}
+	}
+
+	for i, sub := range subs {
+		ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+		defer cancel()
+		_, err := sub.Next(ctx)
+		if err != nil {
+			t.Errorf("failed to receive message: %v at host %d", err, i)
+		}
+	}
+}
+
 func TestSkipPublishingToPeersWithPartialMessageSupport(t *testing.T) {
 	topicName := "test-topic"