refactor(spanner): decouple read-only transactions from async result sets

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

@@ -42,6 +42,7 @@
 import java.util.concurrent.LinkedBlockingDeque;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import javax.annotation.concurrent.GuardedBy;
 
 /** Default implementation for {@link AsyncResultSet}. */
 class AsyncResultSetImpl extends ForwardingStructReader
@@ -95,6 +96,21 @@ private enum State {
   /** Supplies the underlying synchronous {@link ResultSet} that will be producing the rows. */
   private final Supplier<ResultSet> delegateResultSet;
 
+  /**
+   * Callback invoked once when query execution begins or when the result set is closed without
+   * starting, allowing callers (like {@link MultiUseReadOnlyTransaction}) to track query lifecycle.
+   */
+  private final Runnable onStartOrClose;
+
+  @GuardedBy("monitor")
+  private boolean startOrCloseNotified;
+
+  @GuardedBy("monitor")
+  private boolean delegateInitialized;
+
+  @GuardedBy("monitor")
+  private boolean delegateClosed;
+
   /**
    * Any exception that occurs while executing the query and iterating over the result set will be
    * stored in this variable and propagated to the user through {@link #tryNext()}.
@@ -152,30 +168,90 @@ private enum State {
   private volatile CountDownLatch consumingLatch = new CountDownLatch(0);
 
   AsyncResultSetImpl(ExecutorProvider executorProvider, ResultSet delegate, int bufferSize) {
-    this(executorProvider, Suppliers.ofInstance(Preconditions.checkNotNull(delegate)), bufferSize);
+    this(
+        executorProvider,
+        Suppliers.ofInstance(Preconditions.checkNotNull(delegate)),
+        bufferSize,
+        /* onStartOrClose= */ () -> {},
+        /* delegateInitialized= */ true);
   }
 
   AsyncResultSetImpl(
       ExecutorProvider executorProvider, Supplier<ResultSet> delegate, int bufferSize) {
+    this(executorProvider, delegate, bufferSize, /* onStartOrClose= */ () -> {});
+  }
+
+  AsyncResultSetImpl(
+      ExecutorProvider executorProvider,
+      Supplier<ResultSet> delegate,
+      int bufferSize,
+      Runnable onStartOrClose) {
     this(
         executorProvider,
         Suppliers.memoize(Preconditions.checkNotNull(delegate)),
         bufferSize,
-        true);
+        Preconditions.checkNotNull(onStartOrClose),
+        /* delegateInitialized= */ false);
   }
 
   private AsyncResultSetImpl(
       ExecutorProvider executorProvider,
       Supplier<ResultSet> memoizedDelegate,
       int bufferSize,
-      boolean dummy) {
+      Runnable onStartOrClose,
+      boolean delegateInitialized) {
     super(memoizedDelegate);
     this.executorProvider = Preconditions.checkNotNull(executorProvider);
     this.delegateResultSet = memoizedDelegate;
+    this.onStartOrClose = onStartOrClose;
+    this.delegateInitialized = delegateInitialized;
     this.service = MoreExecutors.listeningDecorator(executorProvider.getExecutor());
     this.buffer = new LinkedBlockingDeque<>(bufferSize);
   }
 
