Treat timeouts expanding NestedSet as lost inputs.

NestedSetExpander is now permitted to throw TimeoutException, which it may wish to do if the NestedSet is from remote storage and is not available within a certain threshold. In this case, CppCompileAction will treat this dependency's primary output as a lost input, along with any other dependencies with still-pending NestedSets. Rewinding will allow us to rebuild the nodes with the missing NestedSets.

To facilitate setting a custom NestedSetExpander in tests, it is now set in SkyframeActionExecutor#configure instead of passed into the constructor.

PiperOrigin-RevId: 301382884
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetExpander.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetExpander.java
index fe7473f..2d1b00f 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetExpander.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetExpander.java
@@ -14,17 +14,18 @@
 package com.google.devtools.build.lib.collect.nestedset;
 
 import com.google.common.collect.ImmutableList;
+import java.util.concurrent.TimeoutException;
 
 /**
  * Helper class to expand {@link NestedSet} instances.
  *
- * <p>Implementations besides {@link NO_CALLBACKS} may wish to implement callbacks or timeouts for
+ * <p>Implementations besides {@link #DEFAULT} may wish to implement callbacks or timeouts for
  * dealing with expansions of sets from storage.
  */
 public interface NestedSetExpander {
 
   <T> ImmutableList<? extends T> toListInterruptibly(NestedSet<? extends T> nestedSet)
-      throws InterruptedException;
+      throws InterruptedException, TimeoutException;
 
   /** Simply delegates to {@link NestedSet#toListInterruptibly} without doing anything special. */
   NestedSetExpander DEFAULT = NestedSet::toListInterruptibly;
diff --git a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
index 44964d8..91d4fda 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/cpp/CppCompileAction.java
@@ -30,6 +30,7 @@
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionExecutionException;
 import com.google.devtools.build.lib.actions.ActionInput;
+import com.google.devtools.build.lib.actions.ActionInputDepOwnerMap;
 import com.google.devtools.build.lib.actions.ActionKeyContext;
 import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.ActionResult;
@@ -43,6 +44,7 @@
 import com.google.devtools.build.lib.actions.EnvironmentalExecException;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.actions.ExecutionRequirements;
+import com.google.devtools.build.lib.actions.LostInputsActionExecutionException;
 import com.google.devtools.build.lib.actions.ParameterFile.ParameterFileType;
 import com.google.devtools.build.lib.actions.ResourceSet;
 import com.google.devtools.build.lib.actions.SimpleSpawn;
@@ -98,6 +100,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
 import javax.annotation.Nullable;
 
 /** Action that represents some kind of C++ compilation step. */
@@ -584,7 +587,10 @@
     // used modules. Combining the NestedSets of transitive deps of the top-level modules also
     // gives us an effective way to compute and store discoveredModules.
     Set<Artifact> topLevel = new LinkedHashSet<>(usedModules);
-    for (NestedSet<? extends Artifact> transitive : transitivelyUsedModules.values()) {
+
+    Iterator<Map.Entry<Artifact, NestedSet<? extends Artifact>>> iterator =
+        transitivelyUsedModules.entrySet().iterator();
+    while (iterator.hasNext()) {
       // It is better to iterate over each nested set here instead of creating a joint one and
       // iterating over it, as this makes use of NestedSet's memoization (each of them has likely
       // been iterated over before). Don't use Set.removeAll() here as that iterates over the
@@ -592,8 +598,17 @@
       // (transitive, which is a linear scan).
       // We get a collection view of the NestedSet in a way that can throw an InterruptedException
       // because a NestedSet may contain a future.
-      for (Artifact module :
-          actionExecutionContext.getNestedSetExpander().toListInterruptibly(transitive)) {
+      Map.Entry<Artifact, NestedSet<? extends Artifact>> entry = iterator.next();
+      NestedSet<? extends Artifact> transitive = entry.getValue();
+
+      List<? extends Artifact> modules;
+      try {
+        modules = actionExecutionContext.getNestedSetExpander().toListInterruptibly(transitive);
+      } catch (TimeoutException e) {
+        throw handleTimedOutNestedSetExpansion(entry.getKey(), iterator, e);
+      }
+
+      for (Artifact module : modules) {
         topLevel.remove(module);
       }
     }
@@ -616,6 +631,45 @@
     return additionalInputs;
   }
 
+  /**
+   * Handles a timeout during expansion of transitively used modules.
+   *
+   * <p>A timeout may occur if a nested set of transitively used modules {@linkplain
+   * NestedSet#isFromStorage} but is not {@linkplain NestedSet#isReady ready}. The timeout is
+   * handled by throwing {@link LostInputsActionExecutionException} so that rewinding kicks in and
+   * rebuilds the nodes with the unavailable nested sets.
+   *
+   * <p>As soon as one timeout is seen, any other nested sets of modules which are not ready are
+   * also treated as lost inputs.
+   *
+   * <p>Although the output {@link Artifact} (.pcm file) of dependent modules is not technically
+   * lost, it is treated as a lost input because rewinding will rebuild the corresponding action,
+   * thus reconstituting its nested set of transitively used modules. In lieu of an actual digest,
+   * we use the .pcm file's exec path since rewinding only uses the digest to detect multiple
+   * rewinds of the same input.
+   */
+  private LostInputsActionExecutionException handleTimedOutNestedSetExpansion(
+      Artifact timedOut,
+      Iterator<Map.Entry<Artifact, NestedSet<? extends Artifact>>> remainingModules,
+      TimeoutException e) {
+    ImmutableMap.Builder<String, ActionInput> lostInputsBuilder = ImmutableMap.builder();
+    lostInputsBuilder.put(timedOut.getExecPathString(), timedOut);
+    remainingModules.forEachRemaining(
+        entry -> {
+          if (!entry.getValue().isReady()) {
+            Artifact alsoTimedOut = entry.getKey();
+            lostInputsBuilder.put(alsoTimedOut.getExecPathString(), alsoTimedOut);
+          }
+        });
+    ImmutableMap<String, ActionInput> lostInputs = lostInputsBuilder.build();
+    return new LostInputsActionExecutionException(
+        "Timed out expanding modules",
+        lostInputs,
+        new ActionInputDepOwnerMap(ImmutableList.copyOf(lostInputs.values())),
+        this,
+        e);
+  }
+
   @Override
   public Artifact getPrimaryInput() {
     return getSourceFile();
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutor.java
index 1188368..2323c4e 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SequencedSkyframeExecutor.java
@@ -38,7 +38,6 @@
 import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget;
 import com.google.devtools.build.lib.cmdline.LabelConstants;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
-import com.google.devtools.build.lib.collect.nestedset.NestedSetExpander;
 import com.google.devtools.build.lib.concurrent.Uninterruptibles;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
@@ -191,8 +190,7 @@
         mutableArtifactFactorySupplier,
         new ConfiguredTargetProgressReceiver(),
         /*nonexistentFileReceiver=*/ null,
-        managedDirectoriesKnowledge,
-        NestedSetExpander.DEFAULT);
+        managedDirectoriesKnowledge);
     this.diffAwarenessManager = new DiffAwarenessManager(diffAwarenessFactories);
     this.customDirtinessCheckers = customDirtinessCheckers;
     this.managedDirectoriesKnowledge = managedDirectoriesKnowledge;
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 81b1cbd..6f14413 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
@@ -225,19 +225,17 @@
 
   private boolean bazelRemoteExecutionEnabled;
 
-  private final NestedSetExpander nestedSetExpander;
+  private NestedSetExpander nestedSetExpander;
 
   SkyframeActionExecutor(
       ActionKeyContext actionKeyContext,
       AtomicReference<ActionExecutionStatusReporter> statusReporterRef,
       Supplier<ImmutableList<Root>> sourceRootSupplier,
-      Function<PathFragment, SourceArtifact> sourceArtifactFactory,
-      NestedSetExpander nestedSetExpander) {
+      Function<PathFragment, SourceArtifact> sourceArtifactFactory) {
     this.actionKeyContext = actionKeyContext;
     this.statusReporterRef = statusReporterRef;
     this.sourceRootSupplier = sourceRootSupplier;
     this.sourceArtifactFactory = sourceArtifactFactory;
-    this.nestedSetExpander = nestedSetExpander;
   }
 
   /**
@@ -924,9 +922,13 @@
     return hadExecutionError && !options.getOptions(KeepGoingOption.class).keepGoing;
   }
 
-  void configure(MetadataProvider fileCache, ActionInputPrefetcher actionInputPrefetcher) {
+  public void configure(
+      MetadataProvider fileCache,
+      ActionInputPrefetcher actionInputPrefetcher,
+      NestedSetExpander nestedSetExpander) {
     this.perBuildFileCache = fileCache;
     this.actionInputPrefetcher = actionInputPrefetcher;
+    this.nestedSetExpander = nestedSetExpander;
   }
 
   /**
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
index 2c165e3..f49e829 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/SkyframeExecutor.java
@@ -296,7 +296,7 @@
 
   private final AtomicReference<ActionExecutionStatusReporter> statusReporterRef =
       new AtomicReference<>();
-  private final SkyframeActionExecutor skyframeActionExecutor;
+  protected final SkyframeActionExecutor skyframeActionExecutor;
   private ActionExecutionFunction actionExecutionFunction;
   protected SkyframeProgressReceiver progressReceiver;
   private final AtomicReference<CyclesReporter> cyclesReporter = new AtomicReference<>();
@@ -398,8 +398,7 @@
       MutableArtifactFactorySupplier artifactResolverSupplier,
       @Nullable ConfiguredTargetProgressReceiver configuredTargetProgress,
       @Nullable NonexistentFileReceiver nonexistentFileReceiver,
-      @Nullable ManagedDirectoriesKnowledge managedDirectoriesKnowledge,
-      NestedSetExpander nestedSetExpander) {
+      @Nullable ManagedDirectoriesKnowledge managedDirectoriesKnowledge) {
     // Strictly speaking, these arguments are not required for initialization, but all current
     // callsites have them at hand, so we might as well set them during construction.
     this.skyframeExecutorConsumerOnInit = skyframeExecutorConsumerOnInit;
@@ -429,11 +428,7 @@
     this.defaultBuildOptions = defaultBuildOptions;
     this.skyframeActionExecutor =
         new SkyframeActionExecutor(
-            actionKeyContext,
-            statusReporterRef,
-            this::getPathEntries,
-            this::createSourceArtifact,
-            nestedSetExpander);
+            actionKeyContext, statusReporterRef, this::getPathEntries, this::createSourceArtifact);
     this.skyframeBuildView =
         new SkyframeBuildView(
             directories,
@@ -678,7 +673,7 @@
 
   public void configureActionExecutor(
       MetadataProvider fileCache, ActionInputPrefetcher actionInputPrefetcher) {
-    this.skyframeActionExecutor.configure(fileCache, actionInputPrefetcher);
+    skyframeActionExecutor.configure(fileCache, actionInputPrefetcher, NestedSetExpander.DEFAULT);
   }
 
   public void dump(boolean summarize, PrintStream out) {
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java b/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
index 571dcdcb..dae3824 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/TimestampBuilderTestCase.java
@@ -220,8 +220,7 @@
             actionKeyContext,
             new AtomicReference<>(statusReporter),
             /*sourceRootSupplier=*/ () -> ImmutableList.of(),
-            /*sourceArtifactFactory=*/ unused -> null,
-            NestedSetExpander.DEFAULT);
+            /*sourceArtifactFactory=*/ unused -> null);
 
     Path actionOutputBase = scratch.dir("/usr/local/google/_blaze_jrluser/FAKEMD5/action_out/");
     skyframeActionExecutor.setActionLogBufferPathGenerator(
@@ -229,7 +228,7 @@
 
     MetadataProvider cache =
         new SingleBuildFileCache(rootDirectory.getPathString(), scratch.getFileSystem());
-    skyframeActionExecutor.configure(cache, ActionInputPrefetcher.NONE);
+    skyframeActionExecutor.configure(cache, ActionInputPrefetcher.NONE, NestedSetExpander.DEFAULT);
 
     final InMemoryMemoizingEvaluator evaluator =
         new InMemoryMemoizingEvaluator(