fix(spanner): add closeAsync to ReadContext and make transaction closing non-blocking

Adds `ReadContext.closeAsync()` to allow closing read contexts asynchronously without blocking caller threads. Refactors `MultiUseReadOnlyTransaction` so that background query initializations and `BeginTransaction` RPCs do not block thread execution during close, while preserving legacy synchronous `close()` semantics and guarding against new operations during closing.

Author: olavloite

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

@@ -358,9 +358,40 @@ static Builder newBuilder() {
     @GuardedBy("txnLock")
     private ByteString transactionId;
 
+    /**
+     * Future used to synchronize concurrent operations during transaction initialization:
+     *
+     * <ul>
+     *   <li>When using explicit {@code BeginTransaction} (e.g. via {@link #initTransaction()}), the
+     *       first caller creates this future, releases {@code txnLock}, and executes the RPC
+     *       outside the lock. Subsequent concurrent callers wait on this future outside the lock.
+     *   <li>When using inlined {@code BeginTransaction} options, the first query or read operation
+     *       creates this future and includes {@code Begin} in its {@link TransactionSelector}.
+     *       Concurrent queries wait on this future until the first operation's initial response
+     *       arrives and sets the transaction ID via {@link #onTransactionMetadata}.
+     * </ul>
+     */
     @GuardedBy("txnLock")
     private SettableApiFuture<ByteString> transactionIdFuture;
 
+    /**
+     * Future completed when all in-flight asynchronous query initializations (tracked by {@link
+     * #pendingStarts}) have finished and the read context has been closed.
+     */
+    @GuardedBy("txnLock")
+    private SettableApiFuture<Void> closeFuture;
+
+    /**
+     * Flag indicating whether the transaction has been closed or is currently closing. Once true,
+     * no new reads or queries (synchronous or asynchronous) are permitted to start on this context.
+     */
+    @GuardedBy("txnLock")
+    private boolean isClosedOrClosing;
+
+    /**
+     * Counter tracking the number of asynchronous queries created on this transaction whose
+     * delegate {@link ResultSet} suppliers have not yet been executed or closed.
+     */
     private final AtomicInteger pendingStarts = new AtomicInteger(0);
 
     private static final long WAIT_FOR_INLINE_BEGIN_TIMEOUT_MILLIS = 60_000L;
@@ -468,20 +499,53 @@ TransactionSelector getTransactionSelector() {
       }
     }
 
