Skip to content

Commit

Permalink
chore(spanner): fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@harshachinta
harshachinta committed Dec 17, 2024
1 parent d57fe74 commit dfd3812
Showing 1 changed file with 171 additions and 61 deletions.
232 changes: 171 additions & 61 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assume.assumeFalse;

import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutureCallback;
Expand Down Expand Up @@ -178,9 +179,11 @@ public void asyncTransactionManager_shouldRollbackOnCloseAsync() throws Exceptio
AsyncTransactionManager manager = client().transactionManagerAsync();
TransactionContext txn = manager.beginAsync().get();
txn.executeUpdateAsync(UPDATE_STATEMENT).get();
final TransactionSelector selector =
((TransactionContextImpl) ((SessionPoolTransactionContext) txn).delegate)
.getTransactionSelector();
if (txn instanceof SessionPoolTransactionContext) {
txn = ((SessionPoolTransactionContext) txn).delegate;
}
TransactionContextImpl impl = (TransactionContextImpl) txn;
final TransactionSelector selector = impl.getTransactionSelector();

SpannerApiFutures.get(manager.closeAsync());
// The mock server should already have the Rollback request, as we are waiting for the returned
Expand Down Expand Up @@ -247,6 +250,11 @@ public void asyncTransactionManagerUpdate() throws Exception {

@Test
public void asyncTransactionManagerIsNonBlocking() throws Exception {
// TODO: Remove this condition once DelayedAsyncTransactionManager is made non-blocking with
// multiplexed sessions.
assumeFalse(
"DelayedAsyncTransactionManager is currently blocking with multiplexed sessions.",
spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW());
mockSpanner.freeze();
try (AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync()) {
TransactionContextFuture transactionContextFuture = manager.beginAsync();
Expand Down Expand Up @@ -346,7 +354,7 @@ public void asyncTransactionManagerFireAndForgetInvalidUpdate() throws Exception
}
}
}
ImmutableList<Class<? extends Message>> expectedRequests =
ImmutableList<Class<? extends Message>> expectedRequestsWithRegularSession =
ImmutableList.of(
BatchCreateSessionsRequest.class,
// The first update that fails. This will cause a transaction retry.
Expand All @@ -358,10 +366,24 @@ public void asyncTransactionManagerFireAndForgetInvalidUpdate() throws Exception
ExecuteSqlRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSession =
ImmutableList.of(
CreateSessionRequest.class,
// The first update that fails. This will cause a transaction retry.
ExecuteSqlRequest.class,
// The retry will use an explicit BeginTransaction call.
BeginTransactionRequest.class,
// The first update will again fail, but now there is a transaction id, so the
// transaction can continue.
ExecuteSqlRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSession);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithRegularSession);
}
}

Expand Down Expand Up @@ -501,14 +523,25 @@ public void asyncTransactionManagerUpdateAbortedWithoutGettingResult() throws Ex
// The server may receive 1 or 2 commit requests depending on whether the call to
// commitAsync() already knows that the transaction has aborted. If it does, it will not
// attempt to call the Commit RPC and instead directly propagate the Aborted error.
assertThat(mockSpanner.getRequestTypes())
.containsAtLeast(
BatchCreateSessionsRequest.class,
ExecuteSqlRequest.class,
// The retry will use a BeginTransaction RPC.
BeginTransactionRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsAtLeast(
CreateSessionRequest.class,
ExecuteSqlRequest.class,
// The retry will use a BeginTransaction RPC.
BeginTransactionRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
} else {
assertThat(mockSpanner.getRequestTypes())
.containsAtLeast(
BatchCreateSessionsRequest.class,
ExecuteSqlRequest.class,
// The retry will use a BeginTransaction RPC.
BeginTransactionRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
}
break;
} catch (AbortedException e) {
transactionContextFuture = manager.resetForRetryAsync();
Expand Down Expand Up @@ -556,13 +589,10 @@ public void asyncTransactionManagerWaitsUntilAsyncUpdateHasFinished() throws Exc
executor)
.commitAsync()
.get();
if (isMultiplexedSessionsEnabled()) {
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
CreateSessionRequest.class,
BatchCreateSessionsRequest.class,
ExecuteSqlRequest.class,
CommitRequest.class);
CreateSessionRequest.class, ExecuteSqlRequest.class, CommitRequest.class);
} else {
assertThat(mockSpanner.getRequestTypes())
.containsExactly(
Expand Down Expand Up @@ -600,6 +630,11 @@ public void asyncTransactionManagerBatchUpdate() throws Exception {

@Test
public void asyncTransactionManagerIsNonBlockingWithBatchUpdate() throws Exception {
// TODO: Remove this condition once DelayedAsyncTransactionManager is made non-blocking with
// multiplexed sessions.
assumeFalse(
"DelayedAsyncTransactionManager is currently blocking with multiplexed sessions.",
spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW());
mockSpanner.freeze();
try (AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync()) {
TransactionContextFuture transactionContextFuture = manager.beginAsync();
Expand Down Expand Up @@ -671,16 +706,24 @@ public void asyncTransactionManagerFireAndForgetInvalidBatchUpdate() throws Exce
}
}
}
ImmutableList<Class<? extends Message>> expectedRequests =
ImmutableList<Class<? extends Message>> expectedRequestsWithRegularSession =
ImmutableList.of(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSession =
ImmutableList.of(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSession);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithRegularSession);
}
}

