fix(sqs): fail-closed for unresolved txn read keys

Round 5 introduced fail-closed semantics in ShardRouter.ResolveGroup
(recognised-but-unresolved partitioned keys return ok=false instead
of falling through to the engine), but the read-key path inside
ShardedCoordinator was not audited for the new contract.

engineGroupIDForKey discards the resolver's ok flag and returns 0
for any failure. groupReadKeysByShardID then loops `if gid == 0
{ continue }`, silently dropping unrouted read keys from the
prewrite payload. With the new fail-closed resolver, a partitioned
read key whose queue has drifted out of --sqsFifoPartitionMap
(partial rollout / config drift) gets dropped from OCC validation.
The FSM never sees that key in ReadKeys, a concurrent write to the
same key commits without conflict, and SSI is broken.

Codex round-2 P1 on PR #715 caught this — addressed here.

Fix

groupReadKeysByShardID returns (map, error). Any read key that
fails to route surfaces as ErrInvalidRequest, dispatchTxn
propagates the error, and the transaction aborts before prewrite.
Calls c.router.ResolveGroup directly (rather than via
engineGroupIDForKey) so the (gid, ok) signal is preserved through
the boundary.

Tests

- TestGroupReadKeysByShardID_FailsClosedOnUnroutable replaces
  TestGroupReadKeysByShardID_SkipsUnroutableKeys (which had been
  pinning the buggy skip-silently behaviour). Asserts the new
  fail-closed contract: unroutable keys → error, no partial map.
- TestShardedCoordinator_TxnFailsClosedForUnresolvedReadKey is
  the coordinator-level regression: a transaction with a
  recognised-but-unresolved partitioned read key MUST abort
  before any prewrite. Asserts no group received an RPC.
- The three existing TestGroupReadKeysByShardID_* tests are
  updated for the new (map, error) signature.

Author: bootjp

kv/sharded_coordinator.go

