Let TaskDeduplicator callers decide which executions to join (https://github.com/bazelbuild/bazel/pull/30516) ### Description `TaskDeduplicator` deduplicated executions by key alone, so callers with different requirements either shared an execution that didn't satisfy all of them or had to be separated into different keys, which also gives up sharing in the direction where it is safe. Every execution is now started with caller-provided attributes and every caller supplies a predicate over them, so the join relation can be asymmetric: a caller that needs little joins an execution that produces more, while a caller that needs more starts its own and takes over as the execution that subsequent callers see. The decision is made inside a single `ConcurrentHashMap.compute`, which also lets the retry loop around a concurrently canceled execution go away. `executeIfNew` and `executeUnconditionally` are replaced by a single `execute`. `maybeJoinExecution` is removed as its only use is questionable: Having a `DISCARD` build wait for a `KEEP` build is not a clear performance win since the latter may issue network calls. This also fixes a caller that cancels its future releasing its reference to the shared execution twice: `IndividuallyCancelableFuture` completed itself with `setFuture`, which cancels its argument when it has already been canceled. This was harmless as used today, but not obvious without referring to Guava code and potentially problematic for future changes. ### Motivation Simplify the interface of `TaskDeduplicator` and make its inner working more obvious. At the same time, make the reuse of ongoing executions more explicit and flexible in `MerkleTreeComputer`. ### Build API Changes No ### Checklist - [x] I have added tests for the new use cases (if any). - [ ] I have updated the documentation (if applicable). ### Release Notes RELNOTES: None Closes #30516. PiperOrigin-RevId: 970629545 Change-Id: Ib041dbc815d33349dcf137a3720864d8ac26ea51
diff --git a/src/main/java/com/google/devtools/build/lib/concurrent/BUILD b/src/main/java/com/google/devtools/build/lib/concurrent/BUILD index a1c76c8..246dbe2 100644 --- a/src/main/java/com/google/devtools/build/lib/concurrent/BUILD +++ b/src/main/java/com/google/devtools/build/lib/concurrent/BUILD
@@ -57,6 +57,7 @@ srcs = ["TaskDeduplicator.java"], deps = [ "//third_party/java/error_prone:annotations", + "//third_party/java/guava:annotations", "//third_party/java/guava:concurrent", "//third_party/java/jsr305_annotations", ],
diff --git a/src/main/java/com/google/devtools/build/lib/concurrent/TaskDeduplicator.java b/src/main/java/com/google/devtools/build/lib/concurrent/TaskDeduplicator.java index 849e8c1..ead3df8 100644 --- a/src/main/java/com/google/devtools/build/lib/concurrent/TaskDeduplicator.java +++ b/src/main/java/com/google/devtools/build/lib/concurrent/TaskDeduplicator.java
@@ -15,134 +15,135 @@ import static com.google.common.util.concurrent.MoreExecutors.directExecutor; +import com.google.common.annotations.VisibleForTesting; import com.google.common.util.concurrent.AbstractFuture; +import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.errorprone.annotations.CheckReturnValue; +import java.util.concurrent.CancellationException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Predicate; import java.util.function.Supplier; import javax.annotation.Nullable; /** - * Deduplicates concurrent tasks identified by unique keys. For any given key, only one task is - * actively executed at a time. + * Deduplicates concurrent tasks identified by unique keys. + * + * <p>Every execution is started with caller-provided attributes of type {@code A} that describe + * what it produces. A caller joins an ongoing execution for its key only if its own {@code canJoin} + * predicate accepts that execution's attributes, which lets callers with different needs share work + * asymmetrically: a caller that needs little can join an execution that produces more, while a + * caller that needs more starts its own execution and takes over as the one that subsequent callers + * see. Callers that all need the same thing pass a predicate that is always true. + * + * <p>At most one execution per key is joinable at a time, but an execution that has been taken over + * keeps running for the callers that already joined it. In particular, full deduplication is not + * guaranteed and executions must support being run concurrently. * * <p>Any futures returned by this class can be individually canceled without affecting other * callers. The shared task is only canceled if all callers have canceled their futures and the task * is interrupted if and only if all callers requested interruption. + * + * <p>An execution is forgotten as soon as it completes, so a task that failed or was canceled is + * retried by the next call for the same key. Callers that need results to be remembered across + * calls have to layer their own cache on top. + * + * <p>The {@code taskSupplier} and {@code canJoin} callbacks run while a lock on an internal map is + * held. They must therefore be cheap, must not block and must not call back into this deduplicator, + * which would either deadlock or fail with an {@link IllegalStateException}. Supplying a task + * typically means submitting it to an executor, which satisfies these requirements; running it + * inline does not. */ -public final class TaskDeduplicator<K, V> { - private final ConcurrentMap<K, RefcountedFuture<V>> inFlightTasks = new ConcurrentHashMap<>(); +public final class TaskDeduplicator<K, A, V> { + private final ConcurrentMap<K, RefcountedFuture<A, V>> inFlightTasks = new ConcurrentHashMap<>(); /** - * Returns a future representing either a new or already ongoing execution of the task. + * Returns a future representing either the ongoing execution for the key, if {@code canJoin} + * accepts its attributes, or a new execution started with the given attributes. + * + * <p>A new execution replaces the ongoing one as the execution that subsequent callers can join. + * The replaced execution keeps running for the callers that already joined it. * * <p>The returned future must eventually be completed. The task is only canceled if the futures * returned to all callers for the same key have been canceled. * - * <p>taskSupplier may be called multiple times. It should be inexpensive and free of side - * effects. + * <p>taskSupplier is called at most once and only if no ongoing execution is joined. */ @CheckReturnValue - public ListenableFuture<V> executeIfNew(K key, Supplier<ListenableFuture<V>> taskSupplier) { - while (true) { - var isNewHolder = new boolean[1]; - var future = - inFlightTasks.computeIfAbsent( - key, - unusedKey -> { - isNewHolder[0] = true; - return RefcountedFuture.wrap(taskSupplier.get()); - }); - if (isNewHolder[0]) { - future.addListener(() -> inFlightTasks.remove(key, future), directExecutor()); - } else { - // The shared future may have been canceled between the lookup and the call to retain(). - if (!future.retain()) { - inFlightTasks.remove(key, future); - continue; - } - } - return IndividuallyCancelableFuture.wrap(future); - } - } - - /** - * Returns a future representing either a new or already ongoing execution of the task that is - * guaranteed to happen-after any executions started before the call of this method. - * - * <p>The returned future must eventually be completed. The task is only canceled if the futures - * returned to all callers for the same key have been canceled. - * - * <p>taskSupplier may be called multiple times. It should be inexpensive and free of side - * effects. - */ - @CheckReturnValue - public ListenableFuture<V> executeUnconditionally( - K key, Supplier<ListenableFuture<V>> taskSupplier) { - inFlightTasks.remove(key); - return executeIfNew(key, taskSupplier); - } - - /** - * Returns a future representing an already ongoing execution of the task or null if there is - * none. - * - * <p>The returned future must eventually be completed. The task is only canceled if the futures - * returned to all callers for the same key have been canceled. - */ - @CheckReturnValue - @Nullable - public ListenableFuture<V> maybeJoinExecution(K key) { - var future = inFlightTasks.get(key); - if (future == null) { - return null; - } - if (!future.retain()) { - inFlightTasks.remove(key, future); - return null; + public ListenableFuture<V> execute( + K key, + @Nullable A attributes, + Predicate<? super A> canJoin, + Supplier<ListenableFuture<V>> taskSupplier) { + var isNewHolder = new boolean[1]; + var future = + inFlightTasks.compute( + key, + (unusedKey, ongoingExecution) -> { + if (ongoingExecution != null + && canJoin.test(ongoingExecution.attributes()) + && ongoingExecution.retain()) { + return ongoingExecution; + } + isNewHolder[0] = true; + return new RefcountedFuture<>(attributes, taskSupplier.get()); + }); + if (isNewHolder[0]) { + future.addListener(() -> inFlightTasks.remove(key, future), directExecutor()); } return IndividuallyCancelableFuture.wrap(future); } + @VisibleForTesting + @Nullable + RefcountedFuture<A, V> ongoingExecutionForTesting(K key) { + return inFlightTasks.get(key); + } + /** * A future adapter that is canceled only when {@link #cancel} has been called one more time than * {@link #retain}. */ - private static final class RefcountedFuture<V> extends AbstractFuture<V> implements Runnable { + @VisibleForTesting + static final class RefcountedFuture<A, V> extends AbstractFuture<V> { + @Nullable private final A attributes; private final ListenableFuture<V> delegate; // Initialized to 1 in the constructor and incremented via retain(). Once it drops to 0, it // can never return to 1 or higher (0 is a sticky state). private final AtomicInteger refcount = new AtomicInteger(1); private volatile boolean mayInterruptIfRunning = true; - static <V> RefcountedFuture<V> wrap(ListenableFuture<V> delegate) { - var wrappedFuture = new RefcountedFuture<>(delegate); - delegate.addListener(wrappedFuture, directExecutor()); - return wrappedFuture; - } - - RefcountedFuture(ListenableFuture<V> delegate) { + RefcountedFuture(@Nullable A attributes, ListenableFuture<V> delegate) { + this.attributes = attributes; this.delegate = delegate; + // Completes this future with the delegate's outcome and forwards a cancellation of this + // future to the delegate. setFuture(delegate); } - @Override - public void run() {} + /** Returns the attributes this execution was started with. */ + @Nullable + A attributes() { + return attributes; + } @Override public boolean cancel(boolean mayInterruptIfRunning) { if (!mayInterruptIfRunning) { this.mayInterruptIfRunning = false; } + // The write above happens-before the update below, which in turn happens-before the update + // that drops the refcount to zero. The caller that observes zero thus sees the requests of + // all other callers not to interrupt the task. if (refcount.updateAndGet(oldCount -> oldCount >= 1 ? oldCount - 1 : 0) == 0) { return super.cancel(this.mayInterruptIfRunning); } return false; } - @Nullable @Override protected String pendingToString() { return "delegate=[%s (%d active uses)]".formatted(delegate, refcount.get()); @@ -152,6 +153,12 @@ boolean retain() { return refcount.updateAndGet(oldCount -> oldCount >= 1 ? oldCount + 1 : 0) != 0; } + + /** Returns the number of callers that are still interested in this execution. */ + @VisibleForTesting + int activeUses() { + return refcount.get(); + } } /** @@ -160,21 +167,32 @@ */ private static final class IndividuallyCancelableFuture<V> extends AbstractFuture<V> implements Runnable { - private final RefcountedFuture<V> delegate; + private final RefcountedFuture<?, V> delegate; - static <V> ListenableFuture<V> wrap(RefcountedFuture<V> delegate) { + static <V> ListenableFuture<V> wrap(RefcountedFuture<?, V> delegate) { var wrappedFuture = new IndividuallyCancelableFuture<>(delegate); delegate.addListener(wrappedFuture, directExecutor()); return wrappedFuture; } - IndividuallyCancelableFuture(RefcountedFuture<V> delegate) { + IndividuallyCancelableFuture(RefcountedFuture<?, V> delegate) { this.delegate = delegate; } @Override public void run() { - setFuture(delegate); + // The outcome is copied over manually rather than with setFuture: if this future has already + // been canceled, setFuture would cancel the delegate a second time and thus decrement its + // refcount twice on behalf of this single caller. + try { + set(Futures.getDone(delegate)); + } catch (ExecutionException e) { + setException(e.getCause()); + } catch (CancellationException e) { + // Either all callers canceled their futures, in which case this future is already canceled, + // or the task completed as canceled on its own, e.g. because its executor was shut down. + super.cancel(/* mayInterruptIfRunning= */ false); + } } @Override
diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteExternalOverlayFileSystem.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteExternalOverlayFileSystem.java index 3813029..49448c3 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteExternalOverlayFileSystem.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteExternalOverlayFileSystem.java
@@ -92,7 +92,7 @@ private final int externalDirectorySegmentCount; private final FileSystem nativeFs; private final RemoteExternalFileSystem externalFs; - private final TaskDeduplicator<String, Void> materializations = new TaskDeduplicator<>(); + private final TaskDeduplicator<String, Void, Void> materializations = new TaskDeduplicator<>(); // The names of the repos whose contents have been fully materialized to nativeFs. private final Set<String> materializedRepos = ConcurrentHashMap.newKeySet(); // As long as a repo name appears as a key in this map, the repo contents are available in @@ -398,8 +398,10 @@ } var unused = getFromFuture( - materializations.executeIfNew( + materializations.execute( repo.getName(), + /* attributes= */ null, + /* canJoin= */ unusedAttributes -> true, () -> materializationExecutor.submit( () -> {
diff --git a/src/main/java/com/google/devtools/build/lib/remote/merkletree/MerkleTreeComputer.java b/src/main/java/com/google/devtools/build/lib/remote/merkletree/MerkleTreeComputer.java index 8fe4090..8cc0e47 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/merkletree/MerkleTreeComputer.java +++ b/src/main/java/com/google/devtools/build/lib/remote/merkletree/MerkleTreeComputer.java
@@ -97,6 +97,7 @@ import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicLong; import java.util.function.Predicate; import java.util.function.Supplier; import javax.annotation.Nullable; @@ -182,8 +183,9 @@ private final String workspaceName; private final Digest emptyDigest; private final MerkleTree.Uploadable emptyTree; - private final TaskDeduplicator<InFlightCacheKey, MerkleTree.RootOnly> inFlightComputations = - new TaskDeduplicator<>(); + private final TaskDeduplicator<InFlightCacheKey, InFlightAttributes, MerkleTree.RootOnly> + inFlightComputations = new TaskDeduplicator<>(); + private final AtomicLong inFlightComputationSequence = new AtomicLong(); public MerkleTreeComputer( DigestUtil digestUtil, @@ -203,7 +205,14 @@ new MerkleTree.RootOnly.BlobsUploaded(emptyDigest, 0, 0), ImmutableSortedMap.of()); } - /** Specifies which blobs should be retained in the Merkle tree. */ + /** + * Specifies which blobs should be retained in the Merkle tree. + * + * <p>The constants are ordered from the weakest to the strongest policy: the result of a + * computation performed with a given policy also satisfies every preceding one. Deduplication of + * ongoing sub-Merkle tree computations relies on this order to decide whether an ongoing + * computation can be joined. + */ public enum BlobPolicy { /** * No blobs are retained and the returned MerkleTree is a {@link MerkleTree.RootOnly}. @@ -252,6 +261,36 @@ @Nullable PathFragment unmappedExecPath) {} /** + * Describes what an ongoing sub-Merkle tree computation produces, which determines whether other + * computations can join it. + * + * @param blobPolicy the policy the computation was started with + * @param sequenceNumber a number assigned before the computation is registered, so that a + * computation started before a given call to {@link #computeIfAbsent} has a lower number than + * the one that call assigns to itself + */ + private record InFlightAttributes(BlobPolicy blobPolicy, long sequenceNumber) {} + + /** + * Whether an ongoing computation with the attributes {@code ongoing} produces a result that also + * satisfies a call with the attributes {@code requested}. + */ + private static boolean canJoin(InFlightAttributes ongoing, InFlightAttributes requested) { + // BlobPolicy constants are ordered from the weakest to the strongest policy, so an ongoing + // computation is reusable if its policy retains at least as much as the requested one. In + // particular, a KEEP_AND_REUPLOAD computation never joins a KEEP one, which wouldn't reupload + // the blobs that the remote cache lost. + if (ongoing.blobPolicy().compareTo(requested.blobPolicy()) < 0) { + return false; + } + // A KEEP_AND_REUPLOAD computation additionally has to reupload the blobs after the request + // discovered that they are missing. An ongoing computation that started earlier may already + // have uploaded them before they were lost, so only a later one will do. + return requested.blobPolicy() != BlobPolicy.KEEP_AND_REUPLOAD + || ongoing.sequenceNumber() > requested.sequenceNumber(); + } + + /** * Builds a Merkle tree for the inputs of a {@link Spawn}. * * @param toolInputs the set of paths of inputs that are considered tools. Note that these paths @@ -935,6 +974,9 @@ return immediateFuture(cachedRoot); } } + // Uploading computations are kept under a separate key so that a DISCARD computation never + // joins one: its future only completes after the upload, whereas building the tree again only + // costs local work. var uploadBlobs = blobPolicy != BlobPolicy.DISCARD; // When the upload of a path mapped tree artifact is shared between two actions that each have // that tree artifact as an input under a different unmmapped exec path, CacheNotFoundExceptions @@ -947,25 +989,14 @@ new InFlightCacheKey(cacheKey, isTool, uploadBlobs, uploadBlobs ? unmappedExecPath : null); AsyncCallable<MerkleTree.RootOnly> buildMerkleTreeTask = () -> { - // There is a window in which a concurrent call may have removed the in-flight cache entry - // while this one had already passed the check above. Recheck the persistent cache to - // avoid unnecessary work. + // A concurrent computation may have completed and populated the persistent cache after + // this one had already passed the check above. Recheck it to avoid unnecessary work. var cachedRoot = persistentCache.getIfPresent(cacheKey); if (cachedRoot != null && (blobPolicy == BlobPolicy.DISCARD || cachedRoot instanceof MerkleTree.RootOnly.BlobsUploaded)) { return immediateFuture(cachedRoot); } - // An ongoing computation with blobs can be reused for one that doesn't require them. - if (blobPolicy == BlobPolicy.DISCARD) { - var inFlightComputation = - inFlightComputations.maybeJoinExecution( - new InFlightCacheKey( - cacheKey, isTool, /* uploadBlobs= */ true, unmappedExecPath)); - if (inFlightComputation != null) { - return inFlightComputation; - } - } ListenableFuture<MerkleTree> merkleTreeFuture; try { // Subtrees either consist entirely of tool inputs or don't contain any. The same @@ -1022,11 +1053,15 @@ }; Supplier<ListenableFuture<MerkleTree.RootOnly>> buildMerkleTreeTaskSupplier = () -> Futures.submitAsync(buildMerkleTreeTask, MERKLE_TREE_BUILD_POOL); - if (blobPolicy == BlobPolicy.KEEP_AND_REUPLOAD) { - return inFlightComputations.executeUnconditionally(key, buildMerkleTreeTaskSupplier); - } else { - return inFlightComputations.executeIfNew(key, buildMerkleTreeTaskSupplier); - } + // The sequence number is claimed before the computation is registered, so every computation + // that is already ongoing at this point has a lower one. + var attributes = + new InFlightAttributes(blobPolicy, inFlightComputationSequence.getAndIncrement()); + return inFlightComputations.execute( + key, + attributes, + /* canJoin= */ ongoing -> canJoin(ongoing, attributes), + buildMerkleTreeTaskSupplier); } private static <T> T getFromFuture(Future<T> future) throws IOException, InterruptedException {
diff --git a/src/test/java/com/google/devtools/build/lib/concurrent/TaskDeduplicatorTest.java b/src/test/java/com/google/devtools/build/lib/concurrent/TaskDeduplicatorTest.java index cb1a74e..7fc2685 100644 --- a/src/test/java/com/google/devtools/build/lib/concurrent/TaskDeduplicatorTest.java +++ b/src/test/java/com/google/devtools/build/lib/concurrent/TaskDeduplicatorTest.java
@@ -22,10 +22,11 @@ import java.util.Random; import java.util.concurrent.CancellationException; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; -import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.Predicate; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -34,12 +35,27 @@ @RunWith(JUnit4.class) public class TaskDeduplicatorTest { + // Executions in the tests below that exercise attributes are described by an integer strength, + // mirroring the use case in MerkleTreeComputer: an execution started with a higher strength + // produces everything a weaker one would. A caller may thus join an ongoing execution that is at + // least as strong as what it needs, but never a weaker one. + private static final int WEAK = 0; + private static final int STRONG = 1; + + // Tests that don't exercise attributes accept any ongoing execution. + private static final Predicate<Object> JOIN_ANY = unusedAttributes -> true; + + private static Predicate<Integer> atLeast(int strength) { + return ongoingStrength -> ongoingStrength >= strength; + } + @Test - public void executeIfNew_taskFinished_completed() throws Exception { - var deduplicator = new TaskDeduplicator<String, String>(); + public void execute_taskFinished_completed() throws Exception { + var deduplicator = new TaskDeduplicator<String, Void, String>(); var taskFuture = SettableFuture.<String>create(); - ListenableFuture<String> result = deduplicator.executeIfNew("key1", () -> taskFuture); + ListenableFuture<String> result = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); taskFuture.set("value1"); @@ -47,12 +63,13 @@ } @Test - public void executeIfNew_taskHasError_propagateError() { - var deduplicator = new TaskDeduplicator<String, String>(); + public void execute_taskHasError_propagateError() { + var deduplicator = new TaskDeduplicator<String, Void, String>(); var taskFuture = SettableFuture.<String>create(); var error = new IllegalStateException("error"); - ListenableFuture<String> result = deduplicator.executeIfNew("key1", () -> taskFuture); + ListenableFuture<String> result = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); taskFuture.setException(error); @@ -61,14 +78,16 @@ } @Test - public void executeIfNew_taskInProgress_noReExecution() throws Exception { - var deduplicator = new TaskDeduplicator<String, String>(); + public void execute_taskInProgress_noReExecution() throws Exception { + var deduplicator = new TaskDeduplicator<String, Void, String>(); var taskFuture = SettableFuture.<String>create(); var executionTimes = new AtomicInteger(0); ListenableFuture<String> result1 = - deduplicator.executeIfNew( + deduplicator.execute( "key1", + /* attributes= */ null, + JOIN_ANY, () -> { executionTimes.incrementAndGet(); return taskFuture; @@ -76,8 +95,10 @@ // Second call with the same key should return the same future, not re-execute ListenableFuture<String> result2 = - deduplicator.executeIfNew( + deduplicator.execute( "key1", + /* attributes= */ null, + JOIN_ANY, () -> { throw new IllegalStateException("should not be called"); }); @@ -93,14 +114,16 @@ } @Test - public void executeIfNew_taskFinished_reExecution() throws Exception { - var deduplicator = new TaskDeduplicator<String, String>(); + public void execute_taskFinished_reExecution() throws Exception { + var deduplicator = new TaskDeduplicator<String, Void, String>(); var executionTimes = new AtomicInteger(0); // First execution ListenableFuture<String> result1 = - deduplicator.executeIfNew( + deduplicator.execute( "key1", + /* attributes= */ null, + JOIN_ANY, () -> { executionTimes.incrementAndGet(); var future = SettableFuture.<String>create(); @@ -113,8 +136,10 @@ // Second execution after first is finished should re-execute ListenableFuture<String> result2 = - deduplicator.executeIfNew( + deduplicator.execute( "key1", + /* attributes= */ null, + JOIN_ANY, () -> { executionTimes.incrementAndGet(); var future = SettableFuture.<String>create(); @@ -127,14 +152,16 @@ } @Test - public void executeIfNew_taskCanceled_reExecution() { - var deduplicator = new TaskDeduplicator<String, String>(); + public void execute_taskCanceled_reExecution() { + var deduplicator = new TaskDeduplicator<String, Void, String>(); var executionTimes = new AtomicInteger(0); // First execution ListenableFuture<String> result1 = - deduplicator.executeIfNew( + deduplicator.execute( "key1", + /* attributes= */ null, + JOIN_ANY, () -> { executionTimes.incrementAndGet(); return Futures.immediateCancelledFuture(); @@ -145,8 +172,10 @@ // Second execution after first is finished should re-execute ListenableFuture<String> result2 = - deduplicator.executeIfNew( + deduplicator.execute( "key1", + /* attributes= */ null, + JOIN_ANY, () -> { executionTimes.incrementAndGet(); return Futures.immediateCancelledFuture(); @@ -157,11 +186,12 @@ } @Test - public void executeIfNew_cancel_cancelled() throws Exception { - var deduplicator = new TaskDeduplicator<String, String>(); + public void execute_cancel_cancelled() throws Exception { + var deduplicator = new TaskDeduplicator<String, Void, String>(); var taskFuture = SettableFuture.<String>create(); - ListenableFuture<String> result = deduplicator.executeIfNew("key1", () -> taskFuture); + ListenableFuture<String> result = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); result.cancel(true); @@ -170,12 +200,14 @@ } @Test - public void executeIfNew_cancelWhenMultipleFutures_notCancelled() throws Exception { - var deduplicator = new TaskDeduplicator<String, String>(); + public void execute_cancelWhenMultipleFutures_notCancelled() throws Exception { + var deduplicator = new TaskDeduplicator<String, Void, String>(); var taskFuture = SettableFuture.<String>create(); - ListenableFuture<String> result1 = deduplicator.executeIfNew("key1", () -> taskFuture); - ListenableFuture<String> result2 = deduplicator.executeIfNew("key1", () -> taskFuture); + ListenableFuture<String> result1 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); + ListenableFuture<String> result2 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); // Cancel one future multiple times result1.cancel(true); @@ -191,12 +223,14 @@ } @Test - public void executeIfNew_cancelWhenMultipleFutures_allCancelled() { - var deduplicator = new TaskDeduplicator<String, String>(); + public void execute_cancelWhenMultipleFutures_allCancelled() { + var deduplicator = new TaskDeduplicator<String, Void, String>(); var taskFuture = SettableFuture.<String>create(); - ListenableFuture<String> result1 = deduplicator.executeIfNew("key1", () -> taskFuture); - ListenableFuture<String> result2 = deduplicator.executeIfNew("key1", () -> taskFuture); + ListenableFuture<String> result1 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); + ListenableFuture<String> result2 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); // Cancel both futures result1.cancel(true); @@ -204,16 +238,19 @@ assertThat(result1.isCancelled()).isTrue(); assertThat(result2.isCancelled()).isTrue(); + assertThat(taskFuture.isCancelled()).isTrue(); } @Test - public void executeIfNew_multipleTasks_completeOne() throws Exception { - var deduplicator = new TaskDeduplicator<String, String>(); + public void execute_multipleTasks_completeOne() throws Exception { + var deduplicator = new TaskDeduplicator<String, Void, String>(); var taskFuture1 = SettableFuture.<String>create(); var taskFuture2 = SettableFuture.<String>create(); - ListenableFuture<String> result1 = deduplicator.executeIfNew("key1", () -> taskFuture1); - ListenableFuture<String> result2 = deduplicator.executeIfNew("key2", () -> taskFuture2); + ListenableFuture<String> result1 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture1); + ListenableFuture<String> result2 = + deduplicator.execute("key2", /* attributes= */ null, JOIN_ANY, () -> taskFuture2); taskFuture1.set("value1"); @@ -226,11 +263,274 @@ } @Test - public void executeIfNeeded_executeAndCancelLoop_noErrors() { + public void execute_errorInAsyncCallable_propagated() { + var deduplicator = new TaskDeduplicator<String, Void, String>(); + var expectedException = new RuntimeException("task creation failed"); + + var actualException = + assertThrows( + RuntimeException.class, + () -> + deduplicator.execute( + "key1", + /* attributes= */ null, + JOIN_ANY, + () -> { + throw expectedException; + })); + assertThat(actualException).isSameInstanceAs(expectedException); + } + + @Test + public void execute_taskCanceledIndependently_allCallersCanceled() { + var deduplicator = new TaskDeduplicator<String, Void, String>(); + var taskFuture = SettableFuture.<String>create(); + + ListenableFuture<String> result1 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); + ListenableFuture<String> result2 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); + + // The task is canceled by whoever runs it rather than by any of its callers, e.g. because its + // executor was shut down. + taskFuture.cancel(/* mayInterruptIfRunning= */ true); + + assertThat(result1.isCancelled()).isTrue(); + assertThat(result2.isCancelled()).isTrue(); + assertThrows(CancellationException.class, result1::get); + } + + @Test + public void execute_cancelOneFuture_otherStillSeesError() { + var deduplicator = new TaskDeduplicator<String, Void, String>(); + var taskFuture = SettableFuture.<String>create(); + var error = new IllegalStateException("error"); + + ListenableFuture<String> result1 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); + ListenableFuture<String> result2 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); + + result1.cancel(true); + taskFuture.setException(error); + + assertThat(result1.isCancelled()).isTrue(); + var exception = assertThrows(ExecutionException.class, result2::get); + assertThat(exception).hasCauseThat().isSameInstanceAs(error); + } + + /** + * Canceling a future must release exactly one reference to the shared task, also after the task + * has completed. + * + * <p>A stray second release is otherwise only observable through a rare race: a concurrent {@link + * TaskDeduplicator#execute} that looks up the shared task just before it completes and retains it + * just after would be refused a perfectly good result and would redo the work. + */ + @Test + public void execute_cancelOneFuture_releasesExactlyOneReference() throws Exception { + var deduplicator = new TaskDeduplicator<String, Void, String>(); + var taskFuture = SettableFuture.<String>create(); + + ListenableFuture<String> result1 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); + ListenableFuture<String> result2 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); + ListenableFuture<String> result3 = + deduplicator.execute("key1", /* attributes= */ null, JOIN_ANY, () -> taskFuture); + var sharedTask = deduplicator.ongoingExecutionForTesting("key1"); + assertThat(sharedTask.activeUses()).isEqualTo(3); + + result1.cancel(true); + assertThat(sharedTask.activeUses()).isEqualTo(2); + + taskFuture.set("value1"); + + assertThat(result2.get()).isEqualTo("value1"); + assertThat(result3.get()).isEqualTo("value1"); + assertThat(sharedTask.activeUses()).isEqualTo(2); + } + + @Test + public void execute_canJoinAcceptsOngoingExecution_joins() throws Exception { + var deduplicator = new TaskDeduplicator<String, Integer, String>(); + var taskFuture = SettableFuture.<String>create(); + + ListenableFuture<String> result1 = + deduplicator.execute("key1", STRONG, atLeast(STRONG), () -> taskFuture); + // A weaker caller is satisfied by the stronger ongoing execution. + ListenableFuture<String> result2 = + deduplicator.execute( + "key1", + WEAK, + atLeast(WEAK), + () -> { + throw new IllegalStateException("should not be called"); + }); + + taskFuture.set("value1"); + + assertThat(result1.get()).isEqualTo("value1"); + assertThat(result2.get()).isEqualTo("value1"); + } + + @Test + public void execute_canJoinRejectsOngoingExecution_startsNewExecution() throws Exception { + var deduplicator = new TaskDeduplicator<String, Integer, String>(); + var weakTaskFuture = SettableFuture.<String>create(); + var strongTaskFuture = SettableFuture.<String>create(); + var executionTimes = new AtomicInteger(0); + + ListenableFuture<String> weakResult = + deduplicator.execute( + "key1", + WEAK, + atLeast(WEAK), + () -> { + executionTimes.incrementAndGet(); + return weakTaskFuture; + }); + // The ongoing execution is too weak, so this caller has to start its own. + ListenableFuture<String> strongResult = + deduplicator.execute( + "key1", + STRONG, + atLeast(STRONG), + () -> { + executionTimes.incrementAndGet(); + return strongTaskFuture; + }); + + assertThat(executionTimes.get()).isEqualTo(2); + + // The replaced execution keeps running for the caller that joined it. + strongTaskFuture.set("strong"); + assertThat(strongResult.get()).isEqualTo("strong"); + assertThat(weakResult.isDone()).isFalse(); + + weakTaskFuture.set("weak"); + assertThat(weakResult.get()).isEqualTo("weak"); + } + + @Test + public void execute_newExecutionTakesOverSlot_laterCallersJoinIt() throws Exception { + var deduplicator = new TaskDeduplicator<String, Integer, String>(); + var weakTaskFuture = SettableFuture.<String>create(); + var strongTaskFuture = SettableFuture.<String>create(); + + var unused = deduplicator.execute("key1", WEAK, atLeast(WEAK), () -> weakTaskFuture); + ListenableFuture<String> strongResult = + deduplicator.execute("key1", STRONG, atLeast(STRONG), () -> strongTaskFuture); + // The weak execution is no longer joinable, but the strong one that replaced it is. + ListenableFuture<String> lateWeakResult = + deduplicator.execute( + "key1", + WEAK, + atLeast(WEAK), + () -> { + throw new IllegalStateException("should not be called"); + }); + + strongTaskFuture.set("strong"); + + assertThat(strongResult.get()).isEqualTo("strong"); + assertThat(lateWeakResult.get()).isEqualTo("strong"); + } + + @Test + public void execute_canJoinRejects_cancelWhenMultipleFutures_notCancelled() throws Exception { + var deduplicator = new TaskDeduplicator<String, Integer, String>(); + var taskFuture = SettableFuture.<String>create(); + + ListenableFuture<String> result1 = + deduplicator.execute("key1", STRONG, atLeast(STRONG), () -> taskFuture); + ListenableFuture<String> result2 = + deduplicator.execute("key1", WEAK, atLeast(WEAK), () -> taskFuture); + + result1.cancel(true); + + assertThat(result1.isCancelled()).isTrue(); + assertThat(result2.isDone()).isFalse(); + + taskFuture.set("value1"); + + assertThat(result2.get()).isEqualTo("value1"); + } + + /** + * A caller must never end up on an execution that its {@code canJoin} predicate rejects, not even + * one that a concurrent call starts while it is looking up the key. + * + * <p>This is what makes {@code MerkleTreeComputer}'s {@code KEEP_AND_REUPLOAD} policy reliable: + * joining a {@code KEEP} computation would skip the reupload of the blobs that the remote cache + * lost. + */ + @Test + public void execute_concurrentWeakerExecution_neverJoinsIt() throws Exception { + int rounds = 2000; + var failures = new ConcurrentLinkedQueue<String>(); + + try (var testExecutorService = Executors.newVirtualThreadPerTaskExecutor()) { + for (int round = 0; round < rounds; ++round) { + var deduplicator = new TaskDeduplicator<String, Integer, Integer>(); + var start = new CountDownLatch(1); + var done = new CountDownLatch(2); + + testExecutorService.execute( + () -> { + try { + start.await(); + // Every execution completes with the strength it was started with. + ListenableFuture<Integer> result = + deduplicator.execute( + "key1", STRONG, atLeast(STRONG), () -> Futures.immediateFuture(STRONG)); + if (result.get() != STRONG) { + failures.add("joined an execution of strength " + result.get()); + } + } catch (Throwable e) { + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } + failures.add(e.toString()); + } finally { + done.countDown(); + } + }); + testExecutorService.execute( + () -> { + try { + start.await(); + var taskFuture = SettableFuture.<Integer>create(); + ListenableFuture<Integer> result = + deduplicator.execute("key1", WEAK, atLeast(WEAK), () -> taskFuture); + // Leave the execution in flight for a moment so that it overlaps with the + // concurrent stronger call. + taskFuture.set(WEAK); + var unused = result.get(); + } catch (Throwable e) { + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } + failures.add(e.toString()); + } finally { + done.countDown(); + } + }); + + start.countDown(); + done.await(); + } + } + + assertThat(failures).isEmpty(); + } + + @Test + public void execute_executeAndCancelLoop_noErrors() { int taskCount = 1000; int maxKey = 20; var random = new Random(); - var deduplicator = new TaskDeduplicator<String, Void>(); + var deduplicator = new TaskDeduplicator<String, Void, Void>(); var throwables = new ConcurrentLinkedQueue<Throwable>(); try (var taskExecutorService = Executors.newFixedThreadPool(50); @@ -240,8 +540,10 @@ () -> { try { ListenableFuture<Void> future = - deduplicator.executeIfNew( + deduplicator.execute( "key" + random.nextInt(maxKey), + /* attributes= */ null, + JOIN_ANY, () -> Futures.submit( () -> { @@ -275,54 +577,55 @@ } @Test - public void executeIfNew_taskCompletedBeforeSecondCall_bothGetSameResult() throws Exception { - var deduplicator = new TaskDeduplicator<String, String>(); - var taskStarted = new AtomicBoolean(false); + public void mixedOperations_executeAndCancelLoop_noErrors() { + int taskCount = 20000; + int maxKey = 20; + var random = new Random(); + var deduplicator = new TaskDeduplicator<String, Integer, Integer>(); + var throwables = new ConcurrentLinkedQueue<Throwable>(); - // First call - task completes immediately in-flight - ListenableFuture<String> result1 = - deduplicator.executeIfNew( - "key1", + try (var taskExecutorService = Executors.newFixedThreadPool(8); + var testExecutorService = Executors.newVirtualThreadPerTaskExecutor()) { + for (int i = 0; i < taskCount; ++i) { + testExecutorService.execute( () -> { - taskStarted.set(true); - var future = SettableFuture.<String>create(); - future.set("value1"); - return future; + try { + String key = "key" + random.nextInt(maxKey); + int strength = random.nextInt(3); + ListenableFuture<Integer> future = + deduplicator.execute( + key, + strength, + atLeast(strength), + () -> + Futures.submit( + () -> { + Thread.sleep(0, 200_000); + return strength; + }, + taskExecutorService)); + if (!future.isDone() && random.nextBoolean()) { + future.cancel(true); + } else { + // A joined execution must always be at least as strong as what was requested. + assertThat(future.get()).isAtLeast(strength); + } + } catch (Throwable e) { + if (e instanceof InterruptedException) { + Thread.currentThread().interrupt(); + } + throwables.add(e); + } }); + } + } - assertThat(result1.get()).isEqualTo("value1"); - assertThat(taskStarted.get()).isTrue(); - taskStarted.set(false); - - // Second call - task should execute again since first is done - ListenableFuture<String> result2 = - deduplicator.executeIfNew( - "key1", - () -> { - taskStarted.set(true); - var future = SettableFuture.<String>create(); - future.set("value2"); - return future; - }); - - assertThat(result2.get()).isEqualTo("value2"); - assertThat(taskStarted.get()).isTrue(); - } - - @Test - public void executeIfNew_errorInAsyncCallable_propagated() { - var deduplicator = new TaskDeduplicator<String, String>(); - var expectedException = new RuntimeException("task creation failed"); - - var actualException = - assertThrows( - RuntimeException.class, - () -> - deduplicator.executeIfNew( - "key1", - () -> { - throw expectedException; - })); - assertThat(actualException).isSameInstanceAs(expectedException); + if (!throwables.isEmpty()) { + var combinedError = new AssertionError(); + for (var throwable : throwables) { + combinedError.addSuppressed(throwable); + } + throw combinedError; + } } }