chore: integrate location aware routing with RPCs

google-cloud-spanner/pom.xml

@@ -25,6 +25,13 @@
 
 
   <build>
+    <extensions>
+      <extension>
+        <groupId>kr.motd.maven</groupId>
+        <artifactId>os-maven-plugin</artifactId>
+        <version>1.7.1</version>
+      </extension>
+    </extensions>
     <plugins>
       <plugin>
         <groupId>org.jacoco</groupId>
@@ -147,6 +154,25 @@
           <excludes>com/google/cloud/spanner/spi/v1/**</excludes>
         </configuration>
       </plugin>
+      <plugin>
+        <groupId>org.xolstice.maven.plugins</groupId>
+        <artifactId>protobuf-maven-plugin</artifactId>
+        <version>0.6.1</version>
+        <configuration>
+          <protocArtifact>com.google.protobuf:protoc:4.33.2:exe:${os.detected.classifier}</protocArtifact>
+          <additionalProtoPathElements>
+            <additionalProtoPathElement>${project.basedir}/../proto-google-cloud-spanner-v1/src/main/proto</additionalProtoPathElement>
+          </additionalProtoPathElements>
+        </configuration>
+        <executions>
+          <execution>
+            <id>test-compile</id>
+            <goals>
+              <goal>test-compile</goal>
+            </goals>
+          </execution>
+        </executions>
+      </plugin>
     </plugins>
     <pluginManagement>
       <plugins>

google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java

@@ -404,6 +404,15 @@ ByteString getTransactionId() {
       }
     }
 
+    @Override
+    public void close() {
+      ByteString id = getTransactionId();
+      if (id != null && !id.isEmpty()) {
+        rpc.clearTransactionAffinity(id);
+      }
+      super.close();
+    }
+
     /**
      * Initializes the transaction with the timestamp specified within MultiUseReadOnlyTransaction.
      * This is used only for fallback of PartitionQueryRequest and PartitionReadRequest with

google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java

@@ -53,6 +53,7 @@
 import com.google.cloud.spanner.admin.instance.v1.InstanceAdminSettings;
 import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStubSettings;
 import com.google.cloud.spanner.spi.SpannerRpcFactory;
+import com.google.cloud.spanner.spi.v1.ChannelEndpointCacheFactory;
 import com.google.cloud.spanner.spi.v1.GapicSpannerRpc;
 import com.google.cloud.spanner.spi.v1.SpannerRpc;
 import com.google.cloud.spanner.v1.SpannerSettings;
@@ -202,6 +203,7 @@ public static GcpChannelPoolOptions createDefaultDynamicChannelPoolOptions() {
   }
 
   private final TransportChannelProvider channelProvider;
+  private final ChannelEndpointCacheFactory channelEndpointCacheFactory;
 
   @SuppressWarnings("rawtypes")
   private final ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> channelConfigurator;
@@ -826,6 +828,7 @@ protected SpannerOptions(Builder builder) {
 
     transportChannelExecutorThreadNameFormat = builder.transportChannelExecutorThreadNameFormat;
     channelProvider = builder.channelProvider;
+    channelEndpointCacheFactory = builder.channelEndpointCacheFactory;
     if (builder.mTLSContext != null) {
       channelConfigurator =
           channelBuilder -> {
@@ -1092,6 +1095,7 @@ public static class Builder
             createCustomClientLibToken(LIQUIBASE_API_CLIENT_LIB_TOKEN),
             createCustomClientLibToken(PG_ADAPTER_CLIENT_LIB_TOKEN));
     private TransportChannelProvider channelProvider;
+    private ChannelEndpointCacheFactory channelEndpointCacheFactory;
 
     @SuppressWarnings("rawtypes")
     private ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> channelConfigurator;
@@ -1232,6 +1236,7 @@ protected Builder() {
       this.asyncExecutorProvider = options.asyncExecutorProvider;
       this.compressorName = options.compressorName;
       this.channelProvider = options.channelProvider;
+      this.channelEndpointCacheFactory = options.channelEndpointCacheFactory;
       this.channelConfigurator = options.channelConfigurator;
       this.interceptorProvider = options.interceptorProvider;
       this.enableDirectAccess = options.enableDirectAccess;
@@ -1288,6 +1293,13 @@ public Builder setChannelProvider(TransportChannelProvider channelProvider) {
       return this;
     }
 
+    @InternalApi
+    public Builder setChannelEndpointCacheFactory(
+        ChannelEndpointCacheFactory channelEndpointCacheFactory) {
+      this.channelEndpointCacheFactory = channelEndpointCacheFactory;
+      return this;
+    }
+
     /**
      * Sets an {@link ApiFunction} that will be used to configure the transport channel. This will
      * only be used if no custom {@link TransportChannelProvider} has been set.
@@ -2141,6 +2153,11 @@ public TransportChannelProvider getChannelProvider() {
     return channelProvider;
   }
 
+  @InternalApi
+  public ChannelEndpointCacheFactory getChannelEndpointCacheFactory() {
+    return channelEndpointCacheFactory;
+  }
+
   @SuppressWarnings("rawtypes")
   public ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder> getChannelConfigurator() {
     return channelConfigurator;

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

@@ -0,0 +1,27 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.spi.v1;
+
+import com.google.api.core.InternalApi;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
+import java.io.IOException;
+
+/** Factory for creating {@link ChannelEndpointCache} instances. */
+@InternalApi
+public interface ChannelEndpointCacheFactory {
+  ChannelEndpointCache create(InstantiatingGrpcChannelProvider baseProvider) throws IOException;
+}

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

