fix(spanner): pin inline read-write transactions to default endpoint (#13562)

Fixes location-aware routing for inline read-write transactions that
begin on the default endpoint.

When an inline read-write transaction (`TransactionSelector.BEGIN`)
starts on the default endpoint, the client must remember that default
endpoint affinity for the returned transaction ID. Without this, later
statements in the same transaction can be routed by CacheUpdate to a
direct endpoint, causing mid-transaction endpoint switches and ABORTED
failures.

This change records default endpoint affinity for inline read-write
begins, matching the existing behavior for explicit `BeginTransaction`
requests.

Author: rahul2393

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java

@@ -785,6 +785,7 @@ private void maybeTrackReadWriteBegin(TransactionSelector selector) {
       if (selector.getSelectorCase() == TransactionSelector.SelectorCase.BEGIN
           && !selector.getBegin().hasReadOnly()) {
         shouldRecordTransactionAffinity = true;
+        allowDefaultAffinity = true;
       }
     }
 

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LocationAwareSharedBackendReplicaHarnessTest.java

@@ -67,6 +67,7 @@ public class LocationAwareSharedBackendReplicaHarnessTest {
   private static final String PROJECT = "fake-project";
   private static final String INSTANCE = "fake-instance";
   private static final String DATABASE = "fake-database";
+  private static final AtomicInteger DATABASE_COUNTER = new AtomicInteger();
   private static final String TABLE = "T";
   private static final String REPLICA_LOCATION = "us-east1";
   private static final Statement SEED_QUERY = Statement.of("SELECT 1");
@@ -85,6 +86,14 @@ public class LocationAwareSharedBackendReplicaHarnessTest {
           .build();
   private static SharedBackendReplicaHarness harness;
 
+  // Assigned a fresh value per test so each test uses a distinct database scope. The location-aware
+  // routing layer keeps a process-wide EndpointLatencyRegistry keyed by (databaseScope,
+  // operationUid, ...). Each test builds a new Spanner whose per-channel operationUid counter
+  // restarts at 1, so without a unique scope a prior test's recorded error/latency penalties (e.g.
+  // an aborted commit on the leader) would leak into a later test and skew its routing, making the
+  // routing assertions flaky.
+  private String database;
+
   @BeforeClass
   public static void enableLocationAwareRouting() throws Exception {
     SpannerOptions.useEnvironment(
@@ -100,6 +109,7 @@ public boolean isEnableLocationApi() {
   @Before
   public void resetHarness() {
     harness.reset();
+    database = DATABASE + "-" + DATABASE_COUNTER.incrementAndGet();
   }
 
   @AfterClass
@@ -118,7 +128,7 @@ public static void restoreEnvironment() throws Exception {
   public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Exception {
     try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"));
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
 
       seedLocationMetadata(client);
       waitForReplicaRoutedRead(client, harness);
@@ -135,7 +145,7 @@ public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Ex
                   KeySet.singleKey(Key.of("b")),
                   Arrays.asList("k"),
                   Options.directedRead(DIRECTED_READ_OPTIONS))) {
-        assertTrue(resultSet.next());
+        assertTrue(drain(resultSet));
       }
 
       String diagnostics = routingDiagnostics(harness);
@@ -154,7 +164,7 @@ public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Ex
   public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() throws Exception {
     try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"));
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
 
       seedLocationMetadata(client);
       waitForReplicaRoutedRead(client, harness);
@@ -172,7 +182,7 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
                   KeySet.singleKey(Key.of("b")),
                   Arrays.asList("k"),
                   Options.directedRead(DIRECTED_READ_OPTIONS))) {
-        assertTrue(firstRead.next());
+        assertTrue(drain(firstRead));
       }
 
       try (ResultSet secondRead =
@@ -183,7 +193,7 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
                   KeySet.singleKey(Key.of("b")),
                   Arrays.asList("k"),
                   Options.directedRead(DIRECTED_READ_OPTIONS))) {
-        assertTrue(secondRead.next());
+        assertTrue(drain(secondRead));
       }
 
       String diagnostics = routingDiagnostics(harness);
@@ -210,7 +220,7 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
   public void singleUseReadReroutesOnUnavailableForBypassTraffic() throws Exception {
     try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"));
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
 
       seedLocationMetadata(client);
       waitForReplicaRoutedRead(client, harness);
@@ -227,7 +237,7 @@ public void singleUseReadReroutesOnUnavailableForBypassTraffic() throws Exceptio
                   KeySet.singleKey(Key.of("b")),
                   Arrays.asList("k"),
                   Options.directedRead(DIRECTED_READ_OPTIONS))) {
-        assertTrue(resultSet.next());
+        assertTrue(drain(resultSet));
       }
 
       String diagnostics = routingDiagnostics(harness);
@@ -247,7 +257,7 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
       throws Exception {
     try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"));
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
 
       seedLocationMetadata(client);
       waitForReplicaRoutedRead(client, harness);
@@ -264,7 +274,7 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
                   KeySet.singleKey(Key.of("b")),
                   Arrays.asList("k"),
                   Options.directedRead(DIRECTED_READ_OPTIONS))) {
-        assertTrue(firstRead.next());
+        assertTrue(drain(firstRead));
       }
 
       try (ResultSet secondRead =
@@ -275,7 +285,7 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
                   KeySet.singleKey(Key.of("b")),
                   Arrays.asList("k"),
                   Options.directedRead(DIRECTED_READ_OPTIONS))) {
-        assertTrue(secondRead.next());
+        assertTrue(drain(secondRead));
       }
 
       String diagnostics = routingDiagnostics(harness);
@@ -303,7 +313,7 @@ public void singleUseReadMidStreamRecvFailureWithoutRetryInfoRetriesForBypassTra
       throws Exception {
     try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, multiRowReadResultSet("b", "c", "d"));
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
 
       seedLocationMetadata(client);
       waitForReplicaRoutedRead(client, harness);
@@ -354,12 +364,75 @@ public void singleUseReadMidStreamRecvFailureWithoutRetryInfoRetriesForBypassTra
     }
   }
 
