Add option to filter for targets in the action graph dump.

RELNOTES: Use bazel dump --action_graph=/path/to/action.proto --action_graph:targets://foo:bar,//foo:foo to filter for certain targets in the action graph dump.
PiperOrigin-RevId: 187608321

3 files changed, 46 insertions, 9 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/runtime/commands/DumpCommand.java b/src/main/java/com/google/devtools/build/lib/runtime/commands/DumpCommand.java
index 35d38003d7..2400cf6d18 100644
--- a/src/main/java/com/google/devtools/build/lib/runtime/commands/DumpCommand.java
+++ b/src/main/java/com/google/devtools/build/lib/runtime/commands/DumpCommand.java
@@ -34,6 +34,7 @@ import com.google.devtools.build.lib.runtime.CommandEnvironment;
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor;
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor.RuleStat;
 import com.google.devtools.build.lib.util.ExitCode;
+import com.google.devtools.common.options.Converters.CommaSeparatedOptionListConverter;
 import com.google.devtools.common.options.EnumConverter;
 import com.google.devtools.common.options.Option;
 import com.google.devtools.common.options.OptionDocumentationCategory;
@@ -114,6 +115,19 @@ public class DumpCommand implements BlazeCommand {
     public String dumpActionGraph;
 
     @Option(
+      name = "action_graph:targets",
+      converter = CommaSeparatedOptionListConverter.class,
+      defaultValue = "...",
+      category = "verbosity",
+      documentationCategory = OptionDocumentationCategory.OUTPUT_SELECTION,
+      effectTags = {OptionEffectTag.BAZEL_MONITORING},
+      help =
+          "Comma separated list of targets to include in action graph dump. "
+              + "Defaults to all attributes. This option does only apply to --action_graph."
+    )
+    public List<String> actionGraphTargets;
+
+    @Option(
       name = "rule_classes",
       defaultValue = "false",
       category = "verbosity",
@@ -232,7 +246,12 @@ public class DumpCommand implements BlazeCommand {
 
       if (dumpOptions.dumpActionGraph != null) {
         try {
-          success &= dumpActionGraph(env.getSkyframeExecutor(), dumpOptions.dumpActionGraph, out);
+          success &=
+              dumpActionGraph(
+                  env.getSkyframeExecutor(),
+                  dumpOptions.dumpActionGraph,
+                  dumpOptions.actionGraphTargets,
+                  out);
         } catch (IOException e) {
           env.getReporter()
               .error(
@@ -283,10 +302,12 @@ public class DumpCommand implements BlazeCommand {
     return true;
   }
 
-  private boolean dumpActionGraph(SkyframeExecutor executor, String path, PrintStream out)
+  private boolean dumpActionGraph(
+      SkyframeExecutor executor, String path, List<String> actionGraphTargets, PrintStream out)
       throws IOException {
     out.println("Dumping action graph to '" + path + "'");
-    ActionGraphContainer actionGraphContainer = executor.getActionGraphContainer();
+    ActionGraphContainer actionGraphContainer =
+        executor.getActionGraphContainer(actionGraphTargets);
     FileOutputStream protoOutputStream = new FileOutputStream(path);
     actionGraphContainer.writeTo(protoOutputStream);
     protoOutputStream.close();
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 1652a764ff..ab98da4641 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
@@ -893,8 +893,16 @@ public final class SequencedSkyframeExecutor extends SkyframeExecutor {
     }
   }
 
+  private static boolean includeInActionGraph(String labelString, List<String> actionGraphTargets) {
+    if (actionGraphTargets.size() == 1
+        && Iterables.getOnlyElement(actionGraphTargets).equals("...")) {
+      return true;
+    }
+    return actionGraphTargets.contains(labelString);
+  }
+
   @Override
-  public ActionGraphContainer getActionGraphContainer() {
+  public ActionGraphContainer getActionGraphContainer(List<String> actionGraphTargets) {
     ActionGraphContainer.Builder actionGraphBuilder = ActionGraphContainer.newBuilder();
     ActionGraphIdCache actionGraphIdCache = new ActionGraphIdCache(actionGraphBuilder);
     for (Map.Entry<SkyKey, ? extends NodeEntry> skyKeyAndNodeEntry :
@@ -904,9 +912,9 @@ public final class SequencedSkyframeExecutor extends SkyframeExecutor {
       SkyFunctionName functionName = key.functionName();
       try {
         if (functionName.equals(SkyFunctions.CONFIGURED_TARGET)) {
-          dumpConfiguredTarget(actionGraphBuilder, actionGraphIdCache, entry);
+          dumpConfiguredTarget(actionGraphBuilder, actionGraphIdCache, entry, actionGraphTargets);
         } else if (functionName.equals(SkyFunctions.ASPECT)) {
-          dumpAspect(actionGraphBuilder, actionGraphIdCache, entry);
+          dumpAspect(actionGraphBuilder, actionGraphIdCache, entry, actionGraphTargets);
         }
       } catch (InterruptedException e) {
         Thread.currentThread().interrupt();
@@ -919,7 +927,8 @@ public final class SequencedSkyframeExecutor extends SkyframeExecutor {
   private void dumpAspect(
       ActionGraphContainer.Builder actionGraphBuilder,
       ActionGraphIdCache actionGraphIdCache,
-      NodeEntry entry)
+      NodeEntry entry,
+      List<String> actionGraphTargets)
       throws InterruptedException {
     AspectValue aspectValue = (AspectValue) entry.getValue();
     AspectKey aspectKey = aspectValue.getKey();
@@ -927,6 +936,9 @@ public final class SequencedSkyframeExecutor extends SkyframeExecutor {
         (ConfiguredTargetValue)
             memoizingEvaluator.getExistingValue(aspectKey.getBaseConfiguredTargetKey());
     ConfiguredTarget configuredTarget = value.getConfiguredTarget();
+    if (!includeInActionGraph(configuredTarget.getLabel().toString(), actionGraphTargets)) {
+      return;
+    }
     for (int i = 0; i < aspectValue.getNumActions(); i++) {
       Action action = aspectValue.getAction(i);
       dumpSingleAction(actionGraphIdCache, actionGraphBuilder, configuredTarget, action);
@@ -936,10 +948,14 @@ public final class SequencedSkyframeExecutor extends SkyframeExecutor {
   private void dumpConfiguredTarget(
       ActionGraphContainer.Builder actionGraphBuilder,
       ActionGraphIdCache actionGraphIdCache,
-      NodeEntry entry)
+      NodeEntry entry,
+      List<String> actionGraphTargets)
       throws InterruptedException {
     ConfiguredTargetValue ctValue = (ConfiguredTargetValue) entry.getValue();
     ConfiguredTarget configuredTarget = ctValue.getConfiguredTarget();
+    if (!includeInActionGraph(configuredTarget.getLabel().toString(), actionGraphTargets)) {
+      return;
+    }
     List<ActionAnalysisMetadata> actions = ctValue.getActions();
     for (ActionAnalysisMetadata action : actions) {
       dumpSingleAction(actionGraphIdCache, actionGraphBuilder, configuredTarget, action);
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 68f74aa7f5..1514290052 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
@@ -593,7 +593,7 @@ public abstract class SkyframeExecutor implements WalkableGraphFactory {
     }
   }
 
-  public abstract ActionGraphContainer getActionGraphContainer();
+  public abstract ActionGraphContainer getActionGraphContainer(List<String> actionGraphTargets);
 
   class BuildViewProvider {
     /**