+    /**
+     * Decrements the count of pending asynchronous query starts. When the count drops to 0, if the
+     * transaction is in the process of closing (i.e. {@link #closeAsync()} was called), this method
+     * closes the read context via {@code super.close()} and completes {@link #closeFuture}.
+     */
     private void decrementPendingStartsAndSignal() {
-      if (pendingStarts.decrementAndGet() == 0) {
-        txnLock.lock();
-        try {
+      SettableApiFuture<Void> futureToComplete = null;
+      Throwable error = null;
+      txnLock.lock();
+      try {
+        if (pendingStarts.decrementAndGet() == 0) {
           hasNoPendingStarts.signalAll();
-        } finally {
-          txnLock.unlock();
+          if (closeFuture != null) {
+            try {
+              super.close();
+            } catch (Throwable throwable) {
+              error = throwable;
+            }
+            futureToComplete = closeFuture;
+          }
+        }
+      } finally {
+        txnLock.unlock();
+      }
+      if (futureToComplete != null) {
+        if (error != null) {
+          futureToComplete.setException(error);
+        } else {
+          futureToComplete.set(null);
         }
       }
     }
 
+    /**
+     * Creates an {@link AsyncResultSetImpl} for an asynchronous read or query. Tracks the query in
+     * {@link #pendingStarts} so that {@link #closeAsync()} will not close the read context until
+     * the query initialization has completed.
+     */
     private ListenableAsyncResultSet createAsyncResultSet(
         Supplier<ResultSet> resultSetSupplier, int bufferRows) {
-      pendingStarts.incrementAndGet();
+      txnLock.lock();
+      try {
+        checkState(!isClosedOrClosing, "Context has been closed");
+        pendingStarts.incrementAndGet();
+      } finally {
+        txnLock.unlock();
+      }
       // Make sure that we decrement the counter exactly once, either
       // when the query is actually executed, or when the result set is closed,
       // or if something goes wrong when creating the result set.
@@ -518,6 +582,41 @@ public void close() {
       }
     }
 
+    /**
+     * Ensures that the transaction is neither closed nor in the process of closing before
+     * initiating a new synchronous read or query.
+     */
+    private void checkClosedOrClosing() {
+      txnLock.lock();
+      try {
+        checkState(!isClosedOrClosing, "Context has been closed");
+      } finally {
+        txnLock.unlock();
+      }
+    }
+
+    @Override
+    ResultSet readInternalWithOptions(
+        String table,
+        @Nullable String index,
+        KeySet keys,
+        Iterable<String> columns,
+        final Options readOptions,
+        ByteString partitionToken) {
+      checkClosedOrClosing();
+      return super.readInternalWithOptions(
+          table, index, keys, columns, readOptions, partitionToken);
+    }
+
+    @Override
+    ResultSet executeQueryInternal(
+        final Statement statement,
+        final com.google.spanner.v1.ExecuteSqlRequest.QueryMode queryMode,
+        final QueryOption... options) {
+      checkClosedOrClosing();
+      return super.executeQueryInternal(statement, queryMode, options);
+    }
+
     @Override
     public ListenableAsyncResultSet readAsync(
         String table, KeySet keys, Iterable<String> columns, ReadOption... options) {
@@ -527,7 +626,7 @@ public ListenableAsyncResultSet readAsync(
               ? readOptions.bufferRows()
               : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE;
       return createAsyncResultSet(
-          () -> readInternal(table, null, keys, columns, options), bufferRows);
+          () -> super.readInternal(table, null, keys, columns, options), bufferRows);
     }
 
     @Override
@@ -539,7 +638,7 @@ public ListenableAsyncResultSet readUsingIndexAsync(
               ? readOptions.bufferRows()
               : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE;
       return createAsyncResultSet(
-          () -> readInternal(table, checkNotNull(index), keys, columns, options), bufferRows);
+          () -> super.readInternal(table, checkNotNull(index), keys, columns, options), bufferRows);
     }
 
     @Override
@@ -551,7 +650,7 @@ public ListenableAsyncResultSet executeQueryAsync(Statement statement, QueryOpti
               : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE;
       return createAsyncResultSet(
           () ->
-              executeQueryInternal(
+              super.executeQueryInternal(
                   statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.NORMAL, options),
           bufferRows);
     }
@@ -650,21 +749,37 @@ ByteString getTransactionId() {
       }
     }
 
+    /**
+     * Closes the transaction asynchronously.
+     *
+     * <p>If there are no in-flight asynchronous query initializations ({@code pendingStarts == 0}),
+     * the read context is closed immediately and a completed future is returned. If there are
+     * pending starts, this method marks {@code isClosedOrClosing = true} and returns a {@link
+     * SettableApiFuture} that will be completed by {@link #decrementPendingStartsAndSignal()} when
+     * all pending query initializations finish.
+     */
     @Override
-    public void close() {
+    public ApiFuture<Void> closeAsync() {
       txnLock.lock();
       try {
-        while (pendingStarts.get() > 0) {
-          try {
-            hasNoPendingStarts.await();
-          } catch (InterruptedException e) {
-            throw SpannerExceptionFactory.propagateInterrupt(e);
-          }
+        if (isClosedOrClosing) {
+          return closeFuture != null ? closeFuture : ApiFutures.immediateFuture(null);
+        }
+        isClosedOrClosing = true;
+        if (pendingStarts.get() == 0) {
+          super.close();
+          return ApiFutures.immediateFuture(null);
         }
+        closeFuture = SettableApiFuture.create();
+        return closeFuture;
       } finally {
         txnLock.unlock();
       }
-      super.close();
+    }
+
+    @Override
+    public void close() {
+      SpannerApiFutures.get(closeAsync());
     }
 
     private TransactionOptions createReadOnlyTransactionOptions() {
@@ -686,47 +801,65 @@ private TransactionOptions createReadOnlyTransactionOptions() {
      * Multiplexed Session.
      */
     void initFallbackTransaction() {
-      txnLock.lock();
-      try {
-        span.addAnnotation("Creating Transaction");
-        final BeginTransactionRequest request =
-            BeginTransactionRequest.newBuilder()
-                .setSession(session.getName())
-                .setOptions(createReadOnlyTransactionOptions())
-                .build();
-        initTransactionInternal(request);
-      } finally {
-        txnLock.unlock();
-      }
+      initTransaction();
     }
 
+    /**
+     * Initializes the transaction by issuing a BeginTransaction RPC.
+     *
+     * <p>To prevent blocking concurrent operations (such as {@link #closeAsync()}) while a network
+     * RPC is in-flight, {@code rpc.beginTransaction} is executed <b>outside</b> {@code txnLock}. A
+     * leader/follower pattern using {@link #transactionIdFuture} is used: the first caller acquires
+     * the lock, creates {@code transactionIdFuture}, releases the lock, and executes the RPC.
+     * Subsequent concurrent callers retrieve {@code transactionIdFuture} and wait on it outside the
+     * lock.
+     */
     void initTransaction() {
       SessionImpl.throwIfTransactionsPending();
 
-      // Since we only support synchronous calls, just block on "txnLock" while the RPC is in
-      // flight. Note that we use the strategy of sending an explicit BeginTransaction() RPC,
-      // rather than using the first read in the transaction to begin it implicitly. The chosen
-      // strategy is sub-optimal in the case of the first read being fast, as it incurs an extra
-      // RTT, but optimal if the first read is slow. As the client library is now using streaming
-      // reads, a possible optimization could be to use the first read in the transaction to begin
-      // it implicitly.
+      ApiFuture<ByteString> futureToWaitFor = null;
+      BeginTransactionRequest request = null;
       txnLock.lock();
       try {
         if (transactionId != null) {
           return;
         }
-        span.addAnnotation("Creating Transaction");
-        final BeginTransactionRequest request =
-            BeginTransactionRequest.newBuilder()
-                .setSession(session.getName())
-                .setOptions(createReadOnlyTransactionOptions())
-                .build();
-        initTransactionInternal(request);
+        if (transactionIdFuture != null) {
+          futureToWaitFor = transactionIdFuture;
+        } else {
+          transactionIdFuture = SettableApiFuture.create();
+          span.addAnnotation("Creating Transaction");
+          request =
+              BeginTransactionRequest.newBuilder()
+                  .setSession(session.getName())
+                  .setOptions(createReadOnlyTransactionOptions())
+                  .build();
+        }
       } finally {
         txnLock.unlock();
       }
+
+      if (futureToWaitFor != null) {
+        try {
+          futureToWaitFor.get();
+          return;
+        } catch (ExecutionException executionException) {
+          throw SpannerExceptionFactory.asSpannerException(executionException.getCause());
+        } catch (InterruptedException interruptedException) {
+          Thread.currentThread().interrupt();
+          throw SpannerExceptionFactory.newSpannerExceptionForCancellation(
+              null, interruptedException);
+        }
+      }
+
+      initTransactionInternal(request);
     }
 
+    /**
+     * Executes the BeginTransaction RPC outside {@code txnLock}, updates transaction state under
+     * {@code txnLock}, and completes or fails {@link #transactionIdFuture} so waiting callers are
+     * notified.
+     */
     private void initTransactionInternal(BeginTransactionRequest request) {
       try {
         Transaction transaction =
@@ -739,20 +872,41 @@ private void initTransactionInternal(BeginTransactionRequest request) {
           throw SpannerExceptionFactory.newSpannerException(
               ErrorCode.INTERNAL, "Missing expected transaction.id metadata field");
         }
+        Timestamp readTimestamp;
         try {
-          timestamp = Timestamp.fromProto(transaction.getReadTimestamp());
-        } catch (IllegalArgumentException e) {
+          readTimestamp = Timestamp.fromProto(transaction.getReadTimestamp());
+        } catch (IllegalArgumentException illegalArgumentException) {
           throw SpannerExceptionFactory.newSpannerException(
-              ErrorCode.INTERNAL, "Bad value in transaction.read_timestamp metadata field", e);
+              ErrorCode.INTERNAL,
+              "Bad value in transaction.read_timestamp metadata field",
+              illegalArgumentException);
+        }
+        txnLock.lock();
+        try {
+          timestamp = readTimestamp;
+          transactionId = transaction.getId();
+          if (transactionIdFuture != null && !transactionIdFuture.isDone()) {
+            transactionIdFuture.set(transactionId);
+          }
+        } finally {
+          txnLock.unlock();
         }
-        transactionId = transaction.getId();
         span.addAnnotation(
             "Transaction Creation Done",
             ImmutableMap.of(
-                "Id", transaction.getId().toStringUtf8(), "Timestamp", timestamp.toString()));
-      } catch (SpannerException e) {
-        span.addAnnotation("Transaction Creation Failed", e);
-        throw e;
+                "Id", transaction.getId().toStringUtf8(), "Timestamp", readTimestamp.toString()));
+      } catch (Throwable throwable) {
+        SpannerException spannerException = SpannerExceptionFactory.asSpannerException(throwable);
+        span.addAnnotation("Transaction Creation Failed", spannerException);
+        txnLock.lock();
+        try {
+          if (transactionIdFuture != null && !transactionIdFuture.isDone()) {
+            transactionIdFuture.setException(spannerException);
+          }
+        } finally {
+          txnLock.unlock();
+        }
+        throw spannerException;
       }
     }
   }
@@ -1206,6 +1360,12 @@ public void close() {
     }
   }
 
+  @Override
+  public ApiFuture<Void> closeAsync() {
+    close();
+    return ApiFutures.immediateFuture(null);
+  }
+
   /**
    * Returns the {@link TransactionSelector} that should be used for a statement that is executed on
    * this read context. This could be a reference to an existing transaction ID, or it could be a

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

@@ -139,6 +139,16 @@ public void close() {
     }
   }
 
+  @Override
+  public ApiFuture<Void> closeAsync() {
+    return ApiFutures.catchingAsync(
+        ApiFutures.transformAsync(
+            this.readContextFuture, ReadContext::closeAsync, MoreExecutors.directExecutor()),
+        Throwable.class,
+        throwable -> ApiFutures.immediateFuture(null),
+        MoreExecutors.directExecutor());
+  }
+
   /**
    * Represents a {@link ReadContext} using a multiplexed session that is not yet ready. The
    * execution will be delayed until the multiplexed session has been created and is ready.