+  @Test
+  public void readWriteTransactionInlineBeginOnDefaultKeepsReadOnDefaultForBypassTraffic()
+      throws Exception {
+    try (Spanner spanner = createSpanner(harness)) {
+      configureBackend(harness, singleRowReadResultSet("b"), /* leaderReplicaIndex= */ 1);
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
+
+      seedLocationMetadata(client);
+      waitForReplicaRoutedStrongRead(client, harness, /* expectedReplicaIndex= */ 1);
+      harness.clearRequests();
+
+      client
+          .readWriteTransaction()
+          .run(
+              transaction -> {
+                try (ResultSet resultSet = transaction.executeQuery(SEED_QUERY)) {
+                  assertTrue(drain(resultSet));
+                }
+                try (ResultSet resultSet =
+                    transaction.read(TABLE, KeySet.singleKey(Key.of("b")), Arrays.asList("k"))) {
+                  assertTrue(drain(resultSet));
+                }
+                return null;
+              });
+
+      String diagnostics = routingDiagnostics(harness);
+      assertEquals(
+          "Read after inline begin on default should stay on default.\n" + diagnostics,
+          1,
+          harness
+              .defaultReplica
+              .getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
+              .size());
+      assertEquals(
+          "Direct replicas should not receive the transaction read.\n" + diagnostics,
+          0,
+          harness
+                  .replicas
+                  .get(0)
+                  .getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
+                  .size()
+              + harness
+                  .replicas
+                  .get(1)
+                  .getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
+                  .size());
+      assertEquals(
+          "Commit should use the default transaction affinity.\n" + diagnostics,
+          1,
+          harness.defaultReplica.getRequests(SharedBackendReplicaHarness.METHOD_COMMIT).size());
+      assertEquals(
+          "Direct replicas should not receive commit for the default-pinned transaction.\n"
+              + diagnostics,
+          0,
+          harness.replicas.get(0).getRequests(SharedBackendReplicaHarness.METHOD_COMMIT).size()
+              + harness
+                  .replicas
+                  .get(1)
+                  .getRequests(SharedBackendReplicaHarness.METHOD_COMMIT)
+                  .size());
+    }
+  }
+
   @Test
   public void readWriteTransactionAbortedCommitUsesReadAffinityReplicaForBypassTraffic()
       throws Exception {
     try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"), /* leaderReplicaIndex= */ 1);
-      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
+      DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, database));
 
       seedLocationMetadata(client);
       waitForReplicaRoutedStrongRead(client, harness, /* expectedReplicaIndex= */ 1);
@@ -374,7 +447,7 @@ public void readWriteTransactionAbortedCommitUsesReadAffinityReplicaForBypassTra
                 int attempt = attempts.incrementAndGet();
                 try (ResultSet resultSet =
                     transaction.read(TABLE, KeySet.singleKey(Key.of("b")), Arrays.asList("k"))) {
-                  assertTrue(resultSet.next());
+                  assertTrue(drain(resultSet));
                 }
 
                 if (attempt == 1) {
@@ -486,6 +559,20 @@ private static void seedLocationMetadata(DatabaseClient client) {
     }
   }
 
+  /**
+   * Fully consumes the result set so the underlying gRPC streaming read closes promptly. Leaving a
+   * stream open keeps the routed endpoint's active-request count above zero, which inflates its
+   * selection cost and can bounce the next request to a different replica, making routing
+   * assertions flaky under load. Returns whether at least one row was seen.
+   */
+  private static boolean drain(ResultSet resultSet) {
+    boolean sawRow = false;
+    while (resultSet.next()) {
+      sawRow = true;
+    }
+    return sawRow;
+  }
+
   private static int waitForReplicaRoutedRead(
       DatabaseClient client, SharedBackendReplicaHarness harness) throws InterruptedException {
     long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
@@ -498,7 +585,7 @@ private static int waitForReplicaRoutedRead(
                   KeySet.singleKey(Key.of("b")),
                   Arrays.asList("k"),
                   Options.directedRead(DIRECTED_READ_OPTIONS))) {
-        if (resultSet.next()) {
+        if (drain(resultSet)) {
           for (int replicaIndex = 0; replicaIndex < harness.replicas.size(); replicaIndex++) {
             if (!harness
                 .replicas
@@ -521,17 +608,18 @@ private static void waitForReplicaRoutedStrongRead(
     long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
     while (System.nanoTime() < deadlineNanos) {
       harness.clearRequests();
+      boolean sawRow;
       try (ResultSet resultSet =
           client.singleUse().read(TABLE, KeySet.singleKey(Key.of("b")), Arrays.asList("k"))) {
-        if (resultSet.next()) {
-          if (!harness
+        sawRow = drain(resultSet);
+      }
+      if (sawRow
+          && !harness
               .replicas
               .get(expectedReplicaIndex)
               .getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
               .isEmpty()) {
-            return;
-          }
-        }
+        return;
       }
       Thread.sleep(50L);
     }