chore: integrate location aware routing with RPCs

Author: rahul2393

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;
+    }
+  }
+}

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

@@ -57,6 +57,7 @@
 import com.google.api.gax.rpc.UnavailableException;
 import com.google.api.gax.rpc.WatchdogProvider;
 import com.google.api.pathtemplate.PathTemplate;
+import com.google.auth.Credentials;
 import com.google.cloud.RetryHelper;
 import com.google.cloud.RetryHelper.RetryHelperException;
 import com.google.cloud.grpc.GcpManagedChannel;
@@ -209,6 +210,7 @@
 import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -223,6 +225,8 @@
 public class GapicSpannerRpc implements SpannerRpc {
   private static final PathTemplate PROJECT_NAME_TEMPLATE =
       PathTemplate.create("projects/{project}");
+  private static final String EXPERIMENTAL_LOCATION_API_ENV_VAR =
+      "GOOGLE_SPANNER_EXPERIMENTAL_LOCATION_API";
   private static final PathTemplate OPERATION_NAME_TEMPLATE =
       PathTemplate.create("{database=projects/*/instances/*/databases/*}/operations/{operation}");
   private static final int MAX_MESSAGE_SIZE = 256 * 1024 * 1024;
@@ -285,6 +289,96 @@ public class GapicSpannerRpc implements SpannerRpc {
 
   private final GrpcCallContext baseGrpcCallContext;
 
+  private static final class KeyAwareTransportChannelProvider implements TransportChannelProvider {
+    private final InstantiatingGrpcChannelProvider baseProvider;
+
+    private KeyAwareTransportChannelProvider(InstantiatingGrpcChannelProvider.Builder builder) {
+      this.baseProvider = builder.build();
+    }
+
+    private KeyAwareTransportChannelProvider(InstantiatingGrpcChannelProvider baseProvider) {
+      this.baseProvider = baseProvider;
+    }
+
+    @Override
+    public GrpcTransportChannel getTransportChannel() throws IOException {
+      return GrpcTransportChannel.newBuilder()
+          .setManagedChannel(KeyAwareChannel.create(baseProvider))
+          .build();
+    }
+
+    @Override
+    public String getTransportName() {
+      return baseProvider.getTransportName();
+    }
+
+    @Override
+    public boolean needsEndpoint() {
+      return baseProvider.needsEndpoint();
+    }
+
+    @Override
+    public boolean needsCredentials() {
+      return baseProvider.needsCredentials();
+    }
+
+    @Override
+    public boolean needsExecutor() {
+      return baseProvider.needsExecutor();
+    }
+
+    @Override
+    public boolean needsHeaders() {
+      return baseProvider.needsHeaders();
+    }
+
+    @Override
+    public boolean shouldAutoClose() {
+      return baseProvider.shouldAutoClose();
+    }
+
+    @Override
+    public TransportChannelProvider withEndpoint(String endpoint) {
+      return new KeyAwareTransportChannelProvider(
+          (InstantiatingGrpcChannelProvider) baseProvider.withEndpoint(endpoint));
+    }
+
+    @Override
+    public TransportChannelProvider withCredentials(Credentials credentials) {
+      return new KeyAwareTransportChannelProvider(
+          (InstantiatingGrpcChannelProvider) baseProvider.withCredentials(credentials));
+    }
+
+    @Override
+    public TransportChannelProvider withHeaders(Map<String, String> headers) {
+      return new KeyAwareTransportChannelProvider(
+          (InstantiatingGrpcChannelProvider) baseProvider.withHeaders(headers));
+    }
+
+    @Override
+    public TransportChannelProvider withPoolSize(int poolSize) {
+      return new KeyAwareTransportChannelProvider(
+          (InstantiatingGrpcChannelProvider) baseProvider.withPoolSize(poolSize));
+    }
+
+    @Override
+    public TransportChannelProvider withExecutor(ScheduledExecutorService executor) {
+      return new KeyAwareTransportChannelProvider(
+          (InstantiatingGrpcChannelProvider) baseProvider.withExecutor(executor));
+    }
+
+    @Override
+    public TransportChannelProvider withExecutor(Executor executor) {
+      return new KeyAwareTransportChannelProvider(
+          (InstantiatingGrpcChannelProvider) baseProvider.withExecutor(executor));
+    }
+
+    @Override
+    public boolean acceptsPoolSize() {
+      return baseProvider.acceptsPoolSize();
+    }
+  }
+
   public static GapicSpannerRpc create(SpannerOptions options) {
     return new GapicSpannerRpc(options);
   }
@@ -393,9 +487,13 @@ public GapicSpannerRpc(final SpannerOptions options) {
       // If it is enabled in options uses the channel pool provided by the gRPC-GCP extension.
       maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options);
 
+      boolean enableLocationApi =
+          Boolean.parseBoolean(System.getenv(EXPERIMENTAL_LOCATION_API_ENV_VAR));
       TransportChannelProvider channelProvider =
-          MoreObjects.firstNonNull(
-              options.getChannelProvider(), defaultChannelProviderBuilder.build());
+          enableLocationApi
+              ? new KeyAwareTransportChannelProvider(defaultChannelProviderBuilder)
+              : MoreObjects.firstNonNull(
+                  options.getChannelProvider(), defaultChannelProviderBuilder.build());
 
       CredentialsProvider credentialsProvider =
           GrpcTransportOptions.setUpCredentialsProvider(options);