Update counters in FileDependencySerializer's NodeDependencyHandler on error.

This nodesWaitingForDeps counter is incremented during setup, but there are
a couple of cases where we are missing a decrement.
* If any dependency fails. Uses QuiescingFutureTask's doneWithError to decrement
  in this case.
* NodeDependencyHandler.call has a couple of short-circuit branches that don't
  decrement correctly. Adds the appropriate decrements for these.

PiperOrigin-RevId: 980735799
Change-Id: Ia5ee7ee7c5753d0857c01b807347092c24344c1a

Author: aoeui

src/main/java/com/google/devtools/build/lib/skyframe/serialization/analysis/BUILD

@@ -232,6 +232,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/actions:file_metadata",
         "//src/main/java/com/google/devtools/build/lib/analysis:analysis_cluster",
         "//src/main/java/com/google/devtools/build/lib/compress:compression_service",
+        "//src/main/java/com/google/devtools/build/lib/concurrent",
         "//src/main/java/com/google/devtools/build/lib/concurrent/safeexecutor:safe_executor",
         "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/profiler:counter_series_task_impl",

src/main/java/com/google/devtools/build/lib/skyframe/serialization/analysis/FileDependencySerializer.java

@@ -44,6 +44,7 @@
 import com.google.devtools.build.lib.actions.FileValue;
 import com.google.devtools.build.lib.analysis.ConfiguredRuleClassProvider.BundledFileSystem;
 import com.google.devtools.build.lib.compress.CompressionService;
+import com.google.devtools.build.lib.concurrent.QuiescingFutureTask;
 import com.google.devtools.build.lib.concurrent.safeexecutor.SafeExecutor;
 import com.google.devtools.build.lib.concurrent.safeexecutor.SafeFutures;
 import com.google.devtools.build.lib.profiler.CounterSeriesCollector;
@@ -88,6 +89,7 @@
 import com.google.devtools.build.lib.vfs.RootedPath;
 import com.google.devtools.build.skyframe.InMemoryGraph;
 import com.google.devtools.build.skyframe.InMemoryNodeEntry;
