From 3bc971714dbf9a621f8f243e0461a59040d3406e Mon Sep 17 00:00:00 2001
From: Xueyang Hu <xueyangh@google.com>
Date: Thu, 27 Aug 2026 06:48:37 +0000
Subject: [PATCH] fix(spanner): prioritize leader replica for read-write
 transactions in location-aware routing

When experimental location-aware routing is enabled, read-write transactions
starting with an inline begin (preferLeader=true) could be routed to follower
replicas because operationUid > 0L on SQL requests caused selectTablet() to
bypass the local leader and take the score-aware replica selection path.

This change reorders the checks in KeyRangeCache.selectTablet() so that when
preferLeader=true and a healthy local Paxos leader exists, the leader is selected
directly. If the leader is unhealthy or in transient failure, it gracefully falls
back to score-aware replica selection.
---
 .../cloud/spanner/spi/v1/KeyRangeCache.java   | 32 ++++++++---------
 .../spanner/spi/v1/KeyRangeCacheTest.java     | 34 ++++++++++++++++++-
 2 files changed, 49 insertions(+), 17 deletions(-)

diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java
index d630f70339ab..b969123d7498 100644
--- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java
+++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyRangeCache.java
@@ -936,22 +936,6 @@ private TabletSnapshot selectTablet(
         List<SkippedTabletDetail> skippedTabletDetails,
         Map<String, ChannelEndpoint> resolvedEndpoints,
         SelectionState selectionStats) {
-      if (!preferLeader || hintBuilder.getOperationUid() > 0L) {
-        TabletSnapshot preferredLeader =
-            preferLeader ? localLeaderForScoreBias(snapshot, hasDirectedReadOptions) : null;
-        return selectScoreAwareTablet(
-            snapshot,
-            preferLeader,
-            directedReadOptions,
-            hintBuilder,
-            excludedEndpoints,
-            skippedTabletUids,
-            skippedTabletDetails,
-            resolvedEndpoints,
-            selectionStats,
-            preferredLeader);
-      }
-
       boolean checkedLeader = false;
       if (preferLeader
           && !hasDirectedReadOptions
@@ -971,6 +955,22 @@ private TabletSnapshot selectTablet(
           return snapshot.leader();
         }
       }
+
+      if (hintBuilder.getOperationUid() > 0L) {
+        TabletSnapshot preferredLeader =
+            preferLeader ? localLeaderForScoreBias(snapshot, hasDirectedReadOptions) : null;
+        return selectScoreAwareTablet(
+            snapshot,
+            preferLeader,
+            directedReadOptions,
+            hintBuilder,
+            excludedEndpoints,
+            skippedTabletUids,
+            skippedTabletDetails,
+            resolvedEndpoints,
+            selectionStats,
+            preferredLeader);
+      }
       for (int index = 0; index < snapshot.tablets.size(); index++) {
         if (checkedLeader && index == snapshot.leaderIndex) {
           continue;
diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java
index 4aeeebadc654..66d6e76666b4 100644
--- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java
+++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyRangeCacheTest.java
@@ -784,7 +784,7 @@ public void preferLeaderFalseUsesLowestLatencyReplicaWhenScoresAvailable() {
   }
 
   @Test
-  public void preferLeaderTrueUsesLatencyScoresWhenOperationUidAvailable() {
+  public void preferLeaderTrueWithOperationUidSelectsLeaderEvenWhenFollowerHasLowerLatency() {
     FakeEndpointCache endpointCache = new FakeEndpointCache();
     KeyRangeCache cache = new KeyRangeCache(endpointCache);
     cache.useDeterministicRandom();
@@ -794,6 +794,38 @@ public void preferLeaderTrueUsesLatencyScoresWhenOperationUidAvailable() {
     endpointCache.get("server2");
     endpointCache.get("server3");
 
+    EndpointLatencyRegistry.recordLatency(
+        null, TEST_OPERATION_UID, true, "server1", Duration.ofNanos(300_000L));
+    EndpointLatencyRegistry.recordLatency(
+        null, TEST_OPERATION_UID, true, "server2", Duration.ofNanos(100_000L));
+    EndpointLatencyRegistry.recordLatency(
+        null, TEST_OPERATION_UID, true, "server3", Duration.ofNanos(200_000L));
+
+    RoutingHint.Builder hint =
+        RoutingHint.newBuilder().setKey(bytes("a")).setOperationUid(TEST_OPERATION_UID);
+    ChannelEndpoint server =
+        cache.fillRoutingHint(
+            true,
+            KeyRangeCache.RangeMode.COVERING_SPLIT,
+            DirectedReadOptions.getDefaultInstance(),
+            hint);
+
+    assertNotNull(server);
+    assertEquals("server1", server.getAddress());
+  }
+
+  @Test
+  public void preferLeaderTrueWithOperationUidFallsBackToScoreAwareWhenLeaderUnhealthy() {
+    FakeEndpointCache endpointCache = new FakeEndpointCache();
+    KeyRangeCache cache = new KeyRangeCache(endpointCache);
+    cache.useDeterministicRandom();
+    cache.addRanges(threeReplicaUpdate());
+
+    endpointCache.get("server1");
+    endpointCache.get("server2");
+    endpointCache.get("server3");
+    endpointCache.setState("server1", EndpointHealthState.TRANSIENT_FAILURE);
+
     EndpointLatencyRegistry.recordLatency(
         null, TEST_OPERATION_UID, true, "server1", Duration.ofNanos(300_000L));
     EndpointLatencyRegistry.recordLatency(