Skip to content

Commit 4745c3e

Browse files
harshachintaharshachinta
committed
chore(spanner): fix AsyncTransactionManagerTest
1 parent 261631c commit 4745c3e

File tree

1 file changed

+160
-37
lines changed

1 file changed

+160
-37
lines changed

‎google-cloud-spanner/src/test/java/com/google/cloud/spanner/AsyncTransactionManagerTest.java‎

Lines changed: 160 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import static org.junit.Assert.assertNotNull;
2929
import static org.junit.Assert.assertThrows;
3030
import static org.junit.Assert.assertTrue;
31+
import static org.junit.Assume.assumeFalse;
3132

3233
import com.google.api.core.ApiFuture;
3334
import com.google.api.core.ApiFutureCallback;
@@ -179,9 +180,11 @@ public void asyncTransactionManager_shouldRollbackOnCloseAsync() throws Exceptio
179180
AsyncTransactionManager manager = client().transactionManagerAsync();
180181
TransactionContext txn = manager.beginAsync().get();
181182
txn.executeUpdateAsync(UPDATE_STATEMENT).get();
182-
final TransactionSelector selector =
183-
((TransactionContextImpl) ((SessionPoolTransactionContext) txn).delegate)
184-
.getTransactionSelector();
183+
if (txn instanceof SessionPoolTransactionContext) {
184+
txn = ((SessionPoolTransactionContext) txn).delegate;
185+
}
186+
TransactionContextImpl impl = (TransactionContextImpl) txn;
187+
final TransactionSelector selector = impl.getTransactionSelector();
185188

