Force synchronous upload and reuse of possibly modified spawn outputs

src/main/java/com/google/devtools/build/lib/actions/ActionExecutionMetadata.java

@@ -114,4 +114,14 @@ default String getProgressMessage(RepositoryMapping mainRepositoryMapping) {
   default boolean mayInsensitivelyPropagateInputs() {
     return false;
   }
+
+  /**
+   * Returns true if the action may modify spawn outputs after the spawn has executed.
+   *
+   * <p>If this returns true, any kind of spawn output caching or reuse needs to happen
+   * synchronously directly after the spawn execution.
+   */
+  default boolean mayModifySpawnOutputsAfterExecution() {
+    return false;
+  }
 }

src/main/java/com/google/devtools/build/lib/analysis/test/TestRunnerAction.java

@@ -309,6 +309,16 @@ private static ImmutableSet<Artifact> nonNullAsSet(Artifact... artifacts) {
     this.isExecutedOnWindows = isExecutedOnWindows;
   }
 
+  @Override
+  public boolean mayModifySpawnOutputsAfterExecution() {
+    // Test actions modify test spawn outputs after execution:
+    // - if there are multiple attempts (unavoidable);
+    // - in all cases due to appending any stray stderr output to the test log in
+    //   StandaloneTestStrategy.
+    // TODO: Get rid of the second case and only return true if there are multiple attempts.
+    return true;
+  }
+
   public boolean isExecutedOnWindows() {
     return isExecutedOnWindows;
   }

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

@@ -1333,6 +1333,14 @@ public InMemoryOutput downloadOutputs(RemoteAction action, RemoteActionResult re
   public static final class LocalExecution {
     private final RemoteAction action;
     private final SettableFuture<SpawnResult> spawnResultFuture;
+    private final Phaser spawnResultConsumers =
+        new Phaser(1) {
+          @Override
+          protected boolean onAdvance(int phase, int registeredParties) {
+            // We only use a single phase.
+            return true;
+          }
+        };
 
     private LocalExecution(RemoteAction action) {
       this.action = action;
@@ -1356,6 +1364,37 @@ public static LocalExecution createIfDeduplicatable(RemoteAction action) {
       return new LocalExecution(action);
     }
 
+    /**
+     * Attempts to register a thread waiting for the {@link #spawnResultFuture} to become available
+     * and returns true if successful.
+     *
+     * <p>Every call to this method must be matched by a call to {@link #unregister()} via
+     * try-finally.
+     *
+     * <p>This always returns true for actions that do not modify their spawns' outputs after
+     * execution.
+     */
+    public boolean registerForOutputReuse() {
+      // We only use a single phase.
+      return spawnResultConsumers.register() == 0;
+    }
+
+    /**
+     * Unregisters a thread waiting for the {@link #spawnResultFuture}, either after successful
+     * reuse of the outputs or upon failure.
+     */
+    public void unregister() {
+      spawnResultConsumers.arriveAndDeregister();
+    }
+
+    /**
+     * Waits for all potential consumers of the {@link #spawnResultFuture} to be done with their
+     * output reuse.
+     */
+    public void awaitAllOutputReuse() {
+      spawnResultConsumers.arriveAndAwaitAdvance();
+    }
+
     /**
      * Signals to all potential consumers of the {@link #spawnResultFuture} that this execution has
      * been cancelled and that the result will not be available.
@@ -1571,7 +1610,8 @@ public void uploadOutputs(RemoteAction action, SpawnResult spawnResult, Runnable
         SpawnResult.Status.SUCCESS.equals(spawnResult.status()) && spawnResult.exitCode() == 0,
         "shouldn't upload outputs of failed local action");
 
-    if (remoteOptions.remoteCacheAsync) {
+    if (remoteOptions.remoteCacheAsync
+        && !action.getSpawn().getResourceOwner().mayModifySpawnOutputsAfterExecution()) {
       Single.using(
               remoteCache::retain,
               remoteCache ->

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

@@ -109,89 +109,110 @@ public CacheHandle lookup(Spawn spawn, SpawnExecutionContext context)
       // results haven't been uploaded to the cache yet and deduplicate all of them against the
       // first one.
       LocalExecution previousExecution = null;
-      thisExecution = LocalExecution.createIfDeduplicatable(action);
-      if (shouldUploadLocalResults && thisExecution != null) {
-        previousExecution = inFlightExecutions.putIfAbsent(action.getActionKey(), thisExecution);
-      }
-      // Metadata will be available in context.current() until we detach.
-      // This is done via a thread-local variable.
       try {
-        RemoteActionResult result;
-        try (SilentCloseable c = prof.profile(ProfilerTask.REMOTE_CACHE_CHECK, "check cache hit")) {
-          result = remoteExecutionService.lookupCache(action);
-        }
-        // In case the remote cache returned a failed action (exit code != 0) we treat it as a
-        // cache miss
-        if (result != null && result.getExitCode() == 0) {
-          Stopwatch fetchTime = Stopwatch.createStarted();
-          InMemoryOutput inMemoryOutput;
-          try (SilentCloseable c = prof.profile(REMOTE_DOWNLOAD, "download outputs")) {
-            inMemoryOutput = remoteExecutionService.downloadOutputs(action, result);
-          }
-          fetchTime.stop();
-          totalTime.stop();
-          spawnMetrics
-              .setFetchTimeInMs((int) fetchTime.elapsed().toMillis())
-              .setTotalTimeInMs((int) totalTime.elapsed().toMillis())
-              .setNetworkTimeInMs((int) action.getNetworkTime().getDuration().toMillis());
-          SpawnResult spawnResult =
-              createSpawnResult(
-                  digestUtil,
+        thisExecution = LocalExecution.createIfDeduplicatable(action);
+        if (shouldUploadLocalResults && thisExecution != null) {
+          LocalExecution previousOrThisExecution =
+              inFlightExecutions.merge(
                   action.getActionKey(),
-                  result.getExitCode(),
-                  /* cacheHit= */ true,
-                  result.cacheName(),
-                  inMemoryOutput,
-                  result.getExecutionMetadata().getExecutionStartTimestamp(),
-                  result.getExecutionMetadata().getExecutionCompletedTimestamp(),
-                  spawnMetrics.build(),
-                  spawn.getMnemonic());
-          return SpawnCache.success(spawnResult);
+                  thisExecution,
+                  (existingExecution, thisExecutionArg) -> {
+                    if (existingExecution.registerForOutputReuse()) {
+                      return existingExecution;
+                    } else {
+                      // The existing execution has completed and its results may have already
+                      // been modified by its action, so we can't deduplicate against it. Instead,
+                      // start a new in-flight execution.
+                      return thisExecutionArg;
+                    }
+                  });
+          previousExecution =
+              previousOrThisExecution == thisExecution ? null : previousOrThisExecution;
         }
