adapter/sqs: align peek partition iteration with effectivePartitionCount

Three r3 findings on PR #794:

1. Codex P2 (sqs_admin_peek.go:268). peekStartPartition and
   walkPeek's partition advance / termination keyed off
   meta.PartitionCount directly. In perQueue throughput mode
   (FifoThroughputLimit=perQueue) the data plane collapses every
   MessageGroupId to partition 0 (see effectivePartitionCount in
   sqs_keys_dispatch.go); peek still rotated over all N partitions,
   wasting up to N-1 guaranteed-empty ScanAt calls per peek (31
   extra scans at PartitionCount=32). Switch to
   effectivePartitionCount(meta) in peekStartPartition, walkPeek's
   partition advance, and preparePeekCursor's bounds check.
   Caller audit: effectivePartitionCount is the same helper
   receiveMessage uses for its fanout loop (sqs_messages.go:885),
   so peek and receive now agree on partition usage.

2. CodeRabbit minor (test order-coupling). HappyPath asserted row
   ordering matched send order, but vis-index entries with the same
   visible_at millisecond tie-break on message_id (random hex) — a
   timing-sensitive flake. Replace ordered loop with set-based
   assertion via assertPeekRowsAsSet.

3. CodeRabbit minor (counter type assertion panic). The fanout
   counter test used `counter, _ := v.(*atomic.Uint32)` then
   counter.Load(), which panics if the assertion fails. Extract
   loadFanoutCounter helper that fails the test explicitly on bad
   type / nil rather than panicking.

Tests: TestAdminPeekQueue_PerQueueFIFOCollapsesToOnePartition pins
the perQueue functional contract end-to-end.
TestPreparePeekCursor_PerQueueCollapse pins the cursor codec's
effective-bounds validation (a Partition=2 cursor against a
perQueue PartitionCount=4 queue is rejected — even though the raw
PartitionCount would otherwise allow it).

Author: bootjp

adapter/sqs_admin_peek.go

