From 0cd7e3aaeb3157d00f8fb5a527daba62490f48ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Fri, 12 Dec 2025 15:30:50 +0100 Subject: [PATCH] feat: include RequestID in requests and errors - Send a RequestID to Spanner for each request - Make sure that the attempt number of the RequestID is increased if the RPC is retried. - Include the RequestID in every error that is thrown due to an error that is returned by Spanner. --- .../clirr-ignored-differences.xml | 57 +- .../cloud/spanner/AbortedException.java | 7 +- .../cloud/spanner/AbstractReadContext.java | 33 +- ...minRequestsPerMinuteExceededException.java | 7 +- .../google/cloud/spanner/BatchClientImpl.java | 10 +- .../cloud/spanner/BuiltInMetricsConstant.java | 8 +- .../spanner/DatabaseNotFoundException.java | 7 +- .../spanner/InstanceNotFoundException.java | 7 +- .../MissingDefaultSequenceKindException.java | 5 +- .../MultiplexedSessionDatabaseClient.java | 2 +- .../spanner/PartitionedDmlTransaction.java | 47 +- .../spanner/ResumableStreamIterator.java | 37 +- .../google/cloud/spanner/SessionClient.java | 41 +- .../com/google/cloud/spanner/SessionImpl.java | 55 +- .../spanner/SessionNotFoundException.java | 7 +- .../cloud/spanner/SpannerApiFutures.java | 5 +- .../spanner/SpannerBatchUpdateException.java | 5 +- .../cloud/spanner/SpannerException.java | 53 +- .../spanner/SpannerExceptionFactory.java | 101 +-- .../com/google/cloud/spanner/SpannerImpl.java | 8 + .../cloud/spanner/SpannerRetryHelper.java | 3 +- ...sactionMutationLimitExceededException.java | 5 +- .../cloud/spanner/TransactionRunnerImpl.java | 40 +- .../cloud/spanner/XGoogSpannerRequestId.java | 81 +- .../cloud/spanner/connection/DdlBatch.java | 2 +- .../cloud/spanner/spi/v1/GapicSpannerRpc.java | 144 ++-- .../spanner/spi/v1/HeaderInterceptor.java | 4 +- .../spanner/spi/v1/RequestIdCreatorImpl.java | 43 ++ .../spanner/spi/v1/RequestIdInterceptor.java | 57 ++ .../spi/v1/SpannerErrorInterceptor.java | 19 + .../spi/v1/SpannerInterceptorProvider.java | 1 + .../cloud/spanner/spi/v1/SpannerRpc.java | 16 +- .../cloud/spanner/BaseSessionPoolTest.java | 2 + .../cloud/spanner/DatabaseClientImplTest.java | 120 +-- .../PartitionedDmlTransactionTest.java | 59 +- .../spanner/RequestIdMockServerTest.java | 727 ++++++++++++++++++ .../spanner/ResumableStreamIteratorTest.java | 6 +- .../cloud/spanner/SessionClientTests.java | 11 +- .../google/cloud/spanner/SessionImplTest.java | 6 +- .../google/cloud/spanner/SessionPoolTest.java | 12 +- .../SpannerCloudMonitoringExporterTest.java | 8 +- .../spanner/SpannerExceptionFactoryTest.java | 13 - .../spanner/TransactionContextImplTest.java | 8 +- .../spanner/TransactionManagerImplTest.java | 4 +- .../spanner/TransactionRunnerImplTest.java | 8 +- .../spanner/XGoogSpannerRequestIdTest.java | 54 +- .../RunTransactionMockServerTest.java | 2 +- 47 files changed, 1392 insertions(+), 565 deletions(-) create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdCreatorImpl.java create mode 100644 google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdInterceptor.java create mode 100644 google-cloud-spanner/src/test/java/com/google/cloud/spanner/RequestIdMockServerTest.java diff --git a/google-cloud-spanner/clirr-ignored-differences.xml b/google-cloud-spanner/clirr-ignored-differences.xml index 21b2eb228cb..a57bf40d1ba 100644 --- a/google-cloud-spanner/clirr-ignored-differences.xml +++ b/google-cloud-spanner/clirr-ignored-differences.xml @@ -1087,9 +1087,56 @@ 7002 com/google/cloud/spanner/CompositeTracer void recordGFELatency(java.lang.Float) - - 7002 - com/google/cloud/spanner/CompositeTracer - void recordGfeHeaderMissingCount(java.lang.Long) - + + + 7002 + com/google/cloud/spanner/CompositeTracer + void recordGfeHeaderMissingCount(java.lang.Long) + + + + 7002 + com/google/cloud/spanner/SpannerException + void setRequestId(com.google.cloud.spanner.XGoogSpannerRequestId) + + + 7002 + com/google/cloud/spanner/SpannerExceptionFactory + com.google.cloud.spanner.SpannerBatchUpdateException newSpannerBatchUpdateException(com.google.cloud.spanner.ErrorCode, java.lang.String, long[], com.google.cloud.spanner.XGoogSpannerRequestId) + + + 7002 + com/google/cloud/spanner/SpannerExceptionFactory + com.google.cloud.spanner.SpannerException newSpannerException(com.google.cloud.spanner.ErrorCode, java.lang.String, java.lang.Throwable, com.google.cloud.spanner.XGoogSpannerRequestId) + + + 7002 + com/google/cloud/spanner/SpannerExceptionFactory + com.google.cloud.spanner.SpannerException newSpannerException(com.google.cloud.spanner.ErrorCode, java.lang.String, com.google.cloud.spanner.XGoogSpannerRequestId) + + + 7002 + com/google/cloud/spanner/SpannerExceptionFactory + com.google.cloud.spanner.SpannerException newSpannerException(java.lang.Throwable, com.google.cloud.spanner.XGoogSpannerRequestId) + + + 7002 + com/google/cloud/spanner/SpannerExceptionFactory + com.google.cloud.spanner.SpannerException newSpannerException(io.grpc.Context, java.lang.Throwable, com.google.cloud.spanner.XGoogSpannerRequestId) + + + 7002 + com/google/cloud/spanner/SpannerExceptionFactory + com.google.cloud.spanner.SpannerException propagateInterrupt(java.lang.InterruptedException, com.google.cloud.spanner.XGoogSpannerRequestId) + + + 6001 + com/google/cloud/spanner/XGoogSpannerRequestId + REQUEST_HEADER_KEY + + + 6001 + com/google/cloud/spanner/XGoogSpannerRequestId + REQUEST_ID + diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbortedException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbortedException.java index 74e45062113..28b5f1fa257 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbortedException.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbortedException.java @@ -38,7 +38,7 @@ public class AbortedException extends SpannerException { /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ AbortedException( DoNotConstructDirectly token, @Nullable String message, @Nullable Throwable cause) { - this(token, message, cause, null, null); + this(token, message, cause, null); } /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ @@ -46,9 +46,8 @@ public class AbortedException extends SpannerException { DoNotConstructDirectly token, @Nullable String message, @Nullable Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId reqId) { - super(token, ErrorCode.ABORTED, IS_RETRYABLE, message, cause, apiException, reqId); + @Nullable ApiException apiException) { + super(token, ErrorCode.ABORTED, IS_RETRYABLE, message, cause, apiException); if (cause instanceof AbortedException) { this.transactionID = ((AbortedException) cause).getTransactionID(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java index ffbdfc1cfd2..289acb1a745 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java @@ -457,30 +457,22 @@ void initTransaction() { } private void initTransactionInternal(BeginTransactionRequest request) { - XGoogSpannerRequestId reqId = - session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); try { Transaction transaction = - rpc.beginTransaction( - request, reqId.withOptions(getTransactionChannelHint()), isRouteToLeader()); + rpc.beginTransaction(request, getTransactionChannelHint(), isRouteToLeader()); if (!transaction.hasReadTimestamp()) { throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, - "Missing expected transaction.read_timestamp metadata field", - reqId); + ErrorCode.INTERNAL, "Missing expected transaction.read_timestamp metadata field"); } if (transaction.getId().isEmpty()) { throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, "Missing expected transaction.id metadata field", reqId); + ErrorCode.INTERNAL, "Missing expected transaction.id metadata field"); } try { timestamp = Timestamp.fromProto(transaction.getReadTimestamp()); } catch (IllegalArgumentException e) { throw SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, - "Bad value in transaction.read_timestamp metadata field", - e, - reqId); + ErrorCode.INTERNAL, "Bad value in transaction.read_timestamp metadata field", e); } transactionId = transaction.getId(); span.addAnnotation( @@ -816,7 +808,8 @@ ResultSet executeQueryInternalWithOptions( @Override CloseableIterator startStream( @Nullable ByteString resumeToken, - AsyncResultSet.StreamMessageListener streamListener) { + AsyncResultSet.StreamMessageListener streamListener, + XGoogSpannerRequestId requestId) { GrpcStreamIterator stream = new GrpcStreamIterator( statement, @@ -839,12 +832,12 @@ CloseableIterator startStream( if (selector != null) { request.setTransaction(selector); } - this.ensureNonNullXGoogRequestId(); SpannerRpc.StreamingCall call = rpc.executeQuery( request.build(), stream.consumer(), - this.xGoogRequestId.withOptions(getTransactionChannelHint()), + getTransactionChannelHint(), + requestId, isRouteToLeader()); session.markUsed(clock.instant()); stream.setCall(call, request.getTransaction().hasBegin()); @@ -860,7 +853,7 @@ boolean prepareIteratorForRetryOnDifferentGrpcChannel() { stream, this, options.hasDecodeMode() ? options.decodeMode() : defaultDecodeMode); } - Map getChannelHintOptions( + static Map getChannelHintOptions( Map channelHintForSession, Long channelHintForTransaction) { if (channelHintForSession != null) { return channelHintForSession; @@ -1030,7 +1023,8 @@ ResultSet readInternalWithOptions( @Override CloseableIterator startStream( @Nullable ByteString resumeToken, - AsyncResultSet.StreamMessageListener streamListener) { + AsyncResultSet.StreamMessageListener streamListener, + XGoogSpannerRequestId requestId) { GrpcStreamIterator stream = new GrpcStreamIterator( lastStatement, prefetchChunks, cancelQueryWhenClientIsClosed); @@ -1048,13 +1042,12 @@ CloseableIterator startStream( builder.setTransaction(selector); } builder.setRequestOptions(buildRequestOptions(readOptions)); - this.incrementXGoogRequestIdAttempt(); - this.xGoogRequestId.setChannelId(session.getChannel()); SpannerRpc.StreamingCall call = rpc.read( builder.build(), stream.consumer(), - this.xGoogRequestId.withOptions(getTransactionChannelHint()), + getTransactionChannelHint(), + requestId, isRouteToLeader()); session.markUsed(clock.instant()); stream.setCall(call, /* withBeginTransaction= */ builder.getTransaction().hasBegin()); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AdminRequestsPerMinuteExceededException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AdminRequestsPerMinuteExceededException.java index b7d76ab9e82..72d8b0ab15d 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AdminRequestsPerMinuteExceededException.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/AdminRequestsPerMinuteExceededException.java @@ -32,7 +32,7 @@ public class AdminRequestsPerMinuteExceededException extends SpannerException { /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ AdminRequestsPerMinuteExceededException( DoNotConstructDirectly token, @Nullable String message, @Nullable Throwable cause) { - this(token, message, cause, null, null); + this(token, message, cause, null); } /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ @@ -40,8 +40,7 @@ public class AdminRequestsPerMinuteExceededException extends SpannerException { DoNotConstructDirectly token, @Nullable String message, @Nullable Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId reqId) { - super(token, ErrorCode.RESOURCE_EXHAUSTED, true, message, cause, apiException, reqId); + @Nullable ApiException apiException) { + super(token, ErrorCode.RESOURCE_EXHAUSTED, true, message, cause, apiException); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java index 6e8340784bd..86aea9ef989 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BatchClientImpl.java @@ -250,11 +250,9 @@ private List partitionReadUsingIndex( } builder.setPartitionOptions(pbuilder.build()); - XGoogSpannerRequestId reqId = - session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); final PartitionReadRequest request = builder.build(); try { - PartitionResponse response = rpc.partitionRead(request, reqId.withOptions(options)); + PartitionResponse response = rpc.partitionRead(request, options); ImmutableList.Builder partitions = ImmutableList.builder(); for (com.google.spanner.v1.Partition p : response.getPartitionsList()) { Partition partition = @@ -274,7 +272,6 @@ private List partitionReadUsingIndex( return partitionReadUsingIndex( partitionOptions, table, index, keys, columns, true, option); } - e.setRequestId(reqId); throw e; } } @@ -316,11 +313,9 @@ private List partitionQuery( } builder.setPartitionOptions(pbuilder.build()); - XGoogSpannerRequestId reqId = - session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); final PartitionQueryRequest request = builder.build(); try { - PartitionResponse response = rpc.partitionQuery(request, reqId.withOptions(options)); + PartitionResponse response = rpc.partitionQuery(request, options); ImmutableList.Builder partitions = ImmutableList.builder(); for (com.google.spanner.v1.Partition p : response.getPartitionsList()) { Partition partition = @@ -333,7 +328,6 @@ private List partitionQuery( if (!isFallback && maybeMarkUnimplementedForPartitionedOps(e)) { return partitionQuery(partitionOptions, statement, true, option); } - e.setRequestId(reqId); throw e; } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java index bd78149465f..d210bd85b79 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/BuiltInMetricsConstant.java @@ -16,7 +16,7 @@ package com.google.cloud.spanner; -import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID; +import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME; import com.google.api.core.InternalApi; import com.google.api.gax.tracing.OpenTelemetryMetricsRecorder; @@ -98,10 +98,12 @@ public class BuiltInMetricsConstant { AttributeKey.stringKey("directpath_enabled"); public static final AttributeKey DIRECT_PATH_USED_KEY = AttributeKey.stringKey("directpath_used"); - public static final AttributeKey REQUEST_ID_KEY = AttributeKey.stringKey(REQUEST_ID); + public static final AttributeKey REQUEST_ID_KEY = + AttributeKey.stringKey(REQUEST_ID_HEADER_NAME); public static final AttributeKey GRPC_XDS_RESOURCE_TYPE_KEY = AttributeKey.stringKey("grpc.xds.resource_type"); - public static Set ALLOWED_EXEMPLARS_ATTRIBUTES = new HashSet<>(Arrays.asList(REQUEST_ID)); + public static Set ALLOWED_EXEMPLARS_ATTRIBUTES = + new HashSet<>(Arrays.asList(REQUEST_ID_HEADER_NAME)); // IP address prefixes allocated for DirectPath backends. public static final String DP_IPV6_PREFIX = "2001:4860:8040"; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseNotFoundException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseNotFoundException.java index 474acb04890..cc4a2e32f0b 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseNotFoundException.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/DatabaseNotFoundException.java @@ -35,7 +35,7 @@ public class DatabaseNotFoundException extends ResourceNotFoundException { @Nullable String message, ResourceInfo resourceInfo, @Nullable Throwable cause) { - this(token, message, resourceInfo, cause, null, null); + this(token, message, resourceInfo, cause, null); } /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ @@ -44,8 +44,7 @@ public class DatabaseNotFoundException extends ResourceNotFoundException { @Nullable String message, ResourceInfo resourceInfo, @Nullable Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId reqId) { - super(token, message, resourceInfo, cause, apiException, reqId); + @Nullable ApiException apiException) { + super(token, message, resourceInfo, cause, apiException); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceNotFoundException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceNotFoundException.java index 5139cd3894e..dc4192e109f 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceNotFoundException.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/InstanceNotFoundException.java @@ -35,7 +35,7 @@ public class InstanceNotFoundException extends ResourceNotFoundException { @Nullable String message, ResourceInfo resourceInfo, @Nullable Throwable cause) { - this(token, message, resourceInfo, cause, null, null); + this(token, message, resourceInfo, cause, null); } /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ @@ -44,8 +44,7 @@ public class InstanceNotFoundException extends ResourceNotFoundException { @Nullable String message, ResourceInfo resourceInfo, @Nullable Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId reqId) { - super(token, message, resourceInfo, cause, apiException, reqId); + @Nullable ApiException apiException) { + super(token, message, resourceInfo, cause, apiException); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MissingDefaultSequenceKindException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MissingDefaultSequenceKindException.java index 80e0ac7efaa..d29a9489c3e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MissingDefaultSequenceKindException.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MissingDefaultSequenceKindException.java @@ -38,9 +38,8 @@ public class MissingDefaultSequenceKindException extends SpannerException { ErrorCode errorCode, String message, Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId reqId) { - super(token, errorCode, /* retryable= */ false, message, cause, apiException, reqId); + @Nullable ApiException apiException) { + super(token, errorCode, /* retryable= */ false, message, cause, apiException); } static boolean isMissingDefaultSequenceKindException(Throwable cause) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java index 4b092bc9866..11e83add517 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/MultiplexedSessionDatabaseClient.java @@ -745,7 +745,7 @@ public AsyncTransactionManager transactionManagerAsync(TransactionOption... opti @Override public long executePartitionedUpdate(Statement stmt, UpdateOption... options) { - return createMultiplexedSessionTransaction(/* singleUse= */ true) + return createMultiplexedSessionTransaction(/* singleUse= */ false) .executePartitionedUpdate(stmt, options); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDmlTransaction.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDmlTransaction.java index 5f0d497c74c..394b8bfbd9e 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDmlTransaction.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/PartitionedDmlTransaction.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner; +import static com.google.cloud.spanner.AbstractReadContext.getChannelHintOptions; import static com.google.common.base.Preconditions.checkState; import com.google.api.core.InternalApi; @@ -42,6 +43,7 @@ import java.time.Duration; import java.time.temporal.ChronoUnit; import java.util.Map; +import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import java.util.logging.Level; import java.util.logging.Logger; @@ -56,12 +58,16 @@ public class PartitionedDmlTransaction implements SessionImpl.SessionTransaction private final Ticker ticker; private final IsRetryableInternalError isRetryableInternalErrorPredicate; private volatile boolean isValid = true; + private final Map channelHintOptions; PartitionedDmlTransaction(SessionImpl session, SpannerRpc rpc, Ticker ticker) { this.session = session; this.rpc = rpc; this.ticker = ticker; this.isRetryableInternalErrorPredicate = new IsRetryableInternalError(); + this.channelHintOptions = + getChannelHintOptions( + session.getOptions(), ThreadLocalRandom.current().nextLong(Long.MAX_VALUE)); } /** @@ -79,15 +85,13 @@ long executeStreamingPartitionedUpdate( boolean foundStats = false; long updateCount = 0L; Stopwatch stopwatch = Stopwatch.createStarted(ticker); - XGoogSpannerRequestId reqId = - session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); - UpdateOption[] allOptions = new UpdateOption[updateOptions.length + 1]; - System.arraycopy(updateOptions, 0, allOptions, 0, updateOptions.length); - allOptions[allOptions.length - 1] = new Options.RequestIdOption(reqId); - Options options = Options.fromUpdateOptions(allOptions); + Options options = Options.fromUpdateOptions(updateOptions); try { ExecuteSqlRequest request = newTransactionRequestFrom(statement, options); + // The channel ID is set to zero here. It will be filled in later by SpannerRpc when it reads + // the channel hint from the options that are passed in. + XGoogSpannerRequestId requestId = this.session.getRequestIdCreator().nextRequestId(0); while (true) { final Duration remainingTimeout = tryUpdateTimeout(timeout, stopwatch); @@ -95,7 +99,7 @@ long executeStreamingPartitionedUpdate( try { ServerStream stream = rpc.executeStreamingPartitionedDml( - request, reqId.withOptions(session.getOptions()), remainingTimeout); + request, channelHintOptions, requestId, remainingTimeout); for (PartialResultSet rs : stream) { if (rs.getResumeToken() != null && !rs.getResumeToken().isEmpty()) { @@ -110,8 +114,12 @@ long executeStreamingPartitionedUpdate( } catch (UnavailableException e) { LOGGER.log( Level.FINER, "Retrying PartitionedDml transaction after UnavailableException", e); - reqId.incrementAttempt(); request = resumeOrRestartRequest(resumeToken, statement, request, options); + if (resumeToken.isEmpty()) { + // Create a new xGoogSpannerRequestId if there is no resume token, as that means that + // the entire transaction will be retried. + requestId = session.getRequestIdCreator().nextRequestId(session.getChannel()); + } } catch (InternalException e) { if (!isRetryableInternalErrorPredicate.apply(e)) { throw e; @@ -119,8 +127,12 @@ long executeStreamingPartitionedUpdate( LOGGER.log( Level.FINER, "Retrying PartitionedDml transaction after InternalException - EOS", e); - reqId.incrementAttempt(); request = resumeOrRestartRequest(resumeToken, statement, request, options); + if (resumeToken.isEmpty()) { + // Create a new xGoogSpannerRequestId if there is no resume token, as that means that + // the entire transaction will be retried. + requestId = session.getRequestIdCreator().nextRequestId(session.getChannel()); + } } catch (AbortedException e) { LOGGER.log(Level.FINER, "Retrying PartitionedDml transaction after AbortedException", e); resumeToken = ByteString.EMPTY; @@ -128,22 +140,18 @@ long executeStreamingPartitionedUpdate( updateCount = 0L; request = newTransactionRequestFrom(statement, options); // Create a new xGoogSpannerRequestId. - reqId = session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); - } catch (SpannerException e) { - e.setRequestId(reqId); - throw e; + requestId = session.getRequestIdCreator().nextRequestId(session.getChannel()); } } if (!foundStats) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.INVALID_ARGUMENT, - "Partitioned DML response missing stats possibly due to non-DML statement as input", - reqId); + "Partitioned DML response missing stats possibly due to non-DML statement as input"); } LOGGER.log(Level.FINER, "Finished PartitionedUpdate statement"); return updateCount; } catch (Exception e) { - throw SpannerExceptionFactory.newSpannerException(e, reqId); + throw SpannerExceptionFactory.asSpannerException(e); } } @@ -223,14 +231,11 @@ private ByteString initTransaction(final Options options) { .setExcludeTxnFromChangeStreams( options.withExcludeTxnFromChangeStreams() == Boolean.TRUE)) .build(); - XGoogSpannerRequestId reqId = - session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); - Transaction tx = rpc.beginTransaction(request, reqId.withOptions(session.getOptions()), true); + Transaction tx = rpc.beginTransaction(request, channelHintOptions, true); if (tx.getId().isEmpty()) { throw SpannerExceptionFactory.newSpannerException( ErrorCode.INTERNAL, - "Failed to init transaction, missing transaction id\n" + session.getName(), - reqId); + "Failed to init transaction, missing transaction id\n" + session.getName()); } return tx.getId(); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java index d292e53fe5b..aac7f63c861 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java @@ -69,10 +69,10 @@ abstract class ResumableStreamIterator extends AbstractIterator stream; + private int attempts; private ByteString resumeToken; private boolean finished; - public XGoogSpannerRequestId xGoogRequestId; - private XGoogSpannerRequestId.RequestIdCreator xGoogRequestIdCreator; + private final XGoogSpannerRequestId requestId; /** * Indicates whether it is currently safe to retry RPCs. This will be {@code false} if we have @@ -119,7 +119,8 @@ protected ResumableStreamIterator( this.errorHandler = errorHandler; this.streamingRetrySettings = Preconditions.checkNotNull(streamingRetrySettings); this.retryableCodes = Preconditions.checkNotNull(retryableCodes); - this.xGoogRequestIdCreator = xGoogRequestIdCreator; + // The channel is automatically updated by the gRPC client when the request is actually sent. + this.requestId = xGoogRequestIdCreator.nextRequestId(0); } private ExponentialBackOff newBackOff() { @@ -187,27 +188,15 @@ private void backoffSleep(Context context, long backoffMillis) throws SpannerExc } if (latch.await(backoffMillis, TimeUnit.MILLISECONDS)) { // Woken by context cancellation. - throw newSpannerExceptionForCancellation(context, null, this.xGoogRequestId); + throw newSpannerExceptionForCancellation(context, null); } } catch (InterruptedException interruptExcept) { - throw newSpannerExceptionForCancellation(context, interruptExcept, this.xGoogRequestId); + throw newSpannerExceptionForCancellation(context, interruptExcept); } finally { context.removeListener(listener); } } - public void ensureNonNullXGoogRequestId() { - if (this.xGoogRequestId == null) { - this.xGoogRequestId = - this.xGoogRequestIdCreator.nextRequestId(1 /*TODO: infer channelId*/, 1 /*attempt*/); - } - } - - public void incrementXGoogRequestIdAttempt() { - this.ensureNonNullXGoogRequestId(); - this.xGoogRequestId.incrementAttempt(); - } - private enum DirectExecutor implements Executor { INSTANCE; @@ -218,7 +207,9 @@ public void execute(Runnable command) { } abstract CloseableIterator startStream( - @Nullable ByteString resumeToken, AsyncResultSet.StreamMessageListener streamMessageListener); + @Nullable ByteString resumeToken, + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId); /** * Prepares the iterator for a retry on a different gRPC channel. Returns true if that is @@ -304,7 +295,6 @@ protected PartialResultSet computeNext() { } assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken); stream = null; - incrementXGoogRequestIdAttempt(); try (IScope s = tracer.withSpan(span)) { long delay = spannerException.getRetryDelayInMillis(); if (delay != -1) { @@ -326,14 +316,12 @@ protected PartialResultSet computeNext() { if (++numAttemptsOnOtherChannel < errorHandler.getMaxAttempts() && prepareIteratorForRetryOnDifferentGrpcChannel()) { stream = null; - xGoogRequestId = null; continue; } } } span.addAnnotation("Stream broken. Not safe to retry", spannerException); span.setStatus(spannerException); - spannerException.setRequestId(this.xGoogRequestId); throw spannerException; } catch (RuntimeException e) { span.addAnnotation("Stream broken. Not safe to retry", e); @@ -352,13 +340,8 @@ private void startGrpcStreaming() { try (IScope scope = tracer.withSpan(span)) { // When start a new stream set the Span as current to make the gRPC Span a child of // this Span. - stream = checkNotNull(startStream(resumeToken, streamMessageListener)); + stream = checkNotNull(startStream(resumeToken, streamMessageListener, requestId)); stream.requestPrefetchChunks(); - if (this.xGoogRequestId == null) { - this.xGoogRequestId = - this.xGoogRequestIdCreator.nextRequestId( - 1 /* channelId shall be replaced by the instantiated class. */, 0); - } } } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java index 6592eb74991..a1852b90881 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionClient.java @@ -36,7 +36,7 @@ import javax.annotation.concurrent.GuardedBy; /** Client for creating single sessions and batches of sessions. */ -class SessionClient implements AutoCloseable, XGoogSpannerRequestId.RequestIdCreator { +class SessionClient implements AutoCloseable { static class SessionId { private static final PathTemplate NAME_TEMPLATE = PathTemplate.create( @@ -110,15 +110,8 @@ Object value() { return ImmutableMap.copyOf(tmp); } - static Map createRequestOptions( - long channelId, XGoogSpannerRequestId requestId) { - return ImmutableMap.of( - Option.CHANNEL_HINT, channelId, - Option.REQUEST_ID, requestId); - } - - static Map createRequestOptions(XGoogSpannerRequestId requestId) { - return ImmutableMap.of(Option.REQUEST_ID, requestId); + static Map createRequestOptions(long channelId) { + return ImmutableMap.of(Option.CHANNEL_HINT, channelId); } private final class BatchCreateSessionsRunnable implements Runnable { @@ -220,12 +213,6 @@ DatabaseId getDatabaseId() { return db; } - @Override - public XGoogSpannerRequestId nextRequestId(long channelId, int attempt) { - return XGoogSpannerRequestId.of( - this.nthId, channelId, this.nthRequest.incrementAndGet(), attempt); - } - /** Create a single session. */ SessionImpl createSession() { // The sessionChannelCounter could overflow, but that will just flip it to Integer.MIN_VALUE, @@ -235,7 +222,6 @@ SessionImpl createSession() { channelId = sessionChannelCounter; sessionChannelCounter++; } - XGoogSpannerRequestId reqId = nextRequestId(channelId, 1); ISpan span = spanner.getTracer().spanBuilder(SpannerImpl.CREATE_SESSION, this.databaseAttributes); try (IScope s = spanner.getTracer().withSpan(span)) { @@ -246,16 +232,14 @@ SessionImpl createSession() { db.getName(), spanner.getOptions().getDatabaseRole(), spanner.getOptions().getSessionLabels(), - createRequestOptions(channelId, reqId)); + createRequestOptions(channelId)); SessionReference sessionReference = new SessionReference( session.getName(), session.getCreateTime(), session.getMultiplexed(), optionMap(SessionOption.channelHint(channelId))); - SessionImpl sessionImpl = new SessionImpl(spanner, sessionReference); - sessionImpl.setRequestIdCreator(this); - return sessionImpl; + return new SessionImpl(spanner, sessionReference); } catch (RuntimeException e) { span.setStatus(e); throw e; @@ -291,9 +275,6 @@ SessionImpl createMultiplexedSession() { spanner .getTracer() .spanBuilder(SpannerImpl.CREATE_MULTIPLEXED_SESSION, this.databaseAttributes); - // MultiplexedSession doesn't use a channelId hence this hard-coded value. - int channelId = 0; - XGoogSpannerRequestId reqId = nextRequestId(channelId, 1); try (IScope s = spanner.getTracer().withSpan(span)) { com.google.spanner.v1.Session session = spanner @@ -302,14 +283,13 @@ SessionImpl createMultiplexedSession() { db.getName(), spanner.getOptions().getDatabaseRole(), spanner.getOptions().getSessionLabels(), - createRequestOptions(reqId), + null, true); SessionImpl sessionImpl = new SessionImpl( spanner, new SessionReference( session.getName(), session.getCreateTime(), session.getMultiplexed(), null)); - sessionImpl.setRequestIdCreator(this); span.addAnnotation( String.format("Request for %d multiplexed session returned %d session", 1, 1)); return sessionImpl; @@ -423,8 +403,6 @@ private List internalBatchCreateSessions( .spanBuilderWithExplicitParent(SpannerImpl.BATCH_CREATE_SESSIONS_REQUEST, parent); span.addAnnotation(String.format("Requesting %d sessions", sessionCount)); try (IScope s = spanner.getTracer().withSpan(span)) { - XGoogSpannerRequestId reqId = - XGoogSpannerRequestId.of(this.nthId, channelHint, this.nthRequest.incrementAndGet(), 1); List sessions = spanner .getRpc() @@ -433,7 +411,7 @@ private List internalBatchCreateSessions( sessionCount, spanner.getOptions().getDatabaseRole(), spanner.getOptions().getSessionLabels(), - createRequestOptions(channelHint, reqId)); + createRequestOptions(channelHint)); span.addAnnotation( String.format( "Request for %d sessions returned %d sessions", sessionCount, sessions.size())); @@ -448,7 +426,6 @@ private List internalBatchCreateSessions( session.getCreateTime(), session.getMultiplexed(), optionMap(SessionOption.channelHint(channelHint)))); - sessionImpl.setRequestIdCreator(this); res.add(sessionImpl); } return res; @@ -465,8 +442,6 @@ SessionImpl sessionWithId(String name) { synchronized (this) { options = optionMap(SessionOption.channelHint(sessionChannelCounter++)); } - SessionImpl sessionImpl = new SessionImpl(spanner, new SessionReference(name, options)); - sessionImpl.setRequestIdCreator(this); - return sessionImpl; + return new SessionImpl(spanner, new SessionReference(name, options)); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java index 0cd42255484..d86ff807afd 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionImpl.java @@ -127,31 +127,18 @@ interface SessionTransaction { private final Clock clock; private final Map options; private final ErrorHandler errorHandler; - private XGoogSpannerRequestId.RequestIdCreator requestIdCreator; SessionImpl(SpannerImpl spanner, SessionReference sessionReference) { this(spanner, sessionReference, NO_CHANNEL_HINT); } SessionImpl(SpannerImpl spanner, SessionReference sessionReference, int channelHint) { - this(spanner, sessionReference, channelHint, new XGoogSpannerRequestId.NoopRequestIdCreator()); - } - - SessionImpl( - SpannerImpl spanner, - SessionReference sessionReference, - int channelHint, - XGoogSpannerRequestId.RequestIdCreator requestIdCreator) { this.spanner = spanner; this.tracer = spanner.getTracer(); this.sessionReference = sessionReference; this.clock = spanner.getOptions().getSessionPoolOptions().getPoolMaintainerClock(); this.options = createOptions(sessionReference, channelHint); this.errorHandler = createErrorHandler(spanner.getOptions()); - this.requestIdCreator = requestIdCreator; - if (this.requestIdCreator == null) { - throw new IllegalStateException("requestIdCreator must be non-null"); - } } static Map createOptions( @@ -307,12 +294,7 @@ public CommitResponse writeAtLeastOnceWithOptions( try (IScope s = tracer.withSpan(span)) { return SpannerRetryHelper.runTxWithRetriesOnAborted( - () -> { - // On Aborted, we have to start a fresh request id. - final XGoogSpannerRequestId reqId = reqIdOrFresh(options); - return new CommitResponse( - spanner.getRpc().commit(request, reqId.withOptions(getOptions()))); - }); + () -> new CommitResponse(spanner.getRpc().commit(request, getOptions()))); } catch (RuntimeException e) { span.setStatus(e); throw e; @@ -321,14 +303,6 @@ public CommitResponse writeAtLeastOnceWithOptions( } } - private XGoogSpannerRequestId reqIdOrFresh(Options options) { - XGoogSpannerRequestId reqId = options.reqId(); - if (reqId == null) { - reqId = this.getRequestIdCreator().nextRequestId(this.getChannel(), 1); - } - return reqId; - } - private RequestOptions getRequestOptions(TransactionOption... transactionOptions) { Options requestOptions = Options.fromTransactionOptions(transactionOptions); if (requestOptions.hasPriority() || requestOptions.hasTag()) { @@ -357,7 +331,6 @@ public ServerStream batchWriteAtLeastOnce( .addAllMutationGroups(mutationGroupsProto); RequestOptions batchWriteRequestOptions = getRequestOptions(transactionOptions); Options allOptions = Options.fromTransactionOptions(transactionOptions); - final XGoogSpannerRequestId reqId = reqIdOrFresh(allOptions); if (batchWriteRequestOptions != null) { requestBuilder.setRequestOptions(batchWriteRequestOptions); } @@ -366,9 +339,7 @@ public ServerStream batchWriteAtLeastOnce( } ISpan span = tracer.spanBuilder(SpannerImpl.BATCH_WRITE); try (IScope s = tracer.withSpan(span)) { - return spanner - .getRpc() - .batchWriteAtLeastOnce(requestBuilder.build(), reqId.withOptions(getOptions())); + return spanner.getRpc().batchWriteAtLeastOnce(requestBuilder.build(), getOptions()); } catch (Throwable e) { span.setStatus(e); throw SpannerExceptionFactory.newSpannerException(e); @@ -472,8 +443,7 @@ public ApiFuture asyncClose() { if (getIsMultiplexed()) { return com.google.api.core.ApiFutures.immediateFuture(Empty.getDefaultInstance()); } - XGoogSpannerRequestId reqId = this.getRequestIdCreator().nextRequestId(this.getChannel(), 1); - return spanner.getRpc().asyncDeleteSession(getName(), reqId.withOptions(getOptions())); + return spanner.getRpc().asyncDeleteSession(getName(), getOptions()); } @Override @@ -483,8 +453,7 @@ public void close() { } ISpan span = tracer.spanBuilder(SpannerImpl.DELETE_SESSION); try (IScope s = tracer.withSpan(span)) { - XGoogSpannerRequestId reqId = this.getRequestIdCreator().nextRequestId(this.getChannel(), 1); - spanner.getRpc().deleteSession(getName(), reqId.withOptions(getOptions())); + spanner.getRpc().deleteSession(getName(), getOptions()); } catch (RuntimeException e) { span.setStatus(e); throw e; @@ -518,12 +487,8 @@ ApiFuture beginTransactionAsync( } final BeginTransactionRequest request = requestBuilder.build(); final ApiFuture requestFuture; - XGoogSpannerRequestId reqId = this.getRequestIdCreator().nextRequestId(this.getChannel(), 1); try (IScope ignore = tracer.withSpan(span)) { - requestFuture = - spanner - .getRpc() - .beginTransactionAsync(request, reqId.withOptions(channelHint), routeToLeader); + requestFuture = spanner.getRpc().beginTransactionAsync(request, channelHint, routeToLeader); } requestFuture.addListener( () -> { @@ -531,7 +496,7 @@ ApiFuture beginTransactionAsync( Transaction txn = requestFuture.get(); if (txn.getId().isEmpty()) { throw newSpannerException( - ErrorCode.INTERNAL, "Missing id in transaction\n" + getName(), reqId); + ErrorCode.INTERNAL, "Missing id in transaction\n" + getName()); } span.end(); res.set(txn); @@ -540,7 +505,7 @@ ApiFuture beginTransactionAsync( span.end(); res.setException( SpannerExceptionFactory.newSpannerException( - e.getCause() == null ? e : e.getCause(), reqId)); + e.getCause() == null ? e : e.getCause())); } catch (InterruptedException e) { span.setStatus(e); span.end(); @@ -602,12 +567,8 @@ TraceWrapper getTracer() { return tracer; } - public void setRequestIdCreator(XGoogSpannerRequestId.RequestIdCreator creator) { - this.requestIdCreator = creator; - } - public XGoogSpannerRequestId.RequestIdCreator getRequestIdCreator() { - return this.requestIdCreator; + return this.spanner.getRpc().getRequestIdCreator(); } int getChannel() { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionNotFoundException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionNotFoundException.java index c17384db3ec..f4a62b1954a 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionNotFoundException.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SessionNotFoundException.java @@ -35,7 +35,7 @@ public class SessionNotFoundException extends ResourceNotFoundException { @Nullable String message, ResourceInfo resourceInfo, @Nullable Throwable cause) { - this(token, message, resourceInfo, cause, null, null); + this(token, message, resourceInfo, cause, null); } /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ @@ -44,8 +44,7 @@ public class SessionNotFoundException extends ResourceNotFoundException { @Nullable String message, ResourceInfo resourceInfo, @Nullable Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId reqId) { - super(token, message, resourceInfo, cause, apiException, reqId); + @Nullable ApiException apiException) { + super(token, message, resourceInfo, cause, apiException); } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerApiFutures.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerApiFutures.java index 7e6acd02f97..88e0b84f0fc 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerApiFutures.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerApiFutures.java @@ -35,10 +35,9 @@ public static T getOrNull(ApiFuture future) throws SpannerException { } throw SpannerExceptionFactory.asSpannerException(e.getCause()); } catch (InterruptedException e) { - throw SpannerExceptionFactory.propagateInterrupt(e, null /*TODO: requestId*/); + throw SpannerExceptionFactory.propagateInterrupt(e); } catch (CancellationException e) { - throw SpannerExceptionFactory.newSpannerExceptionForCancellation( - null, e, null /*TODO: requestId*/); + throw SpannerExceptionFactory.newSpannerExceptionForCancellation(null, e); } } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerBatchUpdateException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerBatchUpdateException.java index 0d841d24463..837a008a833 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerBatchUpdateException.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerBatchUpdateException.java @@ -25,9 +25,8 @@ public class SpannerBatchUpdateException extends SpannerException { ErrorCode code, String message, long[] counts, - Throwable cause, - XGoogSpannerRequestId reqId) { - super(token, code, false, message, cause, null, reqId); + Throwable cause) { + super(token, code, false, message, cause, null); updateCounts = counts; } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerException.java index fbe60e2a1d6..0829cc35d62 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerException.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerException.java @@ -16,11 +16,11 @@ package com.google.cloud.spanner; -import com.google.api.core.InternalApi; import com.google.api.gax.rpc.ApiException; import com.google.api.gax.rpc.ErrorDetails; import com.google.cloud.grpc.BaseGrpcServiceException; import com.google.common.base.Preconditions; +import com.google.common.base.Strings; import com.google.protobuf.util.Durations; import com.google.rpc.ResourceInfo; import com.google.rpc.RetryInfo; @@ -41,9 +41,8 @@ public abstract static class ResourceNotFoundException extends SpannerException @Nullable String message, ResourceInfo resourceInfo, @Nullable Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId reqId) { - super(token, ErrorCode.NOT_FOUND, /* retryable */ false, message, cause, apiException, reqId); + @Nullable ApiException apiException) { + super(token, ErrorCode.NOT_FOUND, /* retryable */ false, message, cause, apiException); this.resourceInfo = resourceInfo; } @@ -59,7 +58,7 @@ public String getResourceName() { private final ErrorCode code; private final ApiException apiException; - private XGoogSpannerRequestId requestId; + private final XGoogSpannerRequestId requestId; private String statement; /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ @@ -80,25 +79,13 @@ public String getResourceName() { @Nullable String message, @Nullable Throwable cause, @Nullable ApiException apiException) { - this(token, code, retryable, message, cause, apiException, null); - } - - /** Private constructor. Use {@link SpannerExceptionFactory} to create instances. */ - SpannerException( - DoNotConstructDirectly token, - ErrorCode code, - boolean retryable, - @Nullable String message, - @Nullable Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId requestId) { super(message, cause, code.getCode(), retryable); if (token != DoNotConstructDirectly.ALLOWED) { throw new AssertionError("Do not construct directly: use SpannerExceptionFactory"); } this.code = Preconditions.checkNotNull(code); this.apiException = apiException; - this.requestId = requestId; + this.requestId = extractRequestId(cause); } @Override @@ -109,6 +96,14 @@ public String getMessage() { return String.format("%s - Statement: '%s'", super.getMessage(), this.statement); } + @Override + public String toString() { + if (this.requestId == null) { + return super.toString(); + } + return super.toString() + " - RequestId: " + this.requestId; + } + /** Returns the error code associated with this exception. */ public ErrorCode getErrorCode() { return code; @@ -150,7 +145,7 @@ static long extractRetryDelay(Throwable cause) { Metadata trailers = Status.trailersFromThrowable(cause); if (trailers != null && trailers.containsKey(KEY_RETRY_INFO)) { RetryInfo retryInfo = trailers.get(KEY_RETRY_INFO); - if (retryInfo.hasRetryDelay()) { + if (retryInfo != null && retryInfo.hasRetryDelay()) { return Durations.toMillis(retryInfo.getRetryDelay()); } } @@ -158,6 +153,20 @@ static long extractRetryDelay(Throwable cause) { return -1L; } + @Nullable + static XGoogSpannerRequestId extractRequestId(Throwable cause) { + if (cause != null) { + Metadata trailers = Status.trailersFromThrowable(cause); + if (trailers != null && trailers.containsKey(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY)) { + String requestId = trailers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY); + if (!Strings.isNullOrEmpty(requestId)) { + return XGoogSpannerRequestId.of(requestId); + } + } + } + return null; + } + /** * Checks the underlying reason of the exception and if it's {@link ApiException} then return the * reason otherwise null. @@ -224,10 +233,4 @@ public ErrorDetails getErrorDetails() { void setStatement(String statement) { this.statement = statement; } - - /** Sets the requestId. */ - @InternalApi - public void setRequestId(XGoogSpannerRequestId reqId) { - this.requestId = reqId; - } } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java index b9cc1cfb8b1..185f98b5433 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerExceptionFactory.java @@ -58,37 +58,17 @@ public final class SpannerExceptionFactory { ProtoUtils.keyForProto(ErrorInfo.getDefaultInstance()); public static SpannerException newSpannerException(ErrorCode code, @Nullable String message) { - return newSpannerException(code, message, (XGoogSpannerRequestId) (null)); - } - - public static SpannerException newSpannerException( - ErrorCode code, - @Nullable String message, - @Nullable Throwable cause, - @Nullable XGoogSpannerRequestId reqId) { - return newSpannerExceptionPreformatted( - code, formatMessage(code, message), cause, (ApiException) (null), reqId); - } - - public static SpannerException newSpannerException( - ErrorCode code, @Nullable String message, @Nullable XGoogSpannerRequestId reqId) { - return newSpannerException(code, message, (Throwable) (null), reqId); + return newSpannerException(code, message, null); } public static SpannerException newSpannerException( ErrorCode code, @Nullable String message, @Nullable Throwable cause) { - return newSpannerException(code, message, cause, null); + return newSpannerExceptionPreformatted(code, formatMessage(code, message), cause, null); } public static SpannerException propagateInterrupt(InterruptedException e) { - return propagateInterrupt(e, null); - } - - public static SpannerException propagateInterrupt( - InterruptedException e, XGoogSpannerRequestId reqId) { Thread.currentThread().interrupt(); - return SpannerExceptionFactory.newSpannerException( - ErrorCode.CANCELLED, "Interrupted", e, reqId); + return SpannerExceptionFactory.newSpannerException(ErrorCode.CANCELLED, "Interrupted", e); } /** @@ -132,27 +112,17 @@ public static SpannerException asSpannerException(Throwable t) { * #newSpannerException(ErrorCode, String)} instead of this method. */ public static SpannerException newSpannerException(Throwable cause) { - return newSpannerException(null, cause, null); - } - - public static SpannerException newSpannerException(Throwable cause, XGoogSpannerRequestId reqId) { - return newSpannerException(null, cause, reqId); + return newSpannerException(null, cause); } public static SpannerBatchUpdateException newSpannerBatchUpdateException( ErrorCode code, String message, long[] updateCounts) { - return newSpannerBatchUpdateException(code, message, updateCounts, null); - } - - public static SpannerBatchUpdateException newSpannerBatchUpdateException( - ErrorCode code, String message, long[] updateCounts, @Nullable XGoogSpannerRequestId reqId) { DoNotConstructDirectly token = DoNotConstructDirectly.ALLOWED; SpannerException cause = null; if (isTransactionMutationLimitException(code, message)) { - cause = - new TransactionMutationLimitExceededException(token, code, message, null, null, reqId); + cause = new TransactionMutationLimitExceededException(token, code, message, null, null); } - return new SpannerBatchUpdateException(token, code, message, updateCounts, cause, reqId); + return new SpannerBatchUpdateException(token, code, message, updateCounts, cause); } /** Constructs a specific error that */ @@ -205,10 +175,6 @@ public static SpannerBatchUpdateException newSpannerBatchUpdateException( cause); } - public static SpannerException newSpannerException(@Nullable Context context, Throwable cause) { - return newSpannerException(context, cause, null); - } - /** * Creates a new exception based on {@code cause}. If {@code cause} indicates cancellation, {@code * context} will be inspected to establish the type of cancellation. @@ -216,22 +182,21 @@ public static SpannerException newSpannerException(@Nullable Context context, Th *

Intended for internal library use; user code should use {@link * #newSpannerException(ErrorCode, String)} instead of this method. */ - public static SpannerException newSpannerException( - @Nullable Context context, Throwable cause, @Nullable XGoogSpannerRequestId reqId) { + public static SpannerException newSpannerException(@Nullable Context context, Throwable cause) { if (cause instanceof SpannerException) { SpannerException e = (SpannerException) cause; - return newSpannerExceptionPreformatted(e.getErrorCode(), e.getMessage(), e, null, reqId); + return newSpannerExceptionPreformatted(e.getErrorCode(), e.getMessage(), e, null); } else if (cause instanceof CancellationException) { - return newSpannerExceptionForCancellation(context, cause, reqId); + return newSpannerExceptionForCancellation(context, cause); } else if (cause instanceof ApiException) { - return fromApiException((ApiException) cause, reqId); + return fromApiException((ApiException) cause); } // Extract gRPC status. This will produce "UNKNOWN" for non-gRPC exceptions. Status status = Status.fromThrowable(cause); if (status.getCode() == Status.Code.CANCELLED) { - return newSpannerExceptionForCancellation(context, cause, reqId); + return newSpannerExceptionForCancellation(context, cause); } - return newSpannerException(ErrorCode.fromGrpcStatus(status), cause.getMessage(), cause, reqId); + return newSpannerException(ErrorCode.fromGrpcStatus(status), cause.getMessage(), cause); } public static RuntimeException causeAsRunTimeException(ExecutionException executionException) { @@ -256,11 +221,6 @@ static SpannerException newRetryOnDifferentGrpcChannelException( static SpannerException newSpannerExceptionForCancellation( @Nullable Context context, @Nullable Throwable cause) { - return newSpannerExceptionForCancellation(context, cause, null); - } - - static SpannerException newSpannerExceptionForCancellation( - @Nullable Context context, @Nullable Throwable cause, @Nullable XGoogSpannerRequestId reqId) { if (context != null && context.isCancelled()) { Throwable cancellationCause = context.cancellationCause(); Throwable throwable = @@ -269,14 +229,13 @@ static SpannerException newSpannerExceptionForCancellation( : MoreObjects.firstNonNull(cause, cancellationCause); if (cancellationCause instanceof TimeoutException) { return newSpannerException( - ErrorCode.DEADLINE_EXCEEDED, "Current context exceeded deadline", throwable, reqId); + ErrorCode.DEADLINE_EXCEEDED, "Current context exceeded deadline", throwable); } else { - return newSpannerException( - ErrorCode.CANCELLED, "Current context was cancelled", throwable, reqId); + return newSpannerException(ErrorCode.CANCELLED, "Current context was cancelled", throwable); } } return newSpannerException( - ErrorCode.CANCELLED, cause == null ? "Cancelled" : cause.getMessage(), cause, reqId); + ErrorCode.CANCELLED, cause == null ? "Cancelled" : cause.getMessage(), cause); } private static String formatMessage(ErrorCode code, @Nullable String message) { @@ -355,13 +314,12 @@ static SpannerException newSpannerExceptionPreformatted( ErrorCode code, @Nullable String message, @Nullable Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId reqId) { + @Nullable ApiException apiException) { // This is the one place in the codebase that is allowed to call constructors directly. DoNotConstructDirectly token = DoNotConstructDirectly.ALLOWED; switch (code) { case ABORTED: - return new AbortedException(token, message, cause, apiException, reqId); + return new AbortedException(token, message, cause, apiException); case RESOURCE_EXHAUSTED: ErrorInfo info = extractErrorInfo(cause, apiException); if (info != null @@ -370,8 +328,7 @@ static SpannerException newSpannerExceptionPreformatted( && AdminRequestsPerMinuteExceededException.ADMIN_REQUESTS_LIMIT_VALUE.equals( info.getMetadataMap() .get(AdminRequestsPerMinuteExceededException.ADMIN_REQUESTS_LIMIT_KEY))) { - return new AdminRequestsPerMinuteExceededException( - token, message, cause, apiException, reqId); + return new AdminRequestsPerMinuteExceededException(token, message, cause, apiException); } case NOT_FOUND: ResourceInfo resourceInfo = extractResourceInfo(cause); @@ -379,39 +336,36 @@ static SpannerException newSpannerExceptionPreformatted( switch (resourceInfo.getResourceType()) { case SESSION_RESOURCE_TYPE: return new SessionNotFoundException( - token, message, resourceInfo, cause, apiException, reqId); + token, message, resourceInfo, cause, apiException); case DATABASE_RESOURCE_TYPE: return new DatabaseNotFoundException( - token, message, resourceInfo, cause, apiException, reqId); + token, message, resourceInfo, cause, apiException); case INSTANCE_RESOURCE_TYPE: return new InstanceNotFoundException( - token, message, resourceInfo, cause, apiException, reqId); + token, message, resourceInfo, cause, apiException); } } case INVALID_ARGUMENT: if (isTransactionMutationLimitException(cause, apiException)) { return new TransactionMutationLimitExceededException( - token, code, message, cause, apiException, reqId); + token, code, message, cause, apiException); } if (isMissingDefaultSequenceKindException(apiException)) { - return new MissingDefaultSequenceKindException( - token, code, message, cause, apiException, reqId); + return new MissingDefaultSequenceKindException(token, code, message, cause, apiException); } // Fall through to the default. default: return new SpannerException( - token, code, isRetryable(code, cause), message, cause, apiException, reqId); + token, code, isRetryable(code, cause), message, cause, apiException); } } static SpannerException newSpannerExceptionPreformatted( ErrorCode code, @Nullable String message, @Nullable Throwable cause) { - return newSpannerExceptionPreformatted( - code, message, cause, null, (XGoogSpannerRequestId) (null)); + return newSpannerExceptionPreformatted(code, message, cause, null); } - private static SpannerException fromApiException( - ApiException exception, @Nullable XGoogSpannerRequestId reqId) { + private static SpannerException fromApiException(ApiException exception) { Status.Code code; if (exception.getStatusCode() instanceof GrpcStatusCode) { code = ((GrpcStatusCode) exception.getStatusCode()).getTransportCode(); @@ -426,8 +380,7 @@ private static SpannerException fromApiException( errorCode, formatMessage(errorCode, exception.getMessage()), exception.getCause(), - exception, - reqId); + exception); } private static boolean isRetryable(ErrorCode code, @Nullable Throwable cause) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java index 16f36144723..b4eef3a0f51 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerImpl.java @@ -459,6 +459,14 @@ public boolean isClosed() { } } + void resetRequestIdCounters() { + gapicRpc.getRequestIdCreator().reset(); + } + + long getRequestIdClientId() { + return gapicRpc.getRequestIdCreator().getClientId(); + } + /** Helper class for gRPC calls that can return paginated results. */ abstract static class PageFetcher implements NextPageFetcher { private String nextPageToken; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerRetryHelper.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerRetryHelper.java index 6ca1a4e02e8..0dabcbd0094 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerRetryHelper.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerRetryHelper.java @@ -116,8 +116,7 @@ public TimedAttemptSettings createNextAttempt( public boolean shouldRetry(Throwable prevThrowable, T prevResponse) throws CancellationException { if (Context.current().isCancelled()) { - throw SpannerExceptionFactory.newSpannerExceptionForCancellation( - Context.current(), null, null); + throw SpannerExceptionFactory.newSpannerExceptionForCancellation(Context.current(), null); } return prevThrowable instanceof AbortedException || prevThrowable instanceof com.google.api.gax.rpc.AbortedException; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionMutationLimitExceededException.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionMutationLimitExceededException.java index f6b0a4efd22..de215c5caee 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionMutationLimitExceededException.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionMutationLimitExceededException.java @@ -37,9 +37,8 @@ public class TransactionMutationLimitExceededException extends SpannerException ErrorCode errorCode, String message, Throwable cause, - @Nullable ApiException apiException, - @Nullable XGoogSpannerRequestId reqId) { - super(token, errorCode, /* retryable= */ false, message, cause, apiException, reqId); + @Nullable ApiException apiException) { + super(token, errorCode, /* retryable= */ false, message, cause, apiException); } static boolean isTransactionMutationLimitException(ErrorCode code, String message) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java index 6ce8cc11aa6..7afccce194c 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/TransactionRunnerImpl.java @@ -449,8 +449,6 @@ private final class CommitRunnable implements Runnable { @Override public void run() { - XGoogSpannerRequestId reqId = - session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); try { prev.get(); if (transactionId == null && transactionIdFuture == null) { @@ -493,8 +491,7 @@ public void run() { final ApiFuture commitFuture; final ISpan opSpan = tracer.spanBuilderWithExplicitParent(SpannerImpl.COMMIT, span); try (IScope ignore = tracer.withSpan(opSpan)) { - commitFuture = - rpc.commitAsync(commitRequest, reqId.withOptions(getTransactionChannelHint())); + commitFuture = rpc.commitAsync(commitRequest, getTransactionChannelHint()); } session.markUsed(clock.instant()); commitFuture.addListener( @@ -505,7 +502,7 @@ public void run() { // future, but we add a result here as well as a safety precaution. res.setException( SpannerExceptionFactory.newSpannerException( - ErrorCode.INTERNAL, "commitFuture is not done", reqId)); + ErrorCode.INTERNAL, "commitFuture is not done")); return; } com.google.spanner.v1.CommitResponse proto = commitFuture.get(); @@ -535,9 +532,7 @@ public void run() { } if (!proto.hasCommitTimestamp()) { throw newSpannerException( - ErrorCode.INTERNAL, - "Missing commitTimestamp:\n" + session.getName(), - reqId); + ErrorCode.INTERNAL, "Missing commitTimestamp:\n" + session.getName()); } span.addAnnotation("Commit Done"); opSpan.end(); @@ -577,8 +572,7 @@ public void run() { res.setException(SpannerExceptionFactory.propagateTimeout(e)); } catch (Throwable e) { res.setException( - SpannerExceptionFactory.newSpannerException( - e.getCause() == null ? e : e.getCause(), reqId)); + SpannerExceptionFactory.newSpannerException(e.getCause() == null ? e : e.getCause())); } } } @@ -697,8 +691,7 @@ options, getPreviousTransactionId()))) } throw se; } catch (InterruptedException e) { - throw SpannerExceptionFactory.newSpannerExceptionForCancellation( - null, e, null /*TODO: requestId*/); + throw SpannerExceptionFactory.newSpannerExceptionForCancellation(null, e); } } // There is already a transactionId available. Include that id as the transaction to use. @@ -938,12 +931,9 @@ private ResultSet internalExecuteUpdate( final ExecuteSqlRequest.Builder builder = getExecuteSqlRequestBuilder( statement, queryMode, options, /* withTransactionSelector= */ true); - XGoogSpannerRequestId reqId = - session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); try { com.google.spanner.v1.ResultSet resultSet = - rpc.executeQuery( - builder.build(), reqId.withOptions(getTransactionChannelHint()), isRouteToLeader()); + rpc.executeQuery(builder.build(), getTransactionChannelHint(), isRouteToLeader()); session.markUsed(clock.instant()); if (resultSet.getMetadata().hasTransaction()) { onTransactionMetadata( @@ -1076,11 +1066,9 @@ public long[] batchUpdate(Iterable statements, UpdateOption... update } final ExecuteBatchDmlRequest.Builder builder = getExecuteBatchDmlRequestBuilder(statements, options); - XGoogSpannerRequestId reqId = - session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); try { com.google.spanner.v1.ExecuteBatchDmlResponse response = - rpc.executeBatchDml(builder.build(), reqId.withOptions(getTransactionChannelHint())); + rpc.executeBatchDml(builder.build(), getTransactionChannelHint()); session.markUsed(clock.instant()); long[] results = new long[response.getResultSetsCount()]; for (int i = 0; i < response.getResultSetsCount(); ++i) { @@ -1104,8 +1092,7 @@ public long[] batchUpdate(Iterable statements, UpdateOption... update throw newSpannerBatchUpdateException( ErrorCode.fromRpcStatus(response.getStatus()), response.getStatus().getMessage(), - results, - reqId); + results); } return results; } catch (Throwable e) { @@ -1140,15 +1127,11 @@ public ApiFuture batchUpdateAsync( final ExecuteBatchDmlRequest.Builder builder = getExecuteBatchDmlRequestBuilder(statements, options); ApiFuture response; - XGoogSpannerRequestId reqId = - session.getRequestIdCreator().nextRequestId(session.getChannel(), 1); try { // Register the update as an async operation that must finish before the transaction may // commit. increaseAsyncOperations(); - response = - rpc.executeBatchDmlAsync( - builder.build(), reqId.withOptions(getTransactionChannelHint())); + response = rpc.executeBatchDmlAsync(builder.build(), getTransactionChannelHint()); session.markUsed(clock.instant()); } catch (Throwable t) { decreaseAsyncOperations(); @@ -1178,8 +1161,7 @@ public ApiFuture batchUpdateAsync( throw newSpannerBatchUpdateException( ErrorCode.fromRpcStatus(batchDmlResponse.getStatus()), batchDmlResponse.getStatus().getMessage(), - results, - reqId); + results); } return results; }, @@ -1233,7 +1215,7 @@ public ListenableAsyncResultSet executeQueryAsync( private final SessionImpl session; private final Options options; private ISpan span; - private TraceWrapper tracer; + private final TraceWrapper tracer; private TransactionContextImpl txn; private volatile boolean isValid = true; diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java index 47ccb05231c..d858fdb9273 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/XGoogSpannerRequestId.java @@ -17,13 +17,11 @@ package com.google.cloud.spanner; import com.google.api.core.InternalApi; -import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.annotations.VisibleForTesting; +import io.grpc.CallOptions; import io.grpc.Metadata; import java.math.BigInteger; import java.security.SecureRandom; -import java.util.HashMap; -import java.util.Map; import java.util.Objects; import java.util.regex.MatchResult; import java.util.regex.Matcher; @@ -31,13 +29,15 @@ @InternalApi public class XGoogSpannerRequestId { - // 1. Generate the random process Id singleton. + // 1. Generate the random process ID singleton. @VisibleForTesting static final String RAND_PROCESS_ID = XGoogSpannerRequestId.generateRandProcessId(); - public static String REQUEST_ID = "x-goog-spanner-request-id"; - public static final Metadata.Key REQUEST_HEADER_KEY = - Metadata.Key.of(REQUEST_ID, Metadata.ASCII_STRING_MARSHALLER); + public static String REQUEST_ID_HEADER_NAME = "x-goog-spanner-request-id"; + public static final Metadata.Key REQUEST_ID_HEADER_KEY = + Metadata.Key.of(REQUEST_ID_HEADER_NAME, Metadata.ASCII_STRING_MARSHALLER); + public static final CallOptions.Key REQUEST_ID_CALL_OPTIONS_KEY = + CallOptions.Key.create("XGoogSpannerRequestId"); @VisibleForTesting static final long VERSION = 1; // The version of the specification being implemented. @@ -59,6 +59,20 @@ public static XGoogSpannerRequestId of( return new XGoogSpannerRequestId(nthClientId, nthChannelId, nthRequest, attempt); } + @VisibleForTesting + long getNthClientId() { + return nthClientId; + } + + @VisibleForTesting + long getNthChannelId() { + return nthChannelId; + } + + boolean hasChannelId() { + return nthChannelId > 0; + } + @VisibleForTesting long getAttempt() { return this.attempt; @@ -95,8 +109,8 @@ private static String generateRandProcessId() { return String.format("%016x", bigInt); } - @Override - public String toString() { + /** Returns the string representation of this RequestId as it should be sent to Spanner. */ + public String getHeaderValue() { return String.format( "%d.%s.%d.%d.%d.%d", XGoogSpannerRequestId.VERSION, @@ -107,6 +121,18 @@ public String toString() { this.attempt); } + @Override + public String toString() { + return String.format( + "%d.%s.%d.%s.%d.%d", + XGoogSpannerRequestId.VERSION, + XGoogSpannerRequestId.RAND_PROCESS_ID, + this.nthClientId, + this.nthChannelId < 0 ? "x" : String.valueOf(this.nthChannelId), + this.nthRequest, + this.attempt); + } + public String debugToString() { return String.format( "%d.%s.nth_client=%d.nth_chan=%d.nth_req=%d.attempt=%d", @@ -151,31 +177,38 @@ public void incrementAttempt() { this.attempt++; } - Map withOptions(Map options) { - Map copyOptions = new HashMap<>(); - if (options != null) { - copyOptions.putAll(options); - } - copyOptions.put(SpannerRpc.Option.REQUEST_ID, this); - return copyOptions; - } - @Override public int hashCode() { return Objects.hash(this.nthClientId, this.nthChannelId, this.nthRequest, this.attempt); } - interface RequestIdCreator { - XGoogSpannerRequestId nextRequestId(long channelId, int attempt); + @InternalApi + public interface RequestIdCreator { + long getClientId(); + + XGoogSpannerRequestId nextRequestId(long channelId); + + void reset(); } - static class NoopRequestIdCreator implements RequestIdCreator { - NoopRequestIdCreator() {} + // TODO: Move this class into test code. + static final class NoopRequestIdCreator implements RequestIdCreator { + static final NoopRequestIdCreator INSTANCE = new NoopRequestIdCreator(); + + private NoopRequestIdCreator() {} @Override - public XGoogSpannerRequestId nextRequestId(long channelId, int attempt) { - return XGoogSpannerRequestId.of(1, 1, 1, 0); + public long getClientId() { + return 1L; } + + @Override + public XGoogSpannerRequestId nextRequestId(long channelId) { + return XGoogSpannerRequestId.of(1, channelId, 1, 0); + } + + @Override + public void reset() {} } public void setChannelId(long channelId) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java index c448695afd2..4a8d643b79c 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DdlBatch.java @@ -266,7 +266,7 @@ public ApiFuture runBatchAsync(CallType callType) { } catch (SpannerException e) { long[] updateCounts = extractUpdateCounts(operationReference.get()); throw SpannerExceptionFactory.newSpannerBatchUpdateException( - e.getErrorCode(), e.getMessage(), updateCounts, null /* TODO: requestId */); + e.getErrorCode(), e.getMessage(), updateCounts); } } catch (Throwable t) { span.setStatus(StatusCode.ERROR); diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index 371d73a0af2..306bd9842ac 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -19,6 +19,7 @@ import static com.google.cloud.spanner.SpannerExceptionFactory.asSpannerException; import static com.google.cloud.spanner.SpannerExceptionFactory.newSpannerException; import static com.google.cloud.spanner.ThreadFactoryUtil.tryCreateVirtualThreadPerTaskExecutor; +import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_CALL_OPTIONS_KEY; import com.google.api.core.ApiFunction; import com.google.api.core.ApiFuture; @@ -73,6 +74,7 @@ import com.google.cloud.spanner.SpannerOptions.CallContextConfigurator; import com.google.cloud.spanner.SpannerOptions.CallCredentialsProvider; import com.google.cloud.spanner.XGoogSpannerRequestId; +import com.google.cloud.spanner.XGoogSpannerRequestId.RequestIdCreator; import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub; import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings; import com.google.cloud.spanner.admin.database.v1.stub.GrpcDatabaseAdminCallableFactory; @@ -88,7 +90,6 @@ import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import com.google.common.io.Resources; import com.google.common.util.concurrent.RateLimiter; @@ -241,6 +242,7 @@ public class GapicSpannerRpc implements SpannerRpc { public static boolean DIRECTPATH_CHANNEL_CREATED = false; private static final String API_FILE = "grpc-gcp-apiconfig.json"; + private final RequestIdCreator requestIdCreator = new RequestIdCreatorImpl(); private boolean rpcIsClosed; private final SpannerStub spannerStub; private final RetrySettings executeQueryRetrySettings; @@ -825,7 +827,7 @@ public OperationFuture call() { isRetry = true; if (operationName == null) { - GrpcCallContext context = newCallContext(null, instanceName, initialRequest, method); + GrpcCallContext context = newAdminCallContext(instanceName, initialRequest, method); return operationCallable.futureCall(initialRequest, context); } else { return operationCallable.resumeFutureCall(operationName); @@ -916,8 +918,7 @@ public Paginated listInstanceConfigs(int pageSize, @Nullable Str ListInstanceConfigsRequest request = requestBuilder.build(); GrpcCallContext context = - newCallContext( - null, projectName, request, InstanceAdminGrpc.getListInstanceConfigsMethod()); + newAdminCallContext(projectName, request, InstanceAdminGrpc.getListInstanceConfigsMethod()); ListInstanceConfigsResponse response = get(instanceAdminStub.listInstanceConfigsCallable().futureCall(request, context)); return new Paginated<>(response.getInstanceConfigsList(), response.getNextPageToken()); @@ -940,7 +941,7 @@ public OperationFuture createInsta } CreateInstanceConfigRequest request = builder.build(); GrpcCallContext context = - newCallContext(null, parent, request, InstanceAdminGrpc.getCreateInstanceConfigMethod()); + newAdminCallContext(parent, request, InstanceAdminGrpc.getCreateInstanceConfigMethod()); return instanceAdminStub.createInstanceConfigOperationCallable().futureCall(request, context); } @@ -957,11 +958,8 @@ public OperationFuture updateInsta } UpdateInstanceConfigRequest request = builder.build(); GrpcCallContext context = - newCallContext( - null, - instanceConfig.getName(), - request, - InstanceAdminGrpc.getUpdateInstanceConfigMethod()); + newAdminCallContext( + instanceConfig.getName(), request, InstanceAdminGrpc.getUpdateInstanceConfigMethod()); return instanceAdminStub.updateInstanceConfigOperationCallable().futureCall(request, context); } @@ -971,7 +969,7 @@ public InstanceConfig getInstanceConfig(String instanceConfigName) throws Spanne GetInstanceConfigRequest.newBuilder().setName(instanceConfigName).build(); GrpcCallContext context = - newCallContext(null, projectName, request, InstanceAdminGrpc.getGetInstanceConfigMethod()); + newAdminCallContext(projectName, request, InstanceAdminGrpc.getGetInstanceConfigMethod()); return get(instanceAdminStub.getInstanceConfigCallable().futureCall(request, context)); } @@ -990,8 +988,8 @@ public void deleteInstanceConfig( } DeleteInstanceConfigRequest request = requestBuilder.build(); GrpcCallContext context = - newCallContext( - null, instanceConfigName, request, InstanceAdminGrpc.getDeleteInstanceConfigMethod()); + newAdminCallContext( + instanceConfigName, request, InstanceAdminGrpc.getDeleteInstanceConfigMethod()); get(instanceAdminStub.deleteInstanceConfigCallable().futureCall(request, context)); } @@ -1012,8 +1010,8 @@ public Paginated listInstanceConfigOperations( final ListInstanceConfigOperationsRequest request = requestBuilder.build(); final GrpcCallContext context = - newCallContext( - null, projectName, request, InstanceAdminGrpc.getListInstanceConfigOperationsMethod()); + newAdminCallContext( + projectName, request, InstanceAdminGrpc.getListInstanceConfigOperationsMethod()); ListInstanceConfigOperationsResponse response = runWithRetryOnAdministrativeRequestsExceeded( () -> @@ -1038,7 +1036,7 @@ public Paginated listInstances( ListInstancesRequest request = requestBuilder.build(); GrpcCallContext context = - newCallContext(null, projectName, request, InstanceAdminGrpc.getListInstancesMethod()); + newAdminCallContext(projectName, request, InstanceAdminGrpc.getListInstancesMethod()); ListInstancesResponse response = get(instanceAdminStub.listInstancesCallable().futureCall(request, context)); return new Paginated<>(response.getInstancesList(), response.getNextPageToken()); @@ -1054,7 +1052,7 @@ public OperationFuture createInstance( .setInstance(instance) .build(); GrpcCallContext context = - newCallContext(null, parent, request, InstanceAdminGrpc.getCreateInstanceMethod()); + newAdminCallContext(parent, request, InstanceAdminGrpc.getCreateInstanceMethod()); return instanceAdminStub.createInstanceOperationCallable().futureCall(request, context); } @@ -1064,8 +1062,8 @@ public OperationFuture updateInstance( UpdateInstanceRequest request = UpdateInstanceRequest.newBuilder().setInstance(instance).setFieldMask(fieldMask).build(); GrpcCallContext context = - newCallContext( - null, instance.getName(), request, InstanceAdminGrpc.getUpdateInstanceMethod()); + newAdminCallContext( + instance.getName(), request, InstanceAdminGrpc.getUpdateInstanceMethod()); return instanceAdminStub.updateInstanceOperationCallable().futureCall(request, context); } @@ -1074,7 +1072,7 @@ public Instance getInstance(String instanceName) throws SpannerException { GetInstanceRequest request = GetInstanceRequest.newBuilder().setName(instanceName).build(); GrpcCallContext context = - newCallContext(null, instanceName, request, InstanceAdminGrpc.getGetInstanceMethod()); + newAdminCallContext(instanceName, request, InstanceAdminGrpc.getGetInstanceMethod()); return get(instanceAdminStub.getInstanceCallable().futureCall(request, context)); } @@ -1084,7 +1082,7 @@ public void deleteInstance(String instanceName) throws SpannerException { DeleteInstanceRequest.newBuilder().setName(instanceName).build(); GrpcCallContext context = - newCallContext(null, instanceName, request, InstanceAdminGrpc.getDeleteInstanceMethod()); + newAdminCallContext(instanceName, request, InstanceAdminGrpc.getDeleteInstanceMethod()); get(instanceAdminStub.deleteInstanceCallable().futureCall(request, context)); } @@ -1103,8 +1101,8 @@ public Paginated listBackupOperations( final ListBackupOperationsRequest request = requestBuilder.build(); final GrpcCallContext context = - newCallContext( - null, instanceName, request, DatabaseAdminGrpc.getListBackupOperationsMethod()); + newAdminCallContext( + instanceName, request, DatabaseAdminGrpc.getListBackupOperationsMethod()); ListBackupOperationsResponse response = runWithRetryOnAdministrativeRequestsExceeded( () -> @@ -1128,8 +1126,8 @@ public Paginated listDatabaseOperations( final ListDatabaseOperationsRequest request = requestBuilder.build(); final GrpcCallContext context = - newCallContext( - null, instanceName, request, DatabaseAdminGrpc.getListDatabaseOperationsMethod()); + newAdminCallContext( + instanceName, request, DatabaseAdminGrpc.getListDatabaseOperationsMethod()); ListDatabaseOperationsResponse response = runWithRetryOnAdministrativeRequestsExceeded( () -> @@ -1154,7 +1152,7 @@ public Paginated listDatabaseRoles( final ListDatabaseRolesRequest request = requestBuilder.build(); final GrpcCallContext context = - newCallContext(null, databaseName, request, DatabaseAdminGrpc.getListDatabaseRolesMethod()); + newAdminCallContext(databaseName, request, DatabaseAdminGrpc.getListDatabaseRolesMethod()); ListDatabaseRolesResponse response = runWithRetryOnAdministrativeRequestsExceeded( () -> get(databaseAdminStub.listDatabaseRolesCallable().futureCall(request, context))); @@ -1178,7 +1176,7 @@ public Paginated listBackups( final ListBackupsRequest request = requestBuilder.build(); final GrpcCallContext context = - newCallContext(null, instanceName, request, DatabaseAdminGrpc.getListBackupsMethod()); + newAdminCallContext(instanceName, request, DatabaseAdminGrpc.getListBackupsMethod()); ListBackupsResponse response = runWithRetryOnAdministrativeRequestsExceeded( () -> get(databaseAdminStub.listBackupsCallable().futureCall(request, context))); @@ -1198,7 +1196,7 @@ public Paginated listDatabases( final ListDatabasesRequest request = requestBuilder.build(); final GrpcCallContext context = - newCallContext(null, instanceName, request, DatabaseAdminGrpc.getListDatabasesMethod()); + newAdminCallContext(instanceName, request, DatabaseAdminGrpc.getListDatabasesMethod()); ListDatabasesResponse response = runWithRetryOnAdministrativeRequestsExceeded( () -> get(databaseAdminStub.listDatabasesCallable().futureCall(request, context))); @@ -1300,8 +1298,7 @@ public OperationFuture updateDatabaseDdl( } final UpdateDatabaseDdlRequest request = requestBuilder.build(); final GrpcCallContext context = - newCallContext( - null, + newAdminCallContext( databaseInfo.getId().getName(), request, DatabaseAdminGrpc.getUpdateDatabaseDdlMethod()); @@ -1341,7 +1338,7 @@ public void dropDatabase(String databaseName) throws SpannerException { DropDatabaseRequest.newBuilder().setDatabase(databaseName).build(); final GrpcCallContext context = - newCallContext(null, databaseName, request, DatabaseAdminGrpc.getDropDatabaseMethod()); + newAdminCallContext(databaseName, request, DatabaseAdminGrpc.getDropDatabaseMethod()); runWithRetryOnAdministrativeRequestsExceeded( () -> { get(databaseAdminStub.dropDatabaseCallable().futureCall(request, context)); @@ -1356,7 +1353,7 @@ public Database getDatabase(String databaseName) throws SpannerException { GetDatabaseRequest.newBuilder().setName(databaseName).build(); final GrpcCallContext context = - newCallContext(null, databaseName, request, DatabaseAdminGrpc.getGetDatabaseMethod()); + newAdminCallContext(databaseName, request, DatabaseAdminGrpc.getGetDatabaseMethod()); return runWithRetryOnAdministrativeRequestsExceeded( () -> get(databaseAdminStub.getDatabaseCallable().futureCall(request, context))); } @@ -1367,8 +1364,8 @@ public OperationFuture updateDatabase( UpdateDatabaseRequest request = UpdateDatabaseRequest.newBuilder().setDatabase(database).setUpdateMask(updateMask).build(); GrpcCallContext context = - newCallContext( - null, database.getName(), request, DatabaseAdminGrpc.getUpdateDatabaseMethod()); + newAdminCallContext( + database.getName(), request, DatabaseAdminGrpc.getUpdateDatabaseMethod()); return databaseAdminStub.updateDatabaseOperationCallable().futureCall(request, context); } @@ -1379,7 +1376,7 @@ public GetDatabaseDdlResponse getDatabaseDdl(String databaseName) throws Spanner GetDatabaseDdlRequest.newBuilder().setDatabase(databaseName).build(); final GrpcCallContext context = - newCallContext(null, databaseName, request, DatabaseAdminGrpc.getGetDatabaseDdlMethod()); + newAdminCallContext(databaseName, request, DatabaseAdminGrpc.getGetDatabaseDdlMethod()); return runWithRetryOnAdministrativeRequestsExceeded( () -> get(databaseAdminStub.getDatabaseDdlCallable().futureCall(request, context))); } @@ -1563,7 +1560,7 @@ public Backup updateBackup(Backup backup, FieldMask updateMask) { final UpdateBackupRequest request = UpdateBackupRequest.newBuilder().setBackup(backup).setUpdateMask(updateMask).build(); final GrpcCallContext context = - newCallContext(null, backup.getName(), request, DatabaseAdminGrpc.getUpdateBackupMethod()); + newAdminCallContext(backup.getName(), request, DatabaseAdminGrpc.getUpdateBackupMethod()); return runWithRetryOnAdministrativeRequestsExceeded( () -> databaseAdminStub.updateBackupCallable().call(request, context)); } @@ -1574,7 +1571,7 @@ public void deleteBackup(String backupName) { final DeleteBackupRequest request = DeleteBackupRequest.newBuilder().setName(backupName).build(); final GrpcCallContext context = - newCallContext(null, backupName, request, DatabaseAdminGrpc.getDeleteBackupMethod()); + newAdminCallContext(backupName, request, DatabaseAdminGrpc.getDeleteBackupMethod()); runWithRetryOnAdministrativeRequestsExceeded( () -> { databaseAdminStub.deleteBackupCallable().call(request, context); @@ -1587,7 +1584,7 @@ public Backup getBackup(String backupName) throws SpannerException { acquireAdministrativeRequestsRateLimiter(); final GetBackupRequest request = GetBackupRequest.newBuilder().setName(backupName).build(); final GrpcCallContext context = - newCallContext(null, backupName, request, DatabaseAdminGrpc.getGetBackupMethod()); + newAdminCallContext(backupName, request, DatabaseAdminGrpc.getGetBackupMethod()); return runWithRetryOnAdministrativeRequestsExceeded( () -> get(databaseAdminStub.getBackupCallable().futureCall(request, context))); } @@ -1597,7 +1594,7 @@ public Operation getOperation(String name) throws SpannerException { acquireAdministrativeRequestsRateLimiter(); final GetOperationRequest request = GetOperationRequest.newBuilder().setName(name).build(); final GrpcCallContext context = - newCallContext(null, name, request, OperationsGrpc.getGetOperationMethod()); + newAdminCallContext(name, request, OperationsGrpc.getGetOperationMethod()); return runWithRetryOnAdministrativeRequestsExceeded( () -> get( @@ -1613,7 +1610,7 @@ public void cancelOperation(String name) throws SpannerException { final CancelOperationRequest request = CancelOperationRequest.newBuilder().setName(name).build(); final GrpcCallContext context = - newCallContext(null, name, request, OperationsGrpc.getCancelOperationMethod()); + newAdminCallContext(name, request, OperationsGrpc.getCancelOperationMethod()); runWithRetryOnAdministrativeRequestsExceeded( () -> { get( @@ -1718,10 +1715,16 @@ public StreamingCall read( ReadRequest request, ResultStreamConsumer consumer, @Nullable Map options, + XGoogSpannerRequestId requestId, boolean routeToLeader) { GrpcCallContext context = newCallContext( - options, request.getSession(), request, SpannerGrpc.getReadMethod(), routeToLeader); + options, + requestId, + request.getSession(), + request, + SpannerGrpc.getReadMethod(), + routeToLeader); SpannerResponseObserver responseObserver = new SpannerResponseObserver(consumer); spannerStub.streamingReadCallable().call(request, responseObserver, context); return new GrpcStreamingCall(context, responseObserver.getController()); @@ -1772,10 +1775,14 @@ public RetrySettings getPartitionedDmlRetrySettings() { @Override public ServerStream executeStreamingPartitionedDml( - ExecuteSqlRequest request, Map options, Duration timeout) { + ExecuteSqlRequest request, + Map options, + XGoogSpannerRequestId requestId, + Duration timeout) { GrpcCallContext context = newCallContext( options, + requestId, request.getSession(), request, SpannerGrpc.getExecuteStreamingSqlMethod(), @@ -1798,10 +1805,12 @@ public StreamingCall executeQuery( ExecuteSqlRequest request, ResultStreamConsumer consumer, @Nullable Map options, + XGoogSpannerRequestId requestId, boolean routeToLeader) { GrpcCallContext context = newCallContext( options, + requestId, request.getSession(), request, SpannerGrpc.getExecuteStreamingSqlMethod(), @@ -1987,7 +1996,7 @@ private static T get(final Future future) throws SpannerException { } catch (ExecutionException e) { throw asSpannerException(e.getCause()); } catch (CancellationException e) { - throw newSpannerException(context, e, null); + throw newSpannerException(context, e); } catch (Exception exception) { throw asSpannerException(exception); } @@ -2018,6 +2027,11 @@ GrpcCallContext newCallContext(@Nullable Map options, String resource return newCallContext(options, resource, null, null); } + private GrpcCallContext newAdminCallContext( + String resource, ReqT request, MethodDescriptor method) { + return newCallContext(null, resource, request, method, false); + } + @VisibleForTesting GrpcCallContext newCallContext( @Nullable Map options, @@ -2034,6 +2048,17 @@ GrpcCallContext newCallContext( ReqT request, MethodDescriptor method, boolean routeToLeader) { + return newCallContext(options, /* requestId= */ null, resource, request, method, routeToLeader); + } + + @VisibleForTesting + GrpcCallContext newCallContext( + @Nullable Map options, + @Nullable XGoogSpannerRequestId requestId, + String resource, + ReqT request, + MethodDescriptor method, + boolean routeToLeader) { GrpcCallContext context = this.baseGrpcCallContext; Long affinity = options == null ? null : Option.CHANNEL_HINT.getLong(options); if (affinity != null) { @@ -2052,10 +2077,15 @@ GrpcCallContext newCallContext( context = context.withChannelAffinity(affinity.intValue()); } } - if (options != null) { - // TODO(@odeke-em): Infer the affinity if it doesn't match up with in the request-id. - context = withRequestId(context, options); + int requestIdChannel = convertToRequestIdChannelNumber(affinity); + if (requestId == null) { + requestId = requestIdCreator.nextRequestId(requestIdChannel); + } else { + requestId.setChannelId(requestIdChannel); } + context = + context.withCallOptions( + context.getCallOptions().withOption(REQUEST_ID_CALL_OPTIONS_KEY, requestId)); context = context.withExtraHeaders(metadataProvider.newExtraHeaders(resource, projectName)); if (routeToLeader && leaderAwareRoutingEnabled) { context = context.withExtraHeaders(metadataProvider.newRouteToLeaderHeader()); @@ -2075,17 +2105,19 @@ GrpcCallContext newCallContext( return (GrpcCallContext) context.merge(apiCallContextFromContext); } - GrpcCallContext withRequestId(GrpcCallContext context, Map options) { - XGoogSpannerRequestId reqId = (XGoogSpannerRequestId) options.get(Option.REQUEST_ID); - if (reqId == null) { - return context; - } + @Override + public RequestIdCreator getRequestIdCreator() { + return this.requestIdCreator; + } - Map> withReqId = - ImmutableMap.of( - XGoogSpannerRequestId.REQUEST_HEADER_KEY.name(), - Collections.singletonList(reqId.toString())); - return context.withExtraHeaders(withReqId); + private int convertToRequestIdChannelNumber(@Nullable Long affinity) { + if (affinity == null) { + return 0; + } + int requestIdChannel = affinity.intValue(); + requestIdChannel = requestIdChannel == Integer.MAX_VALUE ? 0 : Math.abs(requestIdChannel); + // Start counting at 1, to distinguish between '0 == Unknown and >0 == known'. + return requestIdChannel % this.numChannels + 1; } void registerResponseObserver(SpannerResponseObserver responseObserver) { diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java index 985ad68a429..861e839a036 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/HeaderInterceptor.java @@ -188,7 +188,7 @@ private void recordSpan(Span span, String requestId) { if (afeLatency != null) { span.setAttribute("afe_latency", afeLatency.toString()); } - span.setAttribute(XGoogSpannerRequestId.REQUEST_ID, requestId); + span.setAttribute(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME, requestId); } } @@ -260,7 +260,7 @@ private DatabaseName extractDatabaseName(Metadata headers) throws ExecutionExcep } private String extractRequestId(Metadata headers) throws ExecutionException { - return headers.get(XGoogSpannerRequestId.REQUEST_HEADER_KEY); + return headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY); } private TagContext getTagContext(String key, String method, DatabaseName databaseName) diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdCreatorImpl.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdCreatorImpl.java new file mode 100644 index 00000000000..5904fa581fd --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdCreatorImpl.java @@ -0,0 +1,43 @@ +/* + * Copyright 2025 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.spi.v1; + +import com.google.cloud.spanner.XGoogSpannerRequestId; +import com.google.cloud.spanner.XGoogSpannerRequestId.RequestIdCreator; +import java.util.concurrent.atomic.AtomicLong; + +class RequestIdCreatorImpl implements RequestIdCreator { + private static final AtomicLong NEXT_CLIENT_ID = new AtomicLong(); + + private final long clientId = NEXT_CLIENT_ID.incrementAndGet(); + private final AtomicLong requestId = new AtomicLong(); + + @Override + public long getClientId() { + return this.clientId; + } + + @Override + public XGoogSpannerRequestId nextRequestId(long channelId) { + return XGoogSpannerRequestId.of(clientId, channelId, requestId.incrementAndGet(), 0); + } + + @Override + public void reset() { + requestId.set(0); + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdInterceptor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdInterceptor.java new file mode 100644 index 00000000000..c32db2a2704 --- /dev/null +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/RequestIdInterceptor.java @@ -0,0 +1,57 @@ +/* + * Copyright 2025 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.spi.v1; + +import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_CALL_OPTIONS_KEY; +import static com.google.cloud.spanner.XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY; + +import com.google.cloud.spanner.XGoogSpannerRequestId; +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.ForwardingClientCall; +import io.grpc.Metadata; +import io.grpc.Metadata.Key; +import io.grpc.MethodDescriptor; +import java.util.concurrent.atomic.AtomicLong; + +class RequestIdInterceptor implements ClientInterceptor { + static final CallOptions.Key ATTEMPT_KEY = CallOptions.Key.create("Attempt"); + private static final String RESPONSE_ENCODING_KEY_NAME = "x-response-encoding"; + private static final Key RESPONSE_ENCODING_KEY = + Key.of(RESPONSE_ENCODING_KEY_NAME, Metadata.ASCII_STRING_MARSHALLER); + + RequestIdInterceptor() {} + + @Override + public ClientCall interceptCall( + MethodDescriptor method, CallOptions callOptions, Channel next) { + return new ForwardingClientCall.SimpleForwardingClientCall( + next.newCall(method, callOptions)) { + @Override + public void start(Listener responseListener, Metadata headers) { + XGoogSpannerRequestId requestId = callOptions.getOption(REQUEST_ID_CALL_OPTIONS_KEY); + if (requestId != null) { + requestId.incrementAttempt(); + headers.put(REQUEST_ID_HEADER_KEY, requestId.getHeaderValue()); + } + super.start(responseListener, headers); + } + }; + } +} diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerErrorInterceptor.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerErrorInterceptor.java index 549ea18a97f..9c3b2af2b06 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerErrorInterceptor.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerErrorInterceptor.java @@ -17,6 +17,8 @@ package com.google.cloud.spanner.spi.v1; import com.google.cloud.spanner.IsRetryableInternalError; +import com.google.cloud.spanner.XGoogSpannerRequestId; +import com.google.common.base.Strings; import com.google.rpc.BadRequest; import com.google.rpc.Help; import com.google.rpc.LocalizedMessage; @@ -35,6 +37,7 @@ import io.grpc.Status; import io.grpc.Status.Code; import io.grpc.protobuf.ProtoUtils; +import java.util.Objects; import java.util.logging.Level; import java.util.logging.Logger; @@ -70,7 +73,23 @@ public void start(Listener responseListener, Metadata headers) { new SimpleForwardingClientCallListener(responseListener) { @Override public void onClose(Status status, Metadata trailers) { + // Return quickly if there is no error. + if (status.isOk()) { + super.onClose(status, trailers); + return; + } try { + if (headers.containsKey(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY)) { + String requestId = headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY); + if (!Strings.isNullOrEmpty(requestId)) { + if (!trailers.containsKey(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY)) { + trailers.put( + XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY, + Objects.requireNonNull( + headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY))); + } + } + } // Translate INTERNAL errors that should be retried to a retryable error code. if (IsRetryableInternalError.INSTANCE.isRetryableInternalError(status)) { status = diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerInterceptorProvider.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerInterceptorProvider.java index ec7a4e54a03..e8d6c3ebddb 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerInterceptorProvider.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerInterceptorProvider.java @@ -64,6 +64,7 @@ public static SpannerInterceptorProvider createDefault( defaultInterceptorList.add( new LoggingInterceptor(Logger.getLogger(GapicSpannerRpc.class.getName()), Level.FINER)); defaultInterceptorList.add(new HeaderInterceptor(new SpannerRpcMetrics(openTelemetry))); + defaultInterceptorList.add(new RequestIdInterceptor()); return new SpannerInterceptorProvider(ImmutableList.copyOf(defaultInterceptorList)); } diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java index 4b5682bb2b0..304729247f7 100644 --- a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java +++ b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/SpannerRpc.java @@ -27,6 +27,8 @@ import com.google.cloud.spanner.BackupId; import com.google.cloud.spanner.Restore; import com.google.cloud.spanner.SpannerException; +import com.google.cloud.spanner.XGoogSpannerRequestId; +import com.google.cloud.spanner.XGoogSpannerRequestId.RequestIdCreator; import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStub; import com.google.cloud.spanner.admin.database.v1.stub.DatabaseAdminStubSettings; import com.google.cloud.spanner.admin.instance.v1.stub.InstanceAdminStub; @@ -78,8 +80,7 @@ public interface SpannerRpc extends ServiceRpc { /** Options passed in {@link SpannerRpc} methods to control how an RPC is issued. */ enum Option { - CHANNEL_HINT("Channel Hint"), - REQUEST_ID("Request Id"); + CHANNEL_HINT("Channel Hint"); private final String value; @@ -189,6 +190,10 @@ interface StreamingCall { void cancel(@Nullable String message); } + default RequestIdCreator getRequestIdCreator() { + throw new UnsupportedOperationException("Not implemented"); + } + // Instance admin APIs. Paginated listInstanceConfigs(int pageSize, @Nullable String pageToken) throws SpannerException; @@ -389,6 +394,7 @@ StreamingCall read( ReadRequest request, ResultStreamConsumer consumer, @Nullable Map options, + XGoogSpannerRequestId requestId, boolean routeToLeader); /** Returns the retry settings for streaming query operations. */ @@ -428,7 +434,10 @@ ApiFuture executeQueryAsync( RetrySettings getPartitionedDmlRetrySettings(); ServerStream executeStreamingPartitionedDml( - ExecuteSqlRequest request, @Nullable Map options, Duration timeout); + ExecuteSqlRequest request, + @Nullable Map options, + XGoogSpannerRequestId requestId, + Duration timeout); ServerStream batchWriteAtLeastOnce( BatchWriteRequest request, @Nullable Map options); @@ -445,6 +454,7 @@ StreamingCall executeQuery( ExecuteSqlRequest request, ResultStreamConsumer consumer, @Nullable Map options, + XGoogSpannerRequestId requestId, boolean routeToLeader); ExecuteBatchDmlResponse executeBatchDml(ExecuteBatchDmlRequest build, Map options); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java index 939114a7f60..fc177ebb420 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/BaseSessionPoolTest.java @@ -27,6 +27,7 @@ import com.google.api.core.ApiFutures; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; import com.google.cloud.spanner.Options.TransactionOption; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; import com.google.cloud.spanner.spi.v1.SpannerRpc.Option; import com.google.protobuf.Empty; import com.google.protobuf.Timestamp; @@ -72,6 +73,7 @@ SessionImpl mockSession() { final SessionImpl session = mock(SessionImpl.class); Map options = new HashMap<>(); options.put(Option.CHANNEL_HINT, channelHint.getAndIncrement()); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); when(session.getOptions()).thenReturn(options); when(session.getName()) .thenReturn( diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java index f7842fc7a40..7b85c8e772e 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java @@ -42,7 +42,6 @@ import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; -import com.google.api.gax.grpc.testing.LocalChannelProvider; import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.rpc.ApiCallContext; import com.google.api.gax.rpc.ServerStream; @@ -104,16 +103,18 @@ import com.google.spanner.v1.TypeAnnotationCode; import com.google.spanner.v1.TypeCode; import io.grpc.Context; +import io.grpc.ManagedChannelBuilder; import io.grpc.Metadata; import io.grpc.MethodDescriptor; import io.grpc.Server; import io.grpc.ServerInterceptors; import io.grpc.Status; import io.grpc.StatusRuntimeException; -import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder; import io.grpc.protobuf.lite.ProtoLiteUtils; import io.opencensus.trace.Tracing; import io.opentelemetry.api.OpenTelemetry; +import java.net.InetSocketAddress; import java.nio.charset.StandardCharsets; import java.time.Duration; import java.time.Instant; @@ -128,7 +129,6 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; -import java.util.concurrent.ScheduledThreadPoolExecutor; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -159,7 +159,6 @@ public class DatabaseClientImplTest { private static XGoogSpannerRequestIdTest.ServerHeaderEnforcer xGoogReqIdInterceptor; private static MockSpannerServiceImpl mockSpanner; private static Server server; - private static LocalChannelProvider channelProvider; private static final Statement UPDATE_STATEMENT = Statement.of("UPDATE FOO SET BAR=1 WHERE BAZ=2"); private static final Statement INVALID_UPDATE_STATEMENT = @@ -247,15 +246,12 @@ public static void startStaticServer() throws Exception { "google.spanner.v1.Spanner/StreamingRead")); xGoogReqIdInterceptor = new XGoogSpannerRequestIdTest.ServerHeaderEnforcer(checkMethods); executor = Executors.newSingleThreadExecutor(); - String uniqueName = InProcessServerBuilder.generateName(); + InetSocketAddress address = new InetSocketAddress("localhost", 0); server = - InProcessServerBuilder.forName(uniqueName) - // We need to use a real executor for timeouts to occur. - .scheduledExecutorService(new ScheduledThreadPoolExecutor(1)) + NettyServerBuilder.forAddress(address) .addService(ServerInterceptors.intercept(mockSpanner, xGoogReqIdInterceptor)) .build() .start(); - channelProvider = LocalChannelProvider.create(uniqueName); } @AfterClass @@ -267,11 +263,13 @@ public static void stopServer() throws InterruptedException { @Before public void setUp() { + String endpoint = "localhost:" + server.getPort(); spanner = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://" + endpoint) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(SessionPoolOptions.newBuilder().setFailOnSessionLeak().build()) .build() @@ -321,7 +319,8 @@ public void tearDown() { SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -401,7 +400,8 @@ public void tearDown() { SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -473,7 +473,8 @@ public void testPoolMaintainer_whenLongRunningPartitionedUpdateRequest_takeNoAct SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -539,7 +540,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -631,7 +633,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -703,7 +706,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -765,7 +769,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -830,7 +835,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -909,7 +915,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -976,7 +983,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -1047,7 +1055,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -1114,7 +1123,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -1181,7 +1191,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -1250,7 +1261,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -1319,7 +1331,8 @@ public void testPoolMaintainer_whenPDMLFollowedByInactiveTransaction_removeSessi SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) .setDatabaseRole(TEST_DATABASE_ROLE) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(sessionPoolOptions) .build() @@ -2905,7 +2918,8 @@ public void testPartitionedDmlDoesNotTimeout() { SpannerOptions.Builder builder = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()); // Set normal DML timeout value. builder.getSpannerStubSettingsBuilder().executeSqlSettings().setRetrySettings(retrySettings); @@ -2974,7 +2988,8 @@ public void testPartitionedDmlWithLowerTimeout() { SpannerOptions.Builder builder = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()); // Set PDML timeout value. builder.setPartitionedDmlTimeoutDuration(Duration.ofMillis(10L)); @@ -3008,7 +3023,8 @@ public void testPartitionedDmlWithHigherTimeout() { SpannerOptions.Builder builder = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()); // Set PDML timeout value to a value that should allow the statement to be executed. builder.setPartitionedDmlTimeoutDuration(Duration.ofMillis(5000L)); @@ -3088,7 +3104,8 @@ public void testPartitionedDmlRetriesOnUnavailable() { SpannerOptions.Builder builder = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()); try (Spanner spanner = builder.build().getService()) { DatabaseClient client = @@ -3111,7 +3128,8 @@ public void testDatabaseOrInstanceDoesNotExistOnInitialization() throws Exceptio try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .build() .getService()) { @@ -3156,7 +3174,8 @@ public void testDatabaseOrInstanceDoesNotExistOnCreate() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption( SessionPoolOptions.newBuilder() @@ -3227,7 +3246,8 @@ public void testDatabaseOrInstanceDoesNotExistOnReplenish() throws Exception { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .build() .getService()) { @@ -3283,7 +3303,8 @@ public void testDatabaseOrInstanceIsDeletedAndThenRecreated() throws Exception { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .build() .getService()) { @@ -3372,7 +3393,8 @@ public void testGetInvalidatedClientMultipleTimes() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build()) .build() @@ -3511,7 +3533,8 @@ public void testBackendQueryOptions() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build()) .build() @@ -3552,7 +3575,8 @@ public void testBackendQueryOptionsWithAnalyzeQuery() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build()) .build() @@ -3595,7 +3619,8 @@ public void testBackendPartitionQueryOptions() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build()) .setDirectedReadOptions(DIRECTED_READ_OPTIONS2) @@ -3653,7 +3678,8 @@ public void testBackendPartitionQueryOptions() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build()) .setDirectedReadOptions(DIRECTED_READ_OPTIONS2) @@ -3709,7 +3735,8 @@ public void testBackendPartitionReadOptions() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build()) .setDirectedReadOptions(DIRECTED_READ_OPTIONS2) @@ -3759,7 +3786,8 @@ public void testBackendPartitionReadOptions() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId("[PROJECT]") - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption(SessionPoolOptions.newBuilder().setMinSessions(0).build()) .setDirectedReadOptions(DIRECTED_READ_OPTIONS2) @@ -3858,7 +3886,8 @@ public void testClientIdReusedOnDatabaseNotFound() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId("my-project") - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .build() .getService()) { @@ -3894,7 +3923,8 @@ public void testBatchCreateSessionsPermissionDenied() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId("my-project") - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption( SessionPoolOptions.newBuilder() @@ -5271,7 +5301,8 @@ public void testRetryOnResourceExhausted() { SpannerOptions.Builder builder = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()); RetryInfo retryInfo = RetryInfo.newBuilder() @@ -5371,7 +5402,8 @@ public void testSessionPoolExhaustedError_containsStackTraces() { try (Spanner spanner = SpannerOptions.newBuilder() .setProjectId(TEST_PROJECT) - .setChannelProvider(channelProvider) + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) .setCredentials(NoCredentials.getInstance()) .setSessionPoolOption( SessionPoolOptions.newBuilder() diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java index 8f71b1278f3..c6155f0cbb6 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/PartitionedDmlTransactionTest.java @@ -33,6 +33,7 @@ import com.google.api.gax.rpc.ServerStream; import com.google.api.gax.rpc.UnavailableException; import com.google.cloud.spanner.Options.RpcPriority; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.common.base.Ticker; import com.google.common.collect.ImmutableList; @@ -96,11 +97,9 @@ public class PartitionedDmlTransactionTest { public void setup() { MockitoAnnotations.initMocks(this); when(session.getName()).thenReturn(sessionId); - when(session.getRequestIdCreator()) - .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator()); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); when(session.getOptions()).thenReturn(Collections.EMPTY_MAP); - when(session.getRequestIdCreator()) - .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator()); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); when(rpc.beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true))) .thenReturn(Transaction.newBuilder().setId(txId).build()); @@ -115,7 +114,7 @@ public void testExecuteStreamingPartitionedUpdate() { ServerStream stream = mock(ServerStream.class); when(stream.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator()); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream); long count = tx.executeStreamingPartitionedUpdate(Statement.of(sql), Duration.ofMinutes(10)); @@ -124,7 +123,7 @@ public void testExecuteStreamingPartitionedUpdate() { verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)); } @Test @@ -135,7 +134,7 @@ public void testExecuteStreamingPartitionedUpdateWithUpdateOptions() { ServerStream stream = mock(ServerStream.class); when(stream.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator()); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithRequestOptions), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithRequestOptions), anyMap(), any(), any(Duration.class))) .thenReturn(stream); long count = @@ -146,7 +145,7 @@ public void testExecuteStreamingPartitionedUpdateWithUpdateOptions() { verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithRequestOptions), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithRequestOptions), anyMap(), any(), any(Duration.class)); } @Test @@ -166,7 +165,7 @@ public void testExecuteStreamingPartitionedUpdateAborted() { ServerStream stream2 = mock(ServerStream.class); when(stream2.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator()); when(rpc.executeStreamingPartitionedDml( - any(ExecuteSqlRequest.class), anyMap(), any(Duration.class))) + any(ExecuteSqlRequest.class), anyMap(), any(), any(Duration.class))) .thenReturn(stream1, stream2); long count = tx.executeStreamingPartitionedUpdate(Statement.of(sql), Duration.ofMinutes(10)); @@ -175,7 +174,7 @@ public void testExecuteStreamingPartitionedUpdateAborted() { verify(rpc, times(2)).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)); verify(rpc, times(2)) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)); } @Test @@ -195,10 +194,10 @@ public void testExecuteStreamingPartitionedUpdateUnavailable() { ServerStream stream2 = mock(ServerStream.class); when(stream2.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator()); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream1); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream2); long count = tx.executeStreamingPartitionedUpdate(Statement.of(sql), Duration.ofMinutes(10)); @@ -207,10 +206,10 @@ public void testExecuteStreamingPartitionedUpdateUnavailable() { verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class)); } @Test @@ -226,7 +225,7 @@ public void testExecuteStreamingPartitionedUpdateUnavailableAndThenDeadlineExcee "temporary unavailable", null, GrpcStatusCode.of(Code.UNAVAILABLE), true)); when(stream1.iterator()).thenReturn(iterator); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream1); when(ticker.read()).thenReturn(0L, 1L, TimeUnit.NANOSECONDS.convert(10L, TimeUnit.MINUTES)); @@ -238,7 +237,7 @@ public void testExecuteStreamingPartitionedUpdateUnavailableAndThenDeadlineExcee verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)); } @Test @@ -254,7 +253,7 @@ public void testExecuteStreamingPartitionedUpdateAbortedAndThenDeadlineExceeded( "transaction aborted", null, GrpcStatusCode.of(Code.ABORTED), true)); when(stream1.iterator()).thenReturn(iterator); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream1); when(ticker.read()).thenReturn(0L, 1L, TimeUnit.NANOSECONDS.convert(10L, TimeUnit.MINUTES)); @@ -266,7 +265,7 @@ public void testExecuteStreamingPartitionedUpdateAbortedAndThenDeadlineExceeded( verify(rpc, times(2)).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)); } @Test @@ -282,7 +281,7 @@ public void testExecuteStreamingPartitionedUpdateMultipleAbortsUntilDeadlineExce "transaction aborted", null, GrpcStatusCode.of(Code.ABORTED), true)); when(stream1.iterator()).thenReturn(iterator); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream1); when(ticker.read()) .thenAnswer( @@ -306,7 +305,7 @@ public Long answer(InvocationOnMock invocation) { // means that the execute method is only executed 9 times. verify(rpc, times(9)) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)); } @Test @@ -329,10 +328,10 @@ public void testExecuteStreamingPartitionedUpdateUnexpectedEOS() { ServerStream stream2 = mock(ServerStream.class); when(stream2.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator()); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream1); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream2); PartitionedDmlTransaction tx = new PartitionedDmlTransaction(session, rpc, ticker); @@ -342,10 +341,10 @@ public void testExecuteStreamingPartitionedUpdateUnexpectedEOS() { verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class)); } @Test @@ -368,10 +367,10 @@ public void testExecuteStreamingPartitionedUpdateRSTstream() { ServerStream stream2 = mock(ServerStream.class); when(stream2.iterator()).thenReturn(ImmutableList.of(p1, p2).iterator()); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream1); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream2); PartitionedDmlTransaction tx = new PartitionedDmlTransaction(session, rpc, ticker); @@ -381,10 +380,10 @@ public void testExecuteStreamingPartitionedUpdateRSTstream() { verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithResumeToken), anyMap(), any(), any(Duration.class)); } @Test @@ -400,7 +399,7 @@ public void testExecuteStreamingPartitionedUpdateGenericInternalException() { "INTERNAL: Error", null, GrpcStatusCode.of(Code.INTERNAL), false)); when(stream1.iterator()).thenReturn(iterator); when(rpc.executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class))) + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class))) .thenReturn(stream1); PartitionedDmlTransaction tx = new PartitionedDmlTransaction(session, rpc, ticker); @@ -412,7 +411,7 @@ public void testExecuteStreamingPartitionedUpdateGenericInternalException() { verify(rpc).beginTransaction(any(BeginTransactionRequest.class), anyMap(), eq(true)); verify(rpc) .executeStreamingPartitionedDml( - Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(Duration.class)); + Mockito.eq(executeRequestWithoutResumeToken), anyMap(), any(), any(Duration.class)); } @Test diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RequestIdMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RequestIdMockServerTest.java new file mode 100644 index 00000000000..c32a633331b --- /dev/null +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/RequestIdMockServerTest.java @@ -0,0 +1,727 @@ +/* + * Copyright 2025 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.DisableDefaultMtlsProvider.disableDefaultMtlsProvider; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.junit.Assume.assumeTrue; + +import com.google.cloud.NoCredentials; +import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime; +import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult; +import com.google.cloud.spanner.connection.RandomResultSetGenerator; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.ByteString; +import com.google.protobuf.ListValue; +import com.google.protobuf.Value; +import com.google.rpc.RetryInfo; +import com.google.spanner.v1.BeginTransactionRequest; +import com.google.spanner.v1.CommitRequest; +import com.google.spanner.v1.CreateSessionRequest; +import com.google.spanner.v1.ExecuteSqlRequest; +import com.google.spanner.v1.ReadRequest; +import com.google.spanner.v1.ResultSetMetadata; +import com.google.spanner.v1.ResultSetStats; +import com.google.spanner.v1.StructType; +import com.google.spanner.v1.StructType.Field; +import com.google.spanner.v1.Type; +import com.google.spanner.v1.TypeCode; +import io.grpc.Context; +import io.grpc.Contexts; +import io.grpc.ManagedChannelBuilder; +import io.grpc.Metadata; +import io.grpc.Server; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; +import io.grpc.Status; +import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder; +import io.grpc.protobuf.ProtoUtils; +import java.net.InetSocketAddress; +import java.util.List; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.TimeUnit; +import org.junit.AfterClass; +import org.junit.Before; +import org.junit.BeforeClass; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.threeten.bp.Duration; + +@SuppressWarnings({"StatementWithEmptyBody", "resource"}) +@RunWith(JUnit4.class) +public class RequestIdMockServerTest { + private static MockSpannerServiceImpl mockSpanner; + private static Server server; + private static Spanner spanner; + + private static final Statement SELECT1 = Statement.of("SELECT 1"); + private static final com.google.spanner.v1.ResultSet SELECT1_RESULT_SET = + com.google.spanner.v1.ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType( + StructType.newBuilder() + .addFields( + Field.newBuilder() + .setName("c") + .setType(Type.newBuilder().setCode(TypeCode.INT64).build()) + .build()) + .build()) + .build()) + .addRows( + ListValue.newBuilder() + .addValues(Value.newBuilder().setStringValue("1").build()) + .build()) + .build(); + private static final Statement DML = Statement.of("insert into test_table (id) values (1)"); + + private static final ConcurrentLinkedQueue requestIds = + new ConcurrentLinkedQueue<>(); + + @BeforeClass + public static void setup() throws Exception { + assumeTrue(System.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS") == null); + assumeTrue(System.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW") == null); + + disableDefaultMtlsProvider(); + mockSpanner = new MockSpannerServiceImpl(); + mockSpanner.setAbortProbability(0.0D); // We don't want any unpredictable aborted transactions. + + InetSocketAddress address = new InetSocketAddress("localhost", 0); + server = + NettyServerBuilder.forAddress(address) + .addService(mockSpanner) + .intercept( + new ServerInterceptor() { + @Override + public ServerCall.Listener interceptCall( + ServerCall call, + Metadata headers, + ServerCallHandler next) { + try { + String requestId = headers.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY); + if (requestId != null) { + requestIds.add(XGoogSpannerRequestId.of(requestId)); + } else { + requestIds.add(XGoogSpannerRequestId.of(0, 0, 0, 0)); + } + } catch (Throwable t) { + // Ignore and continue + } + return Contexts.interceptCall(Context.current(), call, headers, next); + } + }) + .build() + .start(); + spanner = createSpanner(); + + setupResults(); + } + + private static Spanner createSpanner() { + return SpannerOptions.newBuilder() + .setProjectId("test-project") + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .setHost("http://localhost:" + server.getPort()) + .setCredentials(NoCredentials.getInstance()) + .setSessionPoolOption( + SessionPoolOptions.newBuilder() + .setFailOnSessionLeak() + .setMinSessions(0) + .setSkipVerifyingBeginTransactionForMuxRW(true) + .setWaitForMinSessions(Duration.ofSeconds(5)) + .build()) + .build() + .getService(); + } + + private static void setupResults() { + mockSpanner.putStatementResult(StatementResult.query(SELECT1, SELECT1_RESULT_SET)); + mockSpanner.putStatementResult(StatementResult.update(DML, 1L)); + } + + static Metadata createMinimalRetryInfo() { + Metadata trailers = new Metadata(); + RetryInfo retryInfo = + RetryInfo.newBuilder() + .setRetryDelay( + com.google.protobuf.Duration.newBuilder() + .setNanos((int) TimeUnit.MILLISECONDS.toNanos(1L)) + .setSeconds(0L)) + .build(); + trailers.put(ProtoUtils.keyForProto(RetryInfo.getDefaultInstance()), retryInfo); + return trailers; + } + + @AfterClass + public static void teardown() throws InterruptedException { + if (spanner != null) { + spanner.close(); + } + if (server != null) { + server.shutdown(); + server.awaitTermination(); + } + } + + @Before + public void prepareTest() { + // Call getClient() to make sure the multiplexed session has been created. + // Then clear all requests that were received as part of that so we don't need to include + // that in the test verifications. + getClient(); + mockSpanner.reset(); + requestIds.clear(); + ((SpannerImpl) spanner).resetRequestIdCounters(); + } + + private DatabaseClient getClient() { + return spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + } + + private long getClientId() { + return ((SpannerImpl) spanner).getRequestIdClientId(); + } + + @Test + public void testSingleUseQuery() { + try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + + assertEquals(ImmutableList.of(ExecuteSqlRequest.class), mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds(ImmutableList.of(XGoogSpannerRequestId.of(getClientId(), -1, 1, 1)), actual); + } + + @Test + public void testQueryError() { + Statement query = Statement.of("select * from invalid_table"); + mockSpanner.putStatementResult( + StatementResult.exception( + query, Status.NOT_FOUND.withDescription("Table not found").asRuntimeException())); + + XGoogSpannerRequestId requestIdFromException; + try (ResultSet resultSet = getClient().singleUse().executeQuery(query)) { + SpannerException exception = assertThrows(SpannerException.class, resultSet::next); + assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode()); + assertNotNull(exception.getRequestId()); + assertNotEquals("Request ID should not be empty", "", exception.getRequestId()); + requestIdFromException = XGoogSpannerRequestId.of(exception.getRequestId()); + } + + assertEquals(ImmutableList.of(ExecuteSqlRequest.class), mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds(ImmutableList.of(XGoogSpannerRequestId.of(getClientId(), -1, 1, 1)), actual); + assertEquals(actual.get(0), requestIdFromException); + } + + @Test + public void testMultiUseReadOnlyTransaction() { + try (ReadOnlyTransaction transaction = getClient().readOnlyTransaction()) { + for (int i = 0; i < 2; i++) { + try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + } + } + + assertEquals( + ImmutableList.of( + BeginTransactionRequest.class, ExecuteSqlRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 3, 1)), + actual); + verifySameChannelId(actual); + } + + @Test + public void testDml() { + getClient().readWriteTransaction().run(transaction -> transaction.executeUpdate(DML)); + + assertEquals( + ImmutableList.of(ExecuteSqlRequest.class, CommitRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)), + actual); + verifySameChannelId(actual); + } + + @Test + public void testDmlError() { + Statement invalidDml = Statement.of("insert into invalid_table (id) values (1)"); + mockSpanner.putStatementResult( + StatementResult.exception( + invalidDml, Status.NOT_FOUND.withDescription("Table not found").asRuntimeException())); + + SpannerException exception = + assertThrows( + SpannerException.class, + () -> + getClient() + .readWriteTransaction() + .run(transaction -> transaction.executeUpdate(invalidDml))); + assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode()); + assertNotNull(exception.getRequestId()); + assertNotEquals("Request ID should not be empty", "", exception.getRequestId()); + XGoogSpannerRequestId requestIdFromException = + XGoogSpannerRequestId.of(exception.getRequestId()); + + assertEquals(ImmutableList.of(ExecuteSqlRequest.class), mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds(ImmutableList.of(XGoogSpannerRequestId.of(getClientId(), -1, 1, 1)), actual); + assertEquals(actual.get(0), requestIdFromException); + } + + @Test + public void testAbortedTransaction() { + mockSpanner.setCommitExecutionTime( + SimulatedExecutionTime.ofException( + Status.ABORTED.asRuntimeException(createMinimalRetryInfo()))); + getClient() + .readWriteTransaction() + .run( + transaction -> { + try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + return transaction.executeUpdate(DML); + }); + + assertEquals( + ImmutableList.of( + ExecuteSqlRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class, + ExecuteSqlRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + int requestId = 0; + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1)), + actual); + verifySameChannelId(actual.subList(0, 3)); + verifySameChannelId(actual.subList(3, 6)); + } + + @Test + public void testMix() { + getClient() + .readWriteTransaction() + .run( + transaction -> { + try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + return transaction.executeUpdate(DML); + }); + try (ReadOnlyTransaction transaction = getClient().readOnlyTransaction()) { + for (int i = 0; i < 2; i++) { + try (ResultSet resultSet = transaction.executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + } + } + try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + mockSpanner.putStatementResult( + StatementResult.query( + Statement.of("SELECT my_column FROM my_table WHERE 1=1"), SELECT1_RESULT_SET)); + try (ResultSet resultSet = + getClient().singleUse().read("my_table", KeySet.all(), ImmutableList.of("my_column"))) { + while (resultSet.next()) {} + } + + assertEquals( + ImmutableList.of( + ExecuteSqlRequest.class, + ExecuteSqlRequest.class, + CommitRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + ExecuteSqlRequest.class, + ExecuteSqlRequest.class, + ReadRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + int requestId = 0; + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1), + XGoogSpannerRequestId.of(getClientId(), -1, ++requestId, 1)), + actual); + verifySameChannelId(actual.subList(0, 3)); + verifySameChannelId(actual.subList(3, 6)); + } + + @Test + public void testUnaryUnavailable() { + mockSpanner.setExecuteSqlExecutionTime( + SimulatedExecutionTime.ofException( + Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()))); + + getClient().readWriteTransaction().run(transaction -> transaction.executeUpdate(DML)); + + assertEquals( + ImmutableList.of(ExecuteSqlRequest.class, ExecuteSqlRequest.class, CommitRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 1, 2), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)), + actual); + verifySameChannelId(actual); + } + + @Test + public void testStreamingQueryUnavailable() { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException( + Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()))); + + try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + + assertEquals( + ImmutableList.of(ExecuteSqlRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 1, 2)), + actual); + } + + @Test + public void testStreamingQueryUnavailableHalfway() { + int numRows = 5; + Statement statement = Statement.of("select * from random"); + mockSpanner.putStatementResult( + StatementResult.query(statement, new RandomResultSetGenerator(numRows).generate())); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()), 2)); + + try (ResultSet resultSet = getClient().singleUse().executeQuery(statement)) { + while (resultSet.next()) {} + } + + assertEquals( + ImmutableList.of(ExecuteSqlRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 1, 2)), + actual); + List requests = mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(ByteString.empty(), requests.get(0).getResumeToken()); + assertNotEquals(ByteString.empty(), requests.get(1).getResumeToken()); + } + + @Test + public void testStreamingReadUnavailable() { + mockSpanner.setStreamingReadExecutionTime( + SimulatedExecutionTime.ofException( + Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()))); + + mockSpanner.putStatementResult( + StatementResult.query( + Statement.of("SELECT my_column FROM my_table WHERE 1=1"), SELECT1_RESULT_SET)); + try (ResultSet resultSet = + getClient().singleUse().read("my_table", KeySet.all(), ImmutableList.of("my_column"))) { + while (resultSet.next()) {} + } + + assertEquals( + ImmutableList.of(ReadRequest.class, ReadRequest.class), mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 1, 2)), + actual); + } + + @Test + public void testStreamingReadUnavailableHalfway() { + int numRows = 5; + Statement statement = Statement.of("SELECT my_column FROM my_table WHERE 1=1"); + mockSpanner.putStatementResult( + StatementResult.query(statement, new RandomResultSetGenerator(numRows).generate())); + mockSpanner.setStreamingReadExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()), 2)); + + try (ResultSet resultSet = + getClient().singleUse().read("my_table", KeySet.all(), ImmutableList.of("my_column"))) { + while (resultSet.next()) {} + } + + assertEquals( + ImmutableList.of(ReadRequest.class, ReadRequest.class), mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 1, 2)), + actual); + List requests = mockSpanner.getRequestsOfType(ReadRequest.class); + assertEquals(ByteString.empty(), requests.get(0).getResumeToken()); + assertNotEquals(ByteString.empty(), requests.get(1).getResumeToken()); + } + + @Test + public void testPartitionedDml() { + getClient().executePartitionedUpdate(DML); + + assertEquals( + ImmutableList.of(BeginTransactionRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)), + actual); + verifySameChannelId(actual); + } + + @Test + public void testPartitionedDmlError() { + Statement invalidDml = Statement.of("update invalid_table set col=true where col=false"); + mockSpanner.putStatementResult( + StatementResult.exception( + invalidDml, Status.NOT_FOUND.withDescription("Table not found").asRuntimeException())); + + SpannerException exception = + assertThrows( + SpannerException.class, () -> getClient().executePartitionedUpdate(invalidDml)); + assertEquals(ErrorCode.NOT_FOUND, exception.getErrorCode()); + assertNotNull(exception.getRequestId()); + assertNotEquals("", exception.getRequestId()); + + assertEquals( + ImmutableList.of(BeginTransactionRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)), + actual); + verifySameChannelId(actual); + assertEquals(XGoogSpannerRequestId.of(exception.getRequestId()), actual.get(1)); + } + + @Test + public void testPartitionedDmlAborted() { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException( + Status.ABORTED.asRuntimeException(createMinimalRetryInfo()))); + + getClient().executePartitionedUpdate(DML); + + assertEquals( + ImmutableList.of( + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 3, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 4, 1)), + actual); + verifySameChannelId(actual); + } + + @Test + public void testPartitionedDmlUnavailable() { + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofException( + Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()))); + + getClient().executePartitionedUpdate(DML); + + assertEquals( + ImmutableList.of( + BeginTransactionRequest.class, + ExecuteSqlRequest.class, + BeginTransactionRequest.class, + ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 3, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 4, 1)), + actual); + verifySameChannelId(actual); + } + + @Test + public void testPartitionedDmlUnavailableWithResumeToken() { + Statement update = Statement.of("UPDATE my_table SET active=true where 1=1"); + mockSpanner.putStatementResult( + StatementResult.query( + update, + com.google.spanner.v1.ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setRowType(StructType.newBuilder().build()) + .build()) + .addRows(ListValue.newBuilder().build()) + .addRows(ListValue.newBuilder().build()) + .addRows(ListValue.newBuilder().build()) + .setStats(ResultSetStats.newBuilder().setRowCountLowerBound(100L).build()) + .build())); + mockSpanner.setExecuteStreamingSqlExecutionTime( + SimulatedExecutionTime.ofStreamException( + Status.UNAVAILABLE.asRuntimeException(createMinimalRetryInfo()), 2L)); + + getClient().executePartitionedUpdate(update); + + assertEquals( + ImmutableList.of( + BeginTransactionRequest.class, ExecuteSqlRequest.class, ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 2)), + actual); + verifySameChannelId(actual); + } + + @Test + public void testOtherClientId() { + // Execute a query with the default client from this test class. + try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + // Create a new client and use that to execute a query. This should use a different client ID. + long otherClientId; + try (Spanner spanner = createSpanner()) { + otherClientId = ((SpannerImpl) spanner).getRequestIdClientId(); + DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d")); + try (ResultSet resultSet = client.singleUse().executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + } + // Execute another query with the default client. This should use the original client ID. + try (ResultSet resultSet = getClient().singleUse().executeQuery(SELECT1)) { + while (resultSet.next()) {} + } + assertEquals( + ImmutableList.of( + ExecuteSqlRequest.class, + CreateSessionRequest.class, + ExecuteSqlRequest.class, + ExecuteSqlRequest.class), + mockSpanner.getRequestTypes()); + List actual = ImmutableList.copyOf(requestIds); + verifyRequestIds( + ImmutableList.of( + XGoogSpannerRequestId.of(getClientId(), -1, 1, 1), + // The CreateSession RPC from the initialization of the second client is included in + // the requests that we see. This request does not include a channel hint, hence the + // zero value for the channel number in the request ID. + XGoogSpannerRequestId.of(otherClientId, 0, 1, 1), + XGoogSpannerRequestId.of(otherClientId, -1, 2, 1), + XGoogSpannerRequestId.of(getClientId(), -1, 2, 1)), + actual); + } + + private void verifyRequestIds( + List expectedIds, List actualIds) { + assertEquals(message(expectedIds, actualIds), expectedIds.size(), actualIds.size()); + int i = 0; + for (XGoogSpannerRequestId actual : actualIds) { + XGoogSpannerRequestId expected = expectedIds.get(i); + if (expected.getNthChannelId() > -1) { + assertEquals(expected, actual); + } else { + assertTrue(message(expectedIds, actualIds), equalsIgnoringChannelId(expected, actual)); + assertTrue(message(expectedIds, actualIds), actual.hasChannelId()); + } + i++; + } + } + + private void verifySameChannelId(List requestIds) { + for (int i = 0; i < requestIds.size() - 1; i++) { + XGoogSpannerRequestId requestId = requestIds.get(i); + assertTrue(requestId.hasChannelId()); + assertEquals(requestId.getNthChannelId(), requestIds.get(i + 1).getNthChannelId()); + } + } + + private boolean equalsIgnoringChannelId( + XGoogSpannerRequestId expected, XGoogSpannerRequestId actual) { + return expected.getNthClientId() == actual.getNthClientId() + && expected.getNthRequest() == actual.getNthRequest() + && expected.getAttempt() == actual.getAttempt(); + } + + private String message(List expected, List actual) { + return String.format("\n Got: %s\nWant: %s", actual, expected); + } +} diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java index 14f8df04e97..f13c0bb1237 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/ResumableStreamIteratorTest.java @@ -25,6 +25,7 @@ import com.google.api.client.util.BackOff; import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; import com.google.common.collect.AbstractIterator; import com.google.common.collect.ImmutableList; @@ -168,11 +169,12 @@ private void initWithLimit(int maxBufferSize) { DefaultErrorHandler.INSTANCE, SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings(), SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes(), - new XGoogSpannerRequestId.NoopRequestIdCreator()) { + NoopRequestIdCreator.INSTANCE) { @Override AbstractResultSet.CloseableIterator startStream( @Nullable ByteString resumeToken, - AsyncResultSet.StreamMessageListener streamMessageListener) { + AsyncResultSet.StreamMessageListener streamMessageListener, + XGoogSpannerRequestId requestId) { return starter.startStream(resumeToken, null); } }; diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionClientTests.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionClientTests.java index 5d3ed2bca53..07a76970bfd 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionClientTests.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionClientTests.java @@ -18,7 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doNothing; @@ -163,7 +163,6 @@ public void createAndCloseSession() { assertEquals( deleteOptionsCaptor.getValue().get(SpannerRpc.Option.CHANNEL_HINT), options.getValue().get(SpannerRpc.Option.CHANNEL_HINT)); - assertTrue(deleteOptionsCaptor.getValue().containsKey(SpannerRpc.Option.REQUEST_ID)); } } @@ -207,9 +206,7 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount client.createMultiplexedSession(consumer); } // for multiplexed session there is no channel hint pass in the RPC options - assertNotNull(options.getValue()); - assertEquals(options.getValue().get(Option.CHANNEL_HINT), null); - assertNotNull(options.getValue().get(Option.REQUEST_ID)); + assertNull(options.getValue()); assertEquals(1, returnedSessionCount.get()); } @@ -241,9 +238,7 @@ public void onSessionCreateFailure(Throwable t, int createFailureForSessionCount client.createMultiplexedSession(consumer); } // for multiplexed session there is no channel hint pass in the RPC options - assertNotNull(options.getValue()); - assertEquals(options.getValue().get(Option.CHANNEL_HINT), null); - assertNotNull(options.getValue().get(Option.REQUEST_ID)); + assertNull(options.getValue()); } @SuppressWarnings("unchecked") diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java index df957f75ebf..1ac3b7beaf7 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionImplTest.java @@ -21,6 +21,7 @@ import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyMap; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; @@ -34,6 +35,7 @@ import com.google.cloud.Timestamp; import com.google.cloud.grpc.GrpcTransportOptions; import com.google.cloud.grpc.GrpcTransportOptions.ExecutorFactory; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; import com.google.protobuf.ByteString; @@ -144,8 +146,8 @@ public void setUp() { SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetryableCodes()); when(rpc.getCommitRetrySettings()) .thenReturn(SpannerStubSettings.newBuilder().commitSettings().getRetrySettings()); + when(rpc.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); session = spanner.getSessionClient(db).createSession(); - ((SessionImpl) session).setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); Span oTspan = mock(Span.class); ISpan span = new OpenTelemetrySpan(oTspan); when(oTspan.makeCurrent()).thenReturn(mock(Scope.class)); @@ -446,7 +448,7 @@ public void request(int numMessages) {} private void mockRead(final PartialResultSet myResultSet) { final ArgumentCaptor consumer = ArgumentCaptor.forClass(SpannerRpc.ResultStreamConsumer.class); - Mockito.when(rpc.read(Mockito.any(), consumer.capture(), anyMap(), eq(false))) + Mockito.when(rpc.read(Mockito.any(), consumer.capture(), anyMap(), any(), eq(false))) .then( invocation -> { consumer.getValue().onPartialResultSet(myResultSet); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java index 8d00f0889b8..b9d6943e523 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java @@ -72,6 +72,7 @@ import com.google.cloud.spanner.SpannerImpl.ClosedException; import com.google.cloud.spanner.TransactionRunner.TransactionCallable; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.spi.v1.SpannerRpc.ResultStreamConsumer; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; @@ -1455,6 +1456,7 @@ public void testSessionNotFoundReadWriteTransaction() { any(ExecuteSqlRequest.class), any(ResultStreamConsumer.class), any(Map.class), + any(XGoogSpannerRequestId.class), eq(true))) .thenReturn(closedStreamingCall); when(rpc.executeQuery(any(ExecuteSqlRequest.class), any(Map.class), eq(true))) @@ -1481,8 +1483,7 @@ public void testSessionNotFoundReadWriteTransaction() { when(closedSession.getName()) .thenReturn("projects/dummy/instances/dummy/database/dummy/sessions/session-closed"); when(closedSession.getErrorHandler()).thenReturn(DefaultErrorHandler.INSTANCE); - when(closedSession.getRequestIdCreator()) - .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator()); + when(closedSession.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); Span oTspan = mock(Span.class); ISpan span = new OpenTelemetrySpan(oTspan); @@ -1523,8 +1524,7 @@ public void testSessionNotFoundReadWriteTransaction() { TransactionRunnerImpl openTransactionRunner = new TransactionRunnerImpl(openSession); openTransactionRunner.setSpan(span); when(openSession.readWriteTransaction()).thenReturn(openTransactionRunner); - when(openSession.getRequestIdCreator()) - .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator()); + when(openSession.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); ResultSet openResultSet = mock(ResultSet.class); when(openResultSet.next()).thenReturn(true, false); @@ -1648,14 +1648,12 @@ public void testSessionNotFoundWrite() { SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName); List mutations = Collections.singletonList(Mutation.newInsertBuilder("FOO").build()); final SessionImpl closedSession = mockSession(); - closedSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); when(closedSession.writeWithOptions(eq(mutations), any())).thenThrow(sessionNotFound); final SessionImpl openSession = mockSession(); com.google.cloud.spanner.CommitResponse response = mock(com.google.cloud.spanner.CommitResponse.class); when(response.getCommitTimestamp()).thenReturn(Timestamp.now()); - openSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); when(openSession.writeWithOptions(eq(mutations), any())).thenReturn(response); doAnswer( invocation -> { @@ -1693,7 +1691,6 @@ public void testSessionNotFoundWriteAtLeastOnce() { SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName); List mutations = Collections.singletonList(Mutation.newInsertBuilder("FOO").build()); final SessionImpl closedSession = mockSession(); - closedSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); when(closedSession.writeAtLeastOnceWithOptions(eq(mutations), any())) .thenThrow(sessionNotFound); @@ -1701,7 +1698,6 @@ public void testSessionNotFoundWriteAtLeastOnce() { com.google.cloud.spanner.CommitResponse response = mock(com.google.cloud.spanner.CommitResponse.class); when(response.getCommitTimestamp()).thenReturn(Timestamp.now()); - openSession.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); when(openSession.writeAtLeastOnceWithOptions(eq(mutations), any())).thenReturn(response); doAnswer( invocation -> { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerCloudMonitoringExporterTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerCloudMonitoringExporterTest.java index c52cab0546c..590d62db7b5 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerCloudMonitoringExporterTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerCloudMonitoringExporterTest.java @@ -371,7 +371,7 @@ public void testExportingHistogramDataWithExemplars() { DoubleExemplarData exemplar = ImmutableDoubleExemplarData.create( Attributes.builder() - .put(XGoogSpannerRequestId.REQUEST_ID, "test") + .put(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME, "test") .put("lang", "java") .build(), recordTimeEpoch, @@ -448,8 +448,10 @@ public void testExportingHistogramDataWithExemplars() { // Assert only 1 attachment is there with 1 label for request_id. assertThat(filterAttributes.size()).isEqualTo(1); assertThat(filterAttributes.get(0).getLabelCount()).isEqualTo(1); - assertThat(filterAttributes.get(0).containsLabel(XGoogSpannerRequestId.REQUEST_ID)).isTrue(); - assertThat(filterAttributes.get(0).getLabelOrThrow(XGoogSpannerRequestId.REQUEST_ID)) + assertThat(filterAttributes.get(0).containsLabel(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME)) + .isTrue(); + assertThat( + filterAttributes.get(0).getLabelOrThrow(XGoogSpannerRequestId.REQUEST_ID_HEADER_NAME)) .isEqualTo("test"); } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java index 15af6bb8f1b..55b9523d7db 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpannerExceptionFactoryTest.java @@ -251,17 +251,4 @@ private Metadata createResourceTypeMetadata(String resourceType, String resource return trailers; } - - @Test - public void withRequestId() { - XGoogSpannerRequestId reqIdIn = XGoogSpannerRequestId.of(1, 2, 3, 4); - Status status = Status.fromCodeValue(Status.Code.ABORTED.value()); - Exception exc = new StatusRuntimeException(status); - SpannerException spannerExceptionWithReqId = - SpannerExceptionFactory.newSpannerException(exc, reqIdIn); - assertThat(spannerExceptionWithReqId.getRequestId()).isEqualTo(reqIdIn.toString()); - SpannerException spannerExceptionWithoutReqId = - SpannerExceptionFactory.newSpannerException(exc); - assertThat(spannerExceptionWithoutReqId.getRequestId()).isEqualTo(""); - } } diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java index 4c507b7beae..bbfa7cbd0d5 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionContextImplTest.java @@ -27,6 +27,7 @@ import com.google.api.core.ApiFutures; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; import com.google.protobuf.ByteString; @@ -66,9 +67,9 @@ public void setup() { com.google.spanner.v1.CommitResponse.newBuilder() .setCommitTimestamp(Timestamp.newBuilder().setSeconds(99L).setNanos(10).build()) .build())); + when(rpc.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); when(session.getName()).thenReturn("test"); - when(session.getRequestIdCreator()) - .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator()); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); doNothing().when(span).setStatus(any(Throwable.class)); doNothing().when(span).end(); doNothing().when(span).addAnnotation("Starting Commit"); @@ -212,8 +213,7 @@ public void testReturnCommitStats() { private void batchDml(int status) { SessionImpl session = mock(SessionImpl.class); when(session.getName()).thenReturn("test"); - when(session.getRequestIdCreator()) - .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator()); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); SpannerRpc rpc = mock(SpannerRpc.class); ExecuteBatchDmlResponse response = ExecuteBatchDmlResponse.newBuilder() diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java index 19c03859e66..547f6b70a22 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java @@ -241,7 +241,7 @@ public void usesPreparedTransaction() { Mockito.anyString(), Mockito.anyString(), Mockito.anyMap(), - Mockito.anyMap(), + Mockito.eq(null), Mockito.eq(true))) .thenAnswer( invocation -> @@ -324,7 +324,7 @@ public void inlineBegin() { Mockito.anyString(), Mockito.anyString(), Mockito.anyMap(), - Mockito.anyMap(), + Mockito.eq(null), Mockito.eq(true))) .thenAnswer( invocation -> diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java index 2325f2ac40f..3e3358a53bc 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java @@ -36,6 +36,7 @@ import com.google.cloud.spanner.ErrorHandler.DefaultErrorHandler; import com.google.cloud.spanner.SessionClient.SessionId; import com.google.cloud.spanner.TransactionRunnerImpl.TransactionContextImpl; +import com.google.cloud.spanner.XGoogSpannerRequestId.NoopRequestIdCreator; import com.google.cloud.spanner.spi.v1.SpannerRpc; import com.google.cloud.spanner.v1.stub.SpannerStubSettings; import com.google.common.base.Preconditions; @@ -121,8 +122,8 @@ public void setUp() { when(session.getErrorHandler()).thenReturn(DefaultErrorHandler.INSTANCE); when(session.newTransaction(eq(Options.fromTransactionOptions()), any())).thenReturn(txn); when(session.getTracer()).thenReturn(tracer); - when(session.getRequestIdCreator()) - .thenReturn(new XGoogSpannerRequestId.NoopRequestIdCreator()); + when(session.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); + when(rpc.getRequestIdCreator()).thenReturn(NoopRequestIdCreator.INSTANCE); when(rpc.executeQuery(Mockito.any(ExecuteSqlRequest.class), Mockito.anyMap(), eq(true))) .thenAnswer( invocation -> { @@ -195,7 +196,7 @@ public void usesPreparedTransaction() { Mockito.anyString(), Mockito.anyString(), Mockito.anyMap(), - Mockito.anyMap(), + Mockito.eq(null), Mockito.eq(true))) .thenAnswer( invocation -> @@ -336,7 +337,6 @@ public void inlineBegin() { spanner, new SessionReference( "projects/p/instances/i/databases/d/sessions/s", Collections.EMPTY_MAP)) {}; - session.setRequestIdCreator(new XGoogSpannerRequestId.NoopRequestIdCreator()); session.setCurrentSpan(new OpenTelemetrySpan(mock(io.opentelemetry.api.trace.Span.class))); TransactionRunnerImpl runner = new TransactionRunnerImpl(session); runner.setSpan(span); diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java index 719b94593bf..32d1ac29d2e 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/XGoogSpannerRequestIdTest.java @@ -16,6 +16,7 @@ package com.google.cloud.spanner; +import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotEquals; import static org.junit.Assert.assertNotNull; @@ -67,17 +68,14 @@ public void testEnsureHexadecimalFormatForRandProcessID() { } public static class ServerHeaderEnforcer implements ServerInterceptor { - private Map> unaryResults; - private Map> streamingResults; - private List gotValues; - private Set checkMethods; + private final Map> unaryResults = + new ConcurrentHashMap<>(); + private final Map> streamingResults = + new ConcurrentHashMap<>(); + private final List gotValues = new CopyOnWriteArrayList<>(); + private final Set checkMethods; ServerHeaderEnforcer(Set checkMethods) { - this.gotValues = new CopyOnWriteArrayList(); - this.unaryResults = - new ConcurrentHashMap>(); - this.streamingResults = - new ConcurrentHashMap>(); this.checkMethods = checkMethods; } @@ -88,7 +86,7 @@ public ServerCall.Listener interceptCall( ServerCallHandler next) { boolean isUnary = call.getMethodDescriptor().getType() == MethodType.UNARY; String methodName = call.getMethodDescriptor().getFullMethodName(); - String gotReqIdStr = requestHeaders.get(XGoogSpannerRequestId.REQUEST_HEADER_KEY); + String gotReqIdStr = requestHeaders.get(XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY); if (!this.checkMethods.contains(methodName)) { return next.startCall(call, requestHeaders); } @@ -102,7 +100,7 @@ public ServerCall.Listener interceptCall( Status status = Status.fromCode(Status.Code.INVALID_ARGUMENT) .augmentDescription( - methodName + " lacks " + XGoogSpannerRequestId.REQUEST_HEADER_KEY); + methodName + " lacks " + XGoogSpannerRequestId.REQUEST_ID_HEADER_KEY); call.close(status, requestHeaders); return next.startCall(call, requestHeaders); } @@ -128,14 +126,8 @@ public String[] accumulatedValues() { } public void assertIntegrity() { - this.unaryResults.forEach( - (String method, CopyOnWriteArrayList values) -> { - assertMonotonicityOfIds(method, values); - }); - this.streamingResults.forEach( - (String method, CopyOnWriteArrayList values) -> { - assertMonotonicityOfIds(method, values); - }); + this.unaryResults.forEach(this::assertMonotonicityOfIds); + this.streamingResults.forEach(this::assertMonotonicityOfIds); } private void assertMonotonicityOfIds(String prefix, List reqIds) { @@ -161,22 +153,22 @@ private void assertMonotonicityOfIds(String prefix, List } public MethodAndRequestId[] accumulatedUnaryValues() { - List accumulated = new ArrayList(); + List accumulated = new ArrayList<>(); this.unaryResults.forEach( (String method, CopyOnWriteArrayList values) -> { - for (int i = 0; i < values.size(); i++) { - accumulated.add(new MethodAndRequestId(method, values.get(i))); + for (XGoogSpannerRequestId value : values) { + accumulated.add(new MethodAndRequestId(method, value)); } }); return accumulated.toArray(new MethodAndRequestId[0]); } public MethodAndRequestId[] accumulatedStreamingValues() { - List accumulated = new ArrayList(); + List accumulated = new ArrayList<>(); this.streamingResults.forEach( (String method, CopyOnWriteArrayList values) -> { - for (int i = 0; i < values.size(); i++) { - accumulated.add(new MethodAndRequestId(method, values.get(i))); + for (XGoogSpannerRequestId value : values) { + accumulated.add(new MethodAndRequestId(method, value)); } }); return accumulated.toArray(new MethodAndRequestId[0]); @@ -188,7 +180,7 @@ public void checkExpectedUnaryXGoogRequestIds(MethodAndRequestId... wantUnaryVal for (int i = 0; i < gotUnaryValues.length && false; i++) { System.out.println("\033[33misUnary: #" + i + ":: " + gotUnaryValues[i] + "\033[00m"); } - assertEquals(wantUnaryValues, gotUnaryValues); + assertArrayEquals(wantUnaryValues, gotUnaryValues); } public void checkAtLeastHasExpectedUnaryXGoogRequestIds(MethodAndRequestId... wantUnaryValues) { @@ -200,9 +192,9 @@ public void checkAtLeastHasExpectedUnaryXGoogRequestIds(MethodAndRequestId... wa if (wantUnaryValues.length < gotUnaryValues.length) { MethodAndRequestId[] gotSliced = Arrays.copyOfRange(gotUnaryValues, 0, wantUnaryValues.length); - assertEquals(wantUnaryValues, gotSliced); + assertArrayEquals(wantUnaryValues, gotSliced); } else { - assertEquals(wantUnaryValues, gotUnaryValues); + assertArrayEquals(wantUnaryValues, gotUnaryValues); } } @@ -218,9 +210,9 @@ public void checkExpectedUnaryXGoogRequestIdsAsSuffixes(MethodAndRequestId... wa gotUnaryValues, gotUnaryValues.length - wantUnaryValues.length, gotUnaryValues.length); - assertEquals(wantUnaryValues, gotSliced); + assertArrayEquals(wantUnaryValues, gotSliced); } else { - assertEquals(wantUnaryValues, gotUnaryValues); + assertArrayEquals(wantUnaryValues, gotUnaryValues); } } @@ -236,7 +228,7 @@ public void checkExpectedStreamingXGoogRequestIds(MethodAndRequestId... wantStre "\033[32misStreaming: #" + i + ":: " + gotStreamingValues[i] + "\033[00m"); } sortValues(gotStreamingValues); - assertEquals(wantStreamingValues, gotStreamingValues); + assertArrayEquals(wantStreamingValues, gotStreamingValues); } public void reset() { diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RunTransactionMockServerTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RunTransactionMockServerTest.java index bc0acca3255..d4af5dd4e7e 100644 --- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RunTransactionMockServerTest.java +++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/RunTransactionMockServerTest.java @@ -151,7 +151,7 @@ public void testRunTransaction_rollbacksAfterException() { return null; })); assertEquals(ErrorCode.INVALID_ARGUMENT, exception.getErrorCode()); - assertTrue(exception.getMessage(), exception.getMessage().endsWith("invalid statement")); + assertTrue(exception.getMessage(), exception.getMessage().contains("invalid statement")); } assertEquals(2, mockSpanner.countRequestsOfType(ExecuteSqlRequest.class)); assertEquals(0, mockSpanner.countRequestsOfType(CommitRequest.class));