diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index a4d7614ebd2b..c473d0719f49 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -65,7 +65,6 @@ import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.locks.Condition; import java.util.concurrent.locks.ReentrantLock; import java.util.logging.Logger; import javax.annotation.Nullable; @@ -350,7 +349,6 @@ static Builder newBuilder() { private TimestampBound bound; private final ReentrantLock txnLock = new ReentrantLock(); - private final Condition hasNoPendingStarts = txnLock.newCondition(); @GuardedBy("txnLock") private Timestamp timestamp; @@ -358,9 +356,43 @@ static Builder newBuilder() { @GuardedBy("txnLock") private ByteString transactionId; + /** + * Future used to synchronize concurrent operations during transaction initialization: + * + * + */ @GuardedBy("txnLock") private SettableApiFuture 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 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; + + @GuardedBy("txnLock") + private boolean closed; + + /** + * 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 +500,78 @@ TransactionSelector getTransactionSelector() { } } - private void decrementPendingStartsAndSignal() { - if (pendingStarts.decrementAndGet() == 0) { - txnLock.lock(); + @GuardedBy("txnLock") + private boolean isTransactionInitializingLocked() { + return transactionIdFuture != null && !transactionIdFuture.isDone(); + } + + @GuardedBy("txnLock") + private boolean shouldCloseLocked() { + return isClosedOrClosing + && !closed + && pendingStarts.get() <= 0 + && !isTransactionInitializingLocked(); + } + + /** + * Checks if the transaction can be safely closed, and if so, closes it via {@code + * super.close()} and completes {@link #closeFuture}. + */ + private void checkAndClose() { + SettableApiFuture futureToComplete = null; + boolean shouldClose = false; + txnLock.lock(); + try { + if (shouldCloseLocked()) { + closed = true; + shouldClose = true; + futureToComplete = closeFuture; + } + } finally { + txnLock.unlock(); + } + + if (shouldClose) { + Throwable error = null; try { - hasNoPendingStarts.signalAll(); - } finally { - txnLock.unlock(); + super.close(); + } catch (Throwable throwable) { + error = throwable; + } + if (futureToComplete != null) { + if (error != null) { + futureToComplete.setException(error); + } else { + futureToComplete.set(null); + } } } } + /** + * 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() { + pendingStarts.decrementAndGet(); + checkAndClose(); + } + + /** + * 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 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. @@ -510,14 +600,49 @@ public void close() { } } }; - } catch (Throwable t) { + } catch (Throwable throwable) { if (decremented.compareAndSet(false, true)) { decrementPendingStartsAndSignal(); } - throw t; + throw throwable; + } + } + + /** + * 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 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 columns, ReadOption... options) { @@ -527,7 +652,8 @@ public ListenableAsyncResultSet readAsync( ? readOptions.bufferRows() : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; return createAsyncResultSet( - () -> readInternal(table, null, keys, columns, options), bufferRows); + () -> super.readInternalWithOptions(table, null, keys, columns, readOptions, null), + bufferRows); } @Override @@ -539,7 +665,10 @@ public ListenableAsyncResultSet readUsingIndexAsync( ? readOptions.bufferRows() : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; return createAsyncResultSet( - () -> readInternal(table, checkNotNull(index), keys, columns, options), bufferRows); + () -> + super.readInternalWithOptions( + table, checkNotNull(index), keys, columns, readOptions, null), + bufferRows); } @Override @@ -551,7 +680,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); } @@ -589,6 +718,7 @@ public void onTransactionMetadata(Transaction transaction, boolean shouldInclude } finally { txnLock.unlock(); } + checkAndClose(); } @Override @@ -613,21 +743,22 @@ public void onDone(boolean withBeginTransaction) { } @Override - void onStartFailed(boolean withBeginTransaction, Throwable t) { + void onStartFailed(boolean withBeginTransaction, Throwable throwable) { if (withBeginTransaction) { - failTransactionIdFuture(t); + failTransactionIdFuture(throwable); } } - private void failTransactionIdFuture(Throwable t) { + private void failTransactionIdFuture(Throwable throwable) { txnLock.lock(); try { if (transactionIdFuture != null && !transactionIdFuture.isDone()) { - transactionIdFuture.setException(t); + transactionIdFuture.setException(throwable); } } finally { txnLock.unlock(); } + checkAndClose(); } @Override @@ -650,21 +781,49 @@ ByteString getTransactionId() { } } + /** + * Closes the transaction asynchronously. + * + *

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 closeAsync() { + boolean shouldCloseImmediately = false; txnLock.lock(); try { - while (pendingStarts.get() > 0) { - try { - hasNoPendingStarts.await(); - } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e); - } + if (isClosedOrClosing) { + return closeFuture; + } + isClosedOrClosing = true; + closeFuture = SettableApiFuture.create(); + if (shouldCloseLocked()) { + closed = true; + shouldCloseImmediately = true; + } else { + return closeFuture; } } finally { txnLock.unlock(); } - super.close(); + + if (shouldCloseImmediately) { + try { + super.close(); + closeFuture.set(null); + } catch (Throwable throwable) { + closeFuture.setException(throwable); + } + } + return closeFuture; + } + + @Override + public void close() { + SpannerApiFutures.get(closeAsync()); } private TransactionOptions createReadOnlyTransactionOptions() { @@ -686,47 +845,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. + * + *

To prevent blocking concurrent operations (such as {@link #closeAsync()}) while a network + * RPC is in-flight, {@code rpc.beginTransaction} is executed outside {@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 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 +916,43 @@ 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); } - transactionId = transaction.getId(); + txnLock.lock(); + try { + timestamp = readTimestamp; + transactionId = transaction.getId(); + if (transactionIdFuture != null && !transactionIdFuture.isDone()) { + transactionIdFuture.set(transactionId); + } + } finally { + txnLock.unlock(); + } + checkAndClose(); 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(); + } + checkAndClose(); + throw spannerException; } } } diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedReadContext.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedReadContext.java index 86d6ae079d31..87afe485acef 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedReadContext.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DelayedReadContext.java @@ -132,6 +132,9 @@ public ResultSet analyzeQuery(Statement statement, QueryAnalyzeMode queryMode) { @Override public void close() { try { + if (this.readContextFuture.cancel(true) || this.readContextFuture.isCancelled()) { + return; + } this.readContextFuture.get().close(); } catch (Throwable ignore) { // Ignore any errors during close, as this error has already propagated to the user through @@ -139,6 +142,19 @@ public void close() { } } + @Override + public ApiFuture closeAsync() { + if (this.readContextFuture.cancel(true) || this.readContextFuture.isCancelled()) { + return ApiFutures.immediateFuture(null); + } + 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. diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java index 4b5ba8620eed..7508bfabb8a1 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ReadContext.java @@ -17,6 +17,7 @@ package com.google.cloud.spanner; import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import javax.annotation.Nullable; @@ -221,7 +222,29 @@ ApiFuture readRowUsingIndexAsync( */ ResultSet analyzeQuery(Statement statement, QueryAnalyzeMode queryMode); - /** Closes this read context and frees up the underlying resources. */ + /** Closes this read context. */ @Override void close(); + + /** + * Closes this read context asynchronously. + * + *

For read contexts with pending asynchronous operations (such as multi-use read-only + * transactions with background queries initializing), the returned future will complete once all + * pending starts have completed and the read context has been closed. + * + *

Callers must either consume or close all {@link AsyncResultSet} instances created on this + * read context; otherwise, unconsumed and unclosed result sets may prevent {@code closeAsync()} + * from completing. + * + * @return an {@link ApiFuture} that is done when the read context has been closed. + */ + default ApiFuture closeAsync() { + try { + close(); + return ApiFutures.immediateFuture(null); + } catch (Throwable throwable) { + return ApiFutures.immediateFailedFuture(throwable); + } + } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java index f92d63dad175..034af24f8cdd 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncReadOnlyTransactionTest.java @@ -16,21 +16,28 @@ package com.google.cloud.spanner; +import static com.google.cloud.spanner.MockSpannerServiceImpl.NO_EXECUTION_TIME; import static com.google.cloud.spanner.MockSpannerTestUtil.READ_ONE_KEY_VALUE_STATEMENT; import static com.google.cloud.spanner.MockSpannerTestUtil.TEST_DATABASE; import static com.google.cloud.spanner.MockSpannerTestUtil.TEST_INSTANCE; import static com.google.cloud.spanner.MockSpannerTestUtil.TEST_PROJECT; -import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; +import com.google.api.core.ApiFuture; import com.google.cloud.NoCredentials; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.ExecuteSqlRequest; import io.grpc.ManagedChannelBuilder; +import io.grpc.Status; +import java.util.Arrays; +import java.util.Collections; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import org.junit.Test; @@ -52,27 +59,27 @@ public void asyncReadOnlyTransactionIsNonBlocking() throws Exception { mockSpanner.freeze(); // Call executeQueryAsync. It should not block even though mock server is // frozen! - AsyncResultSet rs = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); // Verify that no requests have been sent yet. assertTrue(mockSpanner.getRequestTypes().isEmpty()); // Now register a callback to start the stream. - final CountDownLatch latch = new CountDownLatch(1); - rs.setCallback( + final CountDownLatch callbackLatch = new CountDownLatch(1); + resultSet.setCallback( executor, - resultSet -> { + ignored -> { try { AsyncResultSet.CursorState state; while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) { // consume } if (state == AsyncResultSet.CursorState.DONE) { - latch.countDown(); + callbackLatch.countDown(); } return AsyncResultSet.CallbackResponse.CONTINUE; - } catch (Throwable t) { - latch.countDown(); + } catch (Throwable throwable) { + callbackLatch.countDown(); return AsyncResultSet.CallbackResponse.DONE; } }); @@ -81,12 +88,13 @@ public void asyncReadOnlyTransactionIsNonBlocking() throws Exception { mockSpanner.unfreeze(); // Wait for the callback to complete. - assertTrue("Timeout waiting for callback", latch.await(10, TimeUnit.SECONDS)); + assertTrue("Timeout waiting for callback", callbackLatch.await(10, TimeUnit.SECONDS)); // Verify that requests were sent on the background thread. // It should contain one BeginTransaction and one ExecuteSql. - assertThat(mockSpanner.getRequestTypes()) - .containsExactly(BeginTransactionRequest.class, ExecuteSqlRequest.class); + assertEquals( + Arrays.asList(BeginTransactionRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); } } @@ -101,8 +109,8 @@ public void testMultipleQueriesOnlyCallsBeginTransactionOnce() throws Exception try (ReadOnlyTransaction transaction = client().readOnlyTransaction()) { mockSpanner.freeze(); // Call executeQueryAsync twice. - AsyncResultSet rs1 = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); - AsyncResultSet rs2 = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + AsyncResultSet resultSet1 = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + AsyncResultSet resultSet2 = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); // Verify that no requests have been sent yet. assertTrue(mockSpanner.getRequestTypes().isEmpty()); @@ -111,50 +119,51 @@ public void testMultipleQueriesOnlyCallsBeginTransactionOnce() throws Exception mockSpanner.unfreeze(); // Now register callbacks to start the streams. - final CountDownLatch latch1 = new CountDownLatch(1); - final CountDownLatch latch2 = new CountDownLatch(1); + final CountDownLatch callbackLatch1 = new CountDownLatch(1); + final CountDownLatch callbackLatch2 = new CountDownLatch(1); - rs1.setCallback( + resultSet1.setCallback( executor, - resultSet -> { + ignored -> { try { AsyncResultSet.CursorState state; - while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) {} + while ((state = resultSet1.tryNext()) == AsyncResultSet.CursorState.OK) {} if (state == AsyncResultSet.CursorState.DONE) { - latch1.countDown(); + callbackLatch1.countDown(); } return AsyncResultSet.CallbackResponse.CONTINUE; - } catch (Throwable t) { - latch1.countDown(); + } catch (Throwable throwable) { + callbackLatch1.countDown(); return AsyncResultSet.CallbackResponse.DONE; } }); - rs2.setCallback( + resultSet2.setCallback( executor, - resultSet -> { + ignored -> { try { AsyncResultSet.CursorState state; - while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) {} + while ((state = resultSet2.tryNext()) == AsyncResultSet.CursorState.OK) {} if (state == AsyncResultSet.CursorState.DONE) { - latch2.countDown(); + callbackLatch2.countDown(); } return AsyncResultSet.CallbackResponse.CONTINUE; - } catch (Throwable t) { - latch2.countDown(); + } catch (Throwable throwable) { + callbackLatch2.countDown(); return AsyncResultSet.CallbackResponse.DONE; } }); // Wait for both callbacks to complete. - assertTrue("Timeout waiting for callback 1", latch1.await(10, TimeUnit.SECONDS)); - assertTrue("Timeout waiting for callback 2", latch2.await(10, TimeUnit.SECONDS)); + assertTrue("Timeout waiting for callback 1", callbackLatch1.await(10, TimeUnit.SECONDS)); + assertTrue("Timeout waiting for callback 2", callbackLatch2.await(10, TimeUnit.SECONDS)); // Verify that requests were sent. // It should contain one BeginTransaction and two ExecuteSql. - assertThat(mockSpanner.getRequestTypes()) - .containsExactly( - BeginTransactionRequest.class, ExecuteSqlRequest.class, ExecuteSqlRequest.class); + assertEquals( + Arrays.asList( + BeginTransactionRequest.class, ExecuteSqlRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); } } @@ -184,12 +193,547 @@ public void createAsyncResultSet_handlesExceptionCorrectly() throws Exception { DatabaseClient client = testSpanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE)); try (ReadOnlyTransaction transaction = client.readOnlyTransaction()) { - RuntimeException e = + RuntimeException exception = assertThrows( RuntimeException.class, () -> transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT)); - assertEquals("Failed to get executor", e.getMessage()); + assertEquals("Failed to get executor", exception.getMessage()); } } } + + @Test + public void closeAsyncIsNonBlockingWhenServerIsFrozen() throws Exception { + // Warm up session pool to avoid CreateSession blocking when server is frozen. + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + mockSpanner.freeze(); + + // Call executeQueryAsync and register callback. + // This starts InitiateStreamingRunnable in the background executor, which calls + // BeginTransaction. + AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + final CountDownLatch callbackLatch = new CountDownLatch(1); + resultSet.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + // Calling closeAsync() must return an ApiFuture immediately without blocking the calling + // thread, + // even though the mock server is frozen and BeginTransaction is in flight. + ApiFuture closeFuture = transaction.closeAsync(); + + // Deterministic assertion: closeFuture is NOT done yet because BeginTransaction is still + // blocked in mockSpanner. + assertFalse(closeFuture.isDone()); + + // Unfreeze mock server so BeginTransaction and query execution can proceed. + mockSpanner.unfreeze(); + + // Wait for callback and closeFuture. + assertTrue("Timeout waiting for callback", callbackLatch.await(10, TimeUnit.SECONDS)); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + + // Verify that requests were sent. + assertEquals( + Arrays.asList(BeginTransactionRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + } + + @Test + public void closeAsyncWithoutQueriesCompletesImmediately() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + try (ReadOnlyTransaction transaction = client().readOnlyTransaction()) { + ApiFuture closeFuture = transaction.closeAsync(); + assertTrue(closeFuture.isDone()); + } + } + + @Test + public void closeAsyncIsIdempotent() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + ApiFuture closeFuture1 = transaction.closeAsync(); + ApiFuture closeFuture2 = transaction.closeAsync(); + assertSame(closeFuture1, closeFuture2); + assertTrue(closeFuture1.isDone()); + assertTrue(closeFuture2.isDone()); + closeFuture1.get(10, TimeUnit.SECONDS); + closeFuture2.get(10, TimeUnit.SECONDS); + } + + @Test + public void queryAfterCloseAsyncThrowsException() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + ApiFuture closeFuture = transaction.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + + IllegalStateException exception = + assertThrows( + IllegalStateException.class, + () -> transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT)); + assertEquals("Context has been closed", exception.getMessage()); + } + + @Test + public void closeAsyncWithUnusedResultSetClosedCompletesWithoutHang() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + // User closes resultSet without attaching callback or starting stream + resultSet.close(); + ApiFuture closeFuture = transaction.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + @Test + public void syncCloseDelegatesToCloseAsyncAndWaitsForPendingStarts() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + final CountDownLatch callbackLatch = new CountDownLatch(1); + resultSet.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + // Synchronous close() should complete successfully once background query finishes. + transaction.close(); + assertTrue("Timeout waiting for callback", callbackLatch.await(10, TimeUnit.SECONDS)); + assertEquals( + Arrays.asList(BeginTransactionRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + } + + @Test + public void multipleConcurrentInFlightQueriesCloseAsyncWaitsForAll() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + mockSpanner.freeze(); + + AsyncResultSet resultSet1 = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + AsyncResultSet resultSet2 = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + AsyncResultSet resultSet3 = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + + final CountDownLatch callbackLatch1 = new CountDownLatch(1); + final CountDownLatch callbackLatch2 = new CountDownLatch(1); + final CountDownLatch callbackLatch3 = new CountDownLatch(1); + + resultSet1.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet1.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch1.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch1.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + resultSet2.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet2.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch2.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch2.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + resultSet3.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet3.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch3.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch3.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + ApiFuture closeFuture = transaction.closeAsync(); + assertFalse(closeFuture.isDone()); + + mockSpanner.unfreeze(); + + assertTrue("Timeout waiting for callback 1", callbackLatch1.await(10, TimeUnit.SECONDS)); + assertTrue("Timeout waiting for callback 2", callbackLatch2.await(10, TimeUnit.SECONDS)); + assertTrue("Timeout waiting for callback 3", callbackLatch3.await(10, TimeUnit.SECONDS)); + + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + + assertEquals( + Arrays.asList( + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + ExecuteSqlRequest.class, + ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + } + + @Test + public void syncQueriesAfterCloseAsyncThrowException() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + ApiFuture closeFuture = transaction.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + + IllegalStateException executeQueryException = + assertThrows( + IllegalStateException.class, + () -> transaction.executeQuery(READ_ONE_KEY_VALUE_STATEMENT)); + assertEquals("Context has been closed", executeQueryException.getMessage()); + + IllegalStateException readException = + assertThrows( + IllegalStateException.class, + () -> + transaction.read( + "TestTable", + KeySet.singleKey(Key.of("k1")), + Collections.singletonList("Value"))); + assertEquals("Context has been closed", readException.getMessage()); + + IllegalStateException readRowException = + assertThrows( + IllegalStateException.class, + () -> + transaction.readRow("TestTable", Key.of("k1"), Collections.singletonList("Value"))); + assertEquals("Context has been closed", readRowException.getMessage()); + + IllegalStateException analyzeQueryException = + assertThrows( + IllegalStateException.class, + () -> + transaction.analyzeQuery( + READ_ONE_KEY_VALUE_STATEMENT, ReadContext.QueryAnalyzeMode.PLAN)); + assertEquals("Context has been closed", analyzeQueryException.getMessage()); + } + + @Test + public void closeAsyncIsIdempotentDuringInFlightQuery() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + mockSpanner.freeze(); + + AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + final CountDownLatch callbackLatch = new CountDownLatch(1); + resultSet.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + ApiFuture closeFuture1 = transaction.closeAsync(); + ApiFuture closeFuture2 = transaction.closeAsync(); + + assertSame(closeFuture1, closeFuture2); + assertFalse(closeFuture1.isDone()); + + mockSpanner.unfreeze(); + + assertTrue("Timeout waiting for callback", callbackLatch.await(10, TimeUnit.SECONDS)); + closeFuture1.get(10, TimeUnit.SECONDS); + closeFuture2.get(10, TimeUnit.SECONDS); + + assertTrue(closeFuture1.isDone()); + assertTrue(closeFuture2.isDone()); + } + + @Test + public void closeAsyncCompletesWhenBeginTransactionFails() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + mockSpanner.setBeginTransactionExecutionTime( + SimulatedExecutionTime.ofException( + Status.UNAVAILABLE.withDescription("Service unavailable").asRuntimeException())); + + try { + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + final CountDownLatch callbackLatch = new CountDownLatch(1); + resultSet.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + ApiFuture closeFuture = transaction.closeAsync(); + + assertTrue("Timeout waiting for callback", callbackLatch.await(10, TimeUnit.SECONDS)); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } finally { + mockSpanner.setBeginTransactionExecutionTime(NO_EXECUTION_TIME); + } + } + + @Test + public void closeAsyncWhileBeginTransactionIsInFlightWaitsForBeginTransactionToFinish() + throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + mockSpanner.freeze(); + + final CountDownLatch initStartedLatch = new CountDownLatch(1); + final CountDownLatch initFinishedLatch = new CountDownLatch(1); + executor.execute( + () -> { + initStartedLatch.countDown(); + try { + ((AbstractReadContext.MultiUseReadOnlyTransaction) transaction).initTransaction(); + } finally { + initFinishedLatch.countDown(); + } + }); + + assertTrue( + "Timeout waiting for init thread to start", initStartedLatch.await(10, TimeUnit.SECONDS)); + + mockSpanner.waitForRequestsToContain(BeginTransactionRequest.class, 10_000); + + ApiFuture closeFuture = transaction.closeAsync(); + + assertFalse(closeFuture.isDone()); + + mockSpanner.unfreeze(); + + assertTrue("Timeout waiting for init to finish", initFinishedLatch.await(10, TimeUnit.SECONDS)); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + + assertEquals( + Collections.singletonList(BeginTransactionRequest.class), mockSpanner.getRequestTypes()); + } + + @Test + public void readAsyncWithCloseAsyncWaitsForQueryToFinish() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + mockSpanner.freeze(); + + AsyncResultSet resultSet = + transaction.readAsync("TestTable", KeySet.all(), Collections.singletonList("Value")); + final CountDownLatch callbackLatch = new CountDownLatch(1); + resultSet.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + ApiFuture closeFuture = transaction.closeAsync(); + assertFalse(closeFuture.isDone()); + + mockSpanner.unfreeze(); + + assertTrue("Timeout waiting for callback", callbackLatch.await(10, TimeUnit.SECONDS)); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + @Test + public void readUsingIndexAsyncWithCloseAsyncWaitsForQueryToFinish() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + mockSpanner.freeze(); + + AsyncResultSet resultSet = + transaction.readUsingIndexAsync( + "TestTable", "TestIndex", KeySet.all(), Collections.singletonList("Value")); + final CountDownLatch callbackLatch = new CountDownLatch(1); + resultSet.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + ApiFuture closeFuture = transaction.closeAsync(); + assertFalse(closeFuture.isDone()); + + mockSpanner.unfreeze(); + + assertTrue("Timeout waiting for callback", callbackLatch.await(10, TimeUnit.SECONDS)); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + @Test + public void closeAsyncThenCloseUnusedAsyncResultSetCompletesCleanly() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + ReadOnlyTransaction transaction = client().readOnlyTransaction(); + AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + ApiFuture closeFuture = transaction.closeAsync(); + assertFalse(closeFuture.isDone()); + resultSet.close(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + @Test + public void closeAsyncWithInlinedBeginTransaction() throws Exception { + try (ResultSet resultSet = client().singleUse().executeQuery(READ_ONE_KEY_VALUE_STATEMENT)) { + while (resultSet.next()) {} + } + mockSpanner.clearRequests(); + + ReadOnlyTransaction transaction = + client() + .readOnlyTransaction( + TimestampBound.strong(), + Options.beginTransactionOption(Options.BeginTransactionOption.INLINE)); + mockSpanner.freeze(); + + AsyncResultSet resultSet = transaction.executeQueryAsync(READ_ONE_KEY_VALUE_STATEMENT); + final CountDownLatch callbackLatch = new CountDownLatch(1); + resultSet.setCallback( + executor, + ignored -> { + try { + AsyncResultSet.CursorState state; + while ((state = resultSet.tryNext()) == AsyncResultSet.CursorState.OK) {} + if (state == AsyncResultSet.CursorState.DONE) { + callbackLatch.countDown(); + } + return AsyncResultSet.CallbackResponse.CONTINUE; + } catch (Throwable throwable) { + callbackLatch.countDown(); + return AsyncResultSet.CallbackResponse.DONE; + } + }); + + ApiFuture closeFuture = transaction.closeAsync(); + assertFalse(closeFuture.isDone()); + + mockSpanner.unfreeze(); + + assertTrue("Timeout waiting for callback", callbackLatch.await(10, TimeUnit.SECONDS)); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DelayedReadContextTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DelayedReadContextTest.java new file mode 100644 index 000000000000..8eaa098b7ebd --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DelayedReadContextTest.java @@ -0,0 +1,103 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.spanner; + +import static org.junit.Assert.assertTrue; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.core.SettableApiFuture; +import java.util.concurrent.TimeUnit; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class DelayedReadContextTest { + + @Test + public void closeAsyncDelegatesToReadContextCloseAsyncWhenFutureAlreadyResolved() + throws Exception { + ReadContext mockReadContext = mock(ReadContext.class); + when(mockReadContext.closeAsync()).thenReturn(ApiFutures.immediateFuture(null)); + + SettableApiFuture readContextFuture = SettableApiFuture.create(); + readContextFuture.set(mockReadContext); + DelayedReadContext delayedReadContext = + new DelayedReadContext<>(readContextFuture); + + ApiFuture closeFuture = delayedReadContext.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + verify(mockReadContext).closeAsync(); + verify(mockReadContext, never()).close(); + } + + @Test + public void closeAsyncCancelsPendingReadContextFutureWhenNotYetDone() throws Exception { + SettableApiFuture readContextFuture = SettableApiFuture.create(); + DelayedReadContext delayedReadContext = + new DelayedReadContext<>(readContextFuture); + + ApiFuture closeFuture = delayedReadContext.closeAsync(); + assertTrue(closeFuture.isDone()); + assertTrue(readContextFuture.isCancelled()); + } + + @Test + public void closeAsyncSuppressesExceptionWhenUnderlyingCloseAsyncFails() throws Exception { + ReadContext mockReadContext = mock(ReadContext.class); + when(mockReadContext.closeAsync()) + .thenReturn(ApiFutures.immediateFailedFuture(new RuntimeException("Close failed"))); + + SettableApiFuture readContextFuture = SettableApiFuture.create(); + readContextFuture.set(mockReadContext); + DelayedReadContext delayedReadContext = + new DelayedReadContext<>(readContextFuture); + + ApiFuture closeFuture = delayedReadContext.closeAsync(); + closeFuture.get(10, TimeUnit.SECONDS); + assertTrue(closeFuture.isDone()); + } + + @Test + public void syncCloseCancelsPendingReadContextFutureWhenNotYetDone() { + SettableApiFuture readContextFuture = SettableApiFuture.create(); + DelayedReadContext delayedReadContext = + new DelayedReadContext<>(readContextFuture); + + delayedReadContext.close(); + assertTrue(readContextFuture.isCancelled()); + } + + @Test + public void syncCloseDelegatesToReadContextCloseWhenFutureAlreadyResolved() { + ReadContext mockReadContext = mock(ReadContext.class); + SettableApiFuture readContextFuture = SettableApiFuture.create(); + readContextFuture.set(mockReadContext); + + DelayedReadContext delayedReadContext = + new DelayedReadContext<>(readContextFuture); + delayedReadContext.close(); + + verify(mockReadContext).close(); + } +}