@@ -262,40 +262,56 @@ func (s *SQSServer) AdminPeekQueue(
 // peek pages between receives lands every receive on the same
 // partition, starving the other three).
 func (s *SQSServer) peekStartPartition(queueName string, cursor *peekCursor, meta *sqsQueueMeta) uint32 {
-	if cursor != nil || meta.PartitionCount <= 1 {
+	if cursor != nil {
 		return 0
 	}
-	return s.nextReceiveFanoutStart(queueName, meta.PartitionCount)
+	effective := effectivePartitionCount(meta)
+	if effective <= 1 {
+		return 0
+	}
+	return s.nextReceiveFanoutStart(queueName, effective)
 }
 
 // preparePeekCursor builds the effective cursor for this call. On the
 // first page (cursor==nil) it stamps Generation + a rotated
-// StartPartition (partitioned queues only — non-partitioned start at
-// 0). On a follow-up page it validates the stored Generation matches
-// the queue's current generation AND that StartPartition / Partition
-// are within [0, max(PartitionCount, 1)); a mismatch on either
+// StartPartition (partitioned queues only — non-partitioned and
+// perQueue-throughput FIFO queues collapse to partition 0). On a
+// follow-up page it validates the stored Generation matches the
+// queue's current generation AND that StartPartition / Partition are
+// within [0, effectivePartitionCount(meta)); a mismatch on either
 // returns ErrAdminSQSValidation.
 //
+// effectivePartitionCount (not meta.PartitionCount) is the
+// authoritative iteration bound. perQueue FIFO mode collapses every
+// MessageGroupId to partition 0 (see partitionFor /
+// effectivePartitionCount), so partitions 1..N-1 are guaranteed empty
+// for those queues; walking them would be pointless read
+// amplification (Codex r3 P2). Keying validation off the effective
+// count keeps peek aligned with the data-plane's actual partition
+// usage.
+//
 // Bounds-check rationale: walkPeek terminates the partitioned
-// rotation when `(Partition + 1) % PartitionCount == StartPartition`.
-// If a client supplies StartPartition outside [0, PartitionCount),
-// that termination condition never fires and the call loops
-// ScanAt-by-ScanAt over guaranteed-empty partitions forever — a
-// request-amplification DoS against the admin endpoint (Codex r1
-// P1 on PR #794). Rejecting bad cursor partition indices up-front
-// closes the vector. Generation mismatch separately forces a
-// front-of-stream refresh after a purge.
+// rotation when `(Partition + 1) % effectivePartitionCount ==
+// StartPartition`. If a client supplies StartPartition outside
+// [0, effectivePartitionCount), that termination condition never
+// fires and the call loops ScanAt-by-ScanAt over guaranteed-empty
+// partitions forever — a request-amplification DoS against the
+// admin endpoint (Codex r1 P1 on PR #794). Rejecting bad cursor
+// partition indices up-front closes the vector. Generation mismatch
+// separately forces a front-of-stream refresh after a purge.
 //
 // startPartition is computed by the caller (the SQSServer method's
-// nextReceiveFanoutStart) so this helper stays method-free for
-// unit-testability. Non-partitioned queues pass 0.
+// peekStartPartition wrapper around nextReceiveFanoutStart) so this
+// helper stays method-free for unit-testability. Non-partitioned and
+// perQueue-collapsed queues pass 0.
 func preparePeekCursor(cursor *peekCursor, meta *sqsQueueMeta, startPartition uint32) (*peekCursor, error) {
+	effective := effectivePartitionCount(meta)
 	if cursor == nil {
 		out := &peekCursor{
 			V:          peekCursorSchemaV1,
 			Generation: meta.Generation,
 		}
-		if meta.PartitionCount > 1 {
+		if effective > 1 {
 			out.StartPartition = startPartition
 			out.Partition = startPartition
 		}
@@ -305,18 +321,10 @@ func preparePeekCursor(cursor *peekCursor, meta *sqsQueueMeta, startPartition ui
 		return nil, errors.Wrap(ErrAdminSQSValidation,
 			"admin peek: cursor is from a prior generation; restart from the front")
 	}
-	maxPartition := meta.PartitionCount
-	if maxPartition <= 1 {
-		// Non-partitioned queues only have partition 0. A non-zero
-		// Partition or StartPartition would also fail the rotation
-		// termination check; reject explicitly so the error is the
-		// documented 400, not a silent O(infty) loop.
-		maxPartition = 1
-	}
-	if cursor.StartPartition >= maxPartition || cursor.Partition >= maxPartition {
+	if cursor.StartPartition >= effective || cursor.Partition >= effective {
 		return nil, errors.Wrapf(ErrAdminSQSValidation,
 			"admin peek: cursor partition index out of range (StartPartition=%d, Partition=%d, max=%d)",
-			cursor.StartPartition, cursor.Partition, maxPartition)
+			cursor.StartPartition, cursor.Partition, effective)
 	}
 	return cursor, nil
 }
@@ -350,11 +358,15 @@ func (s *SQSServer) walkPeek(
 			cursor.LastKey = next
 			return rows, cursor, nil
 		}
-		// Partition exhausted before Limit reached.
-		if meta.PartitionCount <= 1 {
+		// Partition exhausted before Limit reached. effectivePartitionCount
+		// (not meta.PartitionCount) caps the rotation so perQueue-collapsed
+		// FIFO queues stop after partition 0 instead of walking N-1
+		// guaranteed-empty partitions (Codex r3 P2 on PR #794).
+		effective := effectivePartitionCount(meta)
+		if effective <= 1 {
 			return rows, nil, nil
 		}
-		nextPart := (cursor.Partition + 1) % meta.PartitionCount
+		nextPart := (cursor.Partition + 1) % effective
 		if nextPart == cursor.StartPartition {
 			return rows, nil, nil
 		}

adapter/sqs_admin_peek_test.go

@@ -46,30 +46,38 @@ func TestAdminPeekQueue_HappyPath(t *testing.T) {
 	if nextCursor != "" {
 		t.Fatalf("nextCursor=%q want empty (queue drained in one page)", nextCursor)
 	}
-	for i, row := range rows {
-		assertPeekRowMatchesIndexedBody(t, i, row, "body-")
-	}
+	want := map[string]struct{}{"body-0": {}, "body-1": {}, "body-2": {}}
+	assertPeekRowsAsSet(t, rows, want)
 }
 
-// assertPeekRowMatchesIndexedBody pins the per-row invariants the
-// happy path expects: the body matches bodyPrefix+i, no truncation
-// fired (the bodies are short), MessageID is populated, and
-// SentTimestamp is non-zero. Pulled into a helper so
-// TestAdminPeekQueue_HappyPath stays under the cyclop budget.
-func assertPeekRowMatchesIndexedBody(t *testing.T, idx int, row AdminPeekedMessage, bodyPrefix string) {
+// assertPeekRowsAsSet pins the per-row invariants (MessageID
+// populated, SentTimestamp non-zero, no truncation) AND the set of
+// observed bodies, without coupling to row order. Vis-index entries
+// with the same visible_at millisecond tie-break on message_id which
+// is random, so an order-sensitive assertion is flaky under fast
+// sends — CodeRabbit r3 caught the earlier ordered-loop version.
+func assertPeekRowsAsSet(t *testing.T, rows []AdminPeekedMessage, want map[string]struct{}) {
 	t.Helper()
-	want := bodyPrefix + strconv.Itoa(idx)
-	if row.Body != want {
-		t.Fatalf("rows[%d].Body=%q want %q", idx, row.Body, want)
-	}
-	if row.BodyTruncated {
-		t.Fatalf("rows[%d].BodyTruncated=true; want false for %d-byte body", idx, len(row.Body))
+	got := make(map[string]struct{}, len(rows))
+	for i, row := range rows {
+		if row.BodyTruncated {
+			t.Fatalf("rows[%d].BodyTruncated=true; want false for short body", i)
+		}
+		if row.MessageID == "" {
+			t.Fatalf("rows[%d].MessageID empty", i)
+		}
+		if row.SentTimestamp.IsZero() {
+			t.Fatalf("rows[%d].SentTimestamp zero", i)
+		}
+		got[row.Body] = struct{}{}
 	}
-	if row.MessageID == "" {
-		t.Fatalf("rows[%d].MessageID empty", idx)
+	for w := range want {
+		if _, ok := got[w]; !ok {
+			t.Fatalf("missing body %q in observed set %v", w, got)
+		}
 	}
-	if row.SentTimestamp.IsZero() {
-		t.Fatalf("rows[%d].SentTimestamp zero", idx)
+	if len(got) != len(want) {
+		t.Fatalf("observed set %v != expected set %v", got, want)
 	}
 }
 
@@ -826,6 +834,24 @@ func TestPreparePeekCursor_PartitionOutOfRange(t *testing.T) {
 	})
 }
 
+// loadFanoutCounter retrieves the per-queue receive fanout counter
+// the server's receive path uses for partition rotation. Returns the
+// counter or fails the test if it is absent / has an unexpected
+// type. Pulled into a helper so callers that read-then-compare the
+// counter value stay under the cyclop budget.
+func loadFanoutCounter(t *testing.T, node Node, queueName string) *atomic.Uint32 {
+	t.Helper()
+	v, ok := node.sqsServer.receiveFanoutCounters.Load(queueName)
+	if !ok {
+		t.Fatalf("fanout counter missing for queue %q (bump-on-first-page contract broken)", queueName)
+	}
+	counter, ok := v.(*atomic.Uint32)
+	if !ok || counter == nil {
+		t.Fatalf("fanout counter for %q has unexpected type %T (want *atomic.Uint32)", queueName, v)
+	}
+	return counter
+}
+
 // TestAdminPeekQueue_ContinuationDoesNotBumpFanoutCounter pins Codex
 // r2 P1: AdminPeekQueue must NOT advance the shared
 // receiveFanoutCounters counter on continuation pages. The counter
@@ -863,11 +889,7 @@ func TestAdminPeekQueue_ContinuationDoesNotBumpFanoutCounter(t *testing.T) {
 		t.Fatalf("first peek: rows=%d cursor=%q — need both populated to exercise continuation", len(rows), cursor)
 	}
 
-	v, ok := node.sqsServer.receiveFanoutCounters.Load(name)
-	if !ok {
-		t.Fatalf("first peek did not create the fanout counter; the bump-on-first-page contract is broken")
-	}
-	counter, _ := v.(*atomic.Uint32)
+	counter := loadFanoutCounter(t, node, name)
 	before := counter.Load()
 
 	// Drive 5 continuation calls; each must not move the counter.
@@ -887,6 +909,75 @@ func TestAdminPeekQueue_ContinuationDoesNotBumpFanoutCounter(t *testing.T) {
 	}
 }
 
+// TestAdminPeekQueue_PerQueueFIFOCollapsesToOnePartition pins Codex
+// r3 P2: a partitioned FIFO queue in perQueue throughput mode
+// collapses every MessageGroupId to partition 0 (see
+// effectivePartitionCount). Peek must align with the data plane: it
+// walks only the effective partitions, not the raw PartitionCount
+// count, so a PartitionCount=4 perQueue queue does not waste 3
+// guaranteed-empty ScanAt calls on partitions 1..3 per peek call.
+//
+// Functional assertion: peek succeeds end-to-end on a perQueue queue
+// and surfaces every sent message. The optimization itself is
+// internal — a follow-up that re-introduces the meta.PartitionCount
+// rotation would still pass this assertion but the (Partition+1) %
+// effective termination encoded in walkPeek would still cap the
+// scan, so the regression vector is closed by the cursor codec
+// validation (effective bounds) being the chokepoint.
+func TestAdminPeekQueue_PerQueueFIFOCollapsesToOnePartition(t *testing.T) {
+	t.Parallel()
+	nodes, _, _ := createNode(t, 1)
+	defer shutdown(nodes)
+	node := sqsLeaderNode(t, nodes)
+
+	const name = "peek-perqueue.fifo"
+	status, out := callSQS(t, node, sqsCreateQueueTarget, map[string]any{
+		"QueueName":  name,
+		"Attributes": map[string]string{"FifoQueue": "true"},
+	})
+	if status != http.StatusOK {
+		t.Fatalf("create: %d %v", status, out)
+	}
+	queueURL, _ := out["QueueUrl"].(string)
+	installPartitionedMetaForTest(t, node, name, 4, htfifoThroughputPerQueue)
+
+	groups := []string{"alpha", "beta", "gamma"}
+	sent := sendFIFOMessagesForPeek(t, node, queueURL, groups)
+	seen := walkPeekUntilDone(t, context.Background(), node, name, 10)
+	assertEveryMessageSeenOnce(t, sent, seen)
+}
+
+// TestPreparePeekCursor_PerQueueCollapse confirms the cursor codec's
+// validation key off effectivePartitionCount, not the raw
+// meta.PartitionCount. A perQueue queue with PartitionCount=4 must
+// reject any cursor with Partition > 0 (the effective bound is 1)
+// even though meta.PartitionCount itself is 4.
+func TestPreparePeekCursor_PerQueueCollapse(t *testing.T) {
+	t.Parallel()
+
+	t.Run("fresh cursor on perQueue stamps StartPartition=0", func(t *testing.T) {
+		t.Parallel()
+		meta := &sqsQueueMeta{Generation: 1, PartitionCount: 4, FifoThroughputLimit: htfifoThroughputPerQueue}
+		out, err := preparePeekCursor(nil, meta, 2)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if out.StartPartition != 0 || out.Partition != 0 {
+			t.Fatalf("perQueue fresh cursor=%+v want StartPartition=Partition=0 (effective=1)", out)
+		}
+	})
+
+	t.Run("perQueue rejects cursor.Partition >= 1", func(t *testing.T) {
+		t.Parallel()
+		meta := &sqsQueueMeta{Generation: 1, PartitionCount: 4, FifoThroughputLimit: htfifoThroughputPerQueue}
+		cursor := &peekCursor{V: peekCursorSchemaV1, Generation: 1, StartPartition: 0, Partition: 2}
+		_, err := preparePeekCursor(cursor, meta, 0)
+		if !errors.Is(err, ErrAdminSQSValidation) {
+			t.Fatalf("perQueue Partition=2 (raw PartitionCount=4 would allow): want ErrAdminSQSValidation, got %v", err)
+		}
+	})
+}
+
 // TestAdminPeekQueue_HostileCursorBoundedRequest is the end-to-end
 // regression: an attacker crafts a wire-level cursor with an
 // out-of-range partition and submits it. The call MUST return