Don't mutate an action's input map after it has been executed

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

@@ -822,7 +822,7 @@ private ActionExecutionValue checkCacheAndExecuteIfNeeded(
             action);
       }
 
-      addDiscoveredInputs(state, env, action);
+      addDiscoveredInputs(state, env, action, /* afterExecution= */ false);
       if (env.valuesMissing()) {
         return null;
       }
@@ -869,30 +869,43 @@ public void run(
         throws InterruptedException, ActionExecutionException {
       if (action.discoversInputs()) {
         state.discoveredInputs = action.getInputs();
-        addDiscoveredInputs(state, env, action);
+        addDiscoveredInputs(state, env, action, /* afterExecution= */ true);
         if (env.valuesMissing()) {
           return;
         }
       }
       checkState(!env.valuesMissing(), action);
       skyframeActionExecutor.updateActionCache(
-          action, inputMetadataProvider, outputMetadataStore, state.token, clientEnv);
+          action,
+          state.inputMetadataProviderIncludingLateDiscoveredInputs(inputMetadataProvider),
+          outputMetadataStore,
+          state.token,
+          clientEnv);
     }
   }
 
+  /**
+   * Adds the metadata of {@link InputDiscoveryState#discoveredInputs} that isn't known yet to the
+   * state.
+   *
+   * <p>Before the action is executed, the metadata is added to {@link
+   * InputDiscoveryState#inputArtifactData} so that the action can access it. Afterwards, it is
+   * added to {@link InputDiscoveryState#lateDiscoveredInputArtifactData} instead: the action file
+   * system reads {@code inputArtifactData} and outlives the action itself (build event artifact
+   * uploads digest the action's outputs through it asynchronously), but {@link ActionInputMap} is
+   * only thread-compatible, so mutating it at that point would race with those readers.
+   */
   private void addDiscoveredInputs(
-      InputDiscoveryState state, Environment env, Action actionForError)
+      InputDiscoveryState state, Environment env, Action actionForError, boolean afterExecution)
       throws InterruptedException, ActionExecutionException {
     // TODO(janakr): This code's assumptions are wrong in the face of Starlark actions with unused
     //  inputs, since ActionExecutionExceptions can come through here and should be aggregated. Fix.
 
-    ActionInputMap inputData = state.inputArtifactData;
-
     // Filter down to unknown discovered inputs eagerly instead of using a lazy Iterables#filter to
     // reduce iteration cost.
     List<Artifact> unknownDiscoveredInputs = new ArrayList<>();
     for (Artifact input : state.discoveredInputs.toList()) {
-      if (inputData.getInputMetadata(input) == null) {
+      if (state.getKnownDiscoveredInputMetadata(input) == null) {
         unknownDiscoveredInputs.add(input);
       }
     }
@@ -901,6 +914,11 @@ private void addDiscoveredInputs(
       return;
     }
 
+    ActionInputMap inputData =
+        afterExecution
+            ? state.getOrCreateLateDiscoveredInputArtifactData(unknownDiscoveredInputs.size())
+            : state.inputArtifactData;
+
     SkyframeLookupResult nonMandatoryDiscovered =
         env.getValuesAndExceptions(Artifact.keys(unknownDiscoveredInputs));
     for (Artifact input : unknownDiscoveredInputs) {
@@ -1273,6 +1291,18 @@ static class InputDiscoveryState implements SerializableSkyKeyComputeState {
      */
     DelegatingPairInputMetadataProvider compositeInputMetadataProvider = null;
 
+    /**
+     * Metadata for inputs that were only discovered after the action was executed.
+     *
+     * <p>Deliberately not part of {@link #inputArtifactData}: that map is read by the action file
+     * system, which outlives the action itself because build events referencing it are uploaded
+     * asynchronously. {@link ActionInputMap} is only thread-compatible, so mutating it while those
+     * uploads are in flight corrupts it for the reader.
+     */
+    @Nullable private ActionInputMap lateDiscoveredInputArtifactData = null;
+
+    @Nullable private InputMetadataProvider lateDiscoveredInputMetadataProvider = null;
+
     Token token = null;
     NestedSet<Artifact> discoveredInputs = null;
     FileSystem actionFileSystem = null;
@@ -1292,6 +1322,40 @@ boolean hasArtifactData() {
       return inputArtifactData != null;
     }
 
+    /**
+     * Returns the metadata already known for a discovered input, or null if it hasn't been looked
+     * up yet.
+     */
+    @Nullable
+    FileArtifactValue getKnownDiscoveredInputMetadata(Artifact input) {
+      FileArtifactValue metadata = inputArtifactData.getInputMetadata(input);
+      if (metadata != null || lateDiscoveredInputArtifactData == null) {
+        return metadata;
+      }
+      return lateDiscoveredInputArtifactData.getInputMetadata(input);
+    }
+
+    ActionInputMap getOrCreateLateDiscoveredInputArtifactData(int sizeHint) {
+      if (lateDiscoveredInputArtifactData == null) {
+        lateDiscoveredInputArtifactData = new ActionInputMap(sizeHint);
+        lateDiscoveredInputMetadataProvider =
+            new ActionInputMetadataProvider(lateDiscoveredInputArtifactData);
+      }
+      return lateDiscoveredInputArtifactData;
+    }
+
+    /**
+     * Returns a provider that also covers inputs discovered after the action was executed, which
+     * are kept out of {@link #inputArtifactData}.
+     */
+    InputMetadataProvider inputMetadataProviderIncludingLateDiscoveredInputs(
+        InputMetadataProvider inputMetadataProvider) {
+      return lateDiscoveredInputMetadataProvider == null
+          ? inputMetadataProvider
+          : new DelegatingPairInputMetadataProvider(
+              lateDiscoveredInputMetadataProvider, inputMetadataProvider);
+    }
+
     boolean hasCheckedActionCache() {
       // If token is null because there was an action cache hit, this method is never called again
       // because we return immediately.

src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTest.java

@@ -18,9 +18,13 @@
 import static org.junit.Assert.assertThrows;
 
 import com.google.common.collect.ImmutableSet;
+import com.google.devtools.build.lib.actions.ActionExecutionContext;
+import com.google.devtools.build.lib.actions.ActionExecutionException;
+import com.google.devtools.build.lib.actions.ActionResult;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ArtifactRoot;
 import com.google.devtools.build.lib.actions.BuildFailedException;
+import com.google.devtools.build.lib.actions.InputMetadataProvider;
 import com.google.devtools.build.lib.actions.util.TestAction;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
@@ -30,6 +34,7 @@
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.PathFragment;
 import com.google.devtools.build.lib.vfs.Root;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -107,6 +112,88 @@ public void testCachingBuilderCachesUntilReset() throws Exception {
     assertThat(button.pressed).isTrue(); // rebuilt
   }
 
+  /**
+   * The action file system reads the {@link
+   * com.google.devtools.build.lib.actions.ActionInputMap} that backs the {@link
+   * InputMetadataProvider} handed to the action, and it outlives the action: build events
+   * referencing it are uploaded asynchronously. Since the map is only thread-compatible, inputs
+   * discovered after execution must not be added to it. See
+   * https://github.com/bazelbuild/bazel/issues/30683.
+   */
+  @Test
+  public void inputsDiscoveredAfterExecution_notAddedToInputMetadataProviderOfAction()
+      throws Exception {
+    Artifact hello = createSourceArtifact("hello");
+    hello.getPath().getParentDirectory().createDirectoryAndParents();
+    FileSystemUtils.writeContentAsLatin1(hello.getPath(), "content1");
+    Artifact late = createSourceArtifact("late");
+    FileSystemUtils.writeContentAsLatin1(late.getPath(), "late1");
+    Artifact goodbye = createDerivedArtifact("goodbye");
+
+    Button button = new Button();
+    AtomicReference<InputMetadataProvider> inputMetadataProvider = new AtomicReference<>();
+    registerAction(
+        new LateInputDiscoveringAction(
+            button, asNestedSet(hello), ImmutableSet.of(goodbye), late, inputMetadataProvider));
+
+    button.pressed = false;
+    buildArtifacts(cachingBuilder(), goodbye);
+    assertThat(button.pressed).isTrue(); // built
+
+    assertThat(inputMetadataProvider.get().getInput(late.getExecPath())).isNull();
+
+    // The late-discovered input is still tracked by the action cache.
+    button.pressed = false;
+    buildArtifacts(cachingBuilder(), goodbye);
+    assertThat(button.pressed).isFalse(); // not rebuilt
+
+    FileSystemUtils.writeContentAsLatin1(late.getPath(), "late2");
+
+    button.pressed = false;
+    buildArtifacts(cachingBuilder(), goodbye);
+    assertThat(button.pressed).isTrue(); // rebuilt
+  }
+
+  /** A {@link TestAction} that only learns about {@code late} once it has been executed. */
+  private static final class LateInputDiscoveringAction extends TestAction {
+    private final Artifact late;
+    private final AtomicReference<InputMetadataProvider> inputMetadataProvider;
+
+    LateInputDiscoveringAction(
+        Runnable effect,
+        NestedSet<Artifact> inputs,
+        ImmutableSet<Artifact> outputs,
+        Artifact late,
+        AtomicReference<InputMetadataProvider> inputMetadataProvider) {
+      super(effect, inputs, outputs);
+      this.late = late;
+      this.inputMetadataProvider = inputMetadataProvider;
+    }
+
+    @Override
+    public boolean discoversInputs() {
+      return true;
+    }
+
+    @Override
+    public NestedSet<Artifact> discoverInputs(ActionExecutionContext actionExecutionContext) {
+      return NestedSetBuilder.emptySet(Order.STABLE_ORDER);
+    }
+
+    @Override
+    public ActionResult execute(ActionExecutionContext actionExecutionContext)
+        throws ActionExecutionException, InterruptedException {
+      inputMetadataProvider.set(actionExecutionContext.getInputMetadataProvider());
+      ActionResult result = super.execute(actionExecutionContext);
+      updateInputs(
+          NestedSetBuilder.<Artifact>stableOrder()
+              .addTransitive(getMandatoryInputs())
+              .add(late)
+              .build());
+      return result;
+    }
+  }
+
   @Test
   public void testUnneededInputs() throws Exception {
     Artifact hello = createSourceArtifact("hello");