-      } catch (CacheNotFoundException e) {
-        // Intentionally left blank
-      } catch (IOException e) {
-        if (BulkTransferException.allCausedByCacheNotFoundException(e)) {
+        try {
+          RemoteActionResult result;
+          try (SilentCloseable c =
+              prof.profile(ProfilerTask.REMOTE_CACHE_CHECK, "check cache hit")) {
+            result = remoteExecutionService.lookupCache(action);
+          }
+          // In case the remote cache returned a failed action (exit code != 0) we treat it as a
+          // cache miss
+          if (result != null && result.getExitCode() == 0) {
+            Stopwatch fetchTime = Stopwatch.createStarted();
+            InMemoryOutput inMemoryOutput;
+            try (SilentCloseable c = prof.profile(REMOTE_DOWNLOAD, "download outputs")) {
+              inMemoryOutput = remoteExecutionService.downloadOutputs(action, result);
+            }
+            fetchTime.stop();
+            totalTime.stop();
+            spawnMetrics
+                .setFetchTimeInMs((int) fetchTime.elapsed().toMillis())
+                .setTotalTimeInMs((int) totalTime.elapsed().toMillis())
+                .setNetworkTimeInMs((int) action.getNetworkTime().getDuration().toMillis());
+            SpawnResult spawnResult =
+                createSpawnResult(
+                    digestUtil,
+                    action.getActionKey(),
+                    result.getExitCode(),
+                    /* cacheHit= */ true,
+                    result.cacheName(),
+                    inMemoryOutput,
+                    result.getExecutionMetadata().getExecutionStartTimestamp(),
+                    result.getExecutionMetadata().getExecutionCompletedTimestamp(),
+                    spawnMetrics.build(),
+                    spawn.getMnemonic());
+            return SpawnCache.success(spawnResult);
+          }
+        } catch (CacheNotFoundException e) {
           // Intentionally left blank
-        } else {
-          String errorMessage = Utils.grpcAwareErrorMessage(e, verboseFailures);
-          if (isNullOrEmpty(errorMessage)) {
-            errorMessage = e.getClass().getSimpleName();
+        } catch (IOException e) {
+          if (BulkTransferException.allCausedByCacheNotFoundException(e)) {
+            // Intentionally left blank
+          } else {
+            String errorMessage = Utils.grpcAwareErrorMessage(e, verboseFailures);
+            if (isNullOrEmpty(errorMessage)) {
+              errorMessage = e.getClass().getSimpleName();
+            }
+            errorMessage = "Remote Cache: " + errorMessage;
+            remoteExecutionService.report(Event.warn(errorMessage));
           }
-          errorMessage = "Remote Cache: " + errorMessage;
-          remoteExecutionService.report(Event.warn(errorMessage));
         }
-      }
-      if (previousExecution != null) {
-        Stopwatch fetchTime = Stopwatch.createStarted();
-        SpawnResult previousResult;
-        try (SilentCloseable c = prof.profile(REMOTE_DOWNLOAD, "reuse outputs")) {
-          previousResult = remoteExecutionService.waitForAndReuseOutputs(action, previousExecution);
-        }
-        if (previousResult != null) {
-          spawnMetrics
-              .setFetchTimeInMs((int) fetchTime.elapsed().toMillis())
-              .setTotalTimeInMs((int) totalTime.elapsed().toMillis())
-              .setNetworkTimeInMs((int) action.getNetworkTime().getDuration().toMillis());
-          SpawnMetrics buildMetrics = spawnMetrics.build();
-          return SpawnCache.success(
-              new SpawnResult.DelegateSpawnResult(previousResult) {
-                @Override
-                public String getRunnerName() {
-                  return "deduplicated";
-                }
+        if (previousExecution != null) {
+          Stopwatch fetchTime = Stopwatch.createStarted();
+          SpawnResult previousResult;
+          try (SilentCloseable c = prof.profile(REMOTE_DOWNLOAD, "reuse outputs")) {
+            previousResult =
+                remoteExecutionService.waitForAndReuseOutputs(action, previousExecution);
+          }
+          if (previousResult != null) {
+            spawnMetrics
+                .setFetchTimeInMs((int) fetchTime.elapsed().toMillis())
+                .setTotalTimeInMs((int) totalTime.elapsed().toMillis())
+                .setNetworkTimeInMs((int) action.getNetworkTime().getDuration().toMillis());
+            SpawnMetrics buildMetrics = spawnMetrics.build();
+            return SpawnCache.success(
+                new SpawnResult.DelegateSpawnResult(previousResult) {
+                  @Override
+                  public String getRunnerName() {
+                    return "deduplicated";
+                  }
 
-                @Override
-                public SpawnMetrics getMetrics() {
-                  return buildMetrics;
-                }
-              });
+                  @Override
+                  public SpawnMetrics getMetrics() {
+                    return buildMetrics;
+                  }
+                });
+          }
+          // If we reach here, the previous execution was not successful (it encountered an
+          // exception or the spawn had an exit code != 0). Since it isn't possible to accurately
+          // recreate the failure without rerunning the action, we fall back to running the action
+          // locally. This means that we have introduced an unnecessary wait, but that can only
+          // happen in the case of a failing build with --keep_going.
+        }
+      } finally {
+        if (previousExecution != null) {
+          previousExecution.unregister();
         }
-        // If we reach here, the previous execution was not successful (it encountered an exception
-        // or the spawn had an exit code != 0). Since it isn't possible to accurately recreate the
-        // failure without rerunning the action, we fall back to running the action locally. This
-        // means that we have introduced an unnecessary wait, but that can only happen in the case
-        // of a failing build with --keep_going.
       }
     }
 
@@ -239,6 +260,17 @@ public void store(SpawnResult result) throws ExecException, InterruptedException
           // large.
           remoteExecutionService.uploadOutputs(
               action, result, () -> inFlightExecutions.remove(action.getActionKey()));
+          if (thisExecutionFinal != null
+              && action.getSpawn().getResourceOwner().mayModifySpawnOutputsAfterExecution()) {
+            // In this case outputs have been uploaded synchronously and the callback above has run,
+            // so no new executions will be deduplicated against this one. We can safely await all
+            // existing executions finish the reuse.
+            // Note that while this call itself isn't interruptible, all operations it awaits are
+            // interruptible.
+            try (SilentCloseable c = prof.profile(REMOTE_DOWNLOAD, "await output reuse")) {
+              thisExecutionFinal.awaitAllOutputReuse();
+            }
+          }
         }
 
         private void checkForConcurrentModifications()

src/main/java/com/google/devtools/build/lib/rules/java/JavaCompileAction.java

@@ -709,6 +709,14 @@ public NestedSet<Artifact> getPossibleInputsForTesting() {
     return null;
   }
 
+  @Override
+  public boolean mayModifySpawnOutputsAfterExecution() {
+    // Causes of spawn output modification after execution:
+    // - Fallback to the full classpath with --experimental_java_classpath=bazel.
+    // - In-place rewriting of .jdeps files with --experimental_output_paths=strip.
+    return true;
+  }
+
   /**
    * Locally rewrites a .jdeps file to replace missing config prefixes.
    *

src/main/java/com/google/devtools/build/lib/rules/java/JavaHeaderCompileAction.java

@@ -157,6 +157,14 @@ protected void afterExecute(
     }
   }
 
+  @Override
+  public boolean mayModifySpawnOutputsAfterExecution() {
+    // Causes of spawn output modification after execution:
+    // - In-place rewriting of .jdeps files with --experimental_output_paths=strip.
+    // TODO: Use separate files as action and spawn output to avoid in-place modification.
+    return true;
+  }
+
   public static Builder newBuilder(RuleContext ruleContext) {
     return new Builder(ruleContext);
   }