Refactor OpenTransaction interface

.palantir/revapi.yml

@@ -449,9 +449,41 @@ acceptedBreaks:
       justification: "breaks on internal APIs only"
   "0.1184.0":
     com.palantir.atlasdb:atlasdb-api:
+    - code: "java.method.addedToInterface"
+      new: "method void com.palantir.atlasdb.transaction.api.OpenTransaction::close()"
+      justification: "OpenTransaction has a single internal consumer, and it's current\
+        \ interface caused confusion in recent refactors"
     - code: "java.method.addedToInterface"
       new: "method void com.palantir.atlasdb.transaction.api.Transaction::onCommitOrAbort(java.lang.Runnable)"
       justification: "similar to onSuccess, already used in other prod codepaths"
+    - code: "java.method.removed"
+      old: "method <T, E extends java.lang.Exception> T com.palantir.atlasdb.transaction.api.OpenTransaction::finish(com.palantir.atlasdb.transaction.api.TransactionTask<T,\
+        \ E>) throws E, com.palantir.atlasdb.transaction.api.TransactionFailedRetriableException"
+      justification: "OpenTransaction has a single internal consumer, and it's current\
+        \ interface caused confusion in recent refactors"
+    - code: "java.method.removed"
+      old: "method <T, E extends java.lang.Exception> T com.palantir.atlasdb.transaction.api.OpenTransaction::finishWithCallback(com.palantir.atlasdb.transaction.api.TransactionTask<T,\
+        \ E>, java.lang.Runnable) throws E, com.palantir.atlasdb.transaction.api.TransactionFailedRetriableException"
+      justification: "OpenTransaction has a single internal consumer, and it's current\
+        \ interface caused confusion in recent refactors"
+    - code: "java.method.returnTypeTypeParametersChanged"
+      old: "method java.util.List<com.palantir.atlasdb.transaction.api.OpenTransaction>\
+        \ com.palantir.atlasdb.transaction.api.AutoDelegate_TransactionManager::startTransactions(java.util.List<?\
+        \ extends com.palantir.atlasdb.transaction.api.PreCommitCondition>)"
+      new: "method java.util.List<? extends com.palantir.atlasdb.transaction.api.OpenTransaction>\
+        \ com.palantir.atlasdb.transaction.api.AutoDelegate_TransactionManager::startTransactions(java.util.List<?\
+        \ extends com.palantir.atlasdb.transaction.api.PreCommitCondition>)"
+      justification: "OpenTransaction has a single internal consumer, and it's current\
+        \ interface caused confusion in recent refactors"
+    - code: "java.method.returnTypeTypeParametersChanged"
+      old: "method java.util.List<com.palantir.atlasdb.transaction.api.OpenTransaction>\
+        \ com.palantir.atlasdb.transaction.api.TransactionManager::startTransactions(java.util.List<?\
+        \ extends com.palantir.atlasdb.transaction.api.PreCommitCondition>)"
+      new: "method java.util.List<? extends com.palantir.atlasdb.transaction.api.OpenTransaction>\
+        \ com.palantir.atlasdb.transaction.api.TransactionManager::startTransactions(java.util.List<?\
+        \ extends com.palantir.atlasdb.transaction.api.PreCommitCondition>)"
+      justification: "OpenTransaction has a single internal consumer, and it's current\
+        \ interface caused confusion in recent refactors"
   "0.770.0":
     com.palantir.atlasdb:atlasdb-api:
     - code: "java.class.removed"

atlasdb-api/src/main/java/com/palantir/atlasdb/transaction/api/OpenTransaction.java

@@ -16,27 +16,14 @@
 
 package com.palantir.atlasdb.transaction.api;
 
