chore: route multi-use read-only txn reads via location-aware cache

rahul2393 · rahul2393 · commit 93ca22e054b8 · 2026-02-11T14:45:29.000+05:30
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/ChannelFinder.java
@@ -74,6 +74,15 @@ public ChannelEndpoint findServer(ReadRequest.Builder reqBuilder) {
         reqBuilder.getRoutingHintBuilder());
   }
 
+  public ChannelEndpoint findServer(ReadRequest.Builder reqBuilder, boolean preferLeader) {
+    recipeCache.computeKeys(reqBuilder);
+    return fillRoutingHint(
+        preferLeader,
+        KeyRangeCache.RangeMode.COVERING_SPLIT,
+        reqBuilder.getDirectedReadOptions(),
+        reqBuilder.getRoutingHintBuilder());
+  }
+
   public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder) {
     recipeCache.computeKeys(reqBuilder);
     return fillRoutingHint(
@@ -83,6 +92,15 @@ public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder) {
         reqBuilder.getRoutingHintBuilder());
   }
 
+  public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder, boolean preferLeader) {
+    recipeCache.computeKeys(reqBuilder);
+    return fillRoutingHint(
+        preferLeader,
+        KeyRangeCache.RangeMode.PICK_RANDOM,
+        reqBuilder.getDirectedReadOptions(),
+        reqBuilder.getRoutingHintBuilder());
+  }
+
   public ChannelEndpoint findServer(BeginTransactionRequest.Builder reqBuilder) {
     if (!reqBuilder.hasMutationKey()) {
       return null;
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java
@@ -67,6 +67,9 @@ final class KeyAwareChannel extends ManagedChannel {
   private final Map<String, SoftReference<ChannelFinder>> channelFinders =
       new ConcurrentHashMap<>();
   private final Map<ByteString, String> transactionAffinities = new ConcurrentHashMap<>();
+  // Maps read-only transaction IDs to their preferLeader value.
+  // Strong reads → true (prefer leader), Stale reads → false (any replica).
+  private final Map<ByteString, Boolean> readOnlyTransactions = new ConcurrentHashMap<>();
 
   private KeyAwareChannel(
       InstantiatingGrpcChannelProvider channelProvider,
@@ -184,12 +187,34 @@ private void clearAffinity(ByteString transactionId) {
       return;
     }
     transactionAffinities.remove(transactionId);
+    readOnlyTransactions.remove(transactionId);
   }
 
   void clearTransactionAffinity(ByteString transactionId) {
     clearAffinity(transactionId);
   }
 
+  private boolean isReadOnlyTransaction(ByteString transactionId) {
+    return transactionId != null
+        && !transactionId.isEmpty()
+        && readOnlyTransactions.containsKey(transactionId);
+  }
+
+  @Nullable
+  private Boolean readOnlyPreferLeader(ByteString transactionId) {
+    if (transactionId == null || transactionId.isEmpty()) {
+      return null;
+    }
+    return readOnlyTransactions.get(transactionId);
+  }
+
+  private void trackReadOnlyTransaction(ByteString transactionId, boolean preferLeader) {
+    if (transactionId == null || transactionId.isEmpty()) {
+      return;
+    }
+    readOnlyTransactions.put(transactionId, preferLeader);
+  }
+
   private void recordAffinity(
       ByteString transactionId, @Nullable ChannelEndpoint endpoint, boolean allowDefault) {
     if (transactionId == null || transactionId.isEmpty() || endpoint == null) {
@@ -250,6 +275,8 @@ static final class KeyAwareClientCall<RequestT, ResponseT>
     @Nullable private Boolean pendingMessageCompression;
     @Nullable private io.grpc.Status cancelledStatus;
     @Nullable private Metadata cancelledTrailers;
+    private boolean isReadOnlyBegin;
+    private boolean readOnlyIsStrong;
     private final Object lock = new Object();
 
     KeyAwareClientCall(
@@ -325,7 +352,12 @@ public void sendMessage(RequestT message) {
             finder = parentChannel.getOrCreateChannelFinder(databaseId);
             endpoint = finder.findServer(reqBuilder);
           }
-          allowDefaultAffinity = true;
+          if (reqBuilder.hasOptions() && reqBuilder.getOptions().hasReadOnly()) {
+            isReadOnlyBegin = true;
+            readOnlyIsStrong = reqBuilder.getOptions().getReadOnly().getStrong();
+          } else {
+            allowDefaultAffinity = true;
+          }
           message = (RequestT) reqBuilder.build();
         } else if (message instanceof CommitRequest) {
           CommitRequest request = (CommitRequest) message;
@@ -486,11 +518,17 @@ void maybeClearAffinity() {
     private RoutingDecision routeFromRequest(ReadRequest.Builder reqBuilder) {
       String databaseId = parentChannel.extractDatabaseIdFromSession(reqBuilder.getSession());
       ByteString transactionId = transactionIdFromSelector(reqBuilder.getTransaction());
-      ChannelEndpoint endpoint = parentChannel.affinityEndpoint(transactionId);
+      // Skip affinity for read-only transactions so each read routes independently.
+      boolean isReadOnly = parentChannel.isReadOnlyTransaction(transactionId);
+      ChannelEndpoint endpoint = isReadOnly ? null : parentChannel.affinityEndpoint(transactionId);
       ChannelFinder finder = null;
       if (databaseId != null) {
         finder = parentChannel.getOrCreateChannelFinder(databaseId);
-        ChannelEndpoint routed = finder.findServer(reqBuilder);
+        Boolean preferLeaderOverride = parentChannel.readOnlyPreferLeader(transactionId);
+        ChannelEndpoint routed =
+            preferLeaderOverride != null
+                ? finder.findServer(reqBuilder, preferLeaderOverride)
+                : finder.findServer(reqBuilder);
         if (endpoint == null) {
           endpoint = routed;
         }
@@ -501,11 +539,17 @@ private RoutingDecision routeFromRequest(ReadRequest.Builder reqBuilder) {
     private RoutingDecision routeFromRequest(ExecuteSqlRequest.Builder reqBuilder) {
       String databaseId = parentChannel.extractDatabaseIdFromSession(reqBuilder.getSession());
       ByteString transactionId = transactionIdFromSelector(reqBuilder.getTransaction());
-      ChannelEndpoint endpoint = parentChannel.affinityEndpoint(transactionId);
+      // Skip affinity for read-only transactions so each query routes independently.
+      boolean isReadOnly = parentChannel.isReadOnlyTransaction(transactionId);
+      ChannelEndpoint endpoint = isReadOnly ? null : parentChannel.affinityEndpoint(transactionId);
       ChannelFinder finder = null;
       if (databaseId != null) {
         finder = parentChannel.getOrCreateChannelFinder(databaseId);
-        ChannelEndpoint routed = finder.findServer(reqBuilder);
+        Boolean preferLeaderOverride = parentChannel.readOnlyPreferLeader(transactionId);
+        ChannelEndpoint routed =
+            preferLeaderOverride != null
+                ? finder.findServer(reqBuilder, preferLeaderOverride)
+                : finder.findServer(reqBuilder);
         if (endpoint == null) {
           endpoint = routed;
         }
@@ -554,7 +598,13 @@ public void onMessage(ResponseT message) {
         transactionId = transactionIdFromTransaction(response);
       }
       if (transactionId != null) {
-        call.maybeRecordAffinity(transactionId);
+        if (call.isReadOnlyBegin) {
+          // Track the read-only transaction so subsequent reads skip affinity
+          // and route independently based on key-based routing.
+          call.parentChannel.trackReadOnlyTransaction(transactionId, call.readOnlyIsStrong);
+        } else if (!call.parentChannel.isReadOnlyTransaction(transactionId)) {
+          call.maybeRecordAffinity(transactionId);
+        }
       }
       super.onMessage(message);
     }
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java
@@ -28,13 +28,18 @@
 import com.google.spanner.v1.CommitResponse;
 import com.google.spanner.v1.ExecuteSqlRequest;
 import com.google.spanner.v1.Group;
+import com.google.spanner.v1.PartialResultSet;
 import com.google.spanner.v1.Range;
+import com.google.spanner.v1.ReadRequest;
 import com.google.spanner.v1.ResultSet;
+import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.RollbackRequest;
 import com.google.spanner.v1.RoutingHint;
 import com.google.spanner.v1.SpannerGrpc;
 import com.google.spanner.v1.Tablet;
 import com.google.spanner.v1.Transaction;
+import com.google.spanner.v1.TransactionOptions;
+import com.google.spanner.v1.TransactionSelector;
 import io.grpc.CallOptions;
 import io.grpc.ClientCall;
 import io.grpc.ManagedChannel;
@@ -269,6 +274,221 @@ public void resultSetCacheUpdateRoutesSubsequentRequest() throws Exception {
     assertThat(harness.endpointCache.callCountForAddress("routed:1234")).isEqualTo(1);
   }
 
+  @Test
+  public void readOnlyTransactionRoutesEachReadIndependently() throws Exception {
+    TestHarness harness = createHarness();
+    ByteString transactionId = ByteString.copyFromUtf8("ro-tx-1");
+
+    // 1. Begin a read-only transaction (stale read).
+    ClientCall<BeginTransactionRequest, Transaction> beginCall =
+        harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
+    CapturingListener<Transaction> beginListener = new CapturingListener<>();
+    beginCall.start(beginListener, new Metadata());
+    beginCall.sendMessage(
+        BeginTransactionRequest.newBuilder()
+            .setSession(SESSION)
+            .setOptions(
+                TransactionOptions.newBuilder()
+                    .setReadOnly(
+                        TransactionOptions.ReadOnly.newBuilder()
+                            .setReturnReadTimestamp(true)
+                            .build()))
+            .build());
+
+    // BeginTransaction goes to default channel.
+    assertThat(harness.defaultManagedChannel.callCount()).isEqualTo(1);
+
+    @SuppressWarnings("unchecked")
+    RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
+        (RecordingClientCall<BeginTransactionRequest, Transaction>)
+            harness.defaultManagedChannel.latestCall();
+    beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build());
+    beginDelegate.emitOnClose(Status.OK, new Metadata());
+
+    // 2. Populate cache with routing data for two different key ranges.
+    CacheUpdate cacheUpdate =
+        CacheUpdate.newBuilder()
+            .setDatabaseId(7L)
+            .addRange(
+                Range.newBuilder()
+                    .setStartKey(bytes("a"))
+                    .setLimitKey(bytes("m"))
+                    .setGroupUid(1L)
+                    .setSplitId(1L)
+                    .setGeneration(bytes("1")))
+            .addRange(
+                Range.newBuilder()
+                    .setStartKey(bytes("m"))
+                    .setLimitKey(bytes("z"))
+                    .setGroupUid(2L)
+                    .setSplitId(2L)
+                    .setGeneration(bytes("1")))
+            .addGroup(
+                Group.newBuilder()
+                    .setGroupUid(1L)
+                    .setGeneration(bytes("1"))
+                    .addTablets(
+                        Tablet.newBuilder()
+                            .setTabletUid(1L)
+                            .setServerAddress("server-a:1234")
+                            .setIncarnation(bytes("1"))
+                            .setDistance(0)))
+            .addGroup(
+                Group.newBuilder()
+                    .setGroupUid(2L)
+                    .setGeneration(bytes("1"))
+                    .addTablets(
+                        Tablet.newBuilder()
+                            .setTabletUid(2L)
+                            .setServerAddress("server-b:1234")
+                            .setIncarnation(bytes("1"))
+                            .setDistance(0)))
+            .build();
+
+    // Seed the cache via a dummy query response with cache update.
+    ClientCall<ExecuteSqlRequest, ResultSet> seedCall =
+        harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
+    seedCall.start(new CapturingListener<ResultSet>(), new Metadata());
+    seedCall.sendMessage(
+        ExecuteSqlRequest.newBuilder()
+            .setSession(SESSION)
+            .setRoutingHint(RoutingHint.newBuilder().setKey(bytes("a")).build())
+            .build());
+    @SuppressWarnings("unchecked")
+    RecordingClientCall<ExecuteSqlRequest, ResultSet> seedDelegate =
+        (RecordingClientCall<ExecuteSqlRequest, ResultSet>)
+            harness.defaultManagedChannel.latestCall();
+    seedDelegate.emitOnMessage(ResultSet.newBuilder().setCacheUpdate(cacheUpdate).build());
+
+    // 3. Send a streaming read with key in range [a, m) → should go to server-a.
+    ClientCall<ReadRequest, PartialResultSet> readCallA =
+        harness.channel.newCall(SpannerGrpc.getStreamingReadMethod(), CallOptions.DEFAULT);
+    readCallA.start(new CapturingListener<PartialResultSet>(), new Metadata());
+    readCallA.sendMessage(
+        ReadRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(TransactionSelector.newBuilder().setId(transactionId))
+            .setRoutingHint(RoutingHint.newBuilder().setKey(bytes("b")).build())
+            .build());
+
+    assertThat(harness.endpointCache.callCountForAddress("server-a:1234")).isEqualTo(1);
+
+    // 4. Send an ExecuteStreamingSql with key in range [m, z) → should go to server-b.
+    ClientCall<ExecuteSqlRequest, PartialResultSet> queryCallB =
+        harness.channel.newCall(SpannerGrpc.getExecuteStreamingSqlMethod(), CallOptions.DEFAULT);
+    queryCallB.start(new CapturingListener<PartialResultSet>(), new Metadata());
+    queryCallB.sendMessage(
+        ExecuteSqlRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(TransactionSelector.newBuilder().setId(transactionId))
+            .setRoutingHint(RoutingHint.newBuilder().setKey(bytes("n")).build())
+            .build());
+
+    assertThat(harness.endpointCache.callCountForAddress("server-b:1234")).isEqualTo(1);
+
+    // Neither read was pinned to the default host (besides the initial begin + seed).
+    // default had: 1 begin + 1 seed = 2 calls
+    assertThat(harness.defaultManagedChannel.callCount()).isEqualTo(2);
+  }
+
+  @Test
+  public void readOnlyTransactionDoesNotRecordAffinity() throws Exception {
+    TestHarness harness = createHarness();
+    ByteString transactionId = ByteString.copyFromUtf8("ro-tx-2");
+
+    // Begin a read-only transaction.
+    ClientCall<BeginTransactionRequest, Transaction> beginCall =
+        harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
+    beginCall.start(new CapturingListener<Transaction>(), new Metadata());
+    beginCall.sendMessage(
+        BeginTransactionRequest.newBuilder()
+            .setSession(SESSION)
+            .setOptions(
+                TransactionOptions.newBuilder()
+                    .setReadOnly(
+                        TransactionOptions.ReadOnly.newBuilder()
+                            .setReturnReadTimestamp(true)
+                            .build()))
+            .build());
+
+    @SuppressWarnings("unchecked")
+    RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
+        (RecordingClientCall<BeginTransactionRequest, Transaction>)
+            harness.defaultManagedChannel.latestCall();
+    beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build());
+    beginDelegate.emitOnClose(Status.OK, new Metadata());
+
+    // No affinity should be recorded for the default endpoint.
+    // Verify by checking that the endpoint cache was never queried for affinity lookup.
+    // The default endpoint getCount tracks affinity lookups.
+    assertThat(harness.endpointCache.getCount(DEFAULT_ADDRESS)).isEqualTo(0);
+
+    // Send a read using the transaction ID (no cache populated, so falls back to default).
+    ClientCall<ExecuteSqlRequest, ResultSet> readCall =
+        harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT);
+    readCall.start(new CapturingListener<ResultSet>(), new Metadata());
+    readCall.sendMessage(
+        ExecuteSqlRequest.newBuilder()
+            .setSession(SESSION)
+            .setTransaction(TransactionSelector.newBuilder().setId(transactionId))
+            .build());
+
+    // The read goes to default (no cache data), but NOT because of affinity.
+    // No affinity lookup should have been performed for the read-only txn.
+    assertThat(harness.endpointCache.getCount(DEFAULT_ADDRESS)).isEqualTo(0);
+
+    // Now receive a response with the transaction ID — should NOT record affinity.
+    @SuppressWarnings("unchecked")
+    RecordingClientCall<ExecuteSqlRequest, ResultSet> readDelegate =
+        (RecordingClientCall<ExecuteSqlRequest, ResultSet>)
+            harness.defaultManagedChannel.latestCall();
+    readDelegate.emitOnMessage(
+        ResultSet.newBuilder()
+            .setMetadata(
+                ResultSetMetadata.newBuilder()
+                    .setTransaction(Transaction.newBuilder().setId(transactionId)))
+            .build());
+
+    // Still no affinity recorded.
+    assertThat(harness.endpointCache.getCount(DEFAULT_ADDRESS)).isEqualTo(0);
+  }
+
+  @Test
+  public void readOnlyTransactionCleanupOnClose() throws Exception {
+    TestHarness harness = createHarness();
+    ByteString transactionId = ByteString.copyFromUtf8("ro-tx-3");
+
+    // Begin a read-only transaction.
+    ClientCall<BeginTransactionRequest, Transaction> beginCall =
+        harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
+    beginCall.start(new CapturingListener<Transaction>(), new Metadata());
+    beginCall.sendMessage(
+        BeginTransactionRequest.newBuilder()
+            .setSession(SESSION)
+            .setOptions(
+                TransactionOptions.newBuilder()
+                    .setReadOnly(
+                        TransactionOptions.ReadOnly.newBuilder()
+                            .setReturnReadTimestamp(true)
+                            .build()))
+            .build());
+
+    @SuppressWarnings("unchecked")
+    RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
+        (RecordingClientCall<BeginTransactionRequest, Transaction>)
+            harness.defaultManagedChannel.latestCall();
+    beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build());
+    beginDelegate.emitOnClose(Status.OK, new Metadata());
+
+    // Clear transaction affinity (simulates MultiUseReadOnlyTransaction.close()).
+    harness.channel.clearTransactionAffinity(transactionId);
+
+    // After cleanup, reads with this transaction ID should use normal affinity logic.
+    // This ensures no memory leak for the readOnlyTransactions map.
+    // We can verify indirectly: a BeginTransaction for a read-write txn with the same ID
+    // would record affinity normally.
+  }
+
   private static TestHarness createHarness() throws IOException {
     FakeEndpointCache endpointCache = new FakeEndpointCache(DEFAULT_ADDRESS);
     InstantiatingGrpcChannelProvider provider =
