fix keep alives

Author: rahul2393

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

@@ -27,6 +27,7 @@
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
@@ -51,6 +52,9 @@ class GrpcChannelEndpointCache implements ChannelEndpointCache {
   /** Timeout for graceful channel shutdown. */
   private static final long SHUTDOWN_TIMEOUT_SECONDS = 5;
 
+  @VisibleForTesting static final Duration ROUTED_KEEPALIVE_TIME = Duration.ofSeconds(2);
+  @VisibleForTesting static final Duration ROUTED_KEEPALIVE_TIMEOUT = Duration.ofSeconds(20);
+
   private final InstantiatingGrpcChannelProvider baseProvider;
   private final Map<String, GrpcChannelEndpoint> servers = new ConcurrentHashMap<>();
   private final GrpcChannelEndpoint defaultEndpoint;
@@ -129,6 +133,8 @@ InstantiatingGrpcChannelProvider createProviderWithAuthorityOverride(String addr
     }
     Builder builder = endpointProvider.toBuilder();
     builder.setChannelPoolSettings(ChannelPoolSettings.staticallySized(1));
+    builder.setKeepAliveTimeDuration(ROUTED_KEEPALIVE_TIME);
+    builder.setKeepAliveTimeoutDuration(ROUTED_KEEPALIVE_TIMEOUT);
     builder.setKeepAliveWithoutCalls(Boolean.TRUE);
     final com.google.api.core.ApiFunction<ManagedChannelBuilder, ManagedChannelBuilder>
         baseConfigurator =

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

@@ -507,6 +507,7 @@ public void readWriteTransactionAbortedCommitUsesReadAffinityReplicaForBypassTra
       DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
 
       seedLocationMetadata(client);
+      waitForReplicaRoutedStrongRead(client, harness, /* expectedReplicaIndex= */ 1);
       harness.clearRequests();
       AtomicInteger attempts = new AtomicInteger();
       AtomicInteger firstReplicaIndex = new AtomicInteger(-1);
@@ -655,6 +656,30 @@ private static int waitForReplicaRoutedRead(
     throw new AssertionError("Timed out waiting for location-aware read to route to replica");
   }
 
+  private static void waitForReplicaRoutedStrongRead(
+      DatabaseClient client, SharedBackendReplicaHarness harness, int expectedReplicaIndex)
+      throws InterruptedException {
+    long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
+    while (System.nanoTime() < deadlineNanos) {
+      harness.clearRequests();
+      try (ResultSet resultSet =
+          client.singleUse().read(TABLE, KeySet.singleKey(Key.of("b")), Arrays.asList("k"))) {
+        if (resultSet.next()) {
+          if (!harness
+              .replicas
+              .get(expectedReplicaIndex)
+              .getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
+              .isEmpty()) {
+            return;
+          }
+        }
+      }
+      Thread.sleep(50L);
+    }
+    throw new AssertionError(
+        "Timed out waiting for strong read to route to replica " + expectedReplicaIndex);
+  }
+
   private static int findReplicaWithRequest(SharedBackendReplicaHarness harness, String method) {
     for (int replicaIndex = 0; replicaIndex < harness.replicas.size(); replicaIndex++) {
       if (!harness.replicas.get(replicaIndex).getRequests(method).isEmpty()) {

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GrpcChannelEndpointCacheTest.java

@@ -98,12 +98,13 @@ public void routedChannelsUseSingleUnderlyingChannel() throws Exception {
   }
 
   @Test
-  public void routedChannelsEnableKeepAliveWithoutCallsOnlyForEndpointProvider() throws Exception {
+  public void routedChannelsOverrideKeepAliveSettingsOnlyForEndpointProvider() throws Exception {
     InstantiatingGrpcChannelProvider provider =
         InstantiatingGrpcChannelProvider.newBuilder()
             .setEndpoint(DEFAULT_ENDPOINT)
             .setPoolSize(4)
             .setKeepAliveTimeDuration(java.time.Duration.ofSeconds(120))
+            .setKeepAliveTimeoutDuration(java.time.Duration.ofSeconds(60))
             .setKeepAliveWithoutCalls(Boolean.FALSE)
             .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
             .build();
@@ -113,9 +114,14 @@ public void routedChannelsEnableKeepAliveWithoutCallsOnlyForEndpointProvider() t
           cache.createProviderWithAuthorityOverride(ROUTED_ENDPOINT_A);
 
       assertThat(provider.getKeepAliveWithoutCalls()).isFalse();
+      assertThat(provider.getKeepAliveTimeDuration()).isEqualTo(java.time.Duration.ofSeconds(120));
+      assertThat(provider.getKeepAliveTimeoutDuration())
+          .isEqualTo(java.time.Duration.ofSeconds(60));
       assertThat(routedProvider.getKeepAliveWithoutCalls()).isTrue();
       assertThat(routedProvider.getKeepAliveTimeDuration())
-          .isEqualTo(provider.getKeepAliveTimeDuration());
+          .isEqualTo(GrpcChannelEndpointCache.ROUTED_KEEPALIVE_TIME);
+      assertThat(routedProvider.getKeepAliveTimeoutDuration())
+          .isEqualTo(GrpcChannelEndpointCache.ROUTED_KEEPALIVE_TIMEOUT);
     } finally {
       cache.shutdown();
     }