Expand Down Expand Up @@ -714,17 +757,26 @@ public void asyncTransactionManagerBatchUpdateAborted() throws Exception {
assertThat(attempt.get()).isEqualTo(2);
// There should only be 1 CommitRequest, as the first attempt should abort already after the
// ExecuteBatchDmlRequest.
ImmutableList<Class<? extends Message>> expectedRequests =
ImmutableList<Class<? extends Message>> expectedRequestsWithRegularSession =
ImmutableList.of(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSession =
ImmutableList.of(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSession);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithRegularSession);
}
}

Expand Down Expand Up @@ -756,17 +808,30 @@ public void asyncTransactionManagerBatchUpdateAbortedBeforeFirstStatement() thro
assertThat(attempt.get()).isEqualTo(2);
// There should only be 1 CommitRequest, as the first attempt should abort already after the
// ExecuteBatchDmlRequest.
ImmutableList<Class<? extends Message>> expectedRequests =
ImmutableList<Class<? extends Message>> expectedRequestsWithRegularSession =
ImmutableList.of(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
// There should only be 1 CommitRequest, as the first attempt should abort already after the
// ExecuteBatchDmlRequest.
// When requests run using multiplexed session, the BatchCreateSessionsRequest will not be
// triggered because we are creating an empty pool during initialization.
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSession =
ImmutableList.of(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSession);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithRegularSession);
}
}

Expand Down Expand Up @@ -816,18 +881,28 @@ public void asyncTransactionManagerWithBatchUpdateCommitAborted() throws Excepti
} finally {
mockSpanner.putStatementResult(StatementResult.update(UPDATE_STATEMENT, UPDATE_COUNT));
}
ImmutableList<Class<? extends Message>> expectedRequests =
ImmutableList<Class<? extends Message>> expectedRequestsWithRegularSession =
ImmutableList.of(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSession =
ImmutableList.of(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSession);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithRegularSession);
}
}

Expand Down Expand Up @@ -865,28 +940,46 @@ public void asyncTransactionManagerBatchUpdateAbortedWithoutGettingResult() thro
}
assertThat(attempt.get()).isEqualTo(2);
List<Class<? extends AbstractMessage>> requests = mockSpanner.getRequestTypes();
// Remove the CreateSession requests for multiplexed sessions, as those are not relevant for
// this test.
requests.removeIf(request -> request == CreateSessionRequest.class);
int size = Iterables.size(requests);
assertThat(size).isIn(Range.closed(5, 6));
if (size == 5) {
assertThat(requests)
.containsExactly(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(requests)
.containsExactly(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
} else {
assertThat(requests)
.containsExactly(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
}
} else {
assertThat(requests)
.containsExactly(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(requests)
.containsExactly(
CreateSessionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
} else {
assertThat(requests)
.containsExactly(
BatchCreateSessionsRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class,
BeginTransactionRequest.class,
ExecuteBatchDmlRequest.class,
CommitRequest.class);
}
}
}

Expand Down Expand Up @@ -914,13 +1007,18 @@ public void asyncTransactionManagerWithBatchUpdateCommitFails() {
assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
assertThat(e.getMessage()).contains("mutation limit exceeded");
}
ImmutableList<Class<? extends Message>> expectedRequests =
ImmutableList<Class<? extends Message>> expectedRequestsWithRegularSession =
ImmutableList.of(
BatchCreateSessionsRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSession =
ImmutableList.of(
CreateSessionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
if (isMultiplexedSessionsEnabledForRW()) {
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSession);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithRegularSession);
}
}

Expand All @@ -945,13 +1043,18 @@ public void asyncTransactionManagerWaitsUntilAsyncBatchUpdateHasFinished() throw
}
}
}
ImmutableList<Class<? extends Message>> expectedRequests =
ImmutableList<Class<? extends Message>> expectedRequestsWithRegularSession =
ImmutableList.of(
BatchCreateSessionsRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSession =
ImmutableList.of(
CreateSessionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
if (isMultiplexedSessionsEnabled()) {
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSession);
} else {
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
assertThat(mockSpanner.getRequestTypes())
.containsExactlyElementsIn(expectedRequestsWithRegularSession);
}
}

Expand Down Expand Up @@ -1090,4 +1193,11 @@ private boolean isMultiplexedSessionsEnabled() {
}
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession();
}

private boolean isMultiplexedSessionsEnabledForRW() {
if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
return false;
}
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
}
}

0 comments on commit dfd3812

Please sign in to comment.