Skip to content

Commit 60acf51

Browse files
stIncMalenhachicha
stIncMale
authored and
nhachicha
committed
Refactor withTransaction
Now it looks similar to the algorithm described in the specification prose. I also addressed other concerns that I expressed in the PR review.
1 parent 8b24709 commit 60acf51

1 file changed

Lines changed: 57 additions & 52 deletions

File tree

‎driver-sync/src/main/com/mongodb/client/internal/ClientSessionImpl.java‎

Lines changed: 57 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import com.mongodb.MongoException;
2222
import com.mongodb.MongoExecutionTimeoutException;
2323
import com.mongodb.MongoInternalException;
24-
import com.mongodb.MongoOperationTimeoutException;
2524
import com.mongodb.MongoTimeoutException;
2625
import com.mongodb.ReadConcern;
2726
import com.mongodb.TransactionOptions;
@@ -33,7 +32,6 @@
3332
import com.mongodb.internal.observability.micrometer.TransactionSpan;
3433
import com.mongodb.internal.operation.AbortTransactionOperation;
3534
import com.mongodb.internal.operation.CommitTransactionOperation;
36-
import com.mongodb.internal.operation.OperationHelper;
3735
import com.mongodb.internal.operation.ReadOperation;
3836
import com.mongodb.internal.operation.WriteConcernHelper;
3937
import com.mongodb.internal.operation.WriteOperation;
@@ -160,6 +158,12 @@ public void abortTransaction() {
160158
}
161159
}
162160

