chore: handle unary gRPC call ordering in KeyAwareChannel

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

@@ -18,6 +18,7 @@
 
 import static com.google.api.gax.util.TimeConversionUtils.toJavaTimeDuration;
 import static com.google.api.gax.util.TimeConversionUtils.toThreetenDuration;
+import static com.google.cloud.spanner.spi.v1.GapicSpannerRpc.EXPERIMENTAL_LOCATION_API_ENV_VAR;
 
 import com.google.api.core.ApiFunction;
 import com.google.api.core.BetaApi;
@@ -257,6 +258,7 @@ public static GcpChannelPoolOptions createDefaultDynamicChannelPoolOptions() {
   private final OpenTelemetry openTelemetry;
   private final boolean enableApiTracing;
   private final boolean enableBuiltInMetrics;
+  private final boolean enableLocationApi;
   private final boolean enableExtendedTracing;
   private final boolean enableEndToEndTracing;
   private final String monitoringHost;
@@ -926,6 +928,7 @@ protected SpannerOptions(Builder builder) {
     } else {
       enableBuiltInMetrics = builder.enableBuiltInMetrics;
     }
+    enableLocationApi = builder.enableLocationApi;
     enableEndToEndTracing = builder.enableEndToEndTracing;
     monitoringHost = builder.monitoringHost;
     defaultTransactionOptions = builder.defaultTransactionOptions;
@@ -993,6 +996,10 @@ default boolean isEnableEndToEndTracing() {
       return false;
     }
 
+    default boolean isEnableLocationApi() {
+      return false;
+    }
+
     @Deprecated
     @ObsoleteApi(
         "This will be removed in an upcoming version without a major version bump. You should use"
@@ -1084,6 +1091,11 @@ public boolean isEnableEndToEndTracing() {
       return Boolean.parseBoolean(System.getenv(SPANNER_ENABLE_END_TO_END_TRACING));
     }
 
+    @Override
+    public boolean isEnableLocationApi() {
+      return Boolean.parseBoolean(System.getenv(EXPERIMENTAL_LOCATION_API_ENV_VAR));
+    }
+
     @Override
     public String getMonitoringHost() {
       return System.getenv(SPANNER_MONITORING_HOST);
@@ -1164,6 +1176,7 @@ public static class Builder
     private boolean enableExtendedTracing = SpannerOptions.environment.isEnableExtendedTracing();
     private boolean enableEndToEndTracing = SpannerOptions.environment.isEnableEndToEndTracing();
     private boolean enableBuiltInMetrics = SpannerOptions.environment.isEnableBuiltInMetrics();
+    private boolean enableLocationApi = SpannerOptions.environment.isEnableLocationApi();
     private String monitoringHost = SpannerOptions.environment.getMonitoringHost();
     private SslContext mTLSContext = null;
     private String experimentalHost = null;
@@ -1270,6 +1283,7 @@ protected Builder() {
       this.enableApiTracing = options.enableApiTracing;
       this.enableExtendedTracing = options.enableExtendedTracing;
       this.enableBuiltInMetrics = options.enableBuiltInMetrics;
+      this.enableLocationApi = options.enableLocationApi;
       this.enableEndToEndTracing = options.enableEndToEndTracing;
       this.monitoringHost = options.monitoringHost;
       this.defaultTransactionOptions = options.defaultTransactionOptions;
@@ -2434,6 +2448,11 @@ public boolean isEnableBuiltInMetrics() {
     return enableBuiltInMetrics;
   }
 
+  @InternalApi
+  public boolean isEnableLocationApi() {
+    return enableLocationApi;
+  }
+
   /** Returns the override metrics Host. */
   String getMonitoringHost() {
     return monitoringHost;

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

@@ -226,7 +226,7 @@
 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 =
+  public 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}");
@@ -400,7 +400,8 @@ public GapicSpannerRpc(final SpannerOptions options) {
       maybeEnableGrpcGcpExtension(defaultChannelProviderBuilder, options);
 
       boolean enableLocationApi =
-          Boolean.parseBoolean(System.getenv(EXPERIMENTAL_LOCATION_API_ENV_VAR));
+          options.isEnableLocationApi()
+              || Boolean.parseBoolean(System.getenv(EXPERIMENTAL_LOCATION_API_ENV_VAR));
       TransportChannelProvider baseChannelProvider =
           MoreObjects.firstNonNull(
               options.getChannelProvider(), defaultChannelProviderBuilder.build());

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

@@ -245,6 +245,12 @@ static final class KeyAwareClientCall<RequestT, ResponseT>
     @Nullable private ChannelEndpoint selectedEndpoint;
     @Nullable private ByteString transactionIdToClear;
     private boolean allowDefaultAffinity;
+    private long pendingRequests;
+    private boolean pendingHalfClose;
+    @Nullable private Boolean pendingMessageCompression;
+    private boolean cancelled;
+    @Nullable private String cancelMessage;
+    @Nullable private Throwable cancelCause;
 
     KeyAwareClientCall(
         KeyAwareChannel parentChannel,
@@ -268,11 +274,22 @@ protected ClientCall<RequestT, ResponseT> delegate() {
     public void start(Listener<ResponseT> responseListener, Metadata headers) {
       this.responseListener = new KeyAwareClientCallListener<>(responseListener, this);
       this.headers = headers;
+      if (cancelled) {
+        this.responseListener.onClose(
+            io.grpc.Status.CANCELLED.withDescription(cancelMessage).withCause(cancelCause),
+            new Metadata());
+      }
     }
 
     @Override
     @SuppressWarnings("unchecked")
     public void sendMessage(RequestT message) {
+      if (cancelled) {
+        return;
+      }
+      if (responseListener == null || headers == null) {
+        throw new IllegalStateException("start must be called before sendMessage");
+      }
       ChannelEndpoint endpoint = null;
       ChannelFinder finder = null;
 
@@ -326,26 +343,81 @@ public void sendMessage(RequestT message) {
       this.channelFinder = finder;
 
       delegate = endpoint.getChannel().newCall(methodDescriptor, callOptions);
+      if (pendingMessageCompression != null) {
+        delegate.setMessageCompression(pendingMessageCompression);
+      }
       delegate.start(responseListener, headers);
+      drainPendingRequests();
       delegate.sendMessage(message);
+      if (pendingHalfClose) {
+        delegate.halfClose();
+      }
     }
 
     @Override
     public void halfClose() {
       if (delegate != null) {
         delegate.halfClose();
       } else {
-        throw new IllegalStateException("halfClose called before sendMessage");
+        pendingHalfClose = true;
       }
     }
 
     @Override
     public void cancel(@Nullable String message, @Nullable Throwable cause) {
       if (delegate != null) {
         delegate.cancel(message, cause);
-      } else if (responseListener != null) {
-        responseListener.onClose(
-            io.grpc.Status.CANCELLED.withDescription(message).withCause(cause), new Metadata());
+      } else {
+        cancelled = true;
+        cancelMessage = message;
+        cancelCause = cause;
+        if (responseListener != null) {
+          responseListener.onClose(
+              io.grpc.Status.CANCELLED.withDescription(message).withCause(cause), new Metadata());
+        }
+      }
+    }
+
+    @Override
+    public void request(int numMessages) {
+      if (delegate != null) {
+        delegate.request(numMessages);
+        return;
+      }
+      if (numMessages <= 0) {
+        return;
+      }
+      long updated = pendingRequests + numMessages;
+      if (updated < 0L) {
+        updated = Long.MAX_VALUE;
+      }
+      pendingRequests = updated;
+    }
+
+    @Override
+    public boolean isReady() {
+      if (delegate == null) {
+        return false;
+      }
+      return delegate.isReady();
+    }
+
+    @Override
+    public void setMessageCompression(boolean enabled) {
+      if (delegate != null) {
+        delegate.setMessageCompression(enabled);
+      } else {
+        pendingMessageCompression = enabled;
+      }
+    }
+
+    private void drainPendingRequests() {
+      long requests = pendingRequests;
+      pendingRequests = 0L;
+      while (requests > 0) {
+        int batch = requests > Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) requests;
+        delegate.request(batch);
+        requests -= batch;
       }
     }
 
@@ -419,6 +491,9 @@ public void onMessage(ResponseT message) {
         transactionId = transactionIdFromMetadata(response);
       } else if (message instanceof ResultSet) {
         ResultSet response = (ResultSet) message;
+        if (response.hasCacheUpdate() && call.channelFinder != null) {
+          call.channelFinder.update(response.getCacheUpdate());
+        }
         transactionId = transactionIdFromMetadata(response);
       } else if (message instanceof Transaction) {
         Transaction response = (Transaction) message;

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

@@ -0,0 +1,140 @@
+/*
+ * 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;
+
+import static org.junit.Assert.assertEquals;
+
+import com.google.cloud.NoCredentials;
+import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
+import com.google.cloud.spanner.connection.AbstractMockServerTest;
+import com.google.cloud.spanner.connection.RandomResultSetGenerator;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
+import io.grpc.ManagedChannelBuilder;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ThreadLocalRandom;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class LocationAwareTest extends AbstractMockServerTest {
+  private static final Statement SELECT_RANDOM_STATEMENT = Statement.of("select * from random");
+  private static final int RANDOM_RESULT_ROW_COUNT = 20;
+  private static Spanner spanner;
+  private static DatabaseClient client;
+
+  @BeforeClass
+  public static void enableLocationApiAndSetupClient() {
+    SpannerOptions.useEnvironment(
+        new SpannerOptions.SpannerEnvironment() {
+          @Override
+          public boolean isEnableLocationApi() {
+            return true;
+          }
+        });
+    spanner =
+        SpannerOptions.newBuilder()
+            .setProjectId("my-project")
+            .setHost(String.format("http://localhost:%d", getPort()))
+            .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+            .setCredentials(NoCredentials.getInstance())
+            .build()
+            .getService();
+    client = spanner.getDatabaseClient(DatabaseId.of("my-project", "my-instance", "my-database"));
+
+    RandomResultSetGenerator generator = new RandomResultSetGenerator(RANDOM_RESULT_ROW_COUNT);
+    mockSpanner.putStatementResult(
+        StatementResult.query(SELECT_RANDOM_STATEMENT, generator.generate()));
+  }
+
+  @AfterClass
+  public static void cleanup() {
+    SpannerOptions.useDefaultEnvironment();
+    if (spanner != null) {
+      spanner.close();
+    }
+  }
+
+  @Test
+  public void testSingleQuery() {
+    int rowCount = 0;
+    try (ResultSet resultSet = client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
+      while (resultSet.next()) {
+        rowCount++;
+      }
+    }
+    assertEquals(RANDOM_RESULT_ROW_COUNT, rowCount);
+  }
+
+  @Test
+  public void testParallelQueries() throws Exception {
+    int numThreads = 10;
+    ListeningExecutorService executor =
+        MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(numThreads));
+    List<ListenableFuture<Void>> results = new ArrayList<>();
+    for (int i = 0; i < numThreads; i++) {
+      results.add(
+          executor.submit(
+              () -> {
+                try (ResultSet resultSet =
+                    client.singleUse().executeQuery(SELECT_RANDOM_STATEMENT)) {
+                  while (resultSet.next()) {
+                    // Randomly stop consuming results somewhere halfway the results (sometimes).
+                    if (ThreadLocalRandom.current().nextInt(RANDOM_RESULT_ROW_COUNT * 2) == 5) {
+                      break;
+                    }
+                  }
+                }
+                return null;
+              }));
+    }
+    executor.shutdown();
+    Futures.allAsList(results).get();
+  }
+
+  @Test
+  public void testSingleReadWriteTransaction() {
+    client.readWriteTransaction().run(transaction -> transaction.executeUpdate(INSERT_STATEMENT));
+  }
+
+  @Test
+  public void testParallelReadWriteTransactions() throws Exception {
+    int numThreads = 10;
+    ListeningExecutorService executor =
+        MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(numThreads));
+    List<ListenableFuture<Void>> results = new ArrayList<>();
+    for (int i = 0; i < numThreads; i++) {
+      results.add(
+          executor.submit(
+              () -> {
+                client
+                    .readWriteTransaction()
+                    .run(transaction -> transaction.executeUpdate(INSERT_STATEMENT));
+                return null;
+              }));
+    }
+    executor.shutdown();
+    Futures.allAsList(results).get();
+  }
+}

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/AbstractMockServerTest.java

@@ -314,7 +314,7 @@ protected String getBaseUrl() {
         server.getPort());
   }
 
-  protected int getPort() {
+  protected static int getPort() {
     return server.getPort();
   }