partialmessages: add explicit MergePartsMetadata function

This allows a peer's local partsMetadata to merge with their remote
partsMetadata. As long as the underlying PartsMetadata type supports
CRDT-like addition, we can avoid sending the peer duplicates.

Take the following example (credit Sukun):

> Say: We have parts 1, 2, 5, 6
> 
> 1. Peer sends a message saying I have part 1, 2
> 2. We send parts 5, 6 and update the peerstate in PublishMessage to (1, 2, 5, 6)
> 3. Concurrently with 2, we receive message from the peer saying I now have parts (1, 2, 3, 4)

With this change we can now correctly store the fact that this peer
should have parts 1,2,3,4,5,6. And in a way that does not assume a
specific representation of the PartsMetadata.

Author: MarcoPolo

gossipsub_test.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/libp2p/go-libp2p-pubsub/internal/gologshim"
 	"github.com/libp2p/go-libp2p-pubsub/partialmessages"
+	"github.com/libp2p/go-libp2p-pubsub/partialmessages/bitmap"
 	pb "github.com/libp2p/go-libp2p-pubsub/pb"
 	"github.com/libp2p/go-libp2p/core/crypto"
 	"github.com/libp2p/go-msgio"
@@ -4408,15 +4409,14 @@ func (m *minimalTestPartialMessage) complete() bool {
 }
 
 // PartsMetadata implements partialmessages.PartialMessage.
