Dispose of async tasks

Author: fmeum

src/main/java/com/google/devtools/build/lib/remote/RemoteExecutionService.java

@@ -152,6 +152,7 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Phaser;
@@ -1815,7 +1816,8 @@ public void uploadOutputs(RemoteAction action, SpawnResult spawnResult, Runnable
     if (remoteOptions.remoteCacheAsync
         && !action.getSpawn().getResourceOwner().mayModifySpawnOutputsAfterExecution()) {
       AtomicLong startTime = new AtomicLong();
-      var unused =
+      CountDownLatch done = new CountDownLatch(1);
+      var asyncUpload =
           Single.using(
                   () -> {
                     backgroundTaskPhaser.register();
@@ -1837,10 +1839,20 @@ public void uploadOutputs(RemoteAction action, SpawnResult spawnResult, Runnable
                     backgroundTaskPhaser.arriveAndDeregister();
                     onUploadComplete.run();
                     cacheResource.release();
+                    done.countDown();
+                    // Clean up the done task, cancelling it is a no-op at this point.
+                    outputService.cancelTasks(
+                        action.getRemoteActionExecutionContext().getSpawnOwner());
                   },
                   /* eager= */ false)
               .subscribeOn(scheduler)
               .subscribe(result -> {}, this::reportUploadError);
+      outputService.registerCancellableTask(
+          action.getRemoteActionExecutionContext().getSpawnOwner(),
+          () -> {
+            asyncUpload.dispose();
+            done.await();
+          });
     } else {
       try (SilentCloseable c =
           Profiler.instance().profile(ProfilerTask.UPLOAD_TIME, "upload outputs")) {

src/main/java/com/google/devtools/build/lib/remote/RemoteOutputService.java

@@ -49,13 +49,16 @@
 import java.io.IOException;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /** Output service implementation for the remote build without local output service daemon. */
 public class RemoteOutputService implements OutputService {
 
   private final BlazeDirectories directories;
+  private final ConcurrentHashMap<ActionExecutionMetadata, Cancellable> cancellableTasks =
+      new ConcurrentHashMap<>();
 
   @Nullable private RemoteOutputChecker remoteOutputChecker;
   @Nullable private RemoteActionInputFetcher actionInputFetcher;
@@ -240,4 +243,26 @@ public ArtifactPathResolver createPathResolverForArtifactValues(
             actionInputFetcher);
     return ArtifactPathResolver.createPathResolver(remoteFileSystem, fileSystem.getPath(execRoot));
   }
+
+  @Override
+  public void registerCancellableTask(ActionExecutionMetadata action, Cancellable task) {
+    // We don't expect to have multiple cancellable tasks for the same action, so we avoid the
+    // overhead of a multi-valued map.
+    cancellableTasks.merge(
+        action,
+        task,
+        (oldTask, newTask) ->
+            () -> {
+              oldTask.cancel();
+              newTask.cancel();
+            });
+  }
+
+  @Override
+  public void cancelTasks(ActionExecutionMetadata action) throws InterruptedException {
+    Cancellable task = cancellableTasks.remove(action);
+    if (task != null) {
+      task.cancel();
+    }
+  }
 }

src/main/java/com/google/devtools/build/lib/remote/RemoteSpawnCache.java

@@ -116,6 +116,7 @@ public CacheHandle lookup(Spawn spawn, SpawnExecutionContext context)
       try {
         thisExecution = LocalExecution.createIfDeduplicatable(action);
         if (shouldUploadLocalResults && thisExecution != null) {
+          System.err.println("Deduplication");
           LocalExecution previousOrThisExecution =
               inFlightExecutions.merge(
                   action.getActionKey(),
@@ -145,6 +146,7 @@ public CacheHandle lookup(Spawn spawn, SpawnExecutionContext context)
             Stopwatch fetchTime = Stopwatch.createStarted();
             InMemoryOutput inMemoryOutput;
             try (SilentCloseable c = prof.profile(REMOTE_DOWNLOAD, "download outputs")) {
+              System.err.println("RemoteSpawnCache.lookup: download outputs " + result);
               inMemoryOutput = remoteExecutionService.downloadOutputs(action, result);
             }
             fetchTime.stop();

src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java

@@ -1070,23 +1070,27 @@ public ActionStepOrResult run(Environment env)
         boolean unlockRewindingLock = false;
 
         try {
+          ActionStartedEvent event = new ActionStartedEvent(action, actionStartTimeNanos);
+          if (statusReporter != null) {
+            statusReporter.updateStatus(event);
+          }
+          env.getListener().post(event);
           if (mustAcquireExclusiveRewindingLock(action)) {
             try (SilentCloseable d =
                 profiler.profile(ProfilerTask.ACTION_LOCK, "action.acquireWriteLockForRewinding")) {
               rewindingLock.writeLock().lockInterruptibly();
+              System.err.println("Acquired exclusive rewinding lock for " + action);
               unlockRewindingLock = true;
+              outputService.cancelTasks(action);
             }
           }
-          ActionStartedEvent event = new ActionStartedEvent(action, actionStartTimeNanos);
-          if (statusReporter != null) {
-            statusReporter.updateStatus(event);
-          }
-          env.getListener().post(event);
           if (actionFileSystemType().shouldDoEagerActionPrep()) {
             try (SilentCloseable d = profiler.profile(ProfilerTask.INFO, "action.prepare")) {
               // This call generally deletes any files at locations that are declared outputs of the
               // action, although some actions perform additional work, while others intentionally
               // keep previous outputs in place.
+              System.err.println(
+                  "Deleting outputs of " + action + (wasRewound(action) ? " (rewound)" : ""));
               action.prepare(
                   actionExecutionContext.getExecRoot(),
                   actionExecutionContext.getPathResolver(),
@@ -1121,6 +1125,7 @@ public ActionStepOrResult run(Environment env)
           return ActionStepOrResult.of(e);
         } finally {
           if (unlockRewindingLock) {
+            System.err.println("Released exclusive rewinding lock for " + action);
             rewindingLock.writeLock().unlock();
           }
           notifyActionCompletion(env.getListener(), !lostInputs);
@@ -1230,12 +1235,14 @@ private ActionStepOrResult executeAction(ExtendedEventHandler eventHandler, Acti
         boolean unlockRewindingLock = false;
         if (mustAcquireSharedRewindingLock()) {
           rewindingLock.readLock().lockInterruptibly();
+          System.err.println("Acquired shared rewinding lock for " + action);
           unlockRewindingLock = true;
         }
         try {
           result = action.execute(actionExecutionContext);
         } finally {
           if (unlockRewindingLock) {
+            System.err.println("Released shared rewinding lock for " + action);
             rewindingLock.readLock().unlock();
           }
         }

src/main/java/com/google/devtools/build/lib/vfs/BUILD

@@ -103,5 +103,6 @@ java_library(
         "//src/main/protobuf:failure_details_java_proto",
         "//third_party:guava",
         "//third_party:jsr305",
+        "//third_party:rxjava3",
     ],
 )

src/main/java/com/google/devtools/build/lib/vfs/OutputService.java

@@ -265,4 +265,24 @@ default BulkDeleter bulkDeleter() {
   default XattrProvider getXattrProvider(XattrProvider delegate) {
     return delegate;
   }
+
+  /** A task with a cancellation callback. */
+  @FunctionalInterface
+  interface Cancellable {
+    void cancel() throws InterruptedException;
+  }
+
+  /**
+   * Registers a cancellation callback for a task that may still be running after the action has
+   * completed.
+   */
+  default void registerCancellableTask(ActionExecutionMetadata action, Cancellable task) {
+    throw new UnsupportedOperationException();
+  }
+
+  /**
+   * Cancels and awaits the completion of all tasks registered with {@link
+   * #registerCancellableTask}.
+   */
+  default void cancelTasks(ActionExecutionMetadata action) throws InterruptedException {}
 }

src/test/java/com/google/devtools/build/lib/skyframe/rewinding/RewindingTest.java

@@ -144,7 +144,7 @@ public void multipleLostInputsForRewindPlan() throws Exception {
 
   @Test
   public void multipleLostInputsForRewindPlan_standalone() throws Exception {
-    addOptions("--remote_executor=", "--spawn_strategy=standalone");
+    addOptions("--spawn_strategy=standalone", "--remote_executor=");
     helper.runMultipleLostInputsForRewindPlan();
   }