Move action post-processing to a continuation

Action post-processing consists of updating action inputs after
execution as well as writing an entry to the action cache. This
implicitly ensures that we do not even attempt to write shared
(duplicate) actions to the action cache.

Progress on #6394.

PiperOrigin-RevId: 234982739
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 4ba628b..14530b4 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
@@ -14,6 +14,7 @@
 package com.google.devtools.build.lib.skyframe;
 
 import com.google.common.base.Function;
+import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -57,6 +58,7 @@
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.rules.cpp.IncludeScannable;
 import com.google.devtools.build.lib.skyframe.ActionRewindStrategy.RewindPlan;
+import com.google.devtools.build.lib.skyframe.SkyframeActionExecutor.ActionPostprocessing;
 import com.google.devtools.build.lib.util.io.FileOutErr;
 import com.google.devtools.build.lib.util.io.TimestampGranularityMonitor;
 import com.google.devtools.build.lib.vfs.FileSystem;
@@ -684,6 +686,7 @@
       throw new ActionExecutionException(
           "Failed to update filesystem context: ", e, action, /*catastrophe=*/ false);
     }
+
     try (ActionExecutionContext actionExecutionContext =
         skyframeActionExecutor.getContext(
             metadataHandler,
@@ -696,60 +699,68 @@
             topLevelFilesets,
             state.actionFileSystem,
             skyframeDepsResult)) {
-      if (!state.hasExecutedAction()) {
-        state.value =
-            skyframeActionExecutor.executeAction(
-                env,
-                action,
-                metadataHandler,
-                actionStartTime,
-                actionExecutionContext,
-                actionLookupData);
-        // If an action is executed asynchronously, we may not have a result. In that case we return
-        // null here and expect the function to be restarted.
-        if (env.valuesMissing()) {
-          return null;
-        }
-      }
+      return skyframeActionExecutor.executeAction(
+          env,
+          action,
+          metadataHandler,
+          actionStartTime,
+          actionExecutionContext,
+          actionLookupData,
+          new ActionPostprocessingImpl(state));
     } catch (IOException e) {
       throw new ActionExecutionException(
           "Failed to close action output", e, action, /*catastrophe=*/ false);
     }
+  }
 
