diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java index bb4b78d8e016..473abc904788 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java @@ -534,9 +534,11 @@ public GapicSpannerRpc(final SpannerOptions options) { .getStreamWatchdogProvider() .withCheckInterval(pdmlSettings.getStreamWatchdogCheckInterval())); } + ClientContext partitionedDmlClientContext = + createPartitionedDmlClientContext(clientContext, this.keyAwareChannel); this.partitionedDmlStub = GrpcSpannerStubWithStubSettingsAndClientContext.create( - pdmlSettings.build(), clientContext); + pdmlSettings.build(), partitionedDmlClientContext); this.instanceAdminStubSettings = options.getInstanceAdminStubSettings().toBuilder() .setTransportChannelProvider(channelProvider) @@ -677,6 +679,21 @@ private static KeyAwareChannel extractKeyAwareChannel(TransportChannel transport return null; } + private static ClientContext createPartitionedDmlClientContext( + ClientContext clientContext, @Nullable KeyAwareChannel keyAwareChannel) { + if (keyAwareChannel == null) { + return clientContext; + } + GrpcTransportChannel defaultTransportChannel = + GrpcTransportChannel.create(keyAwareChannel.defaultChannel()); + return clientContext.toBuilder() + .setTransportChannel(defaultTransportChannel) + .setDefaultCallContext( + clientContext.getDefaultCallContext().withTransportChannel(defaultTransportChannel)) + .setBackgroundResources(Collections.emptyList()) + .build(); + } + @Nullable private static GcpManagedChannel extractGrpcGcpChannel(TransportChannel transportChannel) { if (!(transportChannel instanceof GrpcTransportChannel)) { diff --git a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java index 0344188c2ff0..91477342827e 100644 --- a/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java +++ b/java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java @@ -327,6 +327,10 @@ public String authority() { return authority; } + ManagedChannel defaultChannel() { + return defaultChannel; + } + @Override public ClientCall newCall( MethodDescriptor methodDescriptor, CallOptions callOptions) { @@ -833,7 +837,9 @@ public void sendMessage(RequestT message) { if (finder != null && reqBuilder.hasMutationKey()) { endpoint = finder.findServerAndFillRoutingHint(reqBuilder, excludedEndpoints); } - if (reqBuilder.hasOptions() && reqBuilder.getOptions().hasReadOnly()) { + if (reqBuilder.hasOptions() + && (reqBuilder.getOptions().hasReadOnly() + || reqBuilder.getOptions().hasPartitionedDml())) { shouldRecordTransactionAffinity = false; } else { shouldRecordTransactionAffinity = true; @@ -1134,7 +1140,8 @@ void maybeClearAffinity() { private void maybeTrackReadWriteBegin(TransactionSelector selector) { if (selector.getSelectorCase() == TransactionSelector.SelectorCase.BEGIN - && !selector.getBegin().hasReadOnly()) { + && !selector.getBegin().hasReadOnly() + && !selector.getBegin().hasPartitionedDml()) { shouldRecordTransactionAffinity = true; allowDefaultAffinity = true; } diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java index 52d5833489b8..c2e4bb601665 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpcTest.java @@ -75,24 +75,30 @@ import com.google.cloud.spanner.spi.v1.SpannerRpc.Option; import com.google.common.collect.ImmutableList; import com.google.common.util.concurrent.Futures; +import com.google.protobuf.ByteString; import com.google.protobuf.ListValue; import com.google.rpc.ErrorInfo; import com.google.spanner.admin.database.v1.DatabaseAdminGrpc; import com.google.spanner.admin.database.v1.ListDatabaseRolesRequest; +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.GetSessionRequest; +import com.google.spanner.v1.PartialResultSet; import com.google.spanner.v1.ReadRequest; import com.google.spanner.v1.ResultSetMetadata; import com.google.spanner.v1.Session; import com.google.spanner.v1.SpannerGrpc; import com.google.spanner.v1.StructType; import com.google.spanner.v1.StructType.Field; +import com.google.spanner.v1.TransactionOptions; +import com.google.spanner.v1.TransactionSelector; import com.google.spanner.v1.TypeCode; import io.grpc.CallOptions; import io.grpc.Context; import io.grpc.Contexts; +import io.grpc.ManagedChannel; import io.grpc.ManagedChannelBuilder; import io.grpc.Metadata; import io.grpc.Metadata.Key; @@ -141,6 +147,7 @@ import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.BooleanSupplier; import javax.annotation.Nullable; import org.junit.After; @@ -2208,6 +2215,154 @@ public boolean isEnableLocationApi() { } } + @Test + public void testPartitionedDmlUsesDefaultChannelWithoutEnteringKeyAwareChannel() + throws Exception { + AtomicInteger keyAwareRoutedCalls = new AtomicInteger(); + ChannelEndpointCacheFactory factory = + baseProvider -> { + GrpcChannelEndpointCache delegateCache = new GrpcChannelEndpointCache(baseProvider); + ChannelEndpoint delegateDefault = delegateCache.defaultChannel(); + ChannelEndpoint countingDefault = + new ChannelEndpoint() { + @Override + public String getAddress() { + return delegateDefault.getAddress(); + } + + @Override + public boolean isHealthy() { + return delegateDefault.isHealthy(); + } + + @Override + public boolean isTransientFailure() { + return delegateDefault.isTransientFailure(); + } + + @Override + public ManagedChannel getChannel() { + return delegateDefault.getChannel(); + } + + @Override + public void incrementActiveRequests() { + keyAwareRoutedCalls.incrementAndGet(); + delegateDefault.incrementActiveRequests(); + } + + @Override + public void decrementActiveRequests() { + delegateDefault.decrementActiveRequests(); + } + + @Override + public int getActiveRequestCount() { + return delegateDefault.getActiveRequestCount(); + } + }; + return new ChannelEndpointCache() { + @Override + public ChannelEndpoint defaultChannel() { + return countingDefault; + } + + @Override + public ChannelEndpoint get(String addr) { + return addr.equals(countingDefault.getAddress()) + ? countingDefault + : delegateCache.get(addr); + } + + @Override + @Nullable + public ChannelEndpoint getIfPresent(String addr) { + return addr.equals(countingDefault.getAddress()) + ? countingDefault + : delegateCache.getIfPresent(addr); + } + + @Override + public void evict(String addr) { + delegateCache.evict(addr); + } + + @Override + public void shutdown() { + delegateCache.shutdown(); + } + }; + }; + + try { + SpannerOptions.useEnvironment( + new SpannerOptions.SpannerEnvironment() { + @Override + public boolean isEnableLocationApi() { + return true; + } + }); + SpannerOptions options = + createSpannerOptions().toBuilder().setChannelEndpointCacheFactory(factory).build(); + try (Spanner testSpanner = options.getService()) { + DatabaseClient client = testSpanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "i", "d")); + try (ResultSet rs = client.singleUse().executeQuery(SELECT1AND2)) { + while (rs.next()) {} + } + assertThat(keyAwareRoutedCalls.get()).isGreaterThan(0); + + mockSpanner.clearRequests(); + keyAwareRoutedCalls.set(0); + + // End-to-end Partitioned DML: BeginTransaction goes through spannerStub (1 call), while + // ExecuteStreamingSql goes through partitionedDmlStub directly on the default channel (0 + // KeyAwareClientCall entries). + long updateCount = client.executePartitionedUpdate(UPDATE_FOO_STATEMENT); + assertEquals(1L, updateCount); + assertEquals(1, keyAwareRoutedCalls.get()); + List pdmlStreamingRequests = + mockSpanner.getRequestsOfType(ExecuteSqlRequest.class); + assertEquals(1, pdmlStreamingRequests.size()); + assertFalse(pdmlStreamingRequests.get(0).hasRoutingHint()); + + // Direct unary and streaming Partitioned DML RPCs must never enter KeyAwareChannel at all. + GapicSpannerRpc rpc = getRpc(testSpanner); + String sessionName = + mockSpanner.getRequestsOfType(BeginTransactionRequest.class).get(0).getSession(); + BeginTransactionRequest beginPdmlRequest = + BeginTransactionRequest.newBuilder() + .setSession(sessionName) + .setOptions( + TransactionOptions.newBuilder() + .setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance())) + .build(); + ByteString unaryTxId = rpc.beginTransaction(beginPdmlRequest, optionsMap, true).getId(); + ByteString streamingTxId = rpc.beginTransaction(beginPdmlRequest, optionsMap, true).getId(); + keyAwareRoutedCalls.set(0); + rpc.executePartitionedDml( + ExecuteSqlRequest.newBuilder() + .setSession(sessionName) + .setSql(UPDATE_FOO_STATEMENT.getSql()) + .setTransaction(TransactionSelector.newBuilder().setId(unaryTxId)) + .build(), + optionsMap); + for (PartialResultSet ignored : + rpc.executeStreamingPartitionedDml( + ExecuteSqlRequest.newBuilder() + .setSession(sessionName) + .setSql(UPDATE_FOO_STATEMENT.getSql()) + .setTransaction(TransactionSelector.newBuilder().setId(streamingTxId)) + .build(), + optionsMap, + rpc.getRequestIdCreator().nextRequestId(0), + Duration.ofSeconds(10))) {} + assertEquals(0, keyAwareRoutedCalls.get()); + } + } finally { + SpannerOptions.useDefaultEnvironment(); + } + } + @Test public void testLocationApiDoesNotOverrideExplicitChannelProvider() { AtomicBoolean factoryCalled = new AtomicBoolean(false); diff --git a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java index 21d6ac10dd9a..90b9f39e48d0 100644 --- a/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java +++ b/java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java @@ -1972,6 +1972,80 @@ public void readOnlyTransactionDoesNotRecordAffinity() throws Exception { assertThat(harness.endpointCache.getCount(DEFAULT_ADDRESS)).isEqualTo(0); } + @Test + public void partitionedDmlTransactionDoesNotRecordAffinity() throws Exception { + TestHarness harness = createHarness(); + ByteString transactionId = ByteString.copyFromUtf8("pdml-tx-no-affinity"); + + // Begin a Partitioned DML transaction via BeginTransaction. + ClientCall beginCall = + harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT); + beginCall.start(new CapturingListener(), new Metadata()); + beginCall.sendMessage( + BeginTransactionRequest.newBuilder() + .setSession(SESSION) + .setOptions( + TransactionOptions.newBuilder() + .setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance())) + .build()); + + @SuppressWarnings("unchecked") + RecordingClientCall beginDelegate = + (RecordingClientCall) + harness.defaultManagedChannel.latestCall(); + beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build()); + beginDelegate.emitOnClose(Status.OK, new Metadata()); + + // Send a request referencing the Partitioned DML transaction ID. Because no affinity was + // recorded on BeginTransaction, the endpoint cache must not be queried for affinity lookup. + ClientCall sqlCall = + harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT); + sqlCall.start(new CapturingListener(), new Metadata()); + sqlCall.sendMessage( + ExecuteSqlRequest.newBuilder() + .setSession(SESSION) + .setTransaction(TransactionSelector.newBuilder().setId(transactionId)) + .build()); + assertThat(harness.endpointCache.getIfPresentCount(DEFAULT_ADDRESS)).isEqualTo(0); + + // Also verify inline TransactionSelector.begin with PartitionedDml does not record affinity. + ByteString inlinePdmlId = ByteString.copyFromUtf8("pdml-inline-no-affinity"); + ClientCall inlineBeginCall = + harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT); + inlineBeginCall.start(new CapturingListener(), new Metadata()); + inlineBeginCall.sendMessage( + ExecuteSqlRequest.newBuilder() + .setSession(SESSION) + .setTransaction( + TransactionSelector.newBuilder() + .setBegin( + TransactionOptions.newBuilder() + .setPartitionedDml( + TransactionOptions.PartitionedDml.getDefaultInstance()))) + .build()); + @SuppressWarnings("unchecked") + RecordingClientCall inlineBeginDelegate = + (RecordingClientCall) + harness.defaultManagedChannel.latestCall(); + inlineBeginDelegate.emitOnMessage( + ResultSet.newBuilder() + .setMetadata( + ResultSetMetadata.newBuilder() + .setTransaction(Transaction.newBuilder().setId(inlinePdmlId))) + .build()); + inlineBeginDelegate.emitOnClose(Status.OK, new Metadata()); + + ClientCall followUpCall = + harness.channel.newCall(SpannerGrpc.getExecuteSqlMethod(), CallOptions.DEFAULT); + followUpCall.start(new CapturingListener(), new Metadata()); + followUpCall.sendMessage( + ExecuteSqlRequest.newBuilder() + .setSession(SESSION) + .setTransaction(TransactionSelector.newBuilder().setId(inlinePdmlId)) + .build()); + assertThat(harness.endpointCache.getIfPresentCount(DEFAULT_ADDRESS)).isEqualTo(0); + } + @Test public void abandonedReadWriteTransactionAffinityExpiresAfterInactivity() throws Exception { FakeTicker ticker = new FakeTicker();