feat(spanner): add settings for gRCP keep-alive

java-spanner/google-cloud-spanner/pom.xml

@@ -15,6 +15,7 @@
   </parent>
   <properties>
     <site.installationModule>google-cloud-spanner</site.installationModule>
+    <awaitility.version>4.3.0</awaitility.version>
     <opencensus.version>0.31.1</opencensus.version>
     <google.cloud.monitoring.version>3.85.0</google.cloud.monitoring.version>
     <spanner.testenv.config.class>com.google.cloud.spanner.GceTestEnvConfig</spanner.testenv.config.class>
@@ -435,6 +436,13 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.awaitility</groupId>
+      <artifactId>awaitility</artifactId>
+      <version>${awaitility.version}</version>
+      <scope>test</scope>
+    </dependency>
+
     <dependency>
       <groupId>com.google.api.grpc</groupId>
       <artifactId>proto-google-cloud-spanner-executor-v1</artifactId>

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

@@ -268,6 +268,8 @@ static GcpChannelPoolOptions mergeWithDefaultChannelPoolOptions(
   private final InstanceAdminStubSettings instanceAdminStubSettings;
   private final DatabaseAdminStubSettings databaseAdminStubSettings;
   private final Duration partitionedDmlTimeout;
+  private final Duration grpcKeepAliveTime;
+  private final Duration grpcKeepAliveTimeout;
   private final boolean grpcGcpExtensionEnabled;
   private final GcpManagedChannelOptions grpcGcpOptions;
   private final boolean dynamicChannelPoolEnabled;
@@ -939,6 +941,8 @@ protected SpannerOptions(Builder builder) {
       throw SpannerExceptionFactory.newSpannerException(e);
     }
     partitionedDmlTimeout = builder.partitionedDmlTimeout;
+    grpcKeepAliveTime = builder.grpcKeepAliveTime;
+    grpcKeepAliveTimeout = builder.grpcKeepAliveTimeout;
     grpcGcpExtensionEnabled = builder.grpcGcpExtensionEnabled;
     grpcGcpOptions = builder.grpcGcpOptions;
 
@@ -1317,6 +1321,8 @@ private static Builder prepareBuilder(Builder builder) {
     private Map<DatabaseId, QueryOptions> defaultQueryOptions = new HashMap<>();
     private boolean enableGrpcGcpOtelMetrics =
         SpannerOptions.environment.isEnableGrpcGcpOtelMetrics();
+    private Duration grpcKeepAliveTime = Duration.ofSeconds(120);
+    private Duration grpcKeepAliveTimeout = Duration.ofSeconds(20);
     private CallCredentialsProvider callCredentialsProvider;
     private CloseableExecutorProvider asyncExecutorProvider;
     private String compressorName;
@@ -1430,6 +1436,8 @@ protected Builder() {
       this.enableGrpcGcpOtelMetrics = options.enableGrpcGcpOtelMetrics;
       this.defaultQueryOptions = options.defaultQueryOptions;
       this.callCredentialsProvider = options.callCredentialsProvider;
+      this.grpcKeepAliveTime = options.grpcKeepAliveTime;
+      this.grpcKeepAliveTimeout = options.grpcKeepAliveTimeout;
       this.asyncExecutorProvider = options.asyncExecutorProvider;
       this.compressorName = options.compressorName;
       this.channelProvider = options.channelProvider;
@@ -1692,6 +1700,32 @@ public Builder setPartitionedDmlTimeoutDuration(Duration timeout) {
       return this;
     }
 
+    /**
+     * Sets the keep-alive time for gRPC connections. The default is 120 seconds. Note that the
+     * client-side keepalive time is clamped to a minimum of 10 seconds by gRPC.
+     */
+    public Builder setGrpcKeepAliveTime(Duration grpcKeepAliveTime) {
+      Preconditions.checkNotNull(grpcKeepAliveTime, "grpcKeepAliveTime cannot be null");
+      Preconditions.checkArgument(
+          !grpcKeepAliveTime.isNegative() && !grpcKeepAliveTime.isZero(),
+          "grpcKeepAliveTime must be positive");
+      this.grpcKeepAliveTime = grpcKeepAliveTime;
+      return this;
+    }
+
+    /**
+     * Sets the keep-alive timeout for gRPC connections. The default is 20 seconds. Note that the
+     * client-side keepalive timeout is clamped to a minimum of 20 milliseconds by gRPC.
+     */
+    public Builder setGrpcKeepAliveTimeout(Duration grpcKeepAliveTimeout) {
+      Preconditions.checkNotNull(grpcKeepAliveTimeout, "grpcKeepAliveTimeout cannot be null");
+      Preconditions.checkArgument(
+          !grpcKeepAliveTimeout.isNegative() && !grpcKeepAliveTimeout.isZero(),
+          "grpcKeepAliveTimeout must be positive");
+      this.grpcKeepAliveTimeout = grpcKeepAliveTimeout;
+      return this;
+    }
+
     /**
      * Instructs the client library to automatically throttle the number of administrative requests
      * if the rate of administrative requests generated by this {@link Spanner} instance will exceed
@@ -2493,6 +2527,14 @@ public Duration getPartitionedDmlTimeoutDuration() {
     return partitionedDmlTimeout;
   }
 
+  public Duration getGrpcKeepAliveTime() {
+    return grpcKeepAliveTime;
+  }
+
+  public Duration getGrpcKeepAliveTimeout() {
+    return grpcKeepAliveTimeout;
+  }
+
   public boolean isGrpcGcpExtensionEnabled() {
     return grpcGcpExtensionEnabled;
   }

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java

@@ -42,6 +42,8 @@
 import static com.google.cloud.spanner.connection.ConnectionProperties.ENCODED_CREDENTIALS;
 import static com.google.cloud.spanner.connection.ConnectionProperties.ENDPOINT;
 import static com.google.cloud.spanner.connection.ConnectionProperties.GRPC_INTERCEPTOR_PROVIDER;
+import static com.google.cloud.spanner.connection.ConnectionProperties.GRPC_KEEPALIVE_TIME;
+import static com.google.cloud.spanner.connection.ConnectionProperties.GRPC_KEEPALIVE_TIMEOUT;
 import static com.google.cloud.spanner.connection.ConnectionProperties.IS_EXPERIMENTAL_HOST;
 import static com.google.cloud.spanner.connection.ConnectionProperties.LENIENT;
 import static com.google.cloud.spanner.connection.ConnectionProperties.MAX_COMMIT_DELAY;
@@ -279,6 +281,12 @@ public class ConnectionOptions {
    */
   public static final String OAUTH_TOKEN_PROPERTY_NAME = "oauthToken";
 
+  /** Name of the 'grpcKeepAliveTime' connection property. */
+  public static final String GRPC_KEEPALIVE_TIME_PROPERTY_NAME = "grpcKeepAliveTime";
+
+  /** Name of the 'grpcKeepAliveTimeout' connection property. */
+  public static final String GRPC_KEEPALIVE_TIMEOUT_PROPERTY_NAME = "grpcKeepAliveTimeout";
+
   /** Name of the 'minSessions' connection property. */
   public static final String MIN_SESSIONS_PROPERTY_NAME = "minSessions";
 
@@ -1083,6 +1091,16 @@ public Integer getNumChannels() {
     return getInitialConnectionPropertyValue(NUM_CHANNELS);
   }
 
+  /** The gRPC keepalive time for this connection. */
+  public Duration getGrpcKeepAliveTime() {
+    return getInitialConnectionPropertyValue(GRPC_KEEPALIVE_TIME);
+  }
+
+  /** The gRPC keepalive timeout for this connection. */
+  public Duration getGrpcKeepAliveTimeout() {
+    return getInitialConnectionPropertyValue(GRPC_KEEPALIVE_TIMEOUT);
+  }
+
   /** Whether dynamic channel pooling is enabled for this connection. */
   public Boolean isEnableDynamicChannelPool() {
     return getInitialConnectionPropertyValue(ENABLE_DYNAMIC_CHANNEL_POOL);

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java

@@ -97,6 +97,8 @@
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENABLE_GRPC_INTERCEPTOR_PROVIDER_SYSTEM_PROPERTY;
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENCODED_CREDENTIALS_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENDPOINT_PROPERTY_NAME;
+import static com.google.cloud.spanner.connection.ConnectionOptions.GRPC_KEEPALIVE_TIMEOUT_PROPERTY_NAME;
+import static com.google.cloud.spanner.connection.ConnectionOptions.GRPC_KEEPALIVE_TIME_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.IS_EXPERIMENTAL_HOST_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.KEEP_TRANSACTION_ALIVE_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.LENIENT_PROPERTY_NAME;
@@ -253,6 +255,24 @@ public class ConnectionProperties {
           BOOLEANS,
           BooleanConverter.INSTANCE,
           Context.STARTUP);
+  static final ConnectionProperty<Duration> GRPC_KEEPALIVE_TIME =
+      create(
+          GRPC_KEEPALIVE_TIME_PROPERTY_NAME,
+          "The keepalive time for gRPC connections (e.g. '120s', '20s'). "
+              + "Setting a lower keep-alive time (minimum 10s enforced by the gRPC library) "
+              + "helps detect disconnected connections faster.",
+          null,
+          DurationConverter.INSTANCE,
+          Context.STARTUP);
+  static final ConnectionProperty<Duration> GRPC_KEEPALIVE_TIMEOUT =
+      create(
+          GRPC_KEEPALIVE_TIMEOUT_PROPERTY_NAME,
+          "The keepalive timeout for gRPC connections (e.g. '20s', '5s'). "
+              + "This determines how long the client waits for a keep-alive ping response before terminating "
+              + "the connection. A lower timeout helps speed up recovery during network failures.",
+          null,
+          DurationConverter.INSTANCE,
+          Context.STARTUP);
 
   /**
    * @deprecated Use {@link #TYPE} with value "omni" instead.

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java

@@ -32,6 +32,7 @@
 import io.grpc.ManagedChannelBuilder;
 import io.opentelemetry.api.OpenTelemetry;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -151,6 +152,8 @@ public boolean equals(Object o) {
   static class SpannerPoolKey {
     private final String host;
     private final String projectId;
+    private final Duration grpcKeepAliveTime;
+    private final Duration grpcKeepAliveTimeout;
     private final CredentialsKey credentialsKey;
     private final SessionPoolOptions sessionPoolOptions;
     private final Integer numChannels;
@@ -222,6 +225,8 @@ private SpannerPoolKey(ConnectionOptions options) throws IOException {
       this.enableDirectAccess = options.isEnableDirectAccess();
       this.universeDomain = options.getUniverseDomain();
       this.grpcInterceptorProvider = options.getGrpcInterceptorProviderName();
+      this.grpcKeepAliveTime = options.getGrpcKeepAliveTime();
+      this.grpcKeepAliveTimeout = options.getGrpcKeepAliveTimeout();
     }
 
     @Override
@@ -259,7 +264,9 @@ public boolean equals(Object o) {
           && Objects.equals(this.instanceType, other.instanceType)
           && Objects.equals(this.enableDirectAccess, other.enableDirectAccess)
           && Objects.equals(this.universeDomain, other.universeDomain)
-          && Objects.equals(this.grpcInterceptorProvider, other.grpcInterceptorProvider);
+          && Objects.equals(this.grpcInterceptorProvider, other.grpcInterceptorProvider)
+          && Objects.equals(this.grpcKeepAliveTime, other.grpcKeepAliveTime)
+          && Objects.equals(this.grpcKeepAliveTimeout, other.grpcKeepAliveTimeout);
     }
 
     @Override
@@ -292,7 +299,9 @@ public int hashCode() {
           this.instanceType,
           this.enableDirectAccess,
           this.universeDomain,
-          this.grpcInterceptorProvider);
+          this.grpcInterceptorProvider,
+          this.grpcKeepAliveTime,
+          this.grpcKeepAliveTimeout);
     }
   }
 
@@ -510,6 +519,12 @@ Spanner createSpanner(SpannerPoolKey key, ConnectionOptions options) {
     if (options.getChannelProvider() != null) {
       builder.setChannelProvider(options.getChannelProvider());
     }
+    if (key.grpcKeepAliveTime != null) {
+      builder.setGrpcKeepAliveTime(key.grpcKeepAliveTime);
+    }
+    if (key.grpcKeepAliveTimeout != null) {
+      builder.setGrpcKeepAliveTimeout(key.grpcKeepAliveTimeout);
+    }
     if (!options.isRouteToLeader()) {
       builder.disableLeaderAwareRouting();
     }

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

@@ -251,7 +251,6 @@ public class GapicSpannerRpc implements SpannerRpc {
       "com.google.cloud.spanner.watchdogPeriodSeconds";
   private static final int DEFAULT_TIMEOUT_SECONDS = 30 * 60;
   private static final int DEFAULT_PERIOD_SECONDS = 10;
-  private static final int GRPC_KEEPALIVE_SECONDS = 2 * 60;
   private static final String USER_AGENT_KEY = "user-agent";
   private static final String CLIENT_LIBRARY_LANGUAGE = "spanner-java";
   public static final String DEFAULT_USER_AGENT =
@@ -740,9 +739,10 @@ private InstantiatingGrpcChannelProvider.Builder createBaseChannelProviderBuilde
             .setMaxInboundMetadataSize(MAX_METADATA_SIZE)
             .setPoolSize(options.getNumChannels())
 
-            // Set a keepalive time of 120 seconds to help long running
+            // Set a keepalive time to help long running
             // commit GRPC calls succeed
-            .setKeepAliveTimeDuration(Duration.ofSeconds(GRPC_KEEPALIVE_SECONDS))
+            .setKeepAliveTimeDuration(options.getGrpcKeepAliveTime())
+            .setKeepAliveTimeoutDuration(options.getGrpcKeepAliveTimeout())
 
             // Then check if SpannerOptions provides an InterceptorProvider. Create a default
             // SpannerInterceptorProvider if none is provided

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

@@ -1452,4 +1452,69 @@ public void testLogin() {
     SpannerOptions options = builder.build();
     assertTrue(options.getCredentials() instanceof SpannerOmniCredentials);
   }
+
+  @Test
+  public void testGrpcKeepAliveTime() {
+    SpannerOptions defaultOptions =
+        SpannerOptions.newBuilder()
+            .setProjectId("test-project")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
+    assertEquals(Duration.ofSeconds(120), defaultOptions.getGrpcKeepAliveTime());
+
+    SpannerOptions customOptions =
+        SpannerOptions.newBuilder()
+            .setProjectId("test-project")
+            .setCredentials(NoCredentials.getInstance())
+            .setGrpcKeepAliveTime(Duration.ofSeconds(20))
+            .build();
+    assertEquals(Duration.ofSeconds(20), customOptions.getGrpcKeepAliveTime());
+
+    SpannerOptions optionsFromBuilder = customOptions.toBuilder().build();
+    assertEquals(Duration.ofSeconds(20), optionsFromBuilder.getGrpcKeepAliveTime());
+
+    assertThrows(
+        NullPointerException.class, () -> SpannerOptions.newBuilder().setGrpcKeepAliveTime(null));
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> SpannerOptions.newBuilder().setGrpcKeepAliveTime(Duration.ZERO));
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> SpannerOptions.newBuilder().setGrpcKeepAliveTime(Duration.ofSeconds(-10)));
+  }
+
+  @Test
+  public void testGrpcKeepAliveTimeout() {
+    SpannerOptions defaultOptions =
+        SpannerOptions.newBuilder()
+            .setProjectId("test-project")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
+    assertEquals(Duration.ofSeconds(20), defaultOptions.getGrpcKeepAliveTimeout());
+
+    SpannerOptions customOptions =
+        SpannerOptions.newBuilder()
+            .setProjectId("test-project")
+            .setCredentials(NoCredentials.getInstance())
+            .setGrpcKeepAliveTimeout(Duration.ofSeconds(5))
+            .build();
+    assertEquals(Duration.ofSeconds(5), customOptions.getGrpcKeepAliveTimeout());
+
+    SpannerOptions optionsFromBuilder = customOptions.toBuilder().build();
+    assertEquals(Duration.ofSeconds(5), optionsFromBuilder.getGrpcKeepAliveTimeout());
+
+    assertThrows(
+        NullPointerException.class,
+        () -> SpannerOptions.newBuilder().setGrpcKeepAliveTimeout(null));
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> SpannerOptions.newBuilder().setGrpcKeepAliveTimeout(Duration.ZERO));
+
+    assertThrows(
+        IllegalArgumentException.class,
+        () -> SpannerOptions.newBuilder().setGrpcKeepAliveTimeout(Duration.ofSeconds(-10)));
+  }
 }

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

@@ -191,6 +191,8 @@ public void getOperation(
             .addService(mockDatabaseAdmin)
             .addService(mockOperations)
             .intercept(interceptor)
+            .permitKeepAliveTime(10, TimeUnit.MILLISECONDS)
+            .permitKeepAliveWithoutCalls(true)
             .build()
             .start();
     mockSpanner.putStatementResult(

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

@@ -1597,4 +1597,44 @@ public void testDcpWithAllOptions() {
     assertEquals(Integer.valueOf(15), options.getDcpMaxChannels());
     assertEquals(Integer.valueOf(5), options.getDcpInitialChannels());
   }
+
+  @Test
+  public void testGrpcKeepAliveTimeOption() {
+    ConnectionOptions options =
+        ConnectionOptions.newBuilder()
+            .setUri(
+                "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database"
+                    + "?grpcKeepAliveTime='20s'")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
+    assertEquals(Duration.ofSeconds(20), options.getGrpcKeepAliveTime());
+
+    ConnectionOptions defaultOptions =
+        ConnectionOptions.newBuilder()
+            .setUri(
+                "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
+    assertNull(defaultOptions.getGrpcKeepAliveTime());
+  }
+
+  @Test
+  public void testGrpcKeepAliveTimeoutOption() {
+    ConnectionOptions options =
+        ConnectionOptions.newBuilder()
+            .setUri(
+                "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database"
+                    + "?grpcKeepAliveTimeout='15s'")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
+    assertEquals(Duration.ofSeconds(15), options.getGrpcKeepAliveTimeout());
+
+    ConnectionOptions defaultOptions =
+        ConnectionOptions.newBuilder()
+            .setUri(
+                "cloudspanner:/projects/test-project-123/instances/test-instance/databases/test-database")
+            .setCredentials(NoCredentials.getInstance())
+            .build();
+    assertNull(defaultOptions.getGrpcKeepAliveTimeout());
+  }
 }