fix getState()

Author: rahul2393

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

@@ -17,10 +17,8 @@
 package com.google.cloud.spanner.spi.v1;
 
 import com.google.api.core.InternalApi;
-import com.google.api.gax.grpc.GrpcTransportChannel;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider.Builder;
-import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.cloud.spanner.ErrorCode;
 import com.google.cloud.spanner.SpannerExceptionFactory;
 import com.google.common.annotations.VisibleForTesting;
@@ -92,13 +90,14 @@ public ChannelEndpoint get(String address) {
           try {
             // Create a new provider with the same config but different endpoint.
             // This is thread-safe as withEndpoint() returns a new provider instance.
-            TransportChannelProvider newProvider = createProviderWithAuthorityOverride(addr);
+            InstantiatingGrpcChannelProvider newProvider =
+                createProviderWithAuthorityOverride(addr);
             GrpcChannelEndpoint endpoint = new GrpcChannelEndpoint(addr, newProvider);
-            logger.log(Level.INFO, "Location-aware endpoint created for address: {0}", addr);
+            logger.log(Level.FINE, "Location-aware endpoint created for address: {0}", addr);
             return endpoint;
           } catch (IOException e) {
             logger.log(
-                Level.WARNING, "Failed to create location-aware endpoint for address: " + addr, e);
+                Level.FINE, "Failed to create location-aware endpoint for address: " + addr, e);
             throw SpannerExceptionFactory.newSpannerException(
                 ErrorCode.INTERNAL, "Failed to create channel for address: " + addr, e);
           }
@@ -111,7 +110,7 @@ public ChannelEndpoint getIfPresent(String address) {
     return servers.get(address);
   }
 
-  private TransportChannelProvider createProviderWithAuthorityOverride(String address) {
+  private InstantiatingGrpcChannelProvider createProviderWithAuthorityOverride(String address) {
     InstantiatingGrpcChannelProvider endpointProvider =
         (InstantiatingGrpcChannelProvider) baseProvider.withEndpoint(address);
     if (Objects.equals(defaultAuthority, address)) {
@@ -191,15 +190,19 @@ static class GrpcChannelEndpoint implements ChannelEndpoint {
      * @param provider the channel provider (must be a gRPC provider)
      * @throws IOException if the channel cannot be created
      */
-    GrpcChannelEndpoint(String address, TransportChannelProvider provider) throws IOException {
+    GrpcChannelEndpoint(String address, InstantiatingGrpcChannelProvider provider)
+        throws IOException {
       this.address = address;
-      TransportChannelProvider readyProvider = provider;
+      // Build a raw ManagedChannel directly instead of going through getTransportChannel(),
+      // which wraps the channel in a ChannelPool that does not support getState().
+      // Location-aware routing needs getState() to check channel connectivity.
+      InstantiatingGrpcChannelProvider readyProvider = provider;
       if (provider.needsHeaders()) {
-        readyProvider = provider.withHeaders(java.util.Collections.emptyMap());
+        readyProvider =
+            (InstantiatingGrpcChannelProvider)
+                provider.withHeaders(java.util.Collections.emptyMap());
       }
-      GrpcTransportChannel transportChannel =
-          (GrpcTransportChannel) readyProvider.getTransportChannel();
-      this.channel = (ManagedChannel) transportChannel.getChannel();
+      this.channel = readyProvider.createDecoratedChannelBuilder().build();
     }
 
     /**
@@ -240,7 +243,7 @@ public boolean isHealthy() {
         return ready;
       } catch (UnsupportedOperationException e) {
         logger.log(
-            Level.WARNING,
+            Level.FINE,
             "getState(false) unsupported for location-aware endpoint {0}, treating as not ready",
             address);
         return false;

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

@@ -159,12 +159,10 @@ private ChannelFinder getOrCreateChannelFinder(String databaseId) {
   }
 
   /** Records real traffic to the selected endpoint for idle eviction tracking. */
-  private void onRequestRouted(
-      @Nullable String session, @Nullable ChannelEndpoint selectedEndpoint) {
+  private void onRequestRouted(@Nullable ChannelEndpoint selectedEndpoint) {
     if (lifecycleManager == null) {
       return;
     }
-    // Record real traffic for idle eviction tracking.
     if (selectedEndpoint != null && !defaultEndpointAddress.equals(selectedEndpoint.getAddress())) {
       lifecycleManager.recordRealTraffic(selectedEndpoint.getAddress());
     }
@@ -545,8 +543,7 @@ public void sendMessage(RequestT message) {
         this.channelFinder = finder;
 
         // Record real traffic for idle eviction tracking.
-        String session = extractSessionFromMessage(message);
-        parentChannel.onRequestRouted(session, endpoint);
+        parentChannel.onRequestRouted(endpoint);
 
         recordRouteSelectionTrace(
             methodDescriptor,
@@ -677,22 +674,6 @@ private void drainPendingRequests() {
       }
     }
 
-    @Nullable
-    private static String extractSessionFromMessage(Object message) {
-      if (message instanceof ReadRequest) {
-        return ((ReadRequest) message).getSession();
-      } else if (message instanceof ExecuteSqlRequest) {
-        return ((ExecuteSqlRequest) message).getSession();
-      } else if (message instanceof BeginTransactionRequest) {
-        return ((BeginTransactionRequest) message).getSession();
-      } else if (message instanceof CommitRequest) {
-        return ((CommitRequest) message).getSession();
-      } else if (message instanceof RollbackRequest) {
-        return ((RollbackRequest) message).getSession();
-      }
-      return null;
-    }
-
     void maybeRecordAffinity(ByteString transactionId) {
       parentChannel.recordAffinity(transactionId, selectedEndpoint, allowDefaultAffinity);
     }

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

@@ -31,11 +31,11 @@
 import java.util.Comparator;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.Set;
 import java.util.List;
 import java.util.Map;
 import java.util.NavigableMap;
 import java.util.Objects;
+import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.atomic.AtomicLong;