feat: add session pool option for modelling a timeout around session acquisition.

README.md

@@ -57,13 +57,13 @@ implementation 'com.google.cloud:google-cloud-spanner'
 If you are using Gradle without BOM, add this to your dependencies:
 
 ```Groovy
-implementation 'com.google.cloud:google-cloud-spanner:6.47.0'
+implementation 'com.google.cloud:google-cloud-spanner:6.48.0'
 ```
 
 If you are using SBT, add this to your dependencies:
 
 ```Scala
-libraryDependencies += "com.google.cloud" % "google-cloud-spanner" % "6.47.0"
+libraryDependencies += "com.google.cloud" % "google-cloud-spanner" % "6.48.0"
 ```
 <!-- {x-version-update-end} -->
 
@@ -431,7 +431,7 @@ Java is a registered trademark of Oracle and/or its affiliates.
 [kokoro-badge-link-5]: http://storage.googleapis.com/cloud-devrel-public/java/badges/java-spanner/java11.html
 [stability-image]: https://img.shields.io/badge/stability-stable-green
 [maven-version-image]: https://img.shields.io/maven-central/v/com.google.cloud/google-cloud-spanner.svg
-[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-spanner/6.47.0
+[maven-version-link]: https://central.sonatype.com/artifact/com.google.cloud/google-cloud-spanner/6.48.0
 [authentication]: https://github.com/googleapis/google-cloud-java#authentication
 [auth-scopes]: https://developers.google.com/identity/protocols/oauth2/scopes
 [predefined-iam-roles]: https://cloud.google.com/iam/docs/understanding-roles#predefined_roles

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

@@ -63,7 +63,6 @@
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
 import com.google.common.util.concurrent.ForwardingListenableFuture;
 import com.google.common.util.concurrent.ForwardingListenableFuture.SimpleForwardingListenableFuture;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -120,17 +119,6 @@ class SessionPool {
   private static final Logger logger = Logger.getLogger(SessionPool.class.getName());
   private static final Tracer tracer = Tracing.getTracer();
   static final String WAIT_FOR_SESSION = "SessionPool.WaitForSession";
-  static final ImmutableSet<ErrorCode> SHOULD_STOP_PREPARE_SESSIONS_ERROR_CODES =
-      ImmutableSet.of(
-          ErrorCode.UNKNOWN,
-          ErrorCode.INVALID_ARGUMENT,
-          ErrorCode.PERMISSION_DENIED,
-          ErrorCode.UNAUTHENTICATED,
-          ErrorCode.RESOURCE_EXHAUSTED,
-          ErrorCode.FAILED_PRECONDITION,
-          ErrorCode.OUT_OF_RANGE,
-          ErrorCode.UNIMPLEMENTED,
-          ErrorCode.INTERNAL);
 
   /**
    * If the {@link SessionPoolOptions#getWaitForMinSessions()} duration is greater than zero, waits
@@ -1675,7 +1663,8 @@ public PooledSession get() {
       while (true) {
         Span span = tracer.spanBuilder(WAIT_FOR_SESSION).startSpan();
         try (Scope waitScope = tracer.withSpan(span)) {
-          PooledSession s = pollUninterruptiblyWithTimeout(currentTimeout);
+          PooledSession s =
+              pollUninterruptiblyWithTimeout(currentTimeout, options.getAcquireSessionTimeout());
           if (s == null) {
             // Set the status to DEADLINE_EXCEEDED and retry.
             numWaiterTimeouts.incrementAndGet();
@@ -1685,6 +1674,11 @@ public PooledSession get() {
             return s;
           }
         } catch (Exception e) {
+          if (e instanceof SpannerException
+              && ErrorCode.RESOURCE_EXHAUSTED.equals(((SpannerException) e).getErrorCode())) {
+            numWaiterTimeouts.incrementAndGet();
+            tracer.getCurrentSpan().setStatus(Status.RESOURCE_EXHAUSTED);
+          }
           TraceUtil.setWithFailure(span, e);
           throw e;
         } finally {
@@ -1693,15 +1687,26 @@ public PooledSession get() {
       }
     }
 
-    private PooledSession pollUninterruptiblyWithTimeout(long timeoutMillis) {
+    private PooledSession pollUninterruptiblyWithTimeout(
+        long timeoutMillis, Duration acquireSessionTimeout) {
       boolean interrupted = false;
       try {
         while (true) {
           try {
-            return waiter.get(timeoutMillis, TimeUnit.MILLISECONDS);
+            return acquireSessionTimeout == null
+                ? waiter.get(timeoutMillis, TimeUnit.MILLISECONDS)
+                : waiter.get(acquireSessionTimeout.toMillis(), TimeUnit.MILLISECONDS);
           } catch (InterruptedException e) {
             interrupted = true;
           } catch (TimeoutException e) {
+            if (acquireSessionTimeout != null) {
+              throw SpannerExceptionFactory.newSpannerException(
+                  ErrorCode.RESOURCE_EXHAUSTED,
+                  "Timed out after waiting "
+                      + acquireSessionTimeout.toMillis()
+                      + "ms for acquiring session. To mitigate error SessionPoolOptions#setAcquireSessionTimeout(Duration) to set a higher timeout"
+                      + " or increase the number of sessions in the session pool.");
+            }
             return null;
           } catch (ExecutionException e) {
             throw SpannerExceptionFactory.newSpannerException(e.getCause());

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

@@ -52,9 +52,16 @@ public class SessionPoolOptions {
   private final ActionOnSessionLeak actionOnSessionLeak;
   private final boolean trackStackTraceOfSessionCheckout;
   private final InactiveTransactionRemovalOptions inactiveTransactionRemovalOptions;
-  private final long initialWaitForSessionTimeoutMillis;
+
+  /**
+   * Use {@link #acquireSessionTimeout} instead to specify the total duration to wait while
+   * acquiring session for a transaction.
+   */
+  @Deprecated private final long initialWaitForSessionTimeoutMillis;
+
   private final boolean autoDetectDialect;
   private final Duration waitForMinSessions;
+  private final Duration acquireSessionTimeout;
 
   /** Property for allowing mocking of session maintenance clock. */
   private final Clock poolMaintainerClock;
@@ -78,6 +85,7 @@ private SessionPoolOptions(Builder builder) {
     this.removeInactiveSessionAfter = builder.removeInactiveSessionAfter;
     this.autoDetectDialect = builder.autoDetectDialect;
     this.waitForMinSessions = builder.waitForMinSessions;
+    this.acquireSessionTimeout = builder.acquireSessionTimeout;
     this.inactiveTransactionRemovalOptions = builder.inactiveTransactionRemovalOptions;
     this.poolMaintainerClock = builder.poolMaintainerClock;
   }
@@ -105,6 +113,7 @@ public boolean equals(Object o) {
         && Objects.equals(this.removeInactiveSessionAfter, other.removeInactiveSessionAfter)
         && Objects.equals(this.autoDetectDialect, other.autoDetectDialect)
         && Objects.equals(this.waitForMinSessions, other.waitForMinSessions)
+        && Objects.equals(this.acquireSessionTimeout, other.acquireSessionTimeout)
         && Objects.equals(
             this.inactiveTransactionRemovalOptions, other.inactiveTransactionRemovalOptions)
         && Objects.equals(this.poolMaintainerClock, other.poolMaintainerClock);
@@ -128,6 +137,7 @@ public int hashCode() {
         this.removeInactiveSessionAfter,
         this.autoDetectDialect,
         this.waitForMinSessions,
+        this.acquireSessionTimeout,
         this.inactiveTransactionRemovalOptions,
         this.poolMaintainerClock);
   }
@@ -239,6 +249,11 @@ Duration getWaitForMinSessions() {
     return waitForMinSessions;
   }
 
+  @VisibleForTesting
+  Duration getAcquireSessionTimeout() {
+    return acquireSessionTimeout;
+  }
+
   public static Builder newBuilder() {
     return new Builder();
   }
@@ -424,6 +439,7 @@ public static class Builder {
     private Duration removeInactiveSessionAfter = Duration.ofMinutes(55L);
     private boolean autoDetectDialect = false;
     private Duration waitForMinSessions = Duration.ZERO;
+    private Duration acquireSessionTimeout = Duration.ofSeconds(60);
 
     private Clock poolMaintainerClock;
 
@@ -446,6 +462,7 @@ private Builder(SessionPoolOptions options) {
       this.removeInactiveSessionAfter = options.removeInactiveSessionAfter;
       this.autoDetectDialect = options.autoDetectDialect;
       this.waitForMinSessions = options.waitForMinSessions;
+      this.acquireSessionTimeout = options.acquireSessionTimeout;
       this.inactiveTransactionRemovalOptions = options.inactiveTransactionRemovalOptions;
       this.poolMaintainerClock = options.poolMaintainerClock;
     }
@@ -538,6 +555,10 @@ public Builder setFailIfPoolExhausted() {
     /**
      * If all sessions are in use and there is no more room for creating new sessions, block for a
      * session to become available. Default behavior is same.
+     *
+     * <p>By default the requests are blocked for 60s and will fail with a `SpannerException` with
+     * error code `ResourceExhausted` if this timeout is exceeded. If you wish to block for a
+     * different period use the option {@link Builder#setAcquireSessionTimeout(Duration)} ()}
      */
     public Builder setBlockIfPoolExhausted() {
       this.actionOnExhaustion = ActionOnExhaustion.BLOCK;
@@ -695,6 +716,25 @@ public Builder setWaitForMinSessions(Duration waitForMinSessions) {
       return this;
     }
 
+    /**
+     * If greater than zero, we wait for said duration when no sessions are available in the {@link
+     * SessionPool}. The default is a 60s timeout. Set the value to null to disable the timeout.
+     */
+    public Builder setAcquireSessionTimeout(Duration acquireSessionTimeout) {
+      try {
+        if (acquireSessionTimeout != null) {
+          Preconditions.checkArgument(
+              acquireSessionTimeout.toMillis() > 0,
+              "acquireSessionTimeout should be greater than 0 ns");
+        }
+      } catch (ArithmeticException ex) {
+        throw new IllegalArgumentException(
+            "acquireSessionTimeout in millis should be lesser than Long.MAX_VALUE");
+      }
+      this.acquireSessionTimeout = acquireSessionTimeout;
+      return this;
+    }
+
     /** Build a SessionPoolOption object */
     public SessionPoolOptions build() {
       validate();