+  private ResultSet getDelegate() {
+    try {
+      ResultSet delegate = delegateResultSet.get();
+      boolean shouldCloseImmediately = false;
+      synchronized (monitor) {
+        if (!delegateInitialized) {
+          delegateInitialized = true;
+          if (this.closed || delegateClosed) {
+            shouldCloseImmediately = true;
+            delegateClosed = true;
+          }
+        }
+      }
+      if (shouldCloseImmediately) {
+        try {
+          delegate.close();
+        } catch (Throwable throwable) {
+          log.log(Level.FINE, "Ignoring error from closing delegate result set", throwable);
+        }
+      }
+      return delegate;
+    } finally {
+      notifyStartOrClose();
+    }
+  }
+
+  private void notifyStartOrClose() {
+    boolean shouldNotify = false;
+    synchronized (monitor) {
+      if (!startOrCloseNotified) {
+        startOrCloseNotified = true;
+        shouldNotify = true;
+      }
+    }
+    if (shouldNotify) {
+      try {
+        onStartOrClose.run();
+      } catch (Throwable throwable) {
+        log.log(Level.WARNING, "Failed to execute onStartOrClose callback", throwable);
+      }
+    }
+  }
+
   boolean isUsed() {
     synchronized (monitor) {
       return state != State.INITIALIZED;
@@ -191,15 +267,33 @@ boolean isUsed() {
    */
   @Override
   public void close() {
+    boolean shouldCloseDelegate = false;
+    boolean shouldShutdownService = false;
+    boolean shouldNotifyClose = false;
     synchronized (monitor) {
       if (this.closed) {
         return;
       }
       if (state == State.INITIALIZED || state == State.SYNC) {
-        delegateResultSet.get().close();
+        shouldCloseDelegate = true;
+        if (executorProvider.shouldAutoClose()) {
+          shouldShutdownService = true;
+        }
+      }
+      if (state == State.INITIALIZED) {
+        shouldNotifyClose = true;
       }
       this.closed = true;
     }
+    if (shouldCloseDelegate) {
+      closeDelegateResultSet();
+    }
+    if (shouldNotifyClose) {
+      notifyStartOrClose();
+    }
+    if (shouldShutdownService) {
+      service.shutdown();
+    }
   }
 
   /**
@@ -258,10 +352,19 @@ public CursorState tryNext() throws SpannerException {
   }
 
   private void closeDelegateResultSet() {
-    try {
-      delegateResultSet.get().close();
-    } catch (Throwable t) {
-      log.log(Level.FINE, "Ignoring error from closing delegate result set", t);
+    boolean shouldClose = false;
+    synchronized (monitor) {
+      if (!delegateClosed) {
+        delegateClosed = true;
+        shouldClose = delegateInitialized;
+      }
+    }
+    if (shouldClose) {
+      try {
+        delegateResultSet.get().close();
+      } catch (Throwable t) {
+        log.log(Level.FINE, "Ignoring error from closing delegate result set", t);
+      }
     }
   }
 
@@ -362,15 +465,17 @@ private class ProduceRowsRunnable implements Runnable {
     public void run() {
       boolean stop = false;
       boolean hasNext = false;
+      ResultSet delegate = null;
       try {
-        hasNext = delegateResultSet.get().next();
+        delegate = getDelegate();
+        hasNext = delegate.next();
       } catch (Throwable e) {
         synchronized (monitor) {
           executionException = SpannerExceptionFactory.asSpannerException(e);
         }
       }
       try {
-        while (!stop && hasNext) {
+        while (!stop && hasNext && delegate != null) {
           try {
             synchronized (monitor) {
               stop = shouldStopProducing();
@@ -395,9 +500,9 @@ public void run() {
               }
             }
             if (!stop) {
-              buffer.put(delegateResultSet.get().getCurrentRowAsStruct());
+              buffer.put(delegate.getCurrentRowAsStruct());
               startCallbackIfNecessary();
-              hasNext = delegateResultSet.get().next();
+              hasNext = delegate.next();
             }
           } catch (Throwable e) {
             synchronized (monitor) {
@@ -511,19 +616,41 @@ public void run() {
   /** Sets the callback for this {@link AsyncResultSet}. */
   @Override
   public ApiFuture<Void> setCallback(Executor exec, ReadyCallback cb) {
-    synchronized (monitor) {
-      Preconditions.checkState(!closed, "This AsyncResultSet has been closed");
-      Preconditions.checkState(
-          this.state == State.INITIALIZED, "callback may not be set multiple times");
-
-      // Start to fetch data and buffer these.
-      this.result = SettableApiFuture.create();
-      this.state = State.STREAMING_INITIALIZED;
-      this.service.execute(new InitiateStreamingRunnable());
-      this.executor = MoreExecutors.newSequentialExecutor(Preconditions.checkNotNull(exec));
-      this.callback = Preconditions.checkNotNull(cb);
-      pausedLatch.countDown();
-      return result;
+    Preconditions.checkNotNull(exec, "exec cannot be null");
+    Preconditions.checkNotNull(cb, "cb cannot be null");
+    boolean shouldNotify = false;
+    boolean shouldShutdownService = false;
+    try {
+      synchronized (monitor) {
+        Preconditions.checkState(!closed, "This AsyncResultSet has been closed");
+        Preconditions.checkState(
+            this.state == State.INITIALIZED, "callback may not be set multiple times");
+
+        // Start to fetch data and buffer these.
+        this.result = SettableApiFuture.create();
+        this.executor = MoreExecutors.newSequentialExecutor(exec);
+        this.callback = cb;
+        this.state = State.STREAMING_INITIALIZED;
+        try {
+          this.service.execute(new InitiateStreamingRunnable());
+        } catch (Throwable throwable) {
+          this.closed = true;
+          shouldNotify = true;
+          if (executorProvider.shouldAutoClose()) {
+            shouldShutdownService = true;
+          }
+          throw throwable;
+        }
+        pausedLatch.countDown();
+        return result;
+      }
+    } finally {
+      if (shouldNotify) {
+        notifyStartOrClose();
+      }
+      if (shouldShutdownService) {
+        service.shutdown();
+      }
     }
   }
 
@@ -627,28 +754,41 @@ public <T> List<T> toList(Function<StructReader, T> transformer) throws SpannerE
   @Override
   public boolean next() throws SpannerException {
     synchronized (monitor) {
+      Preconditions.checkState(!closed, "This AsyncResultSet has been closed");
       Preconditions.checkState(
           this.state == State.INITIALIZED || this.state == State.SYNC,
           "Cannot call next() on a result set with a callback.");
       this.state = State.SYNC;
     }
-    boolean res = delegateResultSet.get().next();
-    currentRow = res ? delegateResultSet.get().getCurrentRowAsStruct() : null;
+    ResultSet delegate = getDelegate();
+    boolean res = delegate.next();
+    currentRow = res ? delegate.getCurrentRowAsStruct() : null;
     return res;
   }
 
+  /**
+   * Returns query execution statistics. Calling this method after the result set has closed is
+   * allowed because callers frequently inspect statistics or query plans after consuming all rows
+   * and finishing the query.
+   */
   @Override
   public ResultSetStats getStats() {
-    return delegateResultSet.get().getStats();
+    return getDelegate().getStats();
   }
 
+  /**
+   * Returns query metadata. Calling this method after the result set has closed is allowed so
+   * callers can inspect column and type information after consuming all rows. If the result set was
+   * closed before execution started, {@link #getDelegate()} ensures the underlying delegate is
+   * closed immediately to prevent resource leaks.
+   */
   @Override
   public ResultSetMetadata getMetadata() {
-    return delegateResultSet.get().getMetadata();
+    return getDelegate().getMetadata();
   }
 
   boolean initiateStreaming(StreamMessageListener streamMessageListener) {
-    return StreamingUtil.initiateStreaming(delegateResultSet.get(), streamMessageListener);
+    return StreamingUtil.initiateStreaming(getDelegate(), streamMessageListener);
   }
 
   @Override

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

@@ -38,7 +38,6 @@
 import java.util.concurrent.ThreadLocalRandom;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.locks.ReentrantLock;
 import javax.annotation.Nullable;
@@ -313,38 +312,11 @@ private ListenableAsyncResultSet createAsyncResultSet(
     } 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.
-    final AtomicBoolean decremented = new AtomicBoolean(false);
     try {
       return new AsyncResultSetImpl(
-          executorProvider,
-          () -> {
-            try {
-              return resultSetSupplier.get();
-            } finally {
-              if (decremented.compareAndSet(false, true)) {
-                decrementPendingStartsAndSignal();
-              }
-            }
-          },
-          bufferRows) {
-        @Override
-        public void close() {
-          try {
-            super.close();
-          } finally {
-            if (!isUsed() && decremented.compareAndSet(false, true)) {
-              decrementPendingStartsAndSignal();
-            }
-          }
-        }
-      };
+          executorProvider, resultSetSupplier, bufferRows, this::decrementPendingStartsAndSignal);
     } catch (Throwable throwable) {
-      if (decremented.compareAndSet(false, true)) {
-        decrementPendingStartsAndSignal();
-      }
+      decrementPendingStartsAndSignal();
       throw throwable;
     }
   }