-func (m *minimalTestPartialMessage) PartsMetadata() []byte {
-	out := byte(0)
-	if len(m.Parts[0]) > 0 {
-		out |= 1
-	}
-	if len(m.Parts[1]) > 0 {
-		out |= 2
+func (m *minimalTestPartialMessage) PartsMetadata() partialmessages.PartsMetadata {
+	out := make(bitmap.Bitmap, 1)
+	for i := range m.Parts {
+		if len(m.Parts[i]) > 0 {
+			out.Set(i)
+		}
 	}
-	return []byte{out}
+	return partialmessages.PartsMetadata(out)
 }
 
 func (m *minimalTestPartialMessage) extendFromEncodedPartialMessage(_ peer.ID, data []byte) (extended bool) {
@@ -4471,34 +4471,30 @@ func (m *minimalTestPartialMessage) GroupID() []byte {
 	return m.Group
 }
 
-func (m *minimalTestPartialMessage) PartialMessageBytes(peerPartsMetadata []byte) ([]byte, []byte, error) {
+func (m *minimalTestPartialMessage) PartialMessageBytes(peerPartsMetadata partialmessages.PartsMetadata) ([]byte, error) {
 	if len(peerPartsMetadata) == 0 {
-		return nil, nil, errors.New("invalid metadata")
+		return nil, errors.New("invalid metadata")
 	}
-	peerWants := ^peerPartsMetadata[0]
+	peerHas := bitmap.Bitmap(peerPartsMetadata)
 
 	var temp minimalTestPartialMessage
 	temp.Group = m.Group
-	if peerWants&1 == 1 && m.Parts[0] != nil {
-		peerWants ^= 1
+	if !peerHas.Get(0) && m.Parts[0] != nil {
 		temp.Parts[0] = m.Parts[0]
 	}
-	if peerWants&2 == 2 && m.Parts[1] != nil {
-		peerWants ^= 2
+	if !peerHas.Get(1) && m.Parts[1] != nil {
 		temp.Parts[1] = m.Parts[1]
 	}
 
-	restMetadata := []byte{^peerWants}
-
 	if temp.Parts[0] == nil && temp.Parts[1] == nil {
-		return nil, restMetadata, nil
+		return nil, nil
 	}
 
 	b, err := json.Marshal(temp)
 	if err != nil {
-		return nil, nil, err
+		return nil, err
 	}
-	return b, restMetadata, nil
+	return b, nil
 }
 
 func (m *minimalTestPartialMessage) shouldRequest(_ peer.ID, peerHasMetadata []byte) bool {
@@ -4516,6 +4512,10 @@ type minimalTestPartialMessageChecker struct {
 	fullMessage *minimalTestPartialMessage
 }
 
+func (m *minimalTestPartialMessageChecker) MergePartsMetadata(left, right partialmessages.PartsMetadata) partialmessages.PartsMetadata {
+	return partialmessages.MergeBitmap(left, right)
+}
+
 // EmptyMessage implements partialmessages.InvariantChecker.
 func (m *minimalTestPartialMessageChecker) EmptyMessage() *minimalTestPartialMessage {
 	return &minimalTestPartialMessage{
@@ -4606,6 +4606,9 @@ func TestPartialMessages(t *testing.T) {
 				// 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)]
@@ -4706,6 +4709,9 @@ func TestSkipPublishingToPeersWithPartialMessageSupport(t *testing.T) {
 			ValidateRPC: func(from peer.ID, rpc *pb.PartialMessagesExtension) error {
 				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 {
 				topicID := rpc.GetTopicID()
 				groupID := rpc.GetGroupID()

partialmessages/invariants.go

@@ -30,13 +30,15 @@ type InvariantChecker[P Message] interface {
 	// in (as determined from the parts metadata)
 	ShouldRequest(a P, from peer.ID, partsMetadata []byte) bool
 
+	MergePartsMetadata(left, right PartsMetadata) PartsMetadata
+
 	Equal(a, b P) bool
 }
 
 func TestPartialMessageInvariants[P Message](t *testing.T, checker InvariantChecker[P]) {
 	extend := func(a, b P) (P, error) {
 		emptyParts := checker.EmptyMessage().PartsMetadata()
-		encodedB, _, err := b.PartialMessageBytes(emptyParts)
+		encodedB, err := b.PartialMessageBytes(emptyParts)
 		if err != nil {
 			var out P
 			return out, err
@@ -64,7 +66,7 @@ func TestPartialMessageInvariants[P Message](t *testing.T, checker InvariantChec
 
 		recombined := checker.EmptyMessage()
 		for _, part := range parts {
-			b, _, err := part.PartialMessageBytes(recombined.PartsMetadata())
+			b, err := part.PartialMessageBytes(recombined.PartsMetadata())
 			if err != nil {
 				t.Fatal(err)
 			}
@@ -86,10 +88,11 @@ func TestPartialMessageInvariants[P Message](t *testing.T, checker InvariantChec
 		emptyMsgPartsMeta := emptyMessage.PartsMetadata()
 
 		// Empty message should not be able to fulfill any request
-		response, rest, err := emptyMessage.PartialMessageBytes(emptyMsgPartsMeta)
+		response, err := emptyMessage.PartialMessageBytes(emptyMsgPartsMeta)
 		if err != nil {
 			t.Fatal(err)
 		}
+		rest := checker.MergePartsMetadata(emptyMsgPartsMeta, emptyMessage.PartsMetadata())
 
 		if len(response) != 0 {
 			t.Error("Empty message should return nil response when requesting parts it doesn't have")
@@ -122,10 +125,11 @@ func TestPartialMessageInvariants[P Message](t *testing.T, checker InvariantChec
 		emptyMsgPartsMeta := emptyMessage.PartsMetadata()
 
 		// Request all parts from the partial message
-		response1, rest1, err := parts[0].PartialMessageBytes(emptyMsgPartsMeta)
+		response1, err := parts[0].PartialMessageBytes(emptyMsgPartsMeta)
 		if err != nil {
 			t.Fatal(err)
 		}
+		rest1 := checker.MergePartsMetadata(emptyMsgPartsMeta, parts[0].PartsMetadata())
 
 		// Should get some response since partial message has at least one part
 		if len(response1) == 0 {
@@ -134,10 +138,10 @@ func TestPartialMessageInvariants[P Message](t *testing.T, checker InvariantChec
 
 		// Rest should be non-zero and different from original request since something was fulfilled
 		if len(rest1) == 0 {
-			t.Error("Rest should be non-zero when partial fulfillment occurred")
+			t.Fatal("Rest should be non-zero when partial fulfillment occurred")
 		}
 		if bytes.Equal(rest1, emptyMsgPartsMeta) {
-			t.Error("Rest should be different from original request since partial fulfillment occurred")
+			t.Fatalf("Rest should be different from original request since partial fulfillment occurred")
 		}
 
 		// Create another partial message with the remaining parts
@@ -150,10 +154,11 @@ func TestPartialMessageInvariants[P Message](t *testing.T, checker InvariantChec
 		}
 
 		// The remaining partial message should be able to fulfill the "rest" request
-		response2, rest2, err := remainingPartial.PartialMessageBytes(rest1)
+		response2, err := remainingPartial.PartialMessageBytes(rest1)
 		if err != nil {
 			t.Fatal(err)
 		}
+		rest2 := checker.MergePartsMetadata(rest1, remainingPartial.PartsMetadata())
 
 		// response2 should be non-empty since we have remaining parts to fulfill
 		if len(response2) == 0 {
@@ -192,7 +197,8 @@ func TestPartialMessageInvariants[P Message](t *testing.T, checker InvariantChec
 
 		// Request with empty metadata should return all available parts
 		emptyMeta := checker.EmptyMessage().PartsMetadata()
-		response, rest, err := fullMessage.PartialMessageBytes(emptyMeta)
+		response, err := fullMessage.PartialMessageBytes(emptyMeta)
+		rest := checker.MergePartsMetadata(emptyMeta, fullMessage.PartsMetadata())
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -244,7 +250,8 @@ func TestPartialMessageInvariants[P Message](t *testing.T, checker InvariantChec
 			// Get the MissingParts() and have the full message fulfill the request
 			msgPartsMeta := testMsg.PartsMetadata()
 
-			response, rest, err := fullMessage.PartialMessageBytes(msgPartsMeta)
+			response, err := fullMessage.PartialMessageBytes(msgPartsMeta)
+			rest := checker.MergePartsMetadata(msgPartsMeta, fullMessage.PartsMetadata())
 			if err != nil {
 				t.Fatal(err)
 			}

partialmessages/partialmsgs.go

@@ -7,6 +7,7 @@ import (
 	"log/slog"
 	"slices"
 
+	"github.com/libp2p/go-libp2p-pubsub/partialmessages/bitmap"
 	pb "github.com/libp2p/go-libp2p-pubsub/pb"
 	"github.com/libp2p/go-libp2p/core/peer"
 )
@@ -21,6 +22,10 @@ import (
 
 const minGroupTTL = 3
 
+// PartsMetadata returns metadata about the parts this partial message
+// contains and, possibly implicitly, the parts it wants.
+type PartsMetadata []byte
+
 // Message is a message that can be broken up into parts. It can be
 // complete, partially complete, or empty. It is up to the application to define
 // how a message is split into parts and recombined, as well as how missing and
@@ -43,20 +48,18 @@ type Message interface {
 	//
 	// If the Partial Message is empty, the implementation MUST return:
 	// nil, metadata, nil.
-	PartialMessageBytes(partsMetadata []byte) (msg []byte, newPartsMetadata []byte, _ error)
+	PartialMessageBytes(partsMetadata PartsMetadata) (msg []byte, _ error)
 
-	// PartsMetadata returns metadata about the parts this partial message
-	// contains and, possibly implicitly, the parts it wants.
-	PartsMetadata() []byte
+	PartsMetadata() PartsMetadata
 }
 
 // peerState is the state we keep per peer. Used to make Publish
 // Idempotent.
 type peerState struct {
 	// The parts metadata the peer has sent us
-	partsMetadata []byte
-	// The parts metadata this endpoint has sent to the peer
-	sentPartsMetadata []byte
+	partsMetadata PartsMetadata
+	// The parts metadata this node has sent to the peer
+	sentPartsMetadata PartsMetadata
 }
 
 func (ps *peerState) IsZero() bool {
@@ -84,9 +87,17 @@ func (s *partialMessageStatePerTopicGroup) clearPeerWants(peerID peer.ID) {
 	}
 }
 
+// MergeBitmap is a helper function for merging parts metadata if they are a
+// bitmap.
+func MergeBitmap(left, right PartsMetadata) PartsMetadata {
+	return PartsMetadata(bitmap.Merge(bitmap.Bitmap(left), bitmap.Bitmap(right)))
+}
+
 type PartialMessageExtension struct {
 	Logger *slog.Logger
 
+	MergePartsMetadata func(topic string, left, right PartsMetadata) PartsMetadata
+
 	// OnIncomingRPC is called whenever we receive an encoded
 	// partial message from a peer. This func MUST be fast and non-blocking.
 	// If you need to do slow work, dispatch the work to your own goroutine.
@@ -147,6 +158,9 @@ func (e *PartialMessageExtension) Init(router Router) error {
 	if e.OnIncomingRPC == nil {
 		return errors.New("field OnIncomingRPC must be set")
 	}
+	if e.MergePartsMetadata == nil {
+		return errors.New("field MergePartsMetadata must be set")
+	}
 	e.statePerTopicPerGroup = make(map[string]map[string]*partialMessageStatePerTopicGroup)
 
 	return nil
@@ -186,14 +200,14 @@ func (e *PartialMessageExtension) PublishPartial(topic string, partial Message,
 		if pState.partsMetadata != nil {
 			// This peer has previously asked for a certain part. We'll give
 			// them what we can.
-			pm, rest, err := partial.PartialMessageBytes(pState.partsMetadata)
+			pm, err := partial.PartialMessageBytes(pState.partsMetadata)
 			if err != nil {
 				log.Warn("partial message extension failed to get partial message bytes", "error", err)
 				// Possibly a bad request, we'll delete the request as we will likely error next time we try to handle it
 				state.clearPeerWants(p)
 				continue
 			}
-			pState.partsMetadata = rest
+			pState.partsMetadata = e.MergePartsMetadata(topic, pState.partsMetadata, myPartsMeta)
 			if len(pm) > 0 {
 				log.Debug("Respond to peer's IWant")
 				sendRPC = true
@@ -283,7 +297,7 @@ func (e *PartialMessageExtension) HandleRPC(from peer.ID, rpc *pb.PartialMessage
 			pState = &peerState{}
 			state.peerState[from] = pState
 		}
-		pState.partsMetadata = rpc.PartsMetadata
+		pState.partsMetadata = e.MergePartsMetadata(rpc.GetTopicID(), pState.partsMetadata, rpc.PartsMetadata)
 	}
 
 	return e.OnIncomingRPC(from, rpc)