-    if (action.discoversInputs()) {
-      Iterable<Artifact> newInputs = filterKnownInputs(action.getInputs(), state.inputArtifactData);
-      Map<SkyKey, SkyValue> metadataFoundDuringActionExecution =
-          env.getValues(toKeys(newInputs, action.getMandatoryInputs()));
-      state.discoveredInputs = newInputs;
-      if (env.valuesMissing()) {
-        return null;
-      }
-      if (!metadataFoundDuringActionExecution.isEmpty()) {
-        // We are in the interesting case of an action that discovered its inputs during
-        // execution, and found some new ones, but the new ones were already present in the graph.
-        // We must therefore cache the metadata for those new ones.
-        for (Map.Entry<SkyKey, SkyValue> entry : metadataFoundDuringActionExecution.entrySet()) {
-          state.inputArtifactData.putWithNoDepOwner(
-              ArtifactSkyKey.artifact(entry.getKey()), (FileArtifactValue) entry.getValue());
-        }
-        // TODO(ulfjack): This causes information loss about omitted and injected outputs. Also see
-        // the documentation on MetadataHandler.artifactOmitted. This works by accident because
-        // markOmitted is only called for remote execution, and this code only gets executed for
-        // local execution.
-        metadataHandler =
-            new ActionMetadataHandler(
-                state.inputArtifactData,
-                /*missingArtifactsAllowed=*/ false,
-                action.getOutputs(),
-                tsgm.get(),
-                pathResolver,
-                state.actionFileSystem == null ? new OutputStore() : new MinimalOutputStore());
-      }
+  /** Implementation of {@link ActionPostprocessing}. */
+  private final class ActionPostprocessingImpl implements ActionPostprocessing {
+    private final ContinuationState state;
+
+    ActionPostprocessingImpl(ContinuationState state) {
+      this.state = state;
     }
-    Preconditions.checkState(!env.valuesMissing(), action);
-    skyframeActionExecutor.updateActionCacheIfReallyExecuted(
-        action, metadataHandler, state.token, clientEnv, actionLookupData);
-    return state.value;
+
+    public void run(
+        Environment env,
+        Action action,
+        ActionMetadataHandler metadataHandler,
+        Map<String, String> clientEnv)
+        throws InterruptedException {
+      if (action.discoversInputs()) {
+        Iterable<Artifact> newInputs =
+            filterKnownInputs(action.getInputs(), state.inputArtifactData);
+        Map<SkyKey, SkyValue> metadataFoundDuringActionExecution =
+            env.getValues(toKeys(newInputs, action.getMandatoryInputs()));
+        state.discoveredInputs = newInputs;
+        if (env.valuesMissing()) {
+          return;
+        }
+        if (!metadataFoundDuringActionExecution.isEmpty()) {
+          // We are in the interesting case of an action that discovered its inputs during
+          // execution, and found some new ones, but the new ones were already present in the graph.
+          // We must therefore cache the metadata for those new ones.
+          for (Map.Entry<SkyKey, SkyValue> entry : metadataFoundDuringActionExecution.entrySet()) {
+            state.inputArtifactData.putWithNoDepOwner(
+                ArtifactSkyKey.artifact(entry.getKey()), (FileArtifactValue) entry.getValue());
+          }
+          // TODO(ulfjack): This causes information loss about omitted and injected outputs. Also
+          // see the documentation on MetadataHandler.artifactOmitted. This works by accident
+          // because markOmitted is only called for remote execution, and this code only gets
+          // executed for local execution.
+          metadataHandler =
+              new ActionMetadataHandler(
+                  state.inputArtifactData,
+                  /*missingArtifactsAllowed=*/ false,
+                  action.getOutputs(),
+                  tsgm.get(),
+                  metadataHandler.getArtifactPathResolver(),
+                  state.actionFileSystem == null ? new OutputStore() : new MinimalOutputStore());
+        }
+      }
+      Preconditions.checkState(!env.valuesMissing(), action);
+      skyframeActionExecutor.updateActionCache(action, metadataHandler, state.token, clientEnv);
+    }
   }
 
   private static final Function<Artifact, SkyKey> TO_NONMANDATORY_SKYKEY =
@@ -876,7 +887,7 @@
    * Declare dependency on all known inputs of action. Throws exception if any are known to be
    * missing. Some inputs may not yet be in the graph, in which case the builder should abort.
    */
-  private CheckInputResults checkInputs(
+  private static CheckInputResults checkInputs(
       Environment env,
       Action action,
       Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps)
@@ -887,7 +898,7 @@
   /**
    * Reconstructs the relationships between lost inputs and the direct deps responsible for them.
    */
-  private ActionInputDepOwners getInputDepOwners(
+  private static ActionInputDepOwners getInputDepOwners(
       Environment env,
       Action action,
       Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps,
@@ -901,7 +912,7 @@
         (actionInputMapSink, expandedArtifacts, expandedFilesets) -> actionInputMapSink);
   }
 
-  private <S extends ActionInputMapSink, R> R accumulateInputs(
+  private static <S extends ActionInputMapSink, R> R accumulateInputs(
       Environment env,
       Action action,
       Map<SkyKey, ValueOrException2<MissingInputFileException, ActionExecutionException>> inputDeps,
@@ -1054,7 +1065,6 @@
     Map<Artifact, ImmutableList<FilesetOutputSymlink>> expandedFilesets = null;
     Token token = null;
     Iterable<Artifact> discoveredInputs = null;
-    ActionExecutionValue value = null;
     FileSystem actionFileSystem = null;
 
     boolean hasCollectedInputs() {
@@ -1073,10 +1083,6 @@
       return token != null;
     }
 
-    boolean hasExecutedAction() {
-      return value != null;
-    }
-
     /** Must be called to assign values to the given variables as they change. */
     void updateFileSystemContext(
         SkyframeActionExecutor executor,
@@ -1092,15 +1098,12 @@
 
     @Override
     public String toString() {
-      return token
-          + ", "
-          + value
-          + ", "
-          + allInputs
-          + ", "
-          + inputArtifactData
-          + ", "
-          + discoveredInputs;
+      return MoreObjects.toStringHelper(this)
+          .add("token", token)
+          .add("allInputs", allInputs)
+          .add("inputArtifactData", inputArtifactData)
+          .add("discoveredInputs", discoveredInputs)
+          .toString();
     }
   }
 
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionState.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionState.java
index 7175049..c02d2c9 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionState.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionExecutionState.java
@@ -66,17 +66,13 @@
     this.state = Preconditions.checkNotNull(state);
   }
 
-  public boolean isOwner(ActionLookupData actionLookupData) {
-    return this.actionLookupData.equals(actionLookupData);
-  }
-
   public ActionExecutionValue getResultOrDependOnFuture(
       SkyFunction.Environment env,
       ActionLookupData actionLookupData,
       Action action,
       ActionCompletedReceiver actionCompletedReceiver)
       throws ActionExecutionException, InterruptedException {
-    if (isOwner(actionLookupData)) {
+    if (this.actionLookupData.equals(actionLookupData)) {
       // This continuation is owned by the Skyframe node executed by the current thread, so we use
       // it to run the state machine.
       return runStateMachine(env);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandler.java b/src/main/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandler.java
index 6b65646..a2c2c06 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandler.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/ActionMetadataHandler.java
@@ -144,6 +144,10 @@
     return value;
   }
 
+  public ArtifactPathResolver getArtifactPathResolver() {
+    return artifactPathResolver;
+  }
+
   @Nullable
   private FileArtifactValue getInputFileArtifactValue(Artifact input) {
     if (isKnownOutput(input)) {
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 edd1427..08a1cc8 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
@@ -501,19 +501,6 @@
     }
   }
 
-  private boolean actionReallyExecuted(Action action, ActionLookupData actionLookupData) {
-    // TODO(b/19539699): This method is used only when the action cache is enabled. It is
-    // incompatible with action rewinding, which removes entries from buildActionMap. Action
-    // rewinding is used only with a disabled action cache.
-    ActionExecutionState cachedRun =
-        Preconditions.checkNotNull(
-            buildActionMap.get(new OwnerlessArtifactWrapper(action.getPrimaryOutput())),
-            "%s %s",
-            action,
-            actionLookupData);
-    return cachedRun.isOwner(actionLookupData);
-  }
-
   void noteActionEvaluationStarted(ActionLookupData actionLookupData, Action action) {
     this.completionReceiver.noteActionEvaluationStarted(actionLookupData, action);
   }
@@ -530,7 +517,8 @@
       ActionMetadataHandler metadataHandler,
       long actionStartTime,
       ActionExecutionContext actionExecutionContext,
-      ActionLookupData actionLookupData)
+      ActionLookupData actionLookupData,
+      ActionPostprocessing postprocessing)
       throws ActionExecutionException, InterruptedException {
     // ActionExecutionFunction may directly call into ActionExecutionState.getResultOrDependOnFuture
     // if a shared action already passed these checks.
@@ -563,7 +551,8 @@
                         metadataHandler,
                         actionStartTime,
                         actionExecutionContext,
-                        actionLookupData)));
+                        actionLookupData,
+                        postprocessing)));
     return activeAction.getResultOrDependOnFuture(
         env, actionLookupData, action, completionReceiver);
   }
