fix: retry transaction if first statements fails and had BeginTransac…

…tion option

@@ -105,8 +105,9 @@ public TransactionContext resetForRetry() {
 "resetForRetry can only be called if the previous attempt" + " aborted");
 }
 try (Scope s = tracer.withSpan(span)) {
+boolean useInlinedBegin = inlineBegin && txn.transactionId != null;
 txn = session.newTransaction();
-if (!inlineBegin) {
+if (!useInlinedBegin) {
 txn.ensureTxn();
 }
 txnState = TransactionState.STARTED;

@@ -55,7 +55,6 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.Callable;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -153,15 +152,13 @@ public void removeListener(Runnable listener) {
 private long retryDelayInMillis = -1L;
 
 /**
-* transactionLock guards that only one request can be beginning the transaction at any time. We
-* only hold on to this lock while a request is creating a transaction. After a transaction has
-* been created, the lock is released and concurrent requests can be executed on the
+* transactionIdFuture will return the transaction id returned by the first statement in the
+* transaction if the BeginTransaction option is included with the first statement of the
 * transaction.
 */
-// private final ReentrantLock transactionLock = new ReentrantLock();
-private volatile CountDownLatch transactionLatch = new CountDownLatch(0);
+private volatile SettableApiFuture<ByteString> transactionIdFuture = null;
 
-private volatile ByteString transactionId;
+volatile ByteString transactionId;
 private Timestamp commitTimestamp;
 
 private TransactionContextImpl(Builder builder) {
@@ -402,39 +399,30 @@ TransactionSelector getTransactionSelector() {
 try {
 // Wait if another request is already beginning, committing or rolling back the
 // transaction.
-
-// transactionLock.lockInterruptibly();
-while (true) {
-CountDownLatch latch;
-synchronized (lock) {
-latch = transactionLatch;
+ApiFuture<ByteString> tx = null;
+synchronized (lock) {
+if (transactionIdFuture == null) {
+transactionIdFuture = SettableApiFuture.create();
+} else {
+tx = transactionIdFuture;
 }
-latch.await();
-
+}
+if (tx == null) {
+return TransactionSelector.newBuilder()
+.setBegin(
+TransactionOptions.newBuilder()
+.setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance()))
+.build();
+} else {
+TransactionSelector.newBuilder().setId(tx.get()).build();
+}
+} catch (ExecutionException e) {
+if (e.getCause() instanceof AbortedException) {
 synchronized (lock) {
-if (transactionLatch.getCount() > 0L) {
-continue;
-}
-// Check again if a transactionId is now available. It could be that the thread that
-// was
-// holding the lock and that had sent a statement with a BeginTransaction request
-// caused
-// an error and did not return a transaction.
-if (transactionId == null) {
-transactionLatch = new CountDownLatch(1);
-// Return a TransactionSelector that will start a new transaction as part of the
-// statement that is being executed.
-return TransactionSelector.newBuilder()
-.setBegin(
-TransactionOptions.newBuilder()
-.setReadWrite(TransactionOptions.ReadWrite.getDefaultInstance()))
-.build();
-} else {
-// transactionLock.unlock();
-break;
-}
+aborted = true;
 }
 }
+throw SpannerExceptionFactory.newSpannerException(e.getCause());
 } catch (InterruptedException e) {
 throw SpannerExceptionFactory.newSpannerExceptionForCancellation(null, e);
 }
@@ -449,8 +437,7 @@ public void onTransactionMetadata(Transaction transaction) {
 // transaction on this instance and release the lock to allow other statements to proceed.
 if (this.transactionId == null && transaction != null && transaction.getId() != null) {
 this.transactionId = transaction.getId();
-transactionLatch.countDown();
-// transactionLock.unlock();
+this.transactionIdFuture.set(transaction.getId());
 }
 }
 
@@ -459,12 +446,15 @@ public void onError(SpannerException e, boolean withBeginTransaction) {
 // Release the transactionLock if that is being held by this thread. That would mean that the
 // statement that was trying to start a transaction caused an error. The next statement should
 // in that case also include a BeginTransaction option.
-
-// if (transactionLock.isHeldByCurrentThread()) {
-// transactionLock.unlock();
-// }
 if (withBeginTransaction) {
-transactionLatch.countDown();
+// Simulate an aborted transaction to force a retry with a new transaction.
+this.transactionIdFuture.setException(
+SpannerExceptionFactory.newSpannerException(
+ErrorCode.ABORTED, "Aborted due to failed initial statement", e));
+// synchronized (lock) {
+// retryDelayInMillis = 0;
+// aborted = true;
+// }
 }
 
 if (e.getErrorCode() == ErrorCode.ABORTED) {
@@ -758,21 +748,25 @@ private <T> T runInternal(final TransactionCallable<T> txCallable) {
 new Callable<T>() {
 @Override
 public T call() {
+boolean useInlinedBegin = inlineBegin;
 if (attempt.get() > 0) {
+if (useInlinedBegin) {
+// Do not inline the BeginTransaction during a retry if the initial attempt did not
+// actually start a transaction.
+useInlinedBegin = txn.transactionId != null;
+}
 txn = session.newTransaction();
 }
 checkState(
 isValid,
 "TransactionRunner has been invalidated by a new operation on the session");
 attempt.incrementAndGet();
-// TODO(user): When using reads, consider using the first read to begin
-// the txn.
 span.addAnnotation(
 "Starting Transaction Attempt",
 ImmutableMap.of("Attempt", AttributeValue.longAttributeValue(attempt.longValue())));
 // Only ensure that there is a transaction if we should not inline the beginTransaction
 // with the first statement.
-if (!inlineBegin) {
+if (!useInlinedBegin) {
 txn.ensureTxn();
 }
 

@@ -280,7 +280,10 @@ public Long run(TransactionContext transaction) throws Exception {
 }
 });
 assertThat(updateCount).isEqualTo(UPDATE_COUNT);
-assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+// The transaction will be retried because the first statement that also tried to include the
+// BeginTransaction statement failed and did not return a transaction. That forces a retry of
+// the entire transaction with an explicit BeginTransaction RPC.
+assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(1);
 // The first update will start a transaction, but then fail the update statement. This will
 // start a transaction on the mock server, but that transaction will never be returned to the
 // client.
@@ -498,12 +501,36 @@ public void testTransactionManagerInlinedBeginTxWithError() {
 }
 }
 }
-assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+// The first statement will fail and not return a transaction id. This will trigger a retry of
+// the entire transaction, and the retry will do an explicit BeginTransaction RPC.
+assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(1);
 // The first statement will start a transaction, but it will never be returned to the client as
 // the update statement fails.
 assertThat(countTransactionsStarted()).isEqualTo(2);
 }
 
+@SuppressWarnings("resource")
+@Test
+public void testTransactionManagerInlinedBeginTxWithUncaughtError() {
+DatabaseClient client =
+spanner.getDatabaseClient(DatabaseId.of("[PROJECT]", "[INSTANCE]", "[DATABASE]"));
+try (TransactionManager txMgr = client.transactionManager()) {
+TransactionContext txn = txMgr.begin();
+while (true) {
+try {
+txn.executeUpdate(INVALID_UPDATE_STATEMENT);
+fail("missing expected exception");
+} catch (AbortedException e) {
+txn = txMgr.resetForRetry();
+}
+}
+} catch (SpannerException e) {
+assertThat(e.getErrorCode()).isEqualTo(ErrorCode.INVALID_ARGUMENT);
+}
+assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+assertThat(countTransactionsStarted()).isEqualTo(1);
+}
+
 @Test
 public void testInlinedBeginAsyncTx() throws InterruptedException, ExecutionException {
 DatabaseClient client =
@@ -631,7 +658,9 @@ public ApiFuture<Long> doWorkAsync(TransactionContext transaction) {
 },
 executor);
 assertThat(updateCount.get()).isEqualTo(UPDATE_COUNT);
-assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(0);
+// The first statement will fail and not return a transaction id. This will trigger a retry of
+// the entire transaction, and the retry will do an explicit BeginTransaction RPC.
+assertThat(countRequests(BeginTransactionRequest.class)).isEqualTo(1);
 // The first update will start a transaction, but then fail the update statement. This will
 // start a transaction on the mock server, but that transaction will never be returned to the
 // client.