Add option to dump the action graph.

Note that this dumps the current state in skyframe (which may contain more nodes than you're interested in):
- bazel build --nobuild //interesting:targets
- bazel dump --action_graph=/path/to/file
- printproto --proto2 --raw_protocol_buffer --message=action_graph.ActionGraphContainer --multiline --proto=third_party/bazel/src/main/protobuf/action_graph.proto /path/to/file

We'll add filtering options in a later CL.

RELNOTES[NEW]: Add option to dump the action graph to a file: 'bazel dump --action_graph=/path/to/file'.

PiperOrigin-RevId: 186597930

1 files changed, 253 insertions, 0 deletions

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 3f271b30c8..de51557e25 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
@@ -27,20 +27,35 @@ import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Range;
 import com.google.common.collect.Sets;
+import com.google.devtools.build.lib.actions.Action;
+import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
+import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
+import com.google.devtools.build.lib.actions.ActionGraphProtos;
+import com.google.devtools.build.lib.actions.ActionGraphProtos.ActionGraphContainer;
 import com.google.devtools.build.lib.actions.ActionKeyContext;
+import com.google.devtools.build.lib.actions.ActionOwner;
+import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.BlazeDirectories;
 import com.google.devtools.build.lib.analysis.BuildView;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.WorkspaceStatusAction.Factory;
+import com.google.devtools.build.lib.analysis.actions.SpawnAction;
 import com.google.devtools.build.lib.analysis.buildinfo.BuildInfoFactory;
+import com.google.devtools.build.lib.analysis.config.BuildConfiguration;
 import com.google.devtools.build.lib.analysis.configuredtargets.RuleConfiguredTarget;
 import com.google.devtools.build.lib.buildtool.BuildRequestOptions;
+import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
+import com.google.devtools.build.lib.collect.nestedset.NestedSet;
+import com.google.devtools.build.lib.collect.nestedset.NestedSetBuilder;
+import com.google.devtools.build.lib.collect.nestedset.NestedSetView;
+import com.google.devtools.build.lib.collect.nestedset.Order;
 import com.google.devtools.build.lib.concurrent.Uninterruptibles;
 import com.google.devtools.build.lib.events.Event;
 import com.google.devtools.build.lib.events.EventHandler;
 import com.google.devtools.build.lib.events.ExtendedEventHandler;
 import com.google.devtools.build.lib.packages.AspectClass;
+import com.google.devtools.build.lib.packages.AspectDescriptor;
 import com.google.devtools.build.lib.packages.BuildFileName;
 import com.google.devtools.build.lib.packages.NoSuchPackageException;
 import com.google.devtools.build.lib.packages.NoSuchTargetException;
@@ -52,6 +67,7 @@ import com.google.devtools.build.lib.packages.SkylarkSemanticsOptions;
 import com.google.devtools.build.lib.pkgcache.PackageCacheOptions;
 import com.google.devtools.build.lib.pkgcache.PathPackageLocator;
 import com.google.devtools.build.lib.profiler.AutoProfiler;
+import com.google.devtools.build.lib.skyframe.AspectValue.AspectKey;
 import com.google.devtools.build.lib.skyframe.DirtinessCheckerUtils.BasicFilesystemDirtinessChecker;
 import com.google.devtools.build.lib.skyframe.DirtinessCheckerUtils.ExternalDirtinessChecker;
 import com.google.devtools.build.lib.skyframe.DirtinessCheckerUtils.MissingDiffDirtinessChecker;
@@ -96,6 +112,7 @@ import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.Callable;
@@ -760,6 +777,242 @@ public final class SequencedSkyframeExecutor extends SkyframeExecutor {
     return new ArrayList<>(ruleStats.values());
   }
 
+  private static class ActionGraphIdCache {
+    private final Map<Artifact, String> knownArtifacts = new HashMap<>();
+    private final Map<BuildConfiguration, String> knownConfigurations = new HashMap<>();
+    private final Map<Label, String> knownTargets = new HashMap<>();
+    private final Map<AspectDescriptor, String> knownAspectDescriptors = new HashMap<>();
+    private final Map<String, String> knownRuleClassStrings = new HashMap<>();
+    // The NestedSet is identified by their raw 'children' object since multiple NestedSetViews
+    // can point to the same object.
+    private final Map<Object, String> knownNestedSets = new HashMap<>();
+
+    private final ActionGraphContainer.Builder actionGraphBuilder;
+    private final ActionKeyContext actionKeyContext = new ActionKeyContext();
+
+    ActionGraphIdCache(ActionGraphContainer.Builder actionGraphBuilder) {
+      this.actionGraphBuilder = actionGraphBuilder;
+    }
+
+    public ActionKeyContext getActionKeyContext() {
+      return actionKeyContext;
+    }
+
+    public String ruleClassStringToId(String ruleClassString) {
+      if (!knownRuleClassStrings.containsKey(ruleClassString)) {
+        String targetId = String.valueOf(knownRuleClassStrings.size());
+        knownRuleClassStrings.put(ruleClassString, targetId);
+        ActionGraphProtos.RuleClass.Builder ruleClassBuilder =
+            ActionGraphProtos.RuleClass.newBuilder().setId(targetId).setName(ruleClassString);
+        actionGraphBuilder.addRuleClasses(ruleClassBuilder.build());
+      }
+      return knownRuleClassStrings.get(ruleClassString);
+    }
+
+    public String targetToId(Label label, String ruleClassString) {
+      if (!knownTargets.containsKey(label)) {
+        String targetId = String.valueOf(knownTargets.size());
+        knownTargets.put(label, targetId);
+        ActionGraphProtos.Target.Builder targetBuilder = ActionGraphProtos.Target.newBuilder();
+        targetBuilder.setId(targetId).setLabel(label.toString());
+        if (ruleClassString != null) {
+          targetBuilder.setRuleClassId(ruleClassStringToId(ruleClassString));
+        }
+        actionGraphBuilder.addTargets(targetBuilder.build());
+      }
+      return knownTargets.get(label);
+    }
+
+    public String configurationToId(BuildConfiguration buildConfiguration) {
+      if (!knownConfigurations.containsKey(buildConfiguration)) {
+        String configurationId = String.valueOf(knownConfigurations.size());
+        knownConfigurations.put(buildConfiguration, configurationId);
+        ActionGraphProtos.Configuration configurationProto =
+            ActionGraphProtos.Configuration.newBuilder()
+                .setMnemonic(buildConfiguration.getMnemonic())
+                .setPlatformName(buildConfiguration.getPlatformName())
+                .setId(configurationId)
+                .build();
+        actionGraphBuilder.addConfiguration(configurationProto);
+      }
+      return knownConfigurations.get(buildConfiguration);
+    }
+
+    public String artifactToId(Artifact artifact) {
+      if (!knownArtifacts.containsKey(artifact)) {
+        String artifactId = String.valueOf(knownArtifacts.size());
+        knownArtifacts.put(artifact, artifactId);
+        ActionGraphProtos.Artifact artifactProto =
+            ActionGraphProtos.Artifact.newBuilder()
+                .setId(artifactId)
+                .setExecPath(artifact.getExecPathString())
+                .setIsTreeArtifact(artifact.isTreeArtifact())
+                .build();
+        actionGraphBuilder.addArtifacts(artifactProto);
+      }
+      return knownArtifacts.get(artifact);
+    }
+
+    public String depSetToId(NestedSetView<Artifact> nestedSetView) {
+      if (!knownNestedSets.containsKey(nestedSetView.identifier())) {
+        String nestedSetId = String.valueOf(knownNestedSets.size());
+        knownNestedSets.put(nestedSetView.identifier(), nestedSetId);
+        ActionGraphProtos.DepSetOfFiles.Builder depSetBuilder =
+            ActionGraphProtos.DepSetOfFiles.newBuilder()
+                .setId(nestedSetId);
+        for (NestedSetView<Artifact> transitiveNestedSet : nestedSetView.transitives()) {
+          depSetBuilder.addTransitiveDepSetIds(depSetToId(transitiveNestedSet));
+        }
+        for (Artifact directArtifact : nestedSetView.directs()) {
+          depSetBuilder.addDirectArtifactIds(artifactToId(directArtifact));
+        }
+        actionGraphBuilder.addDepSetOfFiles(depSetBuilder.build());
+      }
+      return knownNestedSets.get(nestedSetView.identifier());
+    }
+
+    public String aspectDescriptorToId(AspectDescriptor aspectDescriptor) {
+      if (!knownAspectDescriptors.containsKey(aspectDescriptor)) {
+        String aspectDescriptorId = String.valueOf(knownAspectDescriptors.size());
+        knownAspectDescriptors.put(aspectDescriptor, aspectDescriptorId);
+        ActionGraphProtos.AspectDescriptor.Builder aspectDescriptorBuilder =
+            ActionGraphProtos.AspectDescriptor.newBuilder()
+                .setId(aspectDescriptorId)
+                .setName(aspectDescriptor.getAspectClass().getName());
+        for (Entry<String, String> parameter :
+            aspectDescriptor.getParameters().getAttributes().entries()) {
+          ActionGraphProtos.KeyValuePair.Builder keyValuePairBuilder =
+              ActionGraphProtos.KeyValuePair.newBuilder();
+          keyValuePairBuilder
+              .setKey(parameter.getKey())
+              .setValue(parameter.getValue());
+          aspectDescriptorBuilder.addParameters(keyValuePairBuilder.build());
+        }
+        actionGraphBuilder.addAspectDescriptors(aspectDescriptorBuilder.build());
+      }
+      return knownAspectDescriptors.get(aspectDescriptor);
+    }
+  }
+
+  @Override
+  public ActionGraphContainer getActionGraphContainer() {
+    ActionGraphContainer.Builder actionGraphBuilder = ActionGraphContainer.newBuilder();
+    ActionGraphIdCache actionGraphIdCache = new ActionGraphIdCache(actionGraphBuilder);
+    for (Map.Entry<SkyKey, ? extends NodeEntry> skyKeyAndNodeEntry :
+        memoizingEvaluator.getGraphMap().entrySet()) {
+      NodeEntry entry = skyKeyAndNodeEntry.getValue();
+      SkyKey key = skyKeyAndNodeEntry.getKey();
+      SkyFunctionName functionName = key.functionName();
+      try {
+        if (functionName.equals(SkyFunctions.CONFIGURED_TARGET)) {
+          dumpConfiguredTarget(actionGraphBuilder, actionGraphIdCache, entry);
+        } else if (functionName.equals(SkyFunctions.ASPECT)) {
+          dumpAspect(actionGraphBuilder, actionGraphIdCache, entry);
+        }
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new IllegalStateException("No interruption in sequenced evaluation", e);
+      }
+    }
+    return actionGraphBuilder.build();
+  }
+
+  private void dumpAspect(
+      ActionGraphContainer.Builder actionGraphBuilder,
+      ActionGraphIdCache actionGraphIdCache,
+      NodeEntry entry)
+      throws InterruptedException {
+    AspectValue aspectValue = (AspectValue) entry.getValue();
+    AspectKey aspectKey = aspectValue.getKey();
+    ConfiguredTargetValue value =
+        (ConfiguredTargetValue)
+            memoizingEvaluator.getExistingValue(aspectKey.getBaseConfiguredTargetKey());
+    ConfiguredTarget configuredTarget = value.getConfiguredTarget();
+    for (int i = 0; i < aspectValue.getNumActions(); i++) {
+      Action action = aspectValue.getAction(i);
+      dumpSingleAction(actionGraphIdCache, actionGraphBuilder, configuredTarget, action);
+    }
+  }
+
+  private void dumpConfiguredTarget(
+      ActionGraphContainer.Builder actionGraphBuilder,
+      ActionGraphIdCache actionGraphIdCache,
+      NodeEntry entry)
+      throws InterruptedException {
+    ConfiguredTargetValue ctValue = (ConfiguredTargetValue) entry.getValue();
+    ConfiguredTarget configuredTarget = ctValue.getConfiguredTarget();
+    List<ActionAnalysisMetadata> actions = ctValue.getActions();
+    for (ActionAnalysisMetadata action : actions) {
+      dumpSingleAction(actionGraphIdCache, actionGraphBuilder, configuredTarget, action);
+    }
+  }
+
+  private void dumpSingleAction(
+      ActionGraphIdCache actionGraphIdCache,
+      ActionGraphContainer.Builder actionGraphBuilder,
+      ConfiguredTarget configuredTarget,
+      ActionAnalysisMetadata action) {
+    Preconditions.checkState(configuredTarget instanceof RuleConfiguredTarget);
+    Label label = configuredTarget.getLabel();
+    String ruleClassString = ((RuleConfiguredTarget) configuredTarget).getRuleClassString();
+    ActionGraphProtos.Action.Builder actionBuilder =
+        ActionGraphProtos.Action.newBuilder()
+            .setMnemonic(action.getMnemonic())
+            .setTargetId(actionGraphIdCache.targetToId(label, ruleClassString));
+
+    if (action instanceof ActionExecutionMetadata) {
+      ActionExecutionMetadata actionExecutionMetadata = (ActionExecutionMetadata) action;
+      actionBuilder
+          .setActionKey(actionExecutionMetadata.getKey(actionGraphIdCache.getActionKeyContext()))
+          .setDiscoversInputs(actionExecutionMetadata.discoversInputs());
+    }
+
+    // store environment
+    if (action instanceof SpawnAction) {
+      SpawnAction spawnAction = (SpawnAction) action;
+      // TODO(twerth): This handles the fixed environemnt. We probably want to output the inherited
+      // environment as well.
+      ImmutableMap<String, String> fixedEnvironment = spawnAction.getEnvironment();
+      for (Entry<String, String> environmentVariable : fixedEnvironment.entrySet()) {
+        ActionGraphProtos.KeyValuePair.Builder keyValuePairBuilder =
+            ActionGraphProtos.KeyValuePair.newBuilder();
+        keyValuePairBuilder
+            .setKey(environmentVariable.getKey())
+            .setValue(environmentVariable.getValue());
+        actionBuilder.addEnvironmentVariables(keyValuePairBuilder.build());
+      }
+    }
+
+    ActionOwner actionOwner = action.getOwner();
+    if (actionOwner != null) {
+      BuildConfiguration buildConfiguration = (BuildConfiguration) actionOwner.getConfiguration();
+      actionBuilder.setConfigurationId(actionGraphIdCache.configurationToId(buildConfiguration));
+
+      // store aspect
+      for (AspectDescriptor aspectDescriptor : actionOwner.getAspectDescriptors()) {
+        actionBuilder.addAspectDescriptorIds(
+            actionGraphIdCache.aspectDescriptorToId(aspectDescriptor));
+      }
+    }
+
+    // store inputs
+    Iterable<Artifact> inputs = action.getInputs();
+    if (!(inputs instanceof NestedSet)) {
+      inputs = NestedSetBuilder.wrap(Order.STABLE_ORDER, inputs);
+    }
+    NestedSetView<Artifact> nestedSetView = new NestedSetView<>((NestedSet<Artifact>) inputs);
+    if (nestedSetView.directs().size() > 0 || nestedSetView.transitives().size() > 0) {
+      actionBuilder.addInputDepSetIds(actionGraphIdCache.depSetToId(nestedSetView));
+    }
+
+    // store outputs
+    for (Artifact artifact : action.getOutputs()) {
+      actionBuilder.addOutputIds(actionGraphIdCache.artifactToId(artifact));
+    }
+
+    actionGraphBuilder.addActions(actionBuilder.build());
+  }
+
   /**
    * In addition to calling the superclass method, deletes all ConfiguredTarget values from the
    * Skyframe cache. This is done to save memory (e.g. on a configuration change); since the