Move action-cache updating to inside ActionExecutionFunction, in preparation for allowing it to be restarted in case of missing deps.

Note that this means that action-cache writing is no longer part of the ACTION_COMPLETE profiling unit.

--
MOS_MIGRATED_REVID=89702039

3 files changed, 40 insertions, 25 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/actions/ActionCacheChecker.java b/src/main/java/com/google/devtools/build/lib/actions/ActionCacheChecker.java
index cda9cb99ad..e3a8ff8e81 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/ActionCacheChecker.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/ActionCacheChecker.java
@@ -81,6 +81,12 @@ public class ActionCacheChecker {
     return null;
   }
 
+  private void removeCacheEntry(Action action) {
+    for (Artifact output : action.getOutputs()) {
+      actionCache.remove(output.getExecPathString());
+    }
+  }
+
   /**
    * Validate metadata state for action input or output artifacts.
    *
@@ -166,6 +172,9 @@ public class ActionCacheChecker {
       }
     }
     if (mustExecute(action, entry, handler, metadataHandler, actionInputs)) {
+      if (entry != null) {
+        removeCacheEntry(action);
+      }
       return new Token(getKeyString(action));
     }
 
@@ -207,6 +216,10 @@ public class ActionCacheChecker {
       throws IOException {
     Preconditions.checkArgument(token != null);
     String key = token.cacheKey;
+    if (actionCache.get(key) != null) {
+      // This cache entry has already been updated by a shared action. We don't need to do it again.
+      return;
+    }
     ActionCache.Entry entry = actionCache.createEntry(action.getKey());
     for (Artifact output : action.getOutputs()) {
       // Remove old records from the cache if they used different key.
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
index f2a81fc6db..02ef3b334e 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionFunction.java
@@ -180,7 +180,7 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
     // must use that other action's value, provided here, since it is populated with metadata
     // for the outputs.
     if (inputArtifactData == null) {
-      return skyframeActionExecutor.executeAction(action, null, null, -1, null);
+      return skyframeActionExecutor.executeAction(action, null, -1, null);
     }
     ContinuationState state;
     if (action.discoversInputs()) {
@@ -229,6 +229,8 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
         skyframeActionExecutor.constructActionExecutionContext(fileAndMetadataCache,
             metadataHandler);
     boolean inputsDiscoveredDuringActionExecution = false;
+    ActionExecutionValue result;
+    Token token;
     try {
       if (action.discoversInputs()) {
         if (!state.hasDiscoveredInputs()) {
@@ -262,13 +264,13 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
       }
       // Clear state before actual execution of action. It will never be needed again because
       // skyframeActionExecutor is guaranteed to have a result after this.
-      Token token = state.token;
+      token = state.token;
       if (action.discoversInputs()) {
         removeState(action);
       }
       state = null;
-      return skyframeActionExecutor.executeAction(action, fileAndMetadataCache, token,
-          actionStartTime, actionExecutionContext);
+      result = skyframeActionExecutor.executeAction(action,
+          fileAndMetadataCache, actionStartTime, actionExecutionContext);
     } finally {
       try {
         actionExecutionContext.getFileOutErr().close();
@@ -279,6 +281,8 @@ public class ActionExecutionFunction implements SkyFunction, CompletionReceiver
         declareAdditionalDependencies(env, action);
       }
     }
+    skyframeActionExecutor.afterExecution(action, fileAndMetadataCache, token);
+    return result;
   }
 
   private static Map<Artifact, FileArtifactValue> addDiscoveredInputs(
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
index 4a95f679e8..b0a2cfb7b5 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeActionExecutor.java
@@ -470,7 +470,7 @@ public final class SkyframeActionExecutor {
    * <p>For use from {@link ArtifactFunction} only.
    */
   ActionExecutionValue executeAction(Action action, FileAndMetadataCache graphFileCache,
-      Token token, long actionStartTime,
+      long actionStartTime,
       ActionExecutionContext actionExecutionContext)
       throws ActionExecutionException, InterruptedException {
     Exception exception = badActionMap.get(action);
@@ -480,7 +480,7 @@ public final class SkyframeActionExecutor {
     }
     Artifact primaryOutput = action.getPrimaryOutput();
     FutureTask<ActionExecutionValue> actionTask =
-        new FutureTask<>(new ActionRunner(action, graphFileCache, token,
+        new FutureTask<>(new ActionRunner(action, graphFileCache,
             actionStartTime, actionExecutionContext));
     // Check to see if another action is already executing/has executed this value.
     Pair<Action, FutureTask<ActionExecutionValue>> oldAction =
@@ -591,6 +591,18 @@ public final class SkyframeActionExecutor {
     return token;
   }
 
+  void afterExecution(Action action, MetadataHandler metadataHandler, Token token) {
+    try {
+      actionCacheChecker.afterExecution(action, token, metadataHandler);
+    } catch (IOException e) {
+      // Skyframe has already done all the filesystem access needed for outputs, and swallows
+      // IOExceptions for inputs. So an IOException is impossible here.
+      throw new IllegalStateException(
+          "failed to update action cache for " + action.prettyPrint()
+              + ", but all outputs should already have been checked", e);
+    }
+  }
+
   /**
    * Perform dependency discovery for action, which must discover its inputs.
    *
@@ -636,16 +648,14 @@ public final class SkyframeActionExecutor {
   private class ActionRunner implements Callable<ActionExecutionValue> {
     private final Action action;
     private final FileAndMetadataCache graphFileCache;
-    private Token token;
     private long actionStartTime;
     private ActionExecutionContext actionExecutionContext;
 
-    ActionRunner(Action action, FileAndMetadataCache graphFileCache, Token token,
+    ActionRunner(Action action, FileAndMetadataCache graphFileCache,
         long actionStartTime,
         ActionExecutionContext actionExecutionContext) {
       this.action = action;
       this.graphFileCache = graphFileCache;
-      this.token = token;
       this.actionStartTime = actionStartTime;
       this.actionExecutionContext = actionExecutionContext;
     }
@@ -673,8 +683,7 @@ public final class SkyframeActionExecutor {
 
         createOutputDirectories(action);
 
-        prepareScheduleExecuteAndCompleteAction(action, token,
-            actionExecutionContext, actionStartTime);
+        prepareScheduleExecuteAndCompleteAction(action, actionExecutionContext, actionStartTime);
         return new ActionExecutionValue(
             graphFileCache.getOutputData(), graphFileCache.getAdditionalOutputData());
       } finally {
@@ -743,17 +752,15 @@ public final class SkyframeActionExecutor {
    * action cache.
    *
    * @param action  The action to execute
-   * @param token  The non-null token returned by dependencyChecker.getTokenIfNeedToExecute()
    * @param context services in the scope of the action
    * @param actionStartTime time when we started the first phase of the action execution.
    * @throws ActionExecutionException if the execution of the specified action
    *   failed for any reason.
    * @throws InterruptedException if the thread was interrupted.
    */
-  private void prepareScheduleExecuteAndCompleteAction(Action action, Token token,
+  private void prepareScheduleExecuteAndCompleteAction(Action action,
       ActionExecutionContext context, long actionStartTime)
       throws ActionExecutionException, InterruptedException {
-    Preconditions.checkNotNull(token, action);
     // Delete the metadataHandler's cache of the action's outputs, since they are being deleted.
     context.getMetadataHandler().discardMetadata(action.getOutputs());
     // Delete the outputs before executing the action, just to ensure that
@@ -776,7 +783,7 @@ public final class SkyframeActionExecutor {
         resourceManager.acquireResources(action, estimate);
       }
       boolean outputDumped = executeActionTask(action, context);
-      completeAction(action, token, context.getMetadataHandler(),
+      completeAction(action, context.getMetadataHandler(),
           context.getFileOutErr(), outputDumped);
     } finally {
       if (estimate != null) {
@@ -856,7 +863,7 @@ public final class SkyframeActionExecutor {
     return false;
   }
 
-  private void completeAction(Action action, Token token, MetadataHandler metadataHandler,
+  private void completeAction(Action action, MetadataHandler metadataHandler,
       FileOutErr fileOutErr, boolean outputAlreadyDumped) throws ActionExecutionException {
     try {
       Preconditions.checkState(action.inputsKnown(),
@@ -876,15 +883,6 @@ public final class SkyframeActionExecutor {
         } catch (IOException e) {
           reportError("failed to set outputs read-only", e, action, null);
         }
-        try {
-          actionCacheChecker.afterExecution(action, token, metadataHandler);
-        } catch (IOException e) {
-          // Skyframe does all the filesystem access needed during the previous calls, and if those
-          // calls failed, we should already have thrown. So an IOException is impossible here.
-          throw new IllegalStateException(
-              "failed to update action cache for " + action.prettyPrint()
-                  + ", but all outputs should already have been checked", e);
-        }
       } finally {
         profiler.completeTask(ProfilerTask.ACTION_COMPLETE);
       }