@@ -395,7 +395,13 @@ func (c *ShardedCoordinator) dispatchTxn(ctx context.Context, startTS uint64, co
 
 	// Multi-shard path: group read keys by shard now. The result is passed
 	// directly to prewriteTxn to avoid a second iteration inside that function.
-	groupedReadKeys := c.groupReadKeysByShardID(readKeys)
+	// A routing failure here aborts the transaction before any prewrite —
+	// silently dropping unresolvable read keys would let OCC validation run
+	// with an incomplete read set and break SSI.
+	groupedReadKeys, err := c.groupReadKeysByShardID(readKeys)
+	if err != nil {
+		return nil, err
+	}
 	prepared, err := c.prewriteTxn(ctx, startTS, commitTS, primaryKey, grouped, gids, groupedReadKeys)
 	if err != nil {
 		return nil, err
@@ -891,19 +897,38 @@ func (c *ShardedCoordinator) engineGroupIDForKey(key []byte) uint64 {
 	return gid
 }
 
-func (c *ShardedCoordinator) groupReadKeysByShardID(readKeys [][]byte) map[uint64][][]byte {
+// groupReadKeysByShardID groups txn read keys by their owning Raft
+// group. Returns an error when ANY read key cannot be routed —
+// silently skipping unresolvable keys would let a transaction
+// commit with an incomplete OCC read-set, which breaks SSI under
+// the §5 ShardRouter resolver-first dispatch (codex round-2 P1 on
+// PR #715).
+//
+// The fail-closed semantic the resolver gained in PR 4-B-2 makes
+// this path matter: a partitioned-shape read key whose queue is
+// missing from --sqsFifoPartitionMap (drift / partial rollout)
+// returns gid=0 from c.router.ResolveGroup. If we silently dropped
+// those keys, the prewrite Raft entry would carry an empty
+// ReadKeys slice for that key, the FSM's read-write conflict
+// validation would never see it, and a concurrent write could
+// commit alongside a stale read. Surface the routing failure as
+// an error so the transaction aborts before any FSM apply.
+func (c *ShardedCoordinator) groupReadKeysByShardID(readKeys [][]byte) (map[uint64][][]byte, error) {
 	if len(readKeys) == 0 {
-		return nil
+		return nil, nil
 	}
 	grouped := make(map[uint64][][]byte)
 	for _, key := range readKeys {
-		gid := c.engineGroupIDForKey(key)
-		if gid == 0 {
-			continue
+		gid, ok := c.router.ResolveGroup(key)
+		if !ok || gid == 0 {
+			return nil, errors.Wrapf(ErrInvalidRequest,
+				"no route for txn read key %q — recognised-but-"+
+					"unresolved partition keys must fail closed to "+
+					"preserve OCC read-set integrity", key)
 		}
 		grouped[gid] = append(grouped[gid], key)
 	}
-	return grouped
+	return grouped, nil
 }
 
 // validateReadOnlyShards checks read-write conflicts on shards that have

kv/sharded_coordinator_partition_test.go

@@ -195,6 +195,85 @@ func TestShardedCoordinator_DispatchSplitsMutationsByResolverGroup(t *testing.T)
 			"mutations across groups via c.router.ResolveGroup.")
 }
 
+// TestShardedCoordinator_TxnFailsClosedForUnresolvedReadKey pins
+// the codex round-2 P1 fix on PR #715: a transaction whose read
+// set contains a partitioned-shape key the resolver recognises
+// but cannot route (queue missing from --sqsFifoPartitionMap
+// during drift / partial rollout) MUST abort before prewrite.
+//
+// Pre-fix path: groupReadKeysByShardID silently skipped read keys
+// where engineGroupIDForKey returned 0 (which is what the
+// fail-closed resolver returns for recognised-but-unresolved
+// keys). The resulting prewrite Raft entry carried an empty
+// ReadKeys slice for that key, the FSM's read-write conflict
+// validation never saw it, and a concurrent write could commit
+// alongside the stale read — SSI violated.
+//
+// Post-fix path: groupReadKeysByShardID returns an error when
+// any read key cannot be routed; dispatchTxn propagates the
+// error and the transaction aborts.
+func TestShardedCoordinator_TxnFailsClosedForUnresolvedReadKey(t *testing.T) {
+	t.Parallel()
+	engine := distribution.NewEngine()
+	engine.UpdateRoute([]byte(""), nil, 1)
+
+	g1 := &recordingTransactional{
+		responses: []*TransactionResponse{{CommitIndex: 1}},
+	}
+	g42 := &recordingTransactional{
+		responses: []*TransactionResponse{{CommitIndex: 1}},
+	}
+
+	coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{
+		1:  {Txn: g1, Store: store.NewMVCCStore()},
+		42: {Txn: g42, Store: store.NewMVCCStore()},
+	}, 1, NewHLC(), nil)
+
+	// Resolver claims the WRITE key but RECOGNISES the read key as
+	// partitioned-shape and returns ok=false (simulating the
+	// queue-missing-from-map drift scenario). The router pairs
+	// Recognised=true with ok=false to fail closed.
+	writeKey := []byte("!sqs|msg|data|p|claimed-write-key")
+	readKey := []byte("!sqs|msg|data|p|drift-unresolved-read")
+	coord.WithPartitionResolver(&stubResolver{
+		claim:            map[string]uint64{string(writeKey): 42},
+		recognisedPrefix: []byte("!sqs|msg|data|p|"),
+	})
+
+	// Issue a transaction that reads the unresolvable partitioned
+	// key and writes the claimed key. The transaction MUST abort —
+	// pre-fix it would silently drop the read key and proceed.
+	resp, err := coord.Dispatch(context.Background(), &OperationGroup[OP]{
+		IsTxn:    true,
+		StartTS:  100,
+		CommitTS: 200,
+		Elems: []*Elem[OP]{
+			{Op: Put, Key: writeKey, Value: []byte("v")},
+		},
+		ReadKeys: [][]byte{readKey},
+	})
+	require.Error(t, err,
+		"transaction with unresolvable partitioned read key MUST "+
+			"abort — silently dropping the key would let OCC "+
+			"validation run with an incomplete read set and "+
+			"break SSI")
+	require.Nil(t, resp)
+
+	// Neither group should have received a prepare — the routing
+	// failure surfaces before any prewrite.
+	g1.mu.Lock()
+	g42.mu.Lock()
+	defer g1.mu.Unlock()
+	defer g42.mu.Unlock()
+	require.Zero(t, len(g1.requests),
+		"engine default group must not see a prewrite — fail-closed "+
+			"happens before any RPC")
+	require.Zero(t, len(g42.requests),
+		"resolver-claimed group must not see a prewrite either — "+
+			"fail-closed aborts the txn before prewrite even though "+
+			"the write key is routable")
+}
+
 // TestShardedCoordinator_DispatchFallsThroughForUnclaimedKeys pins
 // the inverse: a key the resolver does NOT claim must continue to
 // route via the byte-range engine. Without this guard the resolver-

kv/sharded_coordinator_txn_test.go

@@ -309,15 +309,19 @@ func TestGroupReadKeysByShardID_NilReturnsNil(t *testing.T) {
 	engine := distribution.NewEngine()
 	engine.UpdateRoute([]byte(""), nil, 1)
 	coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{1: {}}, 1, NewHLC(), nil)
-	require.Nil(t, coord.groupReadKeysByShardID(nil))
+	grouped, err := coord.groupReadKeysByShardID(nil)
+	require.NoError(t, err)
+	require.Nil(t, grouped)
 }
 
 func TestGroupReadKeysByShardID_EmptyReturnsNil(t *testing.T) {
 	t.Parallel()
 	engine := distribution.NewEngine()
 	engine.UpdateRoute([]byte(""), nil, 1)
 	coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{1: {}}, 1, NewHLC(), nil)
