fix(p2p): trim the simple priority queue below its limit

internal/p2p/load_router_queue_test.go

@@ -33,9 +33,15 @@ func TestLoadRouterSharedQueueOccupancyAndDrops(t *testing.T) {
 	// queue the router builds, at a size small enough to reach its ceiling.
 	const capacity = 8
 	q := newSimplePriorityQueue(ctx, capacity)
-	// Account for the queue's heap and both channel buffers. Depending on
-	// scheduling, draining can make every one of these retained envelopes
-	// observable before the queue gets another chance to trim its heap.
+	// The queue holds envelopes in three places: waiting in the input buffer,
+	// in the heap, and popped into the output buffer. A push that takes the
+	// heap past maxSize trims it, so it holds at most maxSize once a push
+	// finishes. The run loop holds one envelope in none of these between taking
+	// it from input and pushing it, which frees an input slot for the producer,
+	// but if that push overflows the heap it also trims it, so the envelope
+	// cannot add to the count. A flood that nobody drains can fill all three,
+	// and draining can then hand over every one of them before a later push
+	// trims the heap again.
 	ceiling := q.maxSize + cap(q.input) + cap(q.output)
 
 	// Nothing drains while this runs, which is the case that matters: a

internal/p2p/rqueue.go

@@ -9,6 +9,10 @@ import (
 	"github.com/cosmos/gogoproto/proto"
 )
 
+// simpleQueueTrimFraction sets how far below maxSize an overflowing heap is
+// trimmed: by 1/simpleQueueTrimFraction of maxSize.
+const simpleQueueTrimFraction = 16
+
 type simpleQueue struct {
 	input   chan Envelope
 	output  chan Envelope
@@ -73,26 +77,21 @@ func (q *simpleQueue) run(ctx context.Context) {
 				timestamp: time.Now().UTC(),
 			})
 
+			// Check on every push, so the heap never holds more than
+			// maxSize. Checking only when the signal was already pending
+			// let every push that refilled it skip the trim, so with
+			// nobody draining, the heap ran past maxSize until the next
+			// tick.
+			if len(pq) > q.maxSize {
+				pq = q.trim(pq)
+			}
+
 			select {
 			case signal <- struct{}{}:
 			default:
-				if len(pq) > q.maxSize {
-					sort.Sort(pq)
-					for _, dropped := range pq[q.maxSize:] {
-						dropped.envelope.NotifyDelivery()
-					}
-					pq = pq[:q.maxSize]
-				}
 			}
 
 		case <-ticker.C:
-			if len(pq) > q.maxSize {
-				sort.Sort(pq)
-				for _, dropped := range pq[q.maxSize:] {
-					dropped.envelope.NotifyDelivery()
-				}
-				pq = pq[:q.maxSize]
-			}
 			if len(pq) > 0 {
 				select {
 				case signal <- struct{}{}:
@@ -118,3 +117,21 @@ func (q *simpleQueue) run(ctx context.Context) {
 		}
 	}
 }
+
+// trim drops the lowest-ranked envelopes from a heap that has outgrown maxSize.
+// Sorting takes time in proportion to the whole heap, and the run loop accepts
+// nothing while it sorts, so trimming back to exactly maxSize would sort again
+// on the very next push. Trimming a fraction below it leaves room for the pushes
+// that follow, which spreads the cost of one sort across all of them. The sorted
+// slice it returns is still a valid heap.
+func (q *simpleQueue) trim(pq priorityQueue) priorityQueue {
+	sort.Sort(pq)
+	keep := q.maxSize - q.maxSize/simpleQueueTrimFraction
+	for _, dropped := range pq[keep:] {
+		dropped.envelope.NotifyDelivery()
+	}
+	// Later pushes reclaim these slots only one at a time, so release the
+	// dropped envelopes now.
+	clear(pq[keep:])
+	return pq[:keep]
+}

internal/p2p/rqueue_test.go

@@ -1,9 +1,12 @@
 package p2p
 
 import (
+	"container/heap"
 	"context"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestSimpleQueue(t *testing.T) {
@@ -45,3 +48,61 @@ RETRY:
 	}
 
 }
+
+// An overflowing heap is trimmed far enough below maxSize that the next pushes
+// do not each need another sort. It keeps the highest-ranked envelopes, stays a
+// valid heap, and releases both the senders waiting on what it drops and the
+// memory those envelopes hold.
+func TestSimpleQueueTrimDropsLowestRankedBelowMaxSize(t *testing.T) {
+	testCases := []struct {
+		name    string
+		maxSize int
+		keep    int
+	}{
+		{name: "large enough for a margin", maxSize: 64, keep: 60},
+		{name: "too small for a margin", maxSize: 4, keep: 4},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			q := &simpleQueue{maxSize: tc.maxSize}
+			overflow := tc.maxSize + 1
+
+			// Timestamps this far apart decide the rank on their own, and the
+			// newest ranks highest, so the oldest are the ones to lose.
+			start := time.Now()
+			pq := make(priorityQueue, 0, overflow)
+			delivered := make([]<-chan struct{}, overflow)
+			timestamps := make([]time.Time, overflow)
+			for i := range overflow {
+				pqEnv := &pqEnvelope{timestamp: start.Add(time.Duration(i) * time.Second)}
+				delivered[i] = pqEnv.envelope.EnableDeliveryNotification()
+				timestamps[i] = pqEnv.timestamp
+				heap.Push(&pq, pqEnv)
+			}
+
+			trimmed := q.trim(pq)
+			require.Len(t, trimmed, tc.keep)
+
+			firstKept := overflow - tc.keep
+			for i, done := range delivered {
+				select {
+				case <-done:
+					require.Less(t, i, firstKept, "released the sender of a kept envelope")
+				default:
+					require.GreaterOrEqual(t, i, firstKept, "did not release the sender of a dropped envelope")
+				}
+			}
+			for _, slot := range trimmed[tc.keep:overflow] {
+				require.Nil(t, slot, "a dropped envelope is still reachable")
+			}
+
+			for i := overflow - 1; i >= firstKept; i-- {
+				next, ok := heap.Pop(&trimmed).(*pqEnvelope)
+				require.True(t, ok)
+				require.True(t, timestamps[i].Equal(next.timestamp), "the trimmed heap is out of order")
+			}
+			require.Empty(t, trimmed)
+		})
+	}
+}