161+
private void abortIfInTransaction() {
162+
if (transactionState == TransactionState.IN) {
163+
abortTransaction();
164+
}
165+
}
166+
163167
private void startTransaction(final TransactionOptions transactionOptions, final TimeoutContext timeoutContext) {
164168
Boolean snapshot = getOptions().isSnapshot();
165169
if (snapshot != null && snapshot) {
@@ -258,9 +262,10 @@ public <T> T withTransaction(final TransactionBody<T> transactionBody) {
258262
public <T> T withTransaction(final TransactionBody<T> transactionBody, final TransactionOptions options) {
259263
notNull("transactionBody", transactionBody);
260264
TimeoutContext withTransactionTimeoutContext = createTimeoutContext(options);
261-
final boolean hasTimeoutMS = withTransactionTimeoutContext.hasTimeoutMS();
262-
Timeout withTransactionTimeout = withTransactionTimeoutContext.timeoutOrAlternative(
263-
assertNotNull(TimeoutContext.startTimeout(MAX_RETRY_TIME_LIMIT_MS)));
265+
boolean timeoutMsConfigured = withTransactionTimeoutContext.hasTimeoutMS();
266+
Timeout withTransactionTimeout = assertNotNull(timeoutMsConfigured
267+
? withTransactionTimeoutContext.getTimeout()
268+
: TimeoutContext.startTimeout(MAX_RETRY_TIME_LIMIT_MS));
264269
BooleanSupplier withTransactionTimeoutExpired = () -> withTransactionTimeout.call(TimeUnit.MILLISECONDS,
265270
() -> false, ms -> false, () -> true);
266271
int transactionAttempt = 0;
@@ -270,35 +275,37 @@ public <T> T withTransaction(final TransactionBody<T> transactionBody, final Tra
270275
outer:
271276
while (true) {
272277
if (transactionAttempt > 0) {
273-
backoff(transactionAttempt, withTransactionTimeout, assertNotNull(lastError), hasTimeoutMS);
278+
backoff(transactionAttempt, withTransactionTimeout, assertNotNull(lastError), timeoutMsConfigured);
274279
}
275-
T retVal;
276280
try {
277281
startTransaction(options, withTransactionTimeoutContext);
278282
transactionAttempt++;
279-
280283
if (transactionSpan != null) {
281284
transactionSpan.setIsConvenientTransaction();
282285
}
286+
} catch (Throwable e) {
287+
abortIfInTransaction();
288+
throw e;
289+
}
290+
T retVal;
291+
try {
283292
retVal = transactionBody.execute();
284293
} catch (Throwable e) {
285-
if (transactionState == TransactionState.IN) {
286-
abortTransaction();
287-
}
294+
abortIfInTransaction();
288295
if (e instanceof MongoException) {
289-
lastError = (MongoException) e;
290-
if (!(e instanceof MongoOperationTimeoutException)) {
291-
MongoException exceptionToHandle = OperationHelper.unwrap((MongoException) e);
292-
if (exceptionToHandle.hasErrorLabel(TRANSIENT_TRANSACTION_ERROR_LABEL)) {
293-
if (withTransactionTimeoutExpired.getAsBoolean()) {
294-
throw timeoutException(hasTimeoutMS, e);
295-
} else {
296-
if (transactionSpan != null) {
297-
transactionSpan.spanFinalizing(false);
298-
}
299-
continue;
300-
}
296+
MongoException mongoException = (MongoException) e;
297+
if (mongoException.hasErrorLabel(TRANSIENT_TRANSACTION_ERROR_LABEL)) {
298+
if (transactionSpan != null) {
299+
transactionSpan.spanFinalizing(false);
301300
}
301+
lastError = mongoException;
302+
continue;
303+
} else if (mongoException.hasErrorLabel(UNKNOWN_TRANSACTION_COMMIT_RESULT_LABEL)) {
304+
throw e;
305+
} else {
306+
throw withTransactionTimeoutExpired.getAsBoolean()
307+
? wrapInMongoTimeoutException(mongoException, timeoutMsConfigured)
308+
: mongoException;
302309
}
303310
}
304311
throw e;
@@ -308,28 +315,22 @@ public <T> T withTransaction(final TransactionBody<T> transactionBody, final Tra
308315
try {
309316
commitTransaction(false);
310317
break;
311-
} catch (MongoException e) {
312-
lastError = e;
313-
clearTransactionContextOnError(e);
314-
if (!(e instanceof MongoOperationTimeoutException)) {
315-
if (withTransactionTimeoutExpired.getAsBoolean()) {
316-
throw timeoutException(hasTimeoutMS, e);
317-
318-
} else {
319-
applyMajorityWriteConcernToTransactionOptions();
320-
321-
if (!(e instanceof MongoExecutionTimeoutException)
322-
&& e.hasErrorLabel(UNKNOWN_TRANSACTION_COMMIT_RESULT_LABEL)) {
323-
continue;
324-
} else if (e.hasErrorLabel(TRANSIENT_TRANSACTION_ERROR_LABEL)) {
325-
if (transactionSpan != null) {
326-
transactionSpan.spanFinalizing(true);
327-
}
328-
continue outer;
329-
}
318+
} catch (MongoException mongoException) {
319+
if (mongoException.hasErrorLabel(UNKNOWN_TRANSACTION_COMMIT_RESULT_LABEL)
320+
&& !(mongoException instanceof MongoExecutionTimeoutException)
321+
&& !withTransactionTimeoutExpired.getAsBoolean()) {
322+
applyMajorityWriteConcernToTransactionOptions();
323+
continue;
324+
} else if (mongoException.hasErrorLabel(TRANSIENT_TRANSACTION_ERROR_LABEL)) {
325+
if (transactionSpan != null) {
326+
transactionSpan.spanFinalizing(true);
330327
}
328+
lastError = mongoException;
329+
continue outer;
331330
}
332-
throw e;
331+
throw withTransactionTimeoutExpired.getAsBoolean()
332+
? wrapInMongoTimeoutException(mongoException, timeoutMsConfigured)
333+
: mongoException;
333334
}
334335
}
335336
}
@@ -351,9 +352,7 @@ public TransactionSpan getTransactionSpan() {
351352
@Override
352353
public void close() {
353354
try {
354-
if (transactionState == TransactionState.IN) {
355-
abortTransaction();
356-
}
355+
abortIfInTransaction();
357356
} finally {
358357
clearTransactionContext();
359358
super.close();
@@ -391,10 +390,10 @@ private TimeoutContext createTimeoutContext(final TransactionOptions transaction
391390
}
392391

393392
private static void backoff(final int transactionAttempt,
394-
final Timeout withTransactionTimeout, final MongoException lastError, final boolean hasTimeoutMS) {
393+
final Timeout withTransactionTimeout, final MongoException lastError, final boolean timeoutMsConfigured) {
395394
long backoffMs = ExponentialBackoff.calculateTransactionBackoffMs(transactionAttempt);
396395
withTransactionTimeout.shortenBy(backoffMs, TimeUnit.MILLISECONDS).onExpired(() -> {
397-
throw timeoutException(hasTimeoutMS, lastError);
396+
throw wrapInMongoTimeoutException(lastError, timeoutMsConfigured);
398397
});
399398
try {
400399
if (backoffMs > 0) {
@@ -405,9 +404,15 @@ private static void backoff(final int transactionAttempt,
405404
}
406405
}
407406

408-
private static MongoException timeoutException(final boolean hasTimeoutMS, final Throwable cause) {
409-
return hasTimeoutMS
410-
? createMongoTimeoutException(cause) // CSOT timeout exception
411-
: new MongoTimeoutException("Operation exceeded the timeout limit", cause); // Legacy timeout exception
407+
private static MongoException wrapInMongoTimeoutException(final MongoException cause, final boolean timeoutMsConfigured) {
408+
return timeoutMsConfigured
409+
? createMongoTimeoutException(cause)
410+
: wrapInNonTimeoutMsMongoTimeoutException(cause);
411+
}
412+
413+
private static MongoTimeoutException wrapInNonTimeoutMsMongoTimeoutException(final MongoException cause) {
414+
return cause instanceof MongoTimeoutException
415+
? (MongoTimeoutException) cause
416+
: new MongoTimeoutException("Operation exceeded the timeout limit", cause);
412417
}
413418
}

0 commit comments

Comments
 (0)