186189
SpannerApiFutures.get(manager.closeAsync());
187190
// The mock server should already have the Rollback request, as we are waiting for the returned
@@ -248,6 +251,11 @@ public void asyncTransactionManagerUpdate() throws Exception {
248251

249252
@Test
250253
public void asyncTransactionManagerIsNonBlocking() throws Exception {
254+
// TODO: Remove this condition once DelayedAsyncTransactionManager is made non-blocking with
255+
// multiplexed sessions.
256+
assumeFalse(
257+
"DelayedAsyncTransactionManager is currently blocking with multiplexed sessions.",
258+
spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW());
251259
mockSpanner.freeze();
252260
try (AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync()) {
253261
TransactionContextFuture transactionContextFuture = manager.beginAsync();
@@ -359,7 +367,22 @@ public void asyncTransactionManagerFireAndForgetInvalidUpdate() throws Exception
359367
ExecuteSqlRequest.class,
360368
ExecuteSqlRequest.class,
361369
CommitRequest.class);
362-
if (isMultiplexedSessionsEnabled()) {
370+
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionForRW =
371+
ImmutableList.of(
372+
CreateSessionRequest.class,
373+
// The first update that fails. This will cause a transaction retry.
374+
ExecuteSqlRequest.class,
375+
// The retry will use an explicit BeginTransaction call.
376+
BeginTransactionRequest.class,
377+
// The first update will again fail, but now there is a transaction id, so the
378+
// transaction can continue.
379+
ExecuteSqlRequest.class,
380+
ExecuteSqlRequest.class,
381+
CommitRequest.class);
382+
if (isMultiplexedSessionsEnabledForRW()) {
383+
assertThat(mockSpanner.getRequestTypes())
384+
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionForRW);
385+
} else if (isMultiplexedSessionsEnabled()) {
363386
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
364387
} else {
365388
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
@@ -502,14 +525,25 @@ public void asyncTransactionManagerUpdateAbortedWithoutGettingResult() throws Ex
502525
// The server may receive 1 or 2 commit requests depending on whether the call to
503526
// commitAsync() already knows that the transaction has aborted. If it does, it will not
504527
// attempt to call the Commit RPC and instead directly propagate the Aborted error.
505-
assertThat(mockSpanner.getRequestTypes())
506-
.containsAtLeast(
507-
BatchCreateSessionsRequest.class,
508-
ExecuteSqlRequest.class,
509-
// The retry will use a BeginTransaction RPC.
510-
BeginTransactionRequest.class,
511-
ExecuteSqlRequest.class,
512-
CommitRequest.class);
528+
if (isMultiplexedSessionsEnabledForRW()) {
529+
assertThat(mockSpanner.getRequestTypes())
530+
.containsAtLeast(
531+
CreateSessionRequest.class,
532+
ExecuteSqlRequest.class,
533+
// The retry will use a BeginTransaction RPC.
534+
BeginTransactionRequest.class,
535+
ExecuteSqlRequest.class,
536+
CommitRequest.class);
537+
} else {
538+
assertThat(mockSpanner.getRequestTypes())
539+
.containsAtLeast(
540+
BatchCreateSessionsRequest.class,
541+
ExecuteSqlRequest.class,
542+
// The retry will use a BeginTransaction RPC.
543+
BeginTransactionRequest.class,
544+
ExecuteSqlRequest.class,
545+
CommitRequest.class);
546+
}
513547
break;
514548
} catch (AbortedException e) {
515549
transactionContextFuture = manager.resetForRetryAsync();
@@ -557,7 +591,11 @@ public void asyncTransactionManagerWaitsUntilAsyncUpdateHasFinished() throws Exc
557591
executor)
558592
.commitAsync()
559593
.get();
560-
if (isMultiplexedSessionsEnabled()) {
594+
if (isMultiplexedSessionsEnabledForRW()) {
595+
assertThat(mockSpanner.getRequestTypes())
596+
.containsExactly(
597+
CreateSessionRequest.class, ExecuteSqlRequest.class, CommitRequest.class);
598+
} else if (isMultiplexedSessionsEnabled()) {
561599
assertThat(mockSpanner.getRequestTypes())
562600
.containsExactly(
563601
CreateSessionRequest.class,
@@ -601,6 +639,11 @@ public void asyncTransactionManagerBatchUpdate() throws Exception {
601639

602640
@Test
603641
public void asyncTransactionManagerIsNonBlockingWithBatchUpdate() throws Exception {
642+
// TODO: Remove this condition once DelayedAsyncTransactionManager is made non-blocking with
643+
// multiplexed sessions.
644+
assumeFalse(
645+
"DelayedAsyncTransactionManager is currently blocking with multiplexed sessions.",
646+
spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW());
604647
mockSpanner.freeze();
605648
try (AsyncTransactionManager manager = clientWithEmptySessionPool().transactionManagerAsync()) {
606649
TransactionContextFuture transactionContextFuture = manager.beginAsync();
@@ -678,7 +721,16 @@ public void asyncTransactionManagerFireAndForgetInvalidBatchUpdate() throws Exce
678721
ExecuteBatchDmlRequest.class,
679722
ExecuteBatchDmlRequest.class,
680723
CommitRequest.class);
681-
if (isMultiplexedSessionsEnabled()) {
724+
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
725+
ImmutableList.of(
726+
CreateSessionRequest.class,
727+
ExecuteBatchDmlRequest.class,
728+
ExecuteBatchDmlRequest.class,
729+
CommitRequest.class);
730+
if (isMultiplexedSessionsEnabledForRW()) {
731+
assertThat(mockSpanner.getRequestTypes())
732+
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
733+
} else if (isMultiplexedSessionsEnabled()) {
682734
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
683735
} else {
684736
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
@@ -722,7 +774,17 @@ public void asyncTransactionManagerBatchUpdateAborted() throws Exception {
722774
BeginTransactionRequest.class,
723775
ExecuteBatchDmlRequest.class,
724776
CommitRequest.class);
725-
if (isMultiplexedSessionsEnabled()) {
777+
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
778+
ImmutableList.of(
779+
CreateSessionRequest.class,
780+
ExecuteBatchDmlRequest.class,
781+
BeginTransactionRequest.class,
782+
ExecuteBatchDmlRequest.class,
783+
CommitRequest.class);
784+
if (isMultiplexedSessionsEnabledForRW()) {
785+
assertThat(mockSpanner.getRequestTypes())
786+
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
787+
} else if (isMultiplexedSessionsEnabled()) {
726788
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
727789
} else {
728790
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
@@ -764,7 +826,20 @@ public void asyncTransactionManagerBatchUpdateAbortedBeforeFirstStatement() thro
764826
BeginTransactionRequest.class,
765827
ExecuteBatchDmlRequest.class,
766828
CommitRequest.class);
767-
if (isMultiplexedSessionsEnabled()) {
829+
// When requests run using multiplexed session with read-write enabled, the
830+
// BatchCreateSessionsRequest will not be
831+
// triggered because we are creating an empty pool during initialization.
832+
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
833+
ImmutableList.of(
834+
CreateSessionRequest.class,
835+
ExecuteBatchDmlRequest.class,
836+
BeginTransactionRequest.class,
837+
ExecuteBatchDmlRequest.class,
838+
CommitRequest.class);
839+
if (isMultiplexedSessionsEnabledForRW()) {
840+
assertThat(mockSpanner.getRequestTypes())
841+
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
842+
} else if (isMultiplexedSessionsEnabled()) {
768843
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
769844
} else {
770845
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
@@ -825,7 +900,18 @@ public void asyncTransactionManagerWithBatchUpdateCommitAborted() throws Excepti
825900
BeginTransactionRequest.class,
826901
ExecuteBatchDmlRequest.class,
827902
CommitRequest.class);
828-
if (isMultiplexedSessionsEnabled()) {
903+
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
904+
ImmutableList.of(
905+
CreateSessionRequest.class,
906+
ExecuteBatchDmlRequest.class,
907+
CommitRequest.class,
908+
BeginTransactionRequest.class,
909+
ExecuteBatchDmlRequest.class,
910+
CommitRequest.class);
911+
if (isMultiplexedSessionsEnabledForRW()) {
912+
assertThat(mockSpanner.getRequestTypes())
913+
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
914+
} else if (isMultiplexedSessionsEnabled()) {
829915
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
830916
} else {
831917
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
@@ -866,28 +952,46 @@ public void asyncTransactionManagerBatchUpdateAbortedWithoutGettingResult() thro
866952
}
867953
assertThat(attempt.get()).isEqualTo(2);
868954
List<Class<? extends AbstractMessage>> requests = mockSpanner.getRequestTypes();
869-
// Remove the CreateSession requests for multiplexed sessions, as those are not relevant for
870-
// this test.
871-
requests.removeIf(request -> request == CreateSessionRequest.class);
872955
int size = Iterables.size(requests);
873956
assertThat(size).isIn(Range.closed(5, 6));
874957
if (size == 5) {
875-
assertThat(requests)
876-
.containsExactly(
877-
BatchCreateSessionsRequest.class,
878-
ExecuteBatchDmlRequest.class,
879-
BeginTransactionRequest.class,
880-
ExecuteBatchDmlRequest.class,
881-
CommitRequest.class);
958+
if (isMultiplexedSessionsEnabledForRW()) {
959+
assertThat(requests)
960+
.containsExactly(
961+
CreateSessionRequest.class,
962+
ExecuteBatchDmlRequest.class,
963+
BeginTransactionRequest.class,
964+
ExecuteBatchDmlRequest.class,
965+
CommitRequest.class);
966+
} else {
967+
assertThat(requests)
968+
.containsExactly(
969+
BatchCreateSessionsRequest.class,
970+
ExecuteBatchDmlRequest.class,
971+
BeginTransactionRequest.class,
972+
ExecuteBatchDmlRequest.class,
973+
CommitRequest.class);
974+
}
882975
} else {
883-
assertThat(requests)
884-
.containsExactly(
885-
BatchCreateSessionsRequest.class,
886-
ExecuteBatchDmlRequest.class,
887-
CommitRequest.class,
888-
BeginTransactionRequest.class,
889-
ExecuteBatchDmlRequest.class,
890-
CommitRequest.class);
976+
if (isMultiplexedSessionsEnabledForRW()) {
977+
assertThat(requests)
978+
.containsExactly(
979+
CreateSessionRequest.class,
980+
ExecuteBatchDmlRequest.class,
981+
CommitRequest.class,
982+
BeginTransactionRequest.class,
983+
ExecuteBatchDmlRequest.class,
984+
CommitRequest.class);
985+
} else {
986+
assertThat(requests)
987+
.containsExactly(
988+
BatchCreateSessionsRequest.class,
989+
ExecuteBatchDmlRequest.class,
990+
CommitRequest.class,
991+
BeginTransactionRequest.class,
992+
ExecuteBatchDmlRequest.class,
993+
CommitRequest.class);
994+
}
891995
}
892996
}
893997

@@ -918,7 +1022,13 @@ public void asyncTransactionManagerWithBatchUpdateCommitFails() {
9181022
ImmutableList<Class<? extends Message>> expectedRequests =
9191023
ImmutableList.of(
9201024
BatchCreateSessionsRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
921-
if (isMultiplexedSessionsEnabled()) {
1025+
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
1026+
ImmutableList.of(
1027+
CreateSessionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
1028+
if (isMultiplexedSessionsEnabledForRW()) {
1029+
assertThat(mockSpanner.getRequestTypes())
1030+
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
1031+
} else if (isMultiplexedSessionsEnabled()) {
9221032
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
9231033
} else {
9241034
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
@@ -949,7 +1059,13 @@ public void asyncTransactionManagerWaitsUntilAsyncBatchUpdateHasFinished() throw
9491059
ImmutableList<Class<? extends Message>> expectedRequests =
9501060
ImmutableList.of(
9511061
BatchCreateSessionsRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
952-
if (isMultiplexedSessionsEnabled()) {
1062+
ImmutableList<Class<? extends Message>> expectedRequestsWithMultiplexedSessionsRW =
1063+
ImmutableList.of(
1064+
CreateSessionRequest.class, ExecuteBatchDmlRequest.class, CommitRequest.class);
1065+
if (isMultiplexedSessionsEnabledForRW()) {
1066+
assertThat(mockSpanner.getRequestTypes())
1067+
.containsExactlyElementsIn(expectedRequestsWithMultiplexedSessionsRW);
1068+
} else if (isMultiplexedSessionsEnabled()) {
9531069
assertThat(mockSpanner.getRequestTypes()).containsAtLeastElementsIn(expectedRequests);
9541070
} else {
9551071
assertThat(mockSpanner.getRequestTypes()).containsExactlyElementsIn(expectedRequests);
@@ -1122,4 +1238,11 @@ private boolean isMultiplexedSessionsEnabled() {
11221238
}
11231239
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession();
11241240
}
1241+
1242+
private boolean isMultiplexedSessionsEnabledForRW() {
1243+
if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
1244+
return false;
1245+
}
1246+
return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
1247+
}
11251248
}

0 commit comments

Comments
 (0)