cancell enqueued rpcs on receiving IDontWant

Author: sukunrt

comm.go

@@ -165,7 +165,7 @@ func (p *PubSub) handlePeerDead(s network.Stream) {
 }
 
 func (p *PubSub) handleSendingMessages(ctx context.Context, s network.Stream, outgoing *rpcQueue) {
-	writeRpc := func(rpc *RPC) error {
+	writeRpc := func(rpc *pb.RPC) error {
 		size := uint64(rpc.Size())
 
 		buf := pool.Get(varint.UvarintSize(size) + int(size))
@@ -193,8 +193,11 @@ func (p *PubSub) handleSendingMessages(ctx context.Context, s network.Stream, ou
 			p.logger.Debug("error popping message from the queue to send to peer", "peer", s.Conn().RemotePeer(), "err", err)
 			return
 		}
+		if rpc.Size() == 0 {
+			continue
+		}
 
-		err = writeRpc(rpc)
+		err = writeRpc(&rpc.RPC)
 		if err != nil {
 			s.Reset()
 			p.logger.Debug("error writing message to peer", "peer", s.Conn().RemotePeer(), "err", err)
@@ -215,6 +218,10 @@ func rpcWithMessages(msgs ...*pb.Message) *RPC {
 	return &RPC{RPC: pb.RPC{Publish: msgs}}
 }
 
+func rpcWithMessageAndMsgID(msg *pb.Message, msgID string) *RPC {
+	return &RPC{RPC: pb.RPC{Publish: []*pb.Message{msg}}, MsgIDs: []string{msgID}}
+}
+
 func rpcWithControl(msgs []*pb.Message,
 	ihave []*pb.ControlIHave,
 	iwant []*pb.ControlIWant,

gossipsub.go

@@ -1163,6 +1163,9 @@ func (gs *GossipSubRouter) handleIDontWant(p peer.ID, ctl *pb.ControlMessage) {
 	gs.peerdontwant[p]++
 
 	totalUnwantedIds := 0
+	// Collect message IDs for cancellation
+	var msgIDsToCancel []string
+
 	// Remember all the unwanted message ids
 mainIDWLoop:
 	for _, idontwant := range ctl.GetIdontwant() {
@@ -1175,8 +1178,14 @@ mainIDWLoop:
 
 			totalUnwantedIds++
 			gs.unwanted[p][computeChecksum(mid)] = gs.params.IDontWantMessageTTL
+			msgIDsToCancel = append(msgIDsToCancel, mid)
 		}
 	}
+
+	// Cancel these messages in the RPC queue if it exists
+	if queue, ok := gs.p.peers[p]; ok && len(msgIDsToCancel) > 0 {
+		queue.CancelMessages(msgIDsToCancel)
+	}
 }
 
 func (gs *GossipSubRouter) addBackoff(p peer.ID, topic string, isUnsubscribe bool) {
@@ -1370,7 +1379,7 @@ func (gs *GossipSubRouter) rpcs(msg *Message) iter.Seq2[peer.ID, *RPC] {
 			}
 		}
 
-		out := rpcWithMessages(msg.Message)
+		out := rpcWithMessageAndMsgID(msg.Message, gs.p.idGen.ID(msg))
 		for pid := range tosend {
 			if pid == from || pid == peer.ID(msg.GetFrom()) {
 				continue

gossipsub_spam_test.go

@@ -930,20 +930,25 @@ func TestGossipsubHandleIDontwantSpam(t *testing.T) {
 	rPid := hosts[1].ID()
 	ctrlMessage := &pb.ControlMessage{Idontwant: []*pb.ControlIDontWant{{MessageIDs: idwIds}}}
 	grt := psubs[0].rt.(*GossipSubRouter)
-	grt.handleIDontWant(rPid, ctrlMessage)
+	completed := make(chan struct{})
+	psubs[0].eval <- func() {
+		grt.handleIDontWant(rPid, ctrlMessage)
 
-	if grt.peerdontwant[rPid] != 1 {
-		t.Errorf("Wanted message count of %d but received %d", 1, grt.peerdontwant[rPid])
-	}
-	mid := fmt.Sprintf("idontwant-%d", GossipSubMaxIDontWantLength-1)
-	if _, ok := grt.unwanted[rPid][computeChecksum(mid)]; !ok {
-		t.Errorf("Desired message id was not stored in the unwanted map: %s", mid)
-	}
+		if grt.peerdontwant[rPid] != 1 {
+			t.Errorf("Wanted message count of %d but received %d", 1, grt.peerdontwant[rPid])
+		}
+		mid := fmt.Sprintf("idontwant-%d", GossipSubMaxIDontWantLength-1)
+		if _, ok := grt.unwanted[rPid][computeChecksum(mid)]; !ok {
+			t.Errorf("Desired message id was not stored in the unwanted map: %s", mid)
+		}
 
-	mid = fmt.Sprintf("idontwant-%d", GossipSubMaxIDontWantLength)
-	if _, ok := grt.unwanted[rPid][computeChecksum(mid)]; ok {
-		t.Errorf("Unwanted message id was stored in the unwanted map: %s", mid)
+		mid = fmt.Sprintf("idontwant-%d", GossipSubMaxIDontWantLength)
+		if _, ok := grt.unwanted[rPid][computeChecksum(mid)]; ok {
+			t.Errorf("Unwanted message id was stored in the unwanted map: %s", mid)
+		}
+		close(completed)
 	}
+	<-completed
 }
 
 type mockGSOnRead func(writeMsg func(*pb.RPC), irpc *pb.RPC)

pubsub.go

@@ -257,6 +257,8 @@ func (m *Message) GetFrom() peer.ID {
 
 type RPC struct {
 	pb.RPC
+	// MsgIDs are the ids of the messages in the rpc. MsgID[i] = id(rpc.Publish[i])
+	MsgIDs []string
 
 	// unexported on purpose, not sending this over the wire
 	from peer.ID
@@ -274,6 +276,16 @@ func (rpc *RPC) split(limit int) iter.Seq[RPC] {
 
 			messagesInNextRPC := 0
 			messageSlice := rpc.Publish
+			var msgIDSlice []string
+			if len(rpc.MsgIDs) != len(rpc.Publish) {
+				if len(rpc.MsgIDs) == 0 {
+					msgIDSlice = make([]string, len(rpc.Publish))
+				} else {
+					panic("MsgIDs and Publish have different lengths")
+				}
+			} else {
+				msgIDSlice = rpc.MsgIDs
+			}
 
 			// Merge/Append publish messages. This pattern is optimized compared the
 			// the patterns for other fields because this is the common cause for
@@ -285,7 +297,10 @@ func (rpc *RPC) split(limit int) iter.Seq[RPC] {
 					// The message doesn't fit. Let's set the messages that did fit
 					// into this RPC, yield it, then make a new one
 					nextRPC.Publish = messageSlice[:messagesInNextRPC]
+					nextRPC.MsgIDs = msgIDSlice[:messagesInNextRPC]
 					messageSlice = messageSlice[messagesInNextRPC:]
+					msgIDSlice = msgIDSlice[messagesInNextRPC:]
+
 					if !yield(nextRPC) {
 						return
 					}
@@ -303,6 +318,7 @@ func (rpc *RPC) split(limit int) iter.Seq[RPC] {
 				// packing this RPC, but we avoid successively calling .Size()
 				// on the messages for the next parts.
 				nextRPC.Publish = messageSlice[:messagesInNextRPC]
+				nextRPC.MsgIDs = msgIDSlice[:messagesInNextRPC]
 				if !yield(nextRPC) {
 					return
 				}

rpc_queue.go

@@ -3,6 +3,7 @@ package pubsub
 import (
 	"context"
 	"errors"
+	"slices"
 	"sync"
 )
 
@@ -50,20 +51,41 @@ type rpcQueue struct {
 	dataAvailable  sync.Cond
 	spaceAvailable sync.Cond
 	// Mutex used to access queue
-	queueMu sync.Mutex
-	queue   priorityQueue
+	queueMu      sync.Mutex
+	queue        priorityQueue
+	queuedMsgIDs map[string]int      // message ids in queue
+	cancelledIDs map[string]struct{} // messages ids that'll be dropped before sending
 
 	closed  bool
 	maxSize int
 }
 
 func newRpcQueue(maxSize int) *rpcQueue {
-	q := &rpcQueue{maxSize: maxSize}
+	q := &rpcQueue{
+		maxSize:      maxSize,
+		queuedMsgIDs: make(map[string]int),
+		cancelledIDs: make(map[string]struct{}),
+	}
 	q.dataAvailable.L = &q.queueMu
 	q.spaceAvailable.L = &q.queueMu
 	return q
 }
 
+// CancelMessages marks the given message IDs for cancellation only if they are already in queue.
+func (q *rpcQueue) CancelMessages(msgIDs []string) {
+	q.queueMu.Lock()
+	defer q.queueMu.Unlock()
+
+	for _, id := range msgIDs {
+		if id != "" {
+			// Only cancel messages that are actually in the queue
+			if count := q.queuedMsgIDs[id]; count > 0 {
+				q.cancelledIDs[id] = struct{}{}
+			}
+		}
+	}
+}
+
 func (q *rpcQueue) Push(rpc *RPC, block bool) error {
 	return q.push(rpc, false, block)
 }
@@ -91,11 +113,17 @@ func (q *rpcQueue) push(rpc *RPC, urgent bool, block bool) error {
 			return ErrQueueFull
 		}
 	}
+
 	if urgent {
 		q.queue.PriorityPush(rpc)
 	} else {
 		q.queue.NormalPush(rpc)
 	}
+	for _, id := range rpc.MsgIDs {
+		if id != "" {
+			q.queuedMsgIDs[id]++
+		}
+	}
 
 	q.dataAvailable.Signal()
 	return nil
@@ -133,10 +161,48 @@ func (q *rpcQueue) Pop(ctx context.Context) (*RPC, error) {
 		}
 	}
 	rpc := q.queue.Pop()
+	rpc = q.handleCancellations(rpc)
 	q.spaceAvailable.Signal()
 	return rpc, nil
 }
 
+func (q *rpcQueue) handleCancellations(rpc *RPC) *RPC {
+	hasCancellations := false
+	for _, msgID := range rpc.MsgIDs {
+		q.queuedMsgIDs[msgID]--
+		if q.queuedMsgIDs[msgID] <= 0 {
+			delete(q.queuedMsgIDs, msgID)
+		}
+		if _, ok := q.cancelledIDs[msgID]; ok {
+			hasCancellations = true
+		}
+	}
+	if hasCancellations {
+		// clone the RPC parts that we'll modify. It may be shared with other queues.
+		newRPC := *rpc
+		newRPC.RPC.Publish = slices.Clone(rpc.RPC.Publish)
+		newRPC.MsgIDs = slices.Clone(rpc.MsgIDs)
+		rpc = &newRPC
+		for i, msgID := range slices.Backward(newRPC.MsgIDs) {
+			if msgID == "" {
+				continue
+			}
+			_, ok := q.cancelledIDs[msgID]
+			if !ok {
+				continue
+			}
+			rpc.RPC.Publish[i] = nil
+			rpc.MsgIDs[i] = ""
+			if q.queuedMsgIDs[msgID] <= 0 {
+				delete(q.cancelledIDs, msgID)
+			}
+			rpc.RPC.Publish = slices.Delete(rpc.RPC.Publish, i, i+1) // this slice is small
+			rpc.MsgIDs = slices.Delete(rpc.MsgIDs, i, i+1)
+		}
+	}
+	return rpc
+}
+
 func (q *rpcQueue) Close() {
 	q.queueMu.Lock()
 	defer q.queueMu.Unlock()

rpc_queue_test.go

@@ -2,8 +2,12 @@ package pubsub
 
 import (
 	"context"
+	"fmt"
+	"slices"
 	"testing"
 	"time"
+
+	pb "github.com/libp2p/go-libp2p-pubsub/pb"
 )
 
 func TestNewRpcQueue(t *testing.T) {
@@ -227,3 +231,112 @@ func TestRpcQueueCancelPop(t *testing.T) {
 		t.Fatalf("rpc queue Pop returns wrong error when it's cancelled")
 	}
 }
+
+func TestRPCQueueCancellations(t *testing.T) {
+	maxSize := 32
+	q := newRpcQueue(maxSize)
+
+	getMesssages := func(n int) ([]*pb.Message, []string) {
+		msgs := make([]*pb.Message, n)
+		msgIDs := make([]string, n)
+		for i := range msgs {
+			msgs[i] = &pb.Message{Data: []byte(fmt.Sprintf("message%d", i+1))}
+			msgIDs[i] = fmt.Sprintf("msg%d", i+1)
+		}
+		return msgs, msgIDs
+	}
+
+	t.Run("cancel all", func(t *testing.T) {
+		msgs, msgIDs := getMesssages(10)
+		rpc := &RPC{
+			RPC:    pb.RPC{Publish: slices.Clone(msgs)},
+			MsgIDs: slices.Clone(msgIDs),
+		}
+		q.Push(rpc, true)
+		q.CancelMessages(msgIDs)
+		popped, err := q.Pop(context.Background())
+		if err != nil {
+			t.Fatalf("failed to pop RPC: %v", err)
+		}
+		if len(popped.Publish) != 0 {
+			t.Fatalf("expected popped.Publish to be empty, got %v", popped.Publish)
+		}
+		if len(popped.MsgIDs) != 0 {
+			t.Fatalf("expected popped.MsgIDs to be empty, got %v", popped.MsgIDs)
+		}
+		if len(q.queuedMsgIDs) != 0 {
+			t.Fatalf("expected q.queuedMsgIDs to be empty, got %v", q.queuedMsgIDs)
+		}
+		if len(q.cancelledIDs) != 0 {
+			t.Fatalf("expected q.cancelledIDs to be empty, got %v", q.cancelledIDs)
+		}
+	})
+
+	t.Run("cancel some", func(t *testing.T) {
+		msgs, msgIDs := getMesssages(10)
+		rpc := &RPC{
+			RPC:    pb.RPC{Publish: slices.Clone(msgs)},
+			MsgIDs: slices.Clone(msgIDs),
+		}
+		q.Push(rpc, true)
+		q.CancelMessages(msgIDs[:3])
+		popped, err := q.Pop(context.Background())
+		if err != nil {
+			t.Fatalf("failed to pop RPC: %v", err)
+		}
+		if !slices.Equal(msgs[3:], popped.Publish) {
+			t.Fatalf("expected popped.Publish to be %v, got %v", msgs[3:], popped.Publish)
+		}
+		if !slices.Equal(msgIDs[3:], popped.MsgIDs) {
+			t.Fatalf("expected popped.MsgIDs to be %v, got %v", msgIDs[3:], popped.MsgIDs)
+		}
+		if len(q.queuedMsgIDs) != 0 {
+			t.Fatalf("expected q.queuedMsgIDs to be empty, got %v", q.queuedMsgIDs)
+		}
+		if len(q.cancelledIDs) != 0 {
+			t.Fatalf("expected q.cancelledIDs to be empty, got %v", q.cancelledIDs)
+		}
+	})
+
+	t.Run("cancel duplicate", func(t *testing.T) {
+		msgs, msgIDs := getMesssages(10)
+		rpc := &RPC{
+			RPC:    pb.RPC{Publish: slices.Clone(msgs)},
+			MsgIDs: slices.Clone(msgIDs),
+		}
+		q.Push(rpc, true)
+		rpc2 := &RPC{
+			RPC:    pb.RPC{Publish: slices.Clone(msgs)},
+			MsgIDs: slices.Clone(msgIDs),
+		}
+		q.Push(rpc2, true)
+		q.CancelMessages(msgIDs[:3])
+		popped, err := q.Pop(context.Background())
+		if err != nil {
+			t.Fatalf("failed to pop RPC: %v", err)
+		}
+		if !slices.Equal(msgs[3:], popped.Publish) {
+			t.Fatalf("expected popped.Publish to be %v, got %v", msgs[3:], popped.Publish)
+		}
+		if !slices.Equal(msgIDs[3:], popped.MsgIDs) {
+			t.Fatalf("expected popped.MsgIDs to be %v, got %v", msgIDs[3:], popped.MsgIDs)
+		}
+
+		popped, err = q.Pop(context.Background())
+		if err != nil {
+			t.Fatalf("failed to pop RPC: %v", err)
+		}
+		if !slices.Equal(msgs[3:], popped.Publish) {
+			t.Fatalf("expected popped.Publish to be %v, got %v", msgs[3:], popped.Publish)
+		}
+		if !slices.Equal(msgIDs[3:], popped.MsgIDs) {
+			t.Fatalf("expected popped.MsgIDs to be %v, got %v", msgIDs[3:], popped.MsgIDs)
+		}
+		if len(q.queuedMsgIDs) != 0 {
+			t.Fatalf("expected q.queuedMsgIDs to be empty, got %v", q.queuedMsgIDs)
+		}
+		if len(q.cancelledIDs) != 0 {
+			t.Fatalf("expected q.cancelledIDs to be empty, got %v", q.cancelledIDs)
+		}
+	})
+}