-import com.palantir.atlasdb.metrics.Timed;
-
-public interface OpenTransaction extends Transaction {
-
-    /**
-     * Runs a provided task, commits the transaction, and performs cleanup. If no further work needs to be done with the
-     * transaction, a no-op task can be passed in.
-     *
-     * @return value returned by the task
-     */
-    @Timed
-    <T, E extends Exception> T finish(TransactionTask<T, E> task) throws E, TransactionFailedRetriableException;
+import java.io.Closeable;
 
+public interface OpenTransaction extends Transaction, Closeable {
     /**
-     * Like {@link #finish(TransactionTask)}, except runs a callback after the task has finished. This callback will
-     * not run while the transaction remains in an uncommitted state, but may run afterwards, regardless of whether
-     * the commit was successful, finished, or aborted.
-     *
-     * If the callback is run, it will run before any {@link Transaction#onSuccess(Runnable)} callbacks.
+     * Aborts the transaction if uncommitted and cleanups transaction state.
+     * All open transactions <b>must be closed</b>.
+     * Not closing transactions after they're no longer in use may lead to arbitrary delays elsewhere in the system.
      */
-    @Timed
-    <T, E extends Exception> T finishWithCallback(TransactionTask<T, E> task, Runnable callback)
-            throws E, TransactionFailedRetriableException;
+    @Override
+    void close();
 }

atlasdb-api/src/main/java/com/palantir/atlasdb/transaction/api/Transaction.java

@@ -453,6 +453,10 @@ enum TransactionType {
     /**
      * Allow consumers to register callbacks to be run on {@link #commit()} or {@link #abort()},
      * after a transaction has committed or aborted.
+     * <p>
+     * {@link #onCommitOrAbort(Runnable)} callbacks are added in a stack and run in opposite order they were added to
+     * the stack. I.e. they're FILO (first-in-last-out).
+     * <p>
      * {@link #onCommitOrAbort(Runnable)} callbacks run before {@link #onSuccess(Runnable)} callbacks.
      * <p>
      * Callbacks are usually cleanup tasks, e.g. {@link PreCommitCondition#cleanup()}.

atlasdb-api/src/main/java/com/palantir/atlasdb/transaction/api/TransactionManager.java

@@ -429,20 +429,29 @@ <T, C extends PreCommitCondition, E extends Exception> T runTaskWithConditionRea
     void registerClosingCallback(Runnable closingCallback);
 
     /**
-     * This method can be used for direct control over the lifecycle of a batch of transactions. For example, if the
-     * work done in each given transaction is interactive and cannot be expressed as a {@link TransactionTask} ahead of
-     * time, this method allows for a long lived transaction object. For any data read or written to the transaction to
-     * be valid, the transaction must be committed by calling {@link OpenTransaction#finish(TransactionTask)} to
-     * also perform additional cleanup. Note that this does not clean up the pre commit condition associated with that
-     * task. The order of transactions returned corresponds with the pre commit conditions passed in, however there are
+     * This method can be used for direct control over the lifecycle of a batch of transactions.
+     * For example, if the work done in each given transaction is interactive and cannot be expressed as a
+     * {@link TransactionTask} ahead of time, this method allows for a long lived transaction object.
+     * <p>
+     * For any data read or written to the transaction to be valid, the transaction must be committed explicitly
+     * by the caller with {@link Transaction#commit()}.
+     * <p>
+     * {@link PreCommitCondition#cleanup()} associated with a transaction will be run when the respective transaction
+     * commits or aborts.
+     * All pre-commit conditions are also cleaned in case {@link #startTransactions(List)} throws.
+     * <p>
+     * Note the caller must call {@link OpenTransaction#close()} after the transaction is committed to perform
+     * additional cleanup. Failure to do so might incur in
+     * <p>
+     * The order of transactions returned corresponds with the pre commit conditions passed in, however there are
      * no guarantees on the ordering of transactions returned with respect to their start timestamp.
      *
      * @return a batch of transactions with associated immutable timestamp locks
      * @deprecated Similar functionality will exist, but this method is likely to change in the future
      */
     @Deprecated
     @Timed
-    List<OpenTransaction> startTransactions(List<? extends PreCommitCondition> condition);
+    List<? extends OpenTransaction> startTransactions(List<? extends PreCommitCondition> condition);
 
     /**
      * Frees resources used by this TransactionManager, and invokes any callbacks registered to run on close.

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/transaction/impl/SnapshotTransactionManager.java

@@ -22,6 +22,7 @@
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Streams;
+import com.google.common.io.Closer;
 import com.palantir.atlasdb.cache.TimestampCache;
 import com.palantir.atlasdb.cell.api.DataKeyValueServiceManager;
 import com.palantir.atlasdb.cell.api.DdlManager;
@@ -72,6 +73,7 @@
 import com.palantir.timestamp.TimestampManagementService;
 import com.palantir.timestamp.TimestampService;
 import com.palantir.util.SafeShutdownRunner;
+import java.io.IOException;
 import java.time.Duration;
 import java.util.List;
 import java.util.Optional;
@@ -190,20 +192,30 @@ protected boolean shouldStopRetrying(int numTimesFailed) {
     public <T, C extends PreCommitCondition, E extends Exception> T runTaskWithConditionThrowOnConflict(
             C condition, ConditionAwareTransactionTask<T, C, E> task) throws E, TransactionFailedRetriableException {
         checkOpen();
-        OpenTransaction openTransaction;
+        try (OpenTransactionImpl openTransaction =
+                runTimed(() -> Iterables.getOnlyElement(startTransactions(ImmutableList.of(condition))), "setupTask")) {
+            return openTransaction.execute(txn -> task.execute(txn, condition));
+        }
+    }
+
+    @Override
+    public List<OpenTransactionImpl> startTransactions(List<? extends PreCommitCondition> conditions) {
         try {
-            openTransaction = runTimed(
-                    () -> Iterables.getOnlyElement(startTransactions(ImmutableList.of(condition))), "setupTask");
+            return startTransactionsInternal(conditions);
         } catch (Exception e) {
-            condition.cleanup();
+            Closer closer = Closer.create(); // N.B. using closer to run all cleanup tasks even if one cleanup throws
+            conditions.forEach(condition -> closer.register(condition::cleanup));
+            try {
+                closer.close();
+            } catch (IOException ex) {
+                e.addSuppressed(ex);
+                log.info("Failed to cleanup pre-commit conditions on startTransaction failure", ex);
+            }
             throw e;
         }
-        return openTransaction.finishWithCallback(
-                transaction -> task.execute(transaction, condition), condition::cleanup);
     }
 
-    @Override
-    public List<OpenTransaction> startTransactions(List<? extends PreCommitCondition> conditions) {
+    private List<OpenTransactionImpl> startTransactionsInternal(List<? extends PreCommitCondition> conditions) {
         if (conditions.isEmpty()) {
             return ImmutableList.of();
         }
@@ -224,35 +236,63 @@ public List<OpenTransaction> startTransactions(List<? extends PreCommitCondition
             recordImmutableTimestamp(immutableTs);
             cleaner.punch(responses.get(0).startTimestampAndPartition().timestamp());
 
-            List<OpenTransaction> transactions = Streams.zip(
+            List<OpenTransactionImpl> transactions = Streams.zip(
                             responses.stream(), conditions.stream(), (response, condition) -> {
                                 LockToken immutableTsLock =
                                         response.immutableTimestamp().getLock();
                                 LongSupplier startTimestampSupplier = response.startTimestampAndPartition()::timestamp;
 
                                 ExpectationsAwareTransaction transaction = createTransaction(
                                         immutableTs, startTimestampSupplier, immutableTsLock, condition);
-                                transaction.onSuccess(
-                                        () -> lockWatchManager.onTransactionCommit(transaction.getTimestamp()));
+
+                                transaction.onCommitOrAbort(transaction::reportExpectationsCollectedData);
+                                transaction.onCommitOrAbort(condition::cleanup);
+                                transaction.onCommitOrAbort(
+                                        () -> lockWatchManager.requestTransactionStateRemovalFromCache(
+                                                response.startTimestampAndPartition()
+                                                        .timestamp()));
+                                // N.B. run this before clearing the state from the cache with
+                                // requestTransactionStateRemovalFromCache
+                                transaction.onCommitOrAbort(() -> {
+                                    if (!transaction.isAborted()) {
+                                        // in onCommitOrAbort + not aborted == committed
+                                        lockWatchManager.onTransactionCommit(transaction.getTimestamp());
+                                    }
+                                });
                                 return new OpenTransactionImpl(transaction, immutableTsLock);
                             })
                     .collect(Collectors.toList());
             openTransactionCounter.inc(transactions.size());
             return transactions;
         } catch (Throwable t) {
-            responses.forEach(response -> lockWatchManager.requestTransactionStateRemovalFromCache(
-                    response.startTimestampAndPartition().timestamp()));
-            timelockService.tryUnlock(responses.stream()
+            // In case of failure, we need to remove the transaction from our local lock-watch cache and unlocking
+            // immutable timestamp.
+            // Note that in case we don't throw, the immutable timestamp is registered to be unlocked in the
+            // SnapshotTransaction constructor. But in case of failure, we need to manually unlock it here.
+
+            Closer closer = Closer.create(); // N.B. using closer to run all cleanup tasks even if one cleanup throws
+            responses.forEach(
+                    response -> closer.register(() -> lockWatchManager.requestTransactionStateRemovalFromCache(
+                            response.startTimestampAndPartition().timestamp())));
+            closer.register(() -> timelockService.tryUnlock(responses.stream()
                     .map(response -> response.immutableTimestamp().getLock())
-                    .collect(Collectors.toSet()));
+                    .collect(Collectors.toSet())));
+
+            try {
+                closer.close();
+            } catch (IOException e) {
+                t.addSuppressed(e);
+                log.info("Failed to cleanup startTransaction resources on startTransaction failure", t);
+            }
             throw Throwables.rewrapAndThrowUncheckedException(t);
         }
     }
 
-    private final class OpenTransactionImpl extends ForwardingTransaction implements OpenTransaction {
+    final class OpenTransactionImpl extends ForwardingTransaction implements OpenTransaction {
 
         private final ExpectationsAwareTransaction delegate;
         private final LockToken immutableTsLock;
+        private boolean hasClosed = false;
 
         private OpenTransactionImpl(ExpectationsAwareTransaction delegate, LockToken immutableTsLock) {
             this.delegate = delegate;
@@ -265,33 +305,29 @@ public ExpectationsAwareTransaction delegate() {
         }
 
         @Override
-        public <T, E extends Exception> T finish(TransactionTask<T, E> task)
-                throws E, TransactionFailedRetriableException {
-            return finishWithCallback(task, () -> {});
-        }
-
-        @Override
-        public <T, E extends Exception> T finishWithCallback(TransactionTask<T, E> task, Runnable callback)
-                throws E, TransactionFailedRetriableException {
-            Timer postTaskTimer = getTimer("finishTask");
-            Timer.Context postTaskContext;
-
-            TransactionTask<T, E> wrappedTask = wrapTaskIfNecessary(task, immutableTsLock);
+        public synchronized void close() {
+            if (hasClosed) {
+                // Some operations in close are not idempotent, e.g. openTransactionCounter.dec().
+                // Let's guarantee we run close() only once, and no-op in following runs to respect Closeable interface.
+                return;
+            }
 
             ExpectationsAwareTransaction txn = delegate;
-            T result;
             try {
-                txn.onCommitOrAbort(txn::reportExpectationsCollectedData);
-                txn.onCommitOrAbort(callback);
-                result = runTaskThrowOnConflict(wrappedTask, txn);
+                if (txn.isUncommitted()) {
+                    txn.abort();
+                }
             } finally {
-                lockWatchManager.requestTransactionStateRemovalFromCache(getTimestamp());
-                postTaskContext = postTaskTimer.time();
                 openTransactionCounter.dec();
             }
             scrubForAggressiveHardDelete(extractSnapshotTransaction(txn));
-            postTaskContext.stop();
-            return result;
+
+            hasClosed = true;
+        }
+
+        <T, E extends Exception> T execute(TransactionTask<T, E> task) {
+            TransactionTask<T, E> wrappedTask = wrapTaskIfNecessary(task, immutableTsLock);
+            return runTimed(() -> runTaskThrowOnConflict(wrappedTask, delegate), "runTaskThrowOnConflict");
         }
     }
 

atlasdb-service/src/main/java/com/palantir/atlasdb/impl/AtlasDbServiceImpl.java

@@ -192,7 +192,9 @@ public TransactionToken startTransaction() {
     public void commit(TransactionToken token) {
         OpenTransaction openTxn = transactions.getIfPresent(token);
         if (openTxn != null) {
-            openTxn.finish((TxTask) transaction -> null);
+            try (openTxn) {
+                openTxn.commit();
+            }
             transactions.invalidate(token);
         }
     }
@@ -201,10 +203,9 @@ public void commit(TransactionToken token) {
     public void abort(TransactionToken token) {
         OpenTransaction openTxn = transactions.getIfPresent(token);
         if (openTxn != null) {
-            openTxn.finish((TxTask) transaction -> {
-                transaction.abort();
-                return null;
-            });
+            try (openTxn) {
+                openTxn.abort();
+            }
             transactions.invalidate(token);
         }
     }