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 c473d0719f49..a769758861f4 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 @@ -18,7 +18,6 @@ import static com.google.cloud.spanner.SessionClient.optionMap; import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException; -import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; @@ -27,7 +26,8 @@ import com.google.api.core.ApiFutures; import com.google.api.core.SettableApiFuture; import com.google.api.gax.core.ExecutorProvider; -import com.google.cloud.Timestamp; +import com.google.api.gax.retrying.RetrySettings; +import com.google.api.gax.rpc.StatusCode.Code; import com.google.cloud.grpc.GcpManagedChannel.ChannelAffinityRef; import com.google.cloud.spanner.AbstractResultSet.CloseableIterator; import com.google.cloud.spanner.AsyncResultSet.CallbackResponse; @@ -37,14 +37,10 @@ import com.google.cloud.spanner.SessionClient.SessionOption; import com.google.cloud.spanner.SessionImpl.SessionTransaction; import com.google.cloud.spanner.spi.v1.SpannerRpc; -import com.google.cloud.spanner.spi.v1.SpannerRpc.Option; import com.google.common.annotations.VisibleForTesting; -import com.google.common.base.Preconditions; import com.google.common.base.Supplier; -import com.google.common.collect.ImmutableMap; import com.google.common.util.concurrent.MoreExecutors; import com.google.protobuf.ByteString; -import com.google.spanner.v1.BeginTransactionRequest; import com.google.spanner.v1.DirectedReadOptions; import com.google.spanner.v1.ExecuteBatchDmlRequest; import com.google.spanner.v1.ExecuteSqlRequest; @@ -55,17 +51,11 @@ import com.google.spanner.v1.ReadRequest; import com.google.spanner.v1.RequestOptions; import com.google.spanner.v1.Transaction; -import com.google.spanner.v1.TransactionOptions; import com.google.spanner.v1.TransactionSelector; +import io.opentelemetry.api.common.Attributes; import java.util.Map; -import java.util.concurrent.ExecutionException; -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.Set; import java.util.concurrent.atomic.AtomicLong; -import java.util.concurrent.locks.ReentrantLock; import java.util.logging.Logger; import javax.annotation.Nullable; import javax.annotation.concurrent.GuardedBy; @@ -144,7 +134,7 @@ B setExecutorProvider(ExecutorProvider executorProvider) { } B setClock(Clock clock) { - this.clock = Preconditions.checkNotNull(clock); + this.clock = checkNotNull(clock); return self(); } @@ -167,796 +157,10 @@ interface ListenableAsyncResultSet extends AsyncResultSet { void removeListener(Runnable listener); } - /** - * A {@code ReadContext} for standalone reads. This can only be used for a single operation, since - * each standalone read may see a different timestamp of Cloud Spanner data. - */ - static class SingleReadContext extends AbstractReadContext { - static class Builder extends AbstractReadContext.Builder { - private TimestampBound bound; - - private Builder() {} - - Builder setTimestampBound(TimestampBound bound) { - this.bound = bound; - return self(); - } - - @Override - SingleReadContext build() { - return new SingleReadContext(this); - } - - SingleUseReadOnlyTransaction buildSingleUseReadOnlyTransaction() { - return new SingleUseReadOnlyTransaction(this); - } - } - - static Builder newBuilder() { - return new Builder(); - } - - final TimestampBound bound; - - @GuardedBy("lock") - private boolean used; - - private Map channelHint; - - private SingleReadContext(Builder builder) { - super(builder); - this.bound = builder.bound; - // single use transaction have a single RPC and hence there is no need - // of a channel hint. GAX will automatically choose a hint when used - // with a multiplexed session to perform a round-robin channel selection. We are - // passing a hint here to prefer random channel selection instead of doing GAX round-robin. - this.channelHint = - getChannelHintOptions( - session.getOptions(), - ThreadLocalRandom.current().nextLong(Long.MAX_VALUE), - session.getSpanner().getOptions().isGrpcGcpExtensionEnabled()); - } - - @Override - protected boolean isRouteToLeader() { - return false; - } - - @GuardedBy("lock") - @Override - void beforeReadOrQueryLocked() { - super.beforeReadOrQueryLocked(); - checkState(!used, "Cannot use a single-read ReadContext for multiple reads"); - used = true; - } - - @Override - @Nullable - TransactionSelector getTransactionSelector() { - if (bound.getMode() == TimestampBound.Mode.STRONG) { - // Default mode: no need to specify a transaction. - return null; - } - return TransactionSelector.newBuilder() - .setSingleUse(TransactionOptions.newBuilder().setReadOnly(bound.toProto())) - .build(); - } - - @Override - Map getTransactionChannelHint() { - return channelHint; - } - - @Override - boolean prepareRetryOnDifferentGrpcChannel() { - ChannelAffinityRef channelAffinityRef = - Option.CHANNEL_ID_AFFINITY.getChannelAffinityRef(channelHint); - if (session.getIsMultiplexed() && channelAffinityRef != null) { - channelAffinityRef.useDifferentChannelOnNextCall(); - return true; - } - return super.prepareRetryOnDifferentGrpcChannel(); - } - } - - private static void assertTimestampAvailable(boolean available) { + static void assertTimestampAvailable(boolean available) { checkState(available, "Method can only be called after read has returned data or finished"); } - static class SingleUseReadOnlyTransaction extends SingleReadContext - implements ReadOnlyTransaction { - - @GuardedBy("lock") - private Timestamp timestamp; - - private SingleUseReadOnlyTransaction(SingleReadContext.Builder builder) { - super(builder); - } - - @Override - public Timestamp getReadTimestamp() { - synchronized (lock) { - assertTimestampAvailable(timestamp != null); - return timestamp; - } - } - - @Override - @Nullable - TransactionSelector getTransactionSelector() { - TransactionOptions.Builder options = TransactionOptions.newBuilder(); - bound.applyToBuilder(options.getReadOnlyBuilder()).setReturnReadTimestamp(true); - return TransactionSelector.newBuilder().setSingleUse(options).build(); - } - - @Override - public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) { - synchronized (lock) { - if (!transaction.hasReadTimestamp()) { - throw newSpannerException( - ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field"); - } - try { - timestamp = Timestamp.fromProto(transaction.getReadTimestamp()); - } catch (IllegalArgumentException e) { - throw newSpannerException( - ErrorCode.INTERNAL, "Bad value in transaction.read_timestamp metadata field", e); - } - } - } - } - - static class MultiUseReadOnlyTransaction extends AbstractReadContext - implements ReadOnlyTransaction { - static class Builder extends AbstractReadContext.Builder { - private TimestampBound bound; - private Timestamp timestamp; - private ByteString transactionId; - private Options.BeginTransactionOption beginTransactionOption = - Options.BeginTransactionOption.EXPLICIT; - - private Builder() {} - - Builder setTimestampBound(TimestampBound bound) { - this.bound = bound; - return this; - } - - Builder setTimestamp(Timestamp timestamp) { - this.timestamp = timestamp; - return this; - } - - Builder setTransactionId(ByteString transactionId) { - this.transactionId = transactionId; - return this; - } - - Builder setBeginTransactionOption(Options.BeginTransactionOption beginTransactionOption) { - this.beginTransactionOption = beginTransactionOption; - return this; - } - - @Override - MultiUseReadOnlyTransaction build() { - return new MultiUseReadOnlyTransaction(this); - } - } - - static Builder newBuilder() { - return new Builder(); - } - - private TimestampBound bound; - private final ReentrantLock txnLock = new ReentrantLock(); - - @GuardedBy("txnLock") - private Timestamp timestamp; - - @GuardedBy("txnLock") - private ByteString transactionId; - - /** - * Future used to synchronize concurrent operations during transaction initialization: - * - *
    - *
  • 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. - *
  • 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}. - *
- */ - @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; - - private final Map channelHint; - private final Options.BeginTransactionOption beginTransactionOption; - - MultiUseReadOnlyTransaction(Builder builder) { - super(builder); - checkArgument( - !(builder.bound != null && builder.transactionId != null) - && !(builder.bound == null && builder.transactionId == null), - "Either TimestampBound or TransactionId must be specified"); - if (builder.bound != null) { - checkArgument( - builder.bound.getMode() != TimestampBound.Mode.MAX_STALENESS - && builder.bound.getMode() != TimestampBound.Mode.MIN_READ_TIMESTAMP, - "Bounded staleness mode %s is not supported for multi-use read-only transactions." - + " Create a single-use read or read-only transaction instead.", - builder.bound.getMode()); - this.bound = builder.bound; - } else { - this.timestamp = builder.timestamp; - this.transactionId = builder.transactionId; - } - this.channelHint = - getChannelHintOptions( - session.getOptions(), - ThreadLocalRandom.current().nextLong(Long.MAX_VALUE), - session.getSpanner().getOptions().isGrpcGcpExtensionEnabled()); - this.beginTransactionOption = builder.beginTransactionOption; - } - - @Override - public Map getTransactionChannelHint() { - return channelHint; - } - - @Override - protected boolean isRouteToLeader() { - return false; - } - - private boolean shouldUseInlinedBegin() { - return beginTransactionOption == Options.BeginTransactionOption.INLINE; - } - - @Override - void beforeReadOrQuery() { - super.beforeReadOrQuery(); - if (shouldUseInlinedBegin()) { - // Keep the same nested transaction guard as the explicit BeginTransaction path. This checks - // TransactionRunner's thread-local pending state, not the session's active transaction. - SessionImpl.throwIfTransactionsPending(); - } else { - initTransaction(); - } - } - - @Override - @Nullable - TransactionSelector getTransactionSelector() { - if (!shouldUseInlinedBegin()) { - // No need for synchronization: super.readInternal() is always preceded by a check of - // "transactionId" that provides a happens-before from initialization, and the value is - // never changed afterwards. - @SuppressWarnings("GuardedByChecker") - TransactionSelector selector = - TransactionSelector.newBuilder().setId(transactionId).build(); - return selector; - } - - ApiFuture futureToWaitFor = null; - txnLock.lock(); - try { - if (transactionId != null) { - return TransactionSelector.newBuilder().setId(transactionId).build(); - } - if (transactionIdFuture == null) { - transactionIdFuture = SettableApiFuture.create(); - return TransactionSelector.newBuilder() - .setBegin(createReadOnlyTransactionOptions()) - .build(); - } - futureToWaitFor = transactionIdFuture; - } finally { - txnLock.unlock(); - } - - try { - return TransactionSelector.newBuilder() - .setId(futureToWaitFor.get(WAIT_FOR_INLINE_BEGIN_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS)) - .build(); - } catch (ExecutionException e) { - throw SpannerExceptionFactory.asSpannerException(e.getCause()); - } catch (TimeoutException e) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.DEADLINE_EXCEEDED, - "Timeout while waiting for an inlined read-only transaction to be returned by another" - + " statement.", - e); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw SpannerExceptionFactory.newSpannerExceptionForCancellation(null, e); - } - } - - @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 { - 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) { - 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. - 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(); - } - } - } - }; - } catch (Throwable throwable) { - if (decremented.compareAndSet(false, true)) { - decrementPendingStartsAndSignal(); - } - 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) { - Options readOptions = Options.fromReadOptions(options); - final int bufferRows = - readOptions.hasBufferRows() - ? readOptions.bufferRows() - : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; - return createAsyncResultSet( - () -> super.readInternalWithOptions(table, null, keys, columns, readOptions, null), - bufferRows); - } - - @Override - public ListenableAsyncResultSet readUsingIndexAsync( - String table, String index, KeySet keys, Iterable columns, ReadOption... options) { - Options readOptions = Options.fromReadOptions(options); - final int bufferRows = - readOptions.hasBufferRows() - ? readOptions.bufferRows() - : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; - return createAsyncResultSet( - () -> - super.readInternalWithOptions( - table, checkNotNull(index), keys, columns, readOptions, null), - bufferRows); - } - - @Override - public ListenableAsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { - Options readOptions = Options.fromQueryOptions(options); - final int bufferRows = - readOptions.hasBufferRows() - ? readOptions.bufferRows() - : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; - return createAsyncResultSet( - () -> - super.executeQueryInternal( - statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.NORMAL, options), - bufferRows); - } - - @Override - public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) { - Timestamp readTimestamp = null; - if (transaction.hasReadTimestamp()) { - try { - readTimestamp = Timestamp.fromProto(transaction.getReadTimestamp()); - } catch (IllegalArgumentException e) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, "Bad value in transaction.read_timestamp metadata field", e); - } - } - if (shouldIncludeId && transaction.getId().isEmpty()) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, NO_TRANSACTION_RETURNED_MSG); - } - txnLock.lock(); - try { - if (timestamp == null) { - if (readTimestamp == null) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field"); - } - timestamp = readTimestamp; - } - if (shouldIncludeId && transactionId == null) { - transactionId = transaction.getId(); - if (transactionIdFuture != null && !transactionIdFuture.isDone()) { - transactionIdFuture.set(transactionId); - } - } - } finally { - txnLock.unlock(); - } - checkAndClose(); - } - - @Override - public SpannerException onError( - SpannerException e, boolean withBeginTransaction, boolean lastStatement) { - e = super.onError(e, withBeginTransaction, lastStatement); - if (withBeginTransaction) { - failTransactionIdFuture(e); - } - return e; - } - - @Override - public void onDone(boolean withBeginTransaction) { - if (withBeginTransaction) { - failTransactionIdFuture( - SpannerExceptionFactory.newSpannerException( - ErrorCode.FAILED_PRECONDITION, - "ResultSet was closed before a read-only transaction id was returned")); - } - super.onDone(withBeginTransaction); - } - - @Override - void onStartFailed(boolean withBeginTransaction, Throwable throwable) { - if (withBeginTransaction) { - failTransactionIdFuture(throwable); - } - } - - private void failTransactionIdFuture(Throwable throwable) { - txnLock.lock(); - try { - if (transactionIdFuture != null && !transactionIdFuture.isDone()) { - transactionIdFuture.setException(throwable); - } - } finally { - txnLock.unlock(); - } - checkAndClose(); - } - - @Override - public Timestamp getReadTimestamp() { - txnLock.lock(); - try { - assertTimestampAvailable(timestamp != null); - return timestamp; - } finally { - txnLock.unlock(); - } - } - - ByteString getTransactionId() { - txnLock.lock(); - try { - return transactionId; - } finally { - txnLock.unlock(); - } - } - - /** - * 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 ApiFuture closeAsync() { - boolean shouldCloseImmediately = false; - txnLock.lock(); - try { - if (isClosedOrClosing) { - return closeFuture; - } - isClosedOrClosing = true; - closeFuture = SettableApiFuture.create(); - if (shouldCloseLocked()) { - closed = true; - shouldCloseImmediately = true; - } else { - return closeFuture; - } - } finally { - txnLock.unlock(); - } - - 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() { - TransactionOptions.Builder options = TransactionOptions.newBuilder(); - if (timestamp != null) { - options - .getReadOnlyBuilder() - .setReadTimestamp(timestamp.toProto()) - .setReturnReadTimestamp(true); - } else { - bound.applyToBuilder(options.getReadOnlyBuilder()).setReturnReadTimestamp(true); - } - return options.build(); - } - - /** - * Initializes the transaction with the timestamp specified within MultiUseReadOnlyTransaction. - * This is used only for fallback of PartitionQueryRequest and PartitionReadRequest with - * Multiplexed Session. - */ - void initFallbackTransaction() { - 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(); - - ApiFuture futureToWaitFor = null; - BeginTransactionRequest request = null; - txnLock.lock(); - try { - if (transactionId != null) { - return; - } - 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 = - rpc.beginTransaction(request, getTransactionChannelHint(), isRouteToLeader()); - if (!transaction.hasReadTimestamp()) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field"); - } - if (transaction.getId().isEmpty()) { - throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, "Missing expected transaction.id metadata field"); - } - Timestamp readTimestamp; - try { - readTimestamp = Timestamp.fromProto(transaction.getReadTimestamp()); - } catch (IllegalArgumentException illegalArgumentException) { - throw SpannerExceptionFactory.newSpannerException( - 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(); - } - checkAndClose(); - span.addAnnotation( - "Transaction Creation Done", - ImmutableMap.of( - "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; - } - } - } - final Object lock = new Object(); final SessionImpl session; final boolean cancelQueryWhenClientIsClosed; @@ -1268,71 +472,84 @@ ExecuteBatchDmlRequest.Builder getExecuteBatchDmlRequestBuilder( return builder; } - ResultSet executeQueryInternalWithOptions( - final Statement statement, - final com.google.spanner.v1.ExecuteSqlRequest.QueryMode queryMode, - final Options options, - final ByteString partitionToken) { - beforeReadOrQuery(); - final int prefetchChunks = - options.hasPrefetchChunks() ? options.prefetchChunks() : defaultPrefetchChunks; - final ExecuteSqlRequest.Builder request = - getExecuteSqlRequestBuilder( - statement, queryMode, options, /* withTransactionSelector= */ false); + @FunctionalInterface + private interface StreamStarter { + StreamCallResult start( + @Nullable ByteString resumeToken, + XGoogSpannerRequestId requestId, + SpannerRpc.ResultStreamConsumer consumer); + } + + private static final class StreamCallResult { + final SpannerRpc.StreamingCall call; + final boolean withBeginTransaction; + + StreamCallResult(SpannerRpc.StreamingCall call, boolean withBeginTransaction) { + this.call = call; + this.withBeginTransaction = withBeginTransaction; + } + } + + @FunctionalInterface + private interface StreamCall { + SpannerRpc.StreamingCall execute( + SpannerRpc.ResultStreamConsumer consumer, XGoogSpannerRequestId requestId); + } + + private StreamCallResult executeCall( + boolean withBeginTransaction, + SpannerRpc.ResultStreamConsumer consumer, + XGoogSpannerRequestId requestId, + StreamCall call) { + try { + return new StreamCallResult(call.execute(consumer, requestId), withBeginTransaction); + } catch (RuntimeException | Error t) { + onStartFailed(withBeginTransaction, t); + throw t; + } + } + + @Nullable + private TransactionSelector prepareTransactionSelector( + @Nullable ByteString resumeToken, boolean hasTransaction) { + if (resumeToken != null || !hasTransaction) { + return getTransactionSelector(); + } + return null; + } + + private ResultSet createStreamingResultSet( + String operationName, + Attributes attributes, + RetrySettings retrySettings, + Set retryableCodes, + Options options, + Supplier streamSupplier, + StreamStarter streamStarter) { ResumableStreamIterator stream = new ResumableStreamIterator( MAX_BUFFERED_CHUNKS, - SpannerImpl.QUERY, + operationName, span, tracer, - tracer.createStatementAttributes(statement, options), + attributes, session.getErrorHandler(), - rpc.getExecuteQueryRetrySettings(), - rpc.getExecuteQueryRetryableCodes(), + retrySettings, + retryableCodes, session.getRequestIdCreator()) { @Override CloseableIterator startStream( @Nullable ByteString resumeToken, AsyncResultSet.StreamMessageListener streamListener, XGoogSpannerRequestId requestId) { - GrpcStreamIterator stream = - new GrpcStreamIterator( - statement, - request.getLastStatement(), - prefetchChunks, - cancelQueryWhenClientIsClosed); + GrpcStreamIterator stream = streamSupplier.get(); if (streamListener != null) { stream.registerListener(streamListener); } - if (partitionToken != null) { - request.setPartitionToken(partitionToken); - } - TransactionSelector selector = null; - if (resumeToken != null) { - request.setResumeToken(resumeToken); - selector = getTransactionSelector(); - } else if (!request.hasTransaction()) { - selector = getTransactionSelector(); - } - if (selector != null) { - request.setTransaction(selector); - } - boolean withBeginTransaction = request.getTransaction().hasBegin(); - SpannerRpc.StreamingCall call; - try { - call = - rpc.executeQuery( - request.build(), - stream.consumer(), - getTransactionChannelHint(), - requestId, - isRouteToLeader()); - } catch (RuntimeException | Error t) { - onStartFailed(withBeginTransaction, t); - throw t; - } + StreamCallResult result = + streamStarter.start(resumeToken, requestId, stream.consumer()); session.markUsed(clock.instant()); - stream.setCall(call, withBeginTransaction); + stream.setCall(result.call, result.withBeginTransaction); return stream; } @@ -1345,6 +562,55 @@ boolean prepareIteratorForRetryOnDifferentGrpcChannel() { stream, this, options.hasDecodeMode() ? options.decodeMode() : defaultDecodeMode); } + ResultSet executeQueryInternalWithOptions( + final Statement statement, + final QueryMode queryMode, + final Options options, + final ByteString partitionToken) { + beforeReadOrQuery(); + final int prefetchChunks = + options.hasPrefetchChunks() ? options.prefetchChunks() : defaultPrefetchChunks; + final ExecuteSqlRequest.Builder request = + getExecuteSqlRequestBuilder( + statement, queryMode, options, /* withTransactionSelector= */ false); + if (partitionToken != null) { + request.setPartitionToken(partitionToken); + } + return createStreamingResultSet( + SpannerImpl.QUERY, + tracer.createStatementAttributes(statement, options), + rpc.getExecuteQueryRetrySettings(), + rpc.getExecuteQueryRetryableCodes(), + options, + () -> + new GrpcStreamIterator( + statement, + request.getLastStatement(), + prefetchChunks, + cancelQueryWhenClientIsClosed), + (resumeToken, requestId, consumer) -> { + if (resumeToken != null) { + request.setResumeToken(resumeToken); + } + TransactionSelector selector = + prepareTransactionSelector(resumeToken, request.hasTransaction()); + if (selector != null) { + request.setTransaction(selector); + } + return executeCall( + request.getTransaction().hasBegin(), + consumer, + requestId, + (streamConsumer, reqId) -> + rpc.executeQuery( + request.build(), + streamConsumer, + getTransactionChannelHint(), + reqId, + isRouteToLeader())); + }); + } + static Map getChannelHintOptions( Map channelHintForSession, Long channelHintForTransaction) { return getChannelHintOptions( @@ -1513,69 +779,39 @@ ResultSet readInternalWithOptions( builder.setLockHint(readOptions.lockHint()); } } + builder.setRequestOptions(buildRequestOptions(readOptions)); final int prefetchChunks = readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks; final boolean lastStatement = readOptions.hasLastStatement() ? readOptions.isLastStatement() : false; - ResumableStreamIterator stream = - new ResumableStreamIterator( - MAX_BUFFERED_CHUNKS, - SpannerImpl.READ, - span, - tracer, - tracer.createTableAttributes(table, readOptions), - session.getErrorHandler(), - rpc.getReadRetrySettings(), - rpc.getReadRetryableCodes(), - session.getRequestIdCreator()) { - @Override - CloseableIterator startStream( - @Nullable ByteString resumeToken, - AsyncResultSet.StreamMessageListener streamListener, - XGoogSpannerRequestId requestId) { - GrpcStreamIterator stream = - new GrpcStreamIterator( - lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed); - if (streamListener != null) { - stream.registerListener(streamListener); - } - TransactionSelector selector = null; - if (resumeToken != null) { - builder.setResumeToken(resumeToken); - selector = getTransactionSelector(); - } else if (!builder.hasTransaction()) { - selector = getTransactionSelector(); - } - if (selector != null) { - builder.setTransaction(selector); - } - builder.setRequestOptions(buildRequestOptions(readOptions)); - boolean withBeginTransaction = builder.getTransaction().hasBegin(); - SpannerRpc.StreamingCall call; - try { - call = + return createStreamingResultSet( + SpannerImpl.READ, + tracer.createTableAttributes(table, readOptions), + rpc.getReadRetrySettings(), + rpc.getReadRetryableCodes(), + readOptions, + () -> new GrpcStreamIterator(lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed), + (resumeToken, requestId, consumer) -> { + if (resumeToken != null) { + builder.setResumeToken(resumeToken); + } + TransactionSelector selector = + prepareTransactionSelector(resumeToken, builder.hasTransaction()); + if (selector != null) { + builder.setTransaction(selector); + } + return executeCall( + builder.getTransaction().hasBegin(), + consumer, + requestId, + (streamConsumer, reqId) -> rpc.read( builder.build(), - stream.consumer(), + streamConsumer, getTransactionChannelHint(), - requestId, - isRouteToLeader()); - } catch (RuntimeException | Error t) { - onStartFailed(withBeginTransaction, t); - throw t; - } - session.markUsed(clock.instant()); - stream.setCall(call, withBeginTransaction); - return stream; - } - - @Override - boolean prepareIteratorForRetryOnDifferentGrpcChannel() { - return AbstractReadContext.this.prepareRetryOnDifferentGrpcChannel(); - } - }; - return new GrpcResultSet( - stream, this, readOptions.hasDecodeMode() ? readOptions.decodeMode() : defaultDecodeMode); + reqId, + isRouteToLeader())); + }); } private Struct consumeSingleRow(ResultSet resultSet) { diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java index 5cbe01aa7117..6de6fb6a4b0a 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java @@ -16,13 +16,12 @@ package com.google.cloud.spanner; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; -import com.google.cloud.spanner.AbstractReadContext.MultiUseReadOnlyTransaction; import com.google.cloud.spanner.Options.QueryOption; import com.google.cloud.spanner.Options.ReadOption; import com.google.cloud.spanner.spi.v1.SpannerRpc; -import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.protobuf.Struct; import com.google.spanner.v1.ExecuteSqlRequest.QueryMode; @@ -184,7 +183,7 @@ public List partitionReadUsingIndex( ReadOption... option) throws SpannerException { Options readOptions = Options.fromReadOptions(option); - Preconditions.checkArgument( + checkArgument( !readOptions.hasLimit(), "Limit option not supported by partitionRead|partitionReadUsingIndex"); final PartitionReadRequest.Builder builder = diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java new file mode 100644 index 000000000000..3a35a0be07fa --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiUseReadOnlyTransaction.java @@ -0,0 +1,698 @@ +/* + * 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 com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; +import static com.google.common.base.Preconditions.checkState; + +import com.google.api.core.ApiFuture; +import com.google.api.core.SettableApiFuture; +import com.google.cloud.Timestamp; +import com.google.cloud.spanner.Options.QueryOption; +import com.google.cloud.spanner.Options.ReadOption; +import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.common.base.Supplier; +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.ByteString; +import com.google.spanner.v1.BeginTransactionRequest; +import com.google.spanner.v1.Transaction; +import com.google.spanner.v1.TransactionOptions; +import com.google.spanner.v1.TransactionSelector; +import java.util.Map; +import java.util.concurrent.ExecutionException; +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; +import javax.annotation.concurrent.GuardedBy; + +/** + * A {@link ReadOnlyTransaction} implementation for multi-use transactions. This allows multiple + * read and query operations to be executed against a consistent snapshot of Cloud Spanner data, + * coordinating transaction initialization, statement execution, and lifecycle management. + */ +class MultiUseReadOnlyTransaction extends AbstractReadContext implements ReadOnlyTransaction { + static class Builder extends AbstractReadContext.Builder { + private TimestampBound bound; + private Timestamp timestamp; + private ByteString transactionId; + private Options.BeginTransactionOption beginTransactionOption = + Options.BeginTransactionOption.EXPLICIT; + + private Builder() {} + + Builder setTimestampBound(TimestampBound bound) { + this.bound = bound; + return this; + } + + Builder setTimestamp(Timestamp timestamp) { + this.timestamp = timestamp; + return this; + } + + Builder setTransactionId(ByteString transactionId) { + this.transactionId = transactionId; + return this; + } + + Builder setBeginTransactionOption(Options.BeginTransactionOption beginTransactionOption) { + this.beginTransactionOption = beginTransactionOption; + return this; + } + + @Override + MultiUseReadOnlyTransaction build() { + return new MultiUseReadOnlyTransaction(this); + } + } + + static Builder newBuilder() { + return new Builder(); + } + + private TimestampBound bound; + private final ReentrantLock txnLock = new ReentrantLock(); + + @GuardedBy("txnLock") + private Timestamp timestamp; + + @GuardedBy("txnLock") + private ByteString transactionId; + + /** + * Future used to synchronize concurrent operations during transaction initialization: + * + *

    + *
  • 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. + *
  • 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}. + *
+ */ + @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; + + private final Map channelHint; + private final Options.BeginTransactionOption beginTransactionOption; + + MultiUseReadOnlyTransaction(Builder builder) { + super(builder); + checkArgument( + !(builder.bound != null && builder.transactionId != null) + && !(builder.bound == null && builder.transactionId == null), + "Either TimestampBound or TransactionId must be specified"); + if (builder.bound != null) { + checkArgument( + builder.bound.getMode() != TimestampBound.Mode.MAX_STALENESS + && builder.bound.getMode() != TimestampBound.Mode.MIN_READ_TIMESTAMP, + "Bounded staleness mode %s is not supported for multi-use read-only transactions." + + " Create a single-use read or read-only transaction instead.", + builder.bound.getMode()); + this.bound = builder.bound; + } else { + this.timestamp = builder.timestamp; + this.transactionId = builder.transactionId; + } + this.channelHint = + getChannelHintOptions( + session.getOptions(), + ThreadLocalRandom.current().nextLong(Long.MAX_VALUE), + session.getSpanner().getOptions().isGrpcGcpExtensionEnabled()); + this.beginTransactionOption = builder.beginTransactionOption; + } + + @Override + public Map getTransactionChannelHint() { + return channelHint; + } + + @Override + protected boolean isRouteToLeader() { + return false; + } + + private boolean shouldUseInlinedBegin() { + return beginTransactionOption == Options.BeginTransactionOption.INLINE; + } + + @Override + void beforeReadOrQuery() { + super.beforeReadOrQuery(); + if (shouldUseInlinedBegin()) { + // Keep the same nested transaction guard as the explicit BeginTransaction path. This checks + // TransactionRunner's thread-local pending state, not the session's active transaction. + SessionImpl.throwIfTransactionsPending(); + } else { + initTransaction(); + } + } + + @Override + @Nullable + TransactionSelector getTransactionSelector() { + if (!shouldUseInlinedBegin()) { + // No need for synchronization: super.readInternal() is always preceded by a check of + // "transactionId" that provides a happens-before from initialization, and the value is + // never changed afterwards. + @SuppressWarnings("GuardedByChecker") + TransactionSelector selector = TransactionSelector.newBuilder().setId(transactionId).build(); + return selector; + } + + ApiFuture futureToWaitFor = null; + txnLock.lock(); + try { + if (transactionId != null) { + return TransactionSelector.newBuilder().setId(transactionId).build(); + } + if (transactionIdFuture == null) { + transactionIdFuture = SettableApiFuture.create(); + return TransactionSelector.newBuilder() + .setBegin(createReadOnlyTransactionOptions()) + .build(); + } + futureToWaitFor = transactionIdFuture; + } finally { + txnLock.unlock(); + } + + try { + return TransactionSelector.newBuilder() + .setId(futureToWaitFor.get(WAIT_FOR_INLINE_BEGIN_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS)) + .build(); + } catch (ExecutionException e) { + throw SpannerExceptionFactory.asSpannerException(e.getCause()); + } catch (TimeoutException e) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.DEADLINE_EXCEEDED, + "Timeout while waiting for an inlined read-only transaction to be returned by another" + + " statement.", + e); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw SpannerExceptionFactory.newSpannerExceptionForCancellation(null, e); + } + } + + @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 { + 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) { + 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. + 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(); + } + } + } + }; + } catch (Throwable throwable) { + if (decremented.compareAndSet(false, true)) { + decrementPendingStartsAndSignal(); + } + 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) { + Options readOptions = Options.fromReadOptions(options); + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; + return createAsyncResultSet( + () -> super.readInternalWithOptions(table, null, keys, columns, readOptions, null), + bufferRows); + } + + @Override + public ListenableAsyncResultSet readUsingIndexAsync( + String table, String index, KeySet keys, Iterable columns, ReadOption... options) { + Options readOptions = Options.fromReadOptions(options); + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; + return createAsyncResultSet( + () -> + super.readInternalWithOptions( + table, checkNotNull(index), keys, columns, readOptions, null), + bufferRows); + } + + @Override + public ListenableAsyncResultSet executeQueryAsync(Statement statement, QueryOption... options) { + Options readOptions = Options.fromQueryOptions(options); + final int bufferRows = + readOptions.hasBufferRows() + ? readOptions.bufferRows() + : AsyncResultSetImpl.DEFAULT_BUFFER_SIZE; + return createAsyncResultSet( + () -> + super.executeQueryInternal( + statement, com.google.spanner.v1.ExecuteSqlRequest.QueryMode.NORMAL, options), + bufferRows); + } + + @Override + public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) { + Timestamp readTimestamp = null; + if (transaction.hasReadTimestamp()) { + try { + readTimestamp = Timestamp.fromProto(transaction.getReadTimestamp()); + } catch (IllegalArgumentException e) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INTERNAL, "Bad value in transaction.read_timestamp metadata field", e); + } + } + if (shouldIncludeId && transaction.getId().isEmpty()) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, NO_TRANSACTION_RETURNED_MSG); + } + txnLock.lock(); + try { + if (timestamp == null) { + if (readTimestamp == null) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field"); + } + timestamp = readTimestamp; + } + if (shouldIncludeId && transactionId == null) { + transactionId = transaction.getId(); + if (transactionIdFuture != null && !transactionIdFuture.isDone()) { + transactionIdFuture.set(transactionId); + } + } + } finally { + txnLock.unlock(); + } + checkAndClose(); + } + + @Override + public SpannerException onError( + SpannerException e, boolean withBeginTransaction, boolean lastStatement) { + e = super.onError(e, withBeginTransaction, lastStatement); + if (withBeginTransaction) { + failTransactionIdFuture(e); + } + return e; + } + + @Override + public void onDone(boolean withBeginTransaction) { + if (withBeginTransaction) { + failTransactionIdFuture( + SpannerExceptionFactory.newSpannerException( + ErrorCode.FAILED_PRECONDITION, + "ResultSet was closed before a read-only transaction id was returned")); + } + super.onDone(withBeginTransaction); + } + + @Override + void onStartFailed(boolean withBeginTransaction, Throwable throwable) { + if (withBeginTransaction) { + failTransactionIdFuture(throwable); + } + } + + private void failTransactionIdFuture(Throwable throwable) { + txnLock.lock(); + try { + if (transactionIdFuture != null && !transactionIdFuture.isDone()) { + transactionIdFuture.setException(throwable); + } + } finally { + txnLock.unlock(); + } + checkAndClose(); + } + + @Override + public Timestamp getReadTimestamp() { + txnLock.lock(); + try { + assertTimestampAvailable(timestamp != null); + return timestamp; + } finally { + txnLock.unlock(); + } + } + + ByteString getTransactionId() { + txnLock.lock(); + try { + return transactionId; + } finally { + txnLock.unlock(); + } + } + + /** + * 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 starts finish. + */ + @Override + public ApiFuture closeAsync() { + boolean shouldCloseImmediately = false; + txnLock.lock(); + try { + if (isClosedOrClosing) { + return closeFuture; + } + isClosedOrClosing = true; + closeFuture = SettableApiFuture.create(); + if (shouldCloseLocked()) { + closed = true; + shouldCloseImmediately = true; + } else { + return closeFuture; + } + } finally { + txnLock.unlock(); + } + + 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() { + TransactionOptions.Builder options = TransactionOptions.newBuilder(); + if (timestamp != null) { + options + .getReadOnlyBuilder() + .setReadTimestamp(timestamp.toProto()) + .setReturnReadTimestamp(true); + } else { + bound.applyToBuilder(options.getReadOnlyBuilder()).setReturnReadTimestamp(true); + } + return options.build(); + } + + /** + * Initializes the transaction with the timestamp specified within MultiUseReadOnlyTransaction. + * This is used only for fallback of PartitionQueryRequest and PartitionReadRequest with + * Multiplexed Session. + */ + void initFallbackTransaction() { + 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(); + + ApiFuture futureToWaitFor = null; + BeginTransactionRequest request = null; + txnLock.lock(); + try { + if (transactionId != null) { + return; + } + 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 = + rpc.beginTransaction(request, getTransactionChannelHint(), isRouteToLeader()); + if (!transaction.hasReadTimestamp()) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field"); + } + if (transaction.getId().isEmpty()) { + throw SpannerExceptionFactory.newSpannerException( + ErrorCode.INTERNAL, "Missing expected transaction.id metadata field"); + } + Timestamp readTimestamp; + try { + readTimestamp = Timestamp.fromProto(transaction.getReadTimestamp()); + } catch (IllegalArgumentException illegalArgumentException) { + throw SpannerExceptionFactory.newSpannerException( + 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(); + } + checkAndClose(); + span.addAnnotation( + "Transaction Creation Done", + ImmutableMap.of( + "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/SessionImpl.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index b61492f5efd5..ef882da9ed4a 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -24,9 +24,6 @@ import com.google.api.gax.rpc.ServerStream; import com.google.cloud.Timestamp; import com.google.cloud.grpc.GcpManagedChannel.ChannelAffinityRef; -import com.google.cloud.spanner.AbstractReadContext.MultiUseReadOnlyTransaction; -import com.google.cloud.spanner.AbstractReadContext.SingleReadContext; -import com.google.cloud.spanner.AbstractReadContext.SingleUseReadOnlyTransaction; import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler; import com.google.cloud.spanner.Options.ReadOnlyTransactionOption; import com.google.cloud.spanner.Options.TransactionOption; diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SingleReadContext.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SingleReadContext.java new file mode 100644 index 000000000000..3e42c3e010d9 --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SingleReadContext.java @@ -0,0 +1,121 @@ +/* + * 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 com.google.common.base.Preconditions.checkState; + +import com.google.cloud.grpc.GcpManagedChannel.ChannelAffinityRef; +import com.google.cloud.spanner.spi.v1.SpannerRpc; +import com.google.cloud.spanner.spi.v1.SpannerRpc.Option; +import com.google.spanner.v1.TransactionOptions; +import com.google.spanner.v1.TransactionSelector; +import java.util.Map; +import java.util.concurrent.ThreadLocalRandom; +import javax.annotation.Nullable; +import javax.annotation.concurrent.GuardedBy; + +/** + * A {@code ReadContext} for standalone reads. This can only be used for a single operation, since + * each standalone read may see a different timestamp of Cloud Spanner data. + */ +class SingleReadContext extends AbstractReadContext { + static class Builder extends AbstractReadContext.Builder { + private TimestampBound bound; + + private Builder() {} + + Builder setTimestampBound(TimestampBound bound) { + this.bound = bound; + return self(); + } + + @Override + SingleReadContext build() { + return new SingleReadContext(this); + } + + SingleUseReadOnlyTransaction buildSingleUseReadOnlyTransaction() { + return new SingleUseReadOnlyTransaction(this); + } + } + + static Builder newBuilder() { + return new Builder(); + } + + final TimestampBound bound; + + @GuardedBy("lock") + private boolean used; + + private Map channelHint; + + SingleReadContext(Builder builder) { + super(builder); + this.bound = builder.bound; + // single use transaction have a single RPC and hence there is no need + // of a channel hint. GAX will automatically choose a hint when used + // with a multiplexed session to perform a round-robin channel selection. We are + // passing a hint here to prefer random channel selection instead of doing GAX round-robin. + this.channelHint = + getChannelHintOptions( + session.getOptions(), + ThreadLocalRandom.current().nextLong(Long.MAX_VALUE), + session.getSpanner().getOptions().isGrpcGcpExtensionEnabled()); + } + + @Override + protected boolean isRouteToLeader() { + return false; + } + + @GuardedBy("lock") + @Override + void beforeReadOrQueryLocked() { + super.beforeReadOrQueryLocked(); + checkState(!used, "Cannot use a single-read ReadContext for multiple reads"); + used = true; + } + + @Override + @Nullable + TransactionSelector getTransactionSelector() { + if (bound.getMode() == TimestampBound.Mode.STRONG) { + // Default mode: no need to specify a transaction. + return null; + } + return TransactionSelector.newBuilder() + .setSingleUse(TransactionOptions.newBuilder().setReadOnly(bound.toProto())) + .build(); + } + + @Override + Map getTransactionChannelHint() { + return channelHint; + } + + @Override + boolean prepareRetryOnDifferentGrpcChannel() { + ChannelAffinityRef channelAffinityRef = + Option.CHANNEL_ID_AFFINITY.getChannelAffinityRef(channelHint); + if (session.getIsMultiplexed() && channelAffinityRef != null) { + channelAffinityRef.useDifferentChannelOnNextCall(); + return true; + } + return super.prepareRetryOnDifferentGrpcChannel(); + } +} diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SingleUseReadOnlyTransaction.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SingleUseReadOnlyTransaction.java new file mode 100644 index 000000000000..e3002ba7795a --- /dev/null +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SingleUseReadOnlyTransaction.java @@ -0,0 +1,73 @@ +/* + * 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 com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException; + +import com.google.cloud.Timestamp; +import com.google.spanner.v1.Transaction; +import com.google.spanner.v1.TransactionOptions; +import com.google.spanner.v1.TransactionSelector; +import javax.annotation.Nullable; +import javax.annotation.concurrent.GuardedBy; + +/** + * A {@link ReadOnlyTransaction} implementation for single-use transactions. This can only be used + * for a single read or query operation, and provides access to the read timestamp chosen by Cloud + * Spanner for that operation. + */ +class SingleUseReadOnlyTransaction extends SingleReadContext implements ReadOnlyTransaction { + + @GuardedBy("lock") + private Timestamp timestamp; + + SingleUseReadOnlyTransaction(SingleReadContext.Builder builder) { + super(builder); + } + + @Override + public Timestamp getReadTimestamp() { + synchronized (lock) { + assertTimestampAvailable(timestamp != null); + return timestamp; + } + } + + @Override + @Nullable + TransactionSelector getTransactionSelector() { + TransactionOptions.Builder options = TransactionOptions.newBuilder(); + bound.applyToBuilder(options.getReadOnlyBuilder()).setReturnReadTimestamp(true); + return TransactionSelector.newBuilder().setSingleUse(options).build(); + } + + @Override + public void onTransactionMetadata(Transaction transaction, boolean shouldIncludeId) { + synchronized (lock) { + if (!transaction.hasReadTimestamp()) { + throw newSpannerException( + ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field"); + } + try { + timestamp = Timestamp.fromProto(transaction.getReadTimestamp()); + } catch (IllegalArgumentException e) { + throw newSpannerException( + ErrorCode.INTERNAL, "Bad value in transaction.read_timestamp metadata field", e); + } + } + } +} 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 034af24f8cdd..7a1791386a17 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 @@ -577,7 +577,7 @@ public void closeAsyncWhileBeginTransactionIsInFlightWaitsForBeginTransactionToF () -> { initStartedLatch.countDown(); try { - ((AbstractReadContext.MultiUseReadOnlyTransaction) transaction).initTransaction(); + ((MultiUseReadOnlyTransaction) transaction).initTransaction(); } finally { initFinishedLatch.countDown(); }