@@ -669,21 +658,11 @@
     return token;
   }
 
-  void updateActionCacheIfReallyExecuted(
-      Action action,
-      MetadataHandler metadataHandler,
-      Token token,
-      Map<String, String> clientEnv,
-      ActionLookupData actionLookupData) {
+  void updateActionCache(
+      Action action, MetadataHandler metadataHandler, Token token, Map<String, String> clientEnv) {
     if (!actionCacheChecker.enabled()) {
       return;
     }
-    if (!actionReallyExecuted(action, actionLookupData)) {
-      // If an action shared with this one executed, then we need not update the action cache, since
-      // the other action will do it. Moreover, this action is not aware of metadata acquired
-      // during execution, so its metadata handler is likely unusable anyway.
-      return;
-    }
     try {
       actionCacheChecker.updateActionCache(action, token, metadataHandler, clientEnv);
     } catch (IOException e) {
@@ -796,6 +775,19 @@
     ActionResult execute() throws ActionExecutionException, InterruptedException;
   }
 
+  /**
+   * Temporary interface to allow delegation of action postprocessing to ActionExecutionFunction.
+   * The current implementation requires access to local fields in ActionExecutionFunction.
+   */
+  interface ActionPostprocessing {
+    void run(
+        Environment env,
+        Action action,
+        ActionMetadataHandler metadataHandler,
+        Map<String, String> clientEnv)
+        throws InterruptedException;
+  }
+
   /** Represents an action that needs to be run. */
   private final class ActionRunner extends ActionStep {
     private final Action action;
@@ -804,19 +796,22 @@
     private final ActionExecutionContext actionExecutionContext;
     private final ActionLookupData actionLookupData;
     private final ActionExecutionStatusReporter statusReporter;
+    private final ActionPostprocessing postprocessing;
 
     ActionRunner(
         Action action,
         ActionMetadataHandler metadataHandler,
         long actionStartTime,
         ActionExecutionContext actionExecutionContext,
-        ActionLookupData actionLookupData) {
+        ActionLookupData actionLookupData,
+        ActionPostprocessing postprocessing) {
       this.action = action;
       this.metadataHandler = metadataHandler;
       this.actionStartTime = actionStartTime;
       this.actionExecutionContext = actionExecutionContext;
       this.actionLookupData = actionLookupData;
       this.statusReporter = statusReporterRef.get();
+      this.postprocessing = postprocessing;
     }
 
     @Override
@@ -971,7 +966,7 @@
         } catch (InterruptedException e) {
           return ActionStepOrResult.of(e);
         }
-        return ActionStepOrResult.of(actuallyCompleteAction(eventHandler, actionResult));
+        return new ActionCacheWriteStep(actuallyCompleteAction(eventHandler, actionResult));
       } catch (ActionExecutionException e) {
         return ActionStepOrResult.of(e);
       } finally {
@@ -1088,6 +1083,27 @@
             () -> spawnAction.finishSync(futureSpawn, actionExecutionContext.getVerboseFailures()));
       }
     }
+
+    private class ActionCacheWriteStep extends ActionStep {
+      private final ActionExecutionValue value;
+
+      public ActionCacheWriteStep(ActionExecutionValue value) {
+        this.value = value;
+      }
+
+      @Override
+      public ActionStepOrResult run(Environment env) {
+        try {
+          postprocessing.run(env, action, metadataHandler, actionExecutionContext.getClientEnv());
+          if (env.valuesMissing()) {
+            return this;
+          }
+        } catch (InterruptedException e) {
+          return ActionStepOrResult.of(e);
+        }
+        return ActionStepOrResult.of(value);
+      }
+    }
   }
 
   private void createOutputDirectories(Action action, ActionExecutionContext context)