+import com.google.errorprone.annotations.DoNotCall;
 import com.google.protobuf.CodedOutputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -793,50 +795,9 @@ public ListenableFuture<ListingDataInfo> apply(FileDataInfo info) {
   NodeDataInfoOrFuture populateFutureNodeDataInfo(FutureNodeDataInfo future) {
     counters.nodesWaitingForDeps.incrementAndGet();
 
-    AbstractNestedFileOpNodes node = future.key();
-    var dependencyHandler = new NodeDependencyHandler();
-
-    // Loops through all node dependencies, registering them with the dependencyHandler. The
-    // dependencyHandler triggers recursive registration, keeping track of immediate results and
-    // any futures.
-    for (int i = 0; i < node.analysisDependenciesCount(); i++) {
-      switch (node.getAnalysisDependency(i)) {
-        case FileKey fileKey:
-          dependencyHandler.addFileKey(fileKey);
-          break;
-        case DirectoryListingKey listingKey:
-          dependencyHandler.addListingKey(listingKey);
-          break;
-        case AbstractNestedFileOpNodes nestedKeys:
-          dependencyHandler.addNodeKey(nestedKeys);
-          break;
-        case RemoteFileOpNode remoteNode:
-          dependencyHandler.addRemoteNode(remoteNode);
-          break;
-      }
-    }
-
-    switch (node) {
-      case NestedFileOpNodes plainNodes:
-        break;
-      case NestedFileOpNodesWithSource withSource:
-        dependencyHandler.setSourceFile(withSource.source());
-        break;
-    }
-
-    var allFutures = dependencyHandler.getCombinedFutures();
-    if (allFutures.isEmpty()) {
-      NodeDataInfo result;
-      try {
-        result = dependencyHandler.call();
-      } catch (ExecutionException | IOException e) {
-        // Only thrown when calling Future.get, but none should be present if this is reached.
-        throw new IllegalStateException("unexpected failure", e);
-      }
-      return future.completeWith(result);
-    }
-    return future.completeWith(
-        SafeFutures.call(Futures.whenAllSucceed(allFutures), dependencyHandler, executor));
+    var task = new NodeDependencyHandler(future.key());
+    task.run();
+    return future.completeWith(task);
   }
 
   OutputStream getCompressedOutputStream(OutputStream outputStream) throws IOException {
@@ -849,12 +810,15 @@ OutputStream getCompressedOutputStream(OutputStream outputStream) throws IOExcep
 
   /**
    * Accepts all the dependencies associated with a node, registers their serialization and waits
-   * for processing to complete, signalled through the {@link #call} callback.
+   * for processing to complete.
    *
    * <p>Once processing is complete and all keys are known, uploads the node value. {@link
    * #computeNodeBytes} defines the wire format of nodes.
    */
-  class NodeDependencyHandler implements Callable<NodeDataInfo> {
+  private final class NodeDependencyHandler extends QuiescingFutureTask<NodeDataInfo>
+      implements FutureCallback<Object> {
+    private final AbstractNestedFileOpNodes node;
+
     private final ArrayList<String> fileKeys = new ArrayList<>();
     private final ArrayList<String> listingKeys = new ArrayList<>();
     private final ArrayList<NodeInvalidationDataInfo> nodeDependencies = new ArrayList<>();
@@ -866,21 +830,54 @@ class NodeDependencyHandler implements Callable<NodeDataInfo> {
     private final ArrayList<FutureListingDataInfo> futureListingDataInfo = new ArrayList<>();
     private final ArrayList<FutureNodeDataInfo> futureNodeDataInfo = new ArrayList<>();
 
+    private NodeDependencyHandler(AbstractNestedFileOpNodes node) {
+      super(executor);
+      this.node = node;
+    }
+
     @Override
-    public NodeDataInfo call() throws ExecutionException, IOException {
-      for (FutureFileDataInfo futureInfo : futureFileDataInfo) {
-        addFileInfo(Futures.getDone(futureInfo));
+    protected void arrangeSubtasks() {
+      // Loops through all node dependencies, registering them with this handler. This triggers
+      // recursive registration, keeping track of immediate results and any futures.
+      for (int i = 0; i < node.analysisDependenciesCount(); i++) {
+        switch (node.getAnalysisDependency(i)) {
+          case FileKey fileKey -> addFileKey(fileKey);
+          case DirectoryListingKey listingKey -> addListingKey(listingKey);
+          case AbstractNestedFileOpNodes nestedKeys -> addNodeKey(nestedKeys);
+          case RemoteFileOpNode remoteNode -> addRemoteNode(remoteNode);
+        }
       }
-      for (FutureListingDataInfo futureInfo : futureListingDataInfo) {
-        addListingInfo(Futures.getDone(futureInfo));
+
+      switch (node) {
+        case NestedFileOpNodes plainNodes -> {}
+        case NestedFileOpNodesWithSource withSource -> setSourceFile(withSource.source());
       }
-      for (FutureNodeDataInfo futureInfo : futureNodeDataInfo) {
-        addNodeInfo(Futures.getDone(futureInfo));
+    }
+
+    @Override
+    protected NodeDataInfo getValue() {
+      @Nullable String sourceFileKey;
+      try {
+        for (FutureFileDataInfo futureInfo : futureFileDataInfo) {
+          addFileInfo(Futures.getDone(futureInfo));
+        }
+        for (FutureListingDataInfo futureInfo : futureListingDataInfo) {
+          addListingInfo(Futures.getDone(futureInfo));
+        }
+        for (FutureNodeDataInfo futureInfo : futureNodeDataInfo) {
+          addNodeInfo(Futures.getDone(futureInfo));
+        }
+        sourceFileKey = getSourceFileKey();
+      } catch (ExecutionException e) {
+        // The QuiescingFutureTask setup guarantees that ExecutionException will not be thrown if
+        // getValue is called.
+        throw new AssertionError("unexpected failure", e);
       }
-      @Nullable String sourceFileKey = getSourceFileKey();
 
       if (fileKeys.isEmpty() && listingKeys.isEmpty() && sourceFileKey == null) {
+        // TODO(b/558805781): investigate whether this is always dead code in practice
         if (nodeDependencies.isEmpty()) {
+          counters.nodesWaitingForDeps.decrementAndGet();
           return CONSTANT_NODE; // None of the dependencies are relevant to invalidation.
         }
         // There are multiple ways that result could become unary here, even if `node` always has at
@@ -893,6 +890,7 @@ public NodeDataInfo call() throws ExecutionException, IOException {
           //
           // TODO: b/364831651 - consider additional special casing for unary file or listing
           // dependencies.
+          counters.nodesWaitingForDeps.decrementAndGet();
           return nodeDependencies.get(0);
         }
       }
@@ -956,6 +954,13 @@ public NodeDataInfo call() throws ExecutionException, IOException {
       return new NodeInvalidationDataInfo(key, writeStatusBuilder.build());
     }
 
+    @Override
+    protected void doneWithError(
+        @Nullable Throwable primaryCause, ImmutableList<Throwable> secondaryCauses) {
+      counters.nodesWaitingForDeps.decrementAndGet();
+      counters.nodesWithProcessingErrors.incrementAndGet();
+    }
+
     private void addRemoteNode(RemoteFileOpNode remoteNode) {
       addNodeInfo(registerDependency(remoteNode));
     }
@@ -967,6 +972,7 @@ private void addFileKey(FileKey fileKey) {
           break;
         case FutureFileDataInfo futureInfo:
           futureFileDataInfo.add(futureInfo);
+          trackFuture(futureInfo);
           break;
       }
     }
@@ -989,6 +995,7 @@ private void addListingKey(DirectoryListingKey listingKey) {
           break;
         case FutureListingDataInfo futureInfo:
           futureListingDataInfo.add(futureInfo);
+          trackFuture(futureInfo);
           break;
       }
     }
@@ -1011,6 +1018,7 @@ private void addNodeKey(AbstractNestedFileOpNodes nestedKeys) {
           break;
         case FutureNodeDataInfo futureInfo:
           futureNodeDataInfo.add(futureInfo);
+          trackFuture(futureInfo);
           break;
       }
     }
@@ -1033,28 +1041,39 @@ private void setSourceFile(FileKey sourceFile) {
           sourceFile,
           sourceFileOrFuture);
       this.sourceFileOrFuture = registerDependency(sourceFile);
-    }
-
-    private ImmutableList<ListenableFuture<?>> getCombinedFutures() {
-      var combined =
-          ImmutableList.<ListenableFuture<?>>builder()
-              .addAll(futureFileDataInfo)
-              .addAll(futureListingDataInfo)
-              .addAll(futureNodeDataInfo);
       switch (sourceFileOrFuture) {
-        case null -> {}
-        case FileDataInfo unusedSource -> {}
-        case FutureFileDataInfo futureSource -> combined.add(futureSource);
+        case FileDataInfo immediateSource -> {}
+        case FutureFileDataInfo futureSource -> trackFuture(futureSource);
       }
-      return combined.build();
     }
 
-    private byte[] compressBytes(byte[] nodeBytes) throws IOException {
+    private void trackFuture(ListenableFuture<?> future) {
+      increment();
+      Futures.addCallback(future, this, directExecutor());
+    }
+
+    @Override
+    @DoNotCall("Only called via trackFuture")
+    public void onSuccess(Object ignored) {
+      decrement();
+    }
+
+    @Override
+    @DoNotCall("Only called via trackFuture")
+    public void onFailure(Throwable t) {
+      notifyException(t);
+    }
+
+    private byte[] compressBytes(byte[] nodeBytes) {
       ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
-      MagicBytes.writeMagicBytes(outputStream);
-      try (OutputStream compressedBytesStream =
-          compressionService.newZstdOutputStream(outputStream)) {
-        compressedBytesStream.write(nodeBytes);
+      try {
+        MagicBytes.writeMagicBytes(outputStream);
+        try (OutputStream compressedBytesStream =
+            compressionService.newZstdOutputStream(outputStream)) {
+          compressedBytesStream.write(nodeBytes);
+        }
+      } catch (IOException e) {
+        throw new AssertionError("Unexpected IOException during in-memory compression", e);
       }
       return outputStream.toByteArray();
     }

src/test/java/com/google/devtools/build/lib/skyframe/serialization/analysis/BUILD

@@ -244,6 +244,7 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//src/main/java/com/google/devtools/build/lib/vfs/inmemoryfs",
         "//src/main/java/com/google/devtools/build/skyframe",
+        "//third_party/java/guava:collect",
         "//third_party/java/junit",
         "//third_party/java/mockito",
         "//third_party/java/truth",