Skip to content

test(spanner): unflake ITSqlMusicScriptTest #14814

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sakthivelmanii wants to merge 1 commit into
base: main
Choose a base branch
from
Open

test(spanner): unflake ITSqlMusicScriptTest #14814

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 45 additions & 2 deletions ...r/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/KeyAwareChannel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,13 @@ final class KeyAwareChannel extends ManagedChannel {
// Bound and age out entries in case application code abandons a transaction
// without sending Commit/Rollback or otherwise clearing affinity.
private final Cache<ByteString, String> transactionAffinities;
// Tracks the IDs of Partitioned DML transactions. Requests for these transactions must never
// carry a routing hint, as Spanner would then try to execute them on the fastpath, which only
// supports reads and queries and rejects Partitioned DML transactions with
// "PartitionedDml Transactions may only be used to execute DML statements".
// Partitioned DML transactions are never committed or rolled back, so entries age out in the
// same way as abandoned read-write transaction affinities.
private final Cache<ByteString, Boolean> partitionedDmlTransactions;
private final EndpointOverloadCooldownTracker endpointOverloadCooldowns;

private KeyAwareChannel(
Expand Down Expand Up @@ -158,7 +165,8 @@ private KeyAwareChannel(
this.lifecycleManager =
(endpointCacheFactory == null) ? new EndpointLifecycleManager(endpointCache) : null;
this.endpointOverloadCooldowns = endpointOverloadCooldowns;
this.transactionAffinities = newTransactionAffinities(transactionAffinityTicker);
this.transactionAffinities = newTransactionCache(transactionAffinityTicker);
this.partitionedDmlTransactions = newTransactionCache(transactionAffinityTicker);
}

static KeyAwareChannel create(
Expand Down Expand Up @@ -201,7 +209,7 @@ static KeyAwareChannel create(
transactionAffinityTicker);
}

private static Cache<ByteString, String> newTransactionAffinities(Ticker ticker) {
private static <V> Cache<ByteString, V> newTransactionCache(Ticker ticker) {
return CacheBuilder.newBuilder()
.maximumSize(MAX_TRACKED_TRANSACTION_AFFINITIES)
.expireAfterAccess(TRANSACTION_AFFINITY_TTL_MINUTES, TimeUnit.MINUTES)
Expand Down Expand Up @@ -393,6 +401,21 @@ private void clearAffinity(ByteString transactionId) {
return;
}
transactionAffinities.invalidate(transactionId);
partitionedDmlTransactions.invalidate(transactionId);
}

private void recordPartitionedDmlTransaction(ByteString transactionId) {
if (transactionId == null || transactionId.isEmpty()) {
return;
}
partitionedDmlTransactions.put(transactionId, Boolean.TRUE);
}

private boolean isPartitionedDmlTransaction(ByteString transactionId) {
if (transactionId == null || transactionId.isEmpty()) {
return false;
}
return partitionedDmlTransactions.getIfPresent(transactionId) != null;
}

private void recordEndpointCooldown(
Expand Down Expand Up @@ -737,6 +760,7 @@ static final class KeyAwareClientCall<RequestT, ResponseT>
@Nullable private io.grpc.Status cancelledStatus;
@Nullable private Metadata cancelledTrailers;
private boolean shouldRecordTransactionAffinity;
private boolean beginsPartitionedDmlTransaction;
private boolean genericRequest;
private boolean reroutedToDefault;
private boolean responseMessageReceived;
Expand Down Expand Up @@ -839,6 +863,8 @@ public void sendMessage(RequestT message) {
shouldRecordTransactionAffinity = true;
allowDefaultAffinity = true;
}
beginsPartitionedDmlTransaction =
reqBuilder.hasOptions() && reqBuilder.getOptions().hasPartitionedDml();
message = (RequestT) reqBuilder.build();
} else if (message instanceof CommitRequest) {
CommitRequest request = (CommitRequest) message;
Expand Down Expand Up @@ -1128,6 +1154,12 @@ void maybeRecordAffinity(ByteString transactionId) {
parentChannel.recordAffinity(transactionId, selectedEndpoint, allowDefaultAffinity);
}

void maybeRecordPartitionedDmlTransaction(ByteString transactionId) {
if (beginsPartitionedDmlTransaction) {
parentChannel.recordPartitionedDmlTransaction(transactionId);
}
}

void maybeClearAffinity() {
parentChannel.clearAffinity(transactionIdToClear);
}
Expand Down Expand Up @@ -1184,6 +1216,14 @@ private RoutingDecision routeFromRequest(ExecuteSqlRequest.Builder reqBuilder) {
finder = parentChannel.getOrCreateChannelFinder(databaseId);
}
boolean preferLeader = preferLeader(reqBuilder.getTransaction());
if (parentChannel.isPartitionedDmlTransaction(transactionId)) {
// Never attach a routing hint to a Partitioned DML request. A routing hint makes Spanner
// try to execute the statement on the fastpath, which does not support Partitioned DML
// and fails with "PartitionedDml Transactions may only be used to execute DML
// statements". Route the request to the transaction's affinity endpoint (or the default
// endpoint) without a routing hint instead.
return new RoutingDecision(finder, endpoint, databaseId, 0L, preferLeader);
}
if (databaseId != null && endpoint == null) {
// Transaction IDs only pin routing when a read-write affinity entry exists. Otherwise route
// by this request's routing hint without leader bias, as read-only transaction IDs can run
Expand Down Expand Up @@ -1306,6 +1346,9 @@ public void onMessage(ResponseT message) {
if (transactionId != null && call.shouldRecordTransactionAffinity) {
call.maybeRecordAffinity(transactionId);
}
if (transactionId != null) {
call.maybeRecordPartitionedDmlTransaction(transactionId);
}
super.onMessage(message);
}

Expand Down
51 changes: 51 additions & 0 deletions ...ogle-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1386,6 +1386,57 @@ public void readWriteTransactionQueryAttachesRoutingHintWhilePinnedToAffinityEnd
assertThat(routingHint.getTabletUid()).isEqualTo(2L);
}

@Test
public void partitionedDmlQueryDoesNotAttachRoutingHintWhilePinnedToAffinityEndpoint()
throws Exception {
TestHarness harness = createHarness();
seedCache(harness, createTwoRangeCacheUpdate());
ByteString transactionId = ByteString.copyFromUtf8("pdml-no-routing-hint");
String sql = "DELETE FROM T WHERE K = @k";

// Begin a Partitioned DML transaction. This records a transaction affinity in the same way as
// a read-write transaction.
ClientCall<BeginTransactionRequest, Transaction> beginCall =
harness.channel.newCall(SpannerGrpc.getBeginTransactionMethod(), CallOptions.DEFAULT);
beginCall.start(new CapturingListener<Transaction>(), new Metadata());
beginCall.sendMessage(
BeginTransactionRequest.newBuilder()
.setSession(SESSION)
.setOptions(
TransactionOptions.newBuilder()
.setPartitionedDml(TransactionOptions.PartitionedDml.getDefaultInstance()))
.build());
@SuppressWarnings("unchecked")
RecordingClientCall<BeginTransactionRequest, Transaction> beginDelegate =
(RecordingClientCall<BeginTransactionRequest, Transaction>)
harness.defaultManagedChannel.latestCall();
beginDelegate.emitOnMessage(Transaction.newBuilder().setId(transactionId).build());
beginDelegate.emitOnClose(Status.OK, new Metadata());

// Execute the Partitioned DML statement. Spanner rejects ExecuteSql requests for a Partitioned
// DML transaction that carry a routing hint, because the server fastpath treats them as reads:
// "PartitionedDml Transactions may only be used to execute DML statements".
ClientCall<ExecuteSqlRequest, PartialResultSet> pdmlCall =
harness.channel.newCall(SpannerGrpc.getExecuteStreamingSqlMethod(), CallOptions.DEFAULT);
pdmlCall.start(new CapturingListener<PartialResultSet>(), new Metadata());
pdmlCall.sendMessage(
ExecuteSqlRequest.newBuilder()
.setSession(SESSION)
.setSql(sql)
.setTransaction(TransactionSelector.newBuilder().setId(transactionId))
.build());

@SuppressWarnings("unchecked")
RecordingClientCall<ExecuteSqlRequest, PartialResultSet> pdmlDelegate =
(RecordingClientCall<ExecuteSqlRequest, PartialResultSet>)
harness.defaultManagedChannel.latestCall();
assertNotNull(pdmlDelegate.lastMessage);
assertThat(pdmlDelegate.lastMessage.getSql()).isEqualTo(sql);
assertThat(pdmlDelegate.lastMessage.hasRoutingHint()).isFalse();
assertThat(harness.endpointCache.callCountForAddress("server-a:1234")).isEqualTo(0);
assertThat(harness.endpointCache.callCountForAddress("server-b:1234")).isEqualTo(0);
}

@Test
public void genericUnimplementedReroutesOnceAndLearnsDefaultOnlyMethod() throws Exception {
TestHarness harness = createHarness();
Expand Down
Loading