@@ -0,0 +1,115 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.spanner.spi.v1;
+
+import com.google.api.core.InternalApi;
+import com.google.spanner.v1.CacheUpdate;
+import com.google.spanner.v1.DirectedReadOptions;
+import com.google.spanner.v1.ExecuteSqlRequest;
+import com.google.spanner.v1.ReadRequest;
+import com.google.spanner.v1.RoutingHint;
+import com.google.spanner.v1.TransactionSelector;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Finds a server for a request using location-aware routing metadata.
+ *
+ * <p>This component is per-database and maintains both recipe and range caches.
+ */
+@InternalApi
+public final class ChannelFinder {
+  private final Object updateLock = new Object();
+  private final AtomicLong databaseId = new AtomicLong();
+  private final KeyRecipeCache recipeCache = new KeyRecipeCache();
+  private final KeyRangeCache rangeCache;
+  private final String databaseUri;
+
+  public ChannelFinder(ChannelEndpointCache endpointCache, String databaseUri) {
+    this.rangeCache = new KeyRangeCache(Objects.requireNonNull(endpointCache));
+    this.databaseUri = Objects.requireNonNull(databaseUri);
+  }
+
+  void useDeterministicRandom() {
+    rangeCache.useDeterministicRandom();
+  }
+
+  public void update(CacheUpdate update) {
+    synchronized (updateLock) {
+      long currentId = databaseId.get();
+      if (currentId != update.getDatabaseId()) {
+        if (currentId != 0) {
+          recipeCache.clear();
+          rangeCache.clear();
+        }
+        databaseId.set(update.getDatabaseId());
+      }
+      if (update.hasKeyRecipes()) {
+        recipeCache.addRecipes(update.getKeyRecipes());
+      }
+      rangeCache.addRanges(update);
+    }
+  }
+
+  public ChannelEndpoint findServer(ReadRequest.Builder reqBuilder) {
+    recipeCache.computeKeys(reqBuilder);
+    return fillRoutingHint(
+        reqBuilder.getTransaction(),
+        reqBuilder.getDirectedReadOptions(),
+        KeyRangeCache.RangeMode.COVERING_SPLIT,
+        reqBuilder.getRoutingHintBuilder());
+  }
+
+  public ChannelEndpoint findServer(ExecuteSqlRequest.Builder reqBuilder) {
+    recipeCache.computeKeys(reqBuilder);
+    return fillRoutingHint(
+        reqBuilder.getTransaction(),
+        reqBuilder.getDirectedReadOptions(),
+        KeyRangeCache.RangeMode.PICK_RANDOM,
+        reqBuilder.getRoutingHintBuilder());
+  }
+
+  private ChannelEndpoint fillRoutingHint(
+      TransactionSelector transactionSelector,
+      DirectedReadOptions directedReadOptions,
+      KeyRangeCache.RangeMode rangeMode,
+      RoutingHint.Builder hintBuilder) {
+    long id = databaseId.get();
+    if (id == 0) {
+      return null;
+    }
+    hintBuilder.setDatabaseId(id);
+    return rangeCache.fillRoutingHint(
+        preferLeader(transactionSelector), rangeMode, directedReadOptions, hintBuilder);
+  }
+
+  private static boolean preferLeader(TransactionSelector selector) {
+    switch (selector.getSelectorCase()) {
+      case BEGIN:
+        return !selector.getBegin().hasReadOnly() || selector.getBegin().getReadOnly().getStrong();
+      case SINGLE_USE:
+        if (!selector.getSingleUse().hasReadOnly()) {
+          return true;
+        }
+        return selector.getSingleUse().getReadOnly().getStrong();
+      case ID:
+      case SELECTOR_NOT_SET:
+      default:
+        return true;
+    }
+  }
+}