-	require.Nil(t, coord.groupReadKeysByShardID([][]byte{}))
+	grouped, err := coord.groupReadKeysByShardID([][]byte{})
+	require.NoError(t, err)
+	require.Nil(t, grouped)
 }
 
 func TestGroupReadKeysByShardID_GroupsByShardID(t *testing.T) {
@@ -327,11 +331,12 @@ func TestGroupReadKeysByShardID_GroupsByShardID(t *testing.T) {
 	engine.UpdateRoute([]byte("m"), nil, 2)
 	coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{1: {}, 2: {}}, 1, NewHLC(), nil)
 
-	grouped := coord.groupReadKeysByShardID([][]byte{
+	grouped, err := coord.groupReadKeysByShardID([][]byte{
 		[]byte("b"), // shard 1
 		[]byte("c"), // shard 1
 		[]byte("x"), // shard 2
 	})
+	require.NoError(t, err)
 	require.Len(t, grouped, 2)
 	require.Len(t, grouped[1], 2)
 	require.Equal(t, []byte("b"), grouped[1][0])
@@ -340,20 +345,34 @@ func TestGroupReadKeysByShardID_GroupsByShardID(t *testing.T) {
 	require.Equal(t, []byte("x"), grouped[2][0])
 }
 
-func TestGroupReadKeysByShardID_SkipsUnroutableKeys(t *testing.T) {
+// TestGroupReadKeysByShardID_FailsClosedOnUnroutable pins the
+// codex round-2 P1 fix on PR #715: a read key the resolver cannot
+// route (recognised-but-unresolved partition key during drift, or
+// any key outside the engine's range cover) MUST surface as an
+// error so the transaction aborts before any prewrite. Silently
+// skipping unroutable keys would let OCC validation run with an
+// incomplete read set and break SSI — a concurrent write to that
+// key could commit alongside a stale read.
+//
+// This test was previously TestGroupReadKeysByShardID_SkipsUnroutableKeys
+// and pinned the BUGGY skip-silently behaviour. Renamed and rewritten
+// to pin the new fail-closed contract.
+func TestGroupReadKeysByShardID_FailsClosedOnUnroutable(t *testing.T) {
 	t.Parallel()
 	// Only route "a"-"m" to shard 1. Keys outside this range are unroutable.
 	engine := distribution.NewEngine()
 	engine.UpdateRoute([]byte("a"), []byte("m"), 1)
 	coord := NewShardedCoordinator(engine, map[uint64]*ShardGroup{1: {}}, 1, NewHLC(), nil)
 
-	grouped := coord.groupReadKeysByShardID([][]byte{
+	grouped, err := coord.groupReadKeysByShardID([][]byte{
 		[]byte("b"),   // routable → shard 1
-		[]byte("zzz"), // unroutable → skipped
+		[]byte("zzz"), // unroutable → MUST surface as error
 	})
-	require.Len(t, grouped, 1)
-	require.Len(t, grouped[1], 1)
-	require.Equal(t, []byte("b"), grouped[1][0])
+	require.Error(t, err,
+		"unroutable read key MUST fail closed — silently skipping "+
+			"would drop the key from OCC validation and break SSI")
+	require.Nil(t, grouped)
+	require.ErrorIs(t, err, ErrInvalidRequest)
 }
 
 // ---------------------------------------------------------------------------