Fix JSON format for aquery analysis_v2.proto.

The JSON format doesn't work with the streamed output since by nature it needs
to be grouped by message type. This CL adds a monolithic output handler for
analysis_v2.proto that can be used for JSON format.

RELNOTES: None
PiperOrigin-RevId: 289815268

src/main/java/com/google/devtools/build/lib/query2/aquery/ActionGraphProtoV2OutputFormatterCallback.java

@@ -21,8 +21,10 @@
 import com.google.devtools.build.lib.skyframe.ConfiguredTargetValue;
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor;
 import com.google.devtools.build.lib.skyframe.actiongraph.v2.ActionGraphDump;
+import com.google.devtools.build.lib.skyframe.actiongraph.v2.AqueryOutputHandler;
+import com.google.devtools.build.lib.skyframe.actiongraph.v2.AqueryOutputHandler.OutputType;
+import com.google.devtools.build.lib.skyframe.actiongraph.v2.MonolithicOutputHandler;
 import com.google.devtools.build.lib.skyframe.actiongraph.v2.StreamedOutputHandler;
-import com.google.devtools.build.lib.skyframe.actiongraph.v2.StreamedOutputHandler.OutputType;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
@@ -33,7 +35,7 @@ public class ActionGraphProtoV2OutputFormatterCallback extends AqueryThreadsafeC
   private final OutputType outputType;
   private final ActionGraphDump actionGraphDump;
   private final AqueryActionFilter actionFilters;
-  private StreamedOutputHandler streamedOutputHandler;
+  private AqueryOutputHandler aqueryOutputHandler;
 
   /**
    * Pseudo-arbitrarily chosen buffer size for output. Chosen to be large enough to fit a handful of
@@ -52,20 +54,28 @@ public class ActionGraphProtoV2OutputFormatterCallback extends AqueryThreadsafeC
     super(eventHandler, options, out, skyframeExecutor, accessor);
     this.outputType = outputType;
     this.actionFilters = actionFilters;
-    if (out != null) {
-      this.streamedOutputHandler =
-          new StreamedOutputHandler(
-              this.outputType,
-              CodedOutputStream.newInstance(out, OUTPUT_BUFFER_SIZE),
-              this.printStream);
+
+    switch (outputType) {
+      case BINARY:
+      case TEXT:
+        this.aqueryOutputHandler =
+            new StreamedOutputHandler(
+                this.outputType,
+                CodedOutputStream.newInstance(out, OUTPUT_BUFFER_SIZE),
+                this.printStream);
+        break;
+      case JSON:
+        this.aqueryOutputHandler = new MonolithicOutputHandler(this.printStream);
+        break;
     }
+
     this.actionGraphDump =
         new ActionGraphDump(
             options.includeCommandline,
             options.includeArtifacts,
             this.actionFilters,
             options.includeParamFiles,
-            streamedOutputHandler);
+            aqueryOutputHandler);
   }
 
   @Override
@@ -98,7 +108,7 @@ public void processOutput(Iterable<ConfiguredTargetValue> partialResult)
   @Override
   public void close(boolean failFast) throws IOException {
     if (!failFast) {
-      streamedOutputHandler.close();
+      aqueryOutputHandler.close();
     }
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/ActionGraphDump.java

@@ -64,7 +64,7 @@ public class ActionGraphDump {
   private final boolean includeActionCmdLine;
   private final boolean includeArtifacts;
   private final boolean includeParamFiles;
-  private final StreamedOutputHandler streamedOutputHandler;
+  private final AqueryOutputHandler aqueryOutputHandler;
 
   private Map<String, Iterable<String>> paramFileNameToContentMap;
 
@@ -73,14 +73,14 @@ public ActionGraphDump(
       boolean includeArtifacts,
       AqueryActionFilter actionFilters,
       boolean includeParamFiles,
-      StreamedOutputHandler streamedOutputHandler) {
+      AqueryOutputHandler aqueryOutputHandler) {
     this(
         /* actionGraphTargets= */ ImmutableList.of("..."),
         includeActionCmdLine,
         includeArtifacts,
         actionFilters,
         includeParamFiles,
-        streamedOutputHandler);
+        aqueryOutputHandler);
   }
 
   public ActionGraphDump(
@@ -89,21 +89,21 @@ public ActionGraphDump(
       boolean includeArtifacts,
       AqueryActionFilter actionFilters,
       boolean includeParamFiles,
-      StreamedOutputHandler streamedOutputHandler) {
+      AqueryOutputHandler aqueryOutputHandler) {
     this.actionGraphTargets = ImmutableSet.copyOf(actionGraphTargets);
     this.includeActionCmdLine = includeActionCmdLine;
     this.includeArtifacts = includeArtifacts;
     this.actionFilters = actionFilters;
     this.includeParamFiles = includeParamFiles;
-    this.streamedOutputHandler = streamedOutputHandler;
+    this.aqueryOutputHandler = aqueryOutputHandler;
 
-    knownRuleClassStrings = new KnownRuleClassStrings(streamedOutputHandler);
-    knownArtifacts = new KnownArtifacts(streamedOutputHandler);
-    knownConfigurations = new KnownConfigurations(streamedOutputHandler);
-    knownNestedSets = new KnownNestedSets(streamedOutputHandler, knownArtifacts);
-    knownAspectDescriptors = new KnownAspectDescriptors(streamedOutputHandler);
+    knownRuleClassStrings = new KnownRuleClassStrings(aqueryOutputHandler);
+    knownArtifacts = new KnownArtifacts(aqueryOutputHandler);
+    knownConfigurations = new KnownConfigurations(aqueryOutputHandler);
+    knownNestedSets = new KnownNestedSets(aqueryOutputHandler, knownArtifacts);
+    knownAspectDescriptors = new KnownAspectDescriptors(aqueryOutputHandler);
     knownRuleConfiguredTargets =
-        new KnownRuleConfiguredTargets(streamedOutputHandler, knownRuleClassStrings);
+        new KnownRuleConfiguredTargets(aqueryOutputHandler, knownRuleClassStrings);
   }
 
   public ActionKeyContext getActionKeyContext() {
@@ -232,11 +232,7 @@ private void dumpSingleAction(ConfiguredTarget configuredTarget, ActionAnalysisM
       }
     }
 
-    printToOutput(actionBuilder.build());
-  }
-
-  private void printToOutput(AnalysisProtosV2.Action actionProto) throws IOException {
-    streamedOutputHandler.printAction(actionProto);
+    aqueryOutputHandler.outputAction(actionBuilder.build());
   }
 
   public void dumpAspect(AspectValue aspectValue, ConfiguredTargetValue configuredTargetValue)

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/AqueryOutputHandler.java

@@ -0,0 +1,63 @@
+// Copyright 2020 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+package com.google.devtools.build.lib.skyframe.actiongraph.v2;
+
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.Action;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.Artifact;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.AspectDescriptor;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.Configuration;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.DepSetOfFiles;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.PathFragment;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.RuleClass;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.Target;
+import java.io.IOException;
+
+/** Outputs various messages of analysis_v2.proto. */
+public interface AqueryOutputHandler {
+  /** Defines the types of proto output this class can handle. */
+  enum OutputType {
+    BINARY("proto"),
+    TEXT("textproto"),
+    JSON("jsonproto");
+
+    private final String formatName;
+
+    OutputType(String formatName) {
+      this.formatName = formatName;
+    }
+
+    public String formatName() {
+      return formatName;
+    }
+  }
+
+  void outputArtifact(Artifact message) throws IOException;
+
+  void outputAction(Action message) throws IOException;
+
+  void outputTarget(Target message) throws IOException;
+
+  void outputDepSetOfFiles(DepSetOfFiles message) throws IOException;
+
+  void outputConfiguration(Configuration message) throws IOException;
+
+  void outputAspectDescriptor(AspectDescriptor message) throws IOException;
+
+  void outputRuleClass(RuleClass message) throws IOException;
+
+  void outputPathFragment(PathFragment message) throws IOException;
+
+  /** Called at the end of the query process. */
+  void close() throws IOException;
+}

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/BaseCache.java

@@ -22,10 +22,10 @@
  */
 abstract class BaseCache<K, P> {
   private final Map<K, Integer> cache = new HashMap<>();
-  protected final StreamedOutputHandler streamedOutputHandler;
+  protected final AqueryOutputHandler aqueryOutputHandler;
 
-  BaseCache(StreamedOutputHandler streamedOutputHandler) {
-    this.streamedOutputHandler = streamedOutputHandler;
+  BaseCache(AqueryOutputHandler aqueryOutputHandler) {
+    this.aqueryOutputHandler = aqueryOutputHandler;
   }
 
   private int generateNextId() {
@@ -55,12 +55,12 @@ int dataToIdAndStreamOutputProto(K data) throws IOException {
       id = generateNextId();
       cache.put(key, id);
       P proto = createProto(data, id);
-      streamToOutput(proto);
+      toOutput(proto);
     }
     return id;
   }
 
   abstract P createProto(K key, int id) throws IOException;
 
-  abstract void streamToOutput(P proto) throws IOException;
+  abstract void toOutput(P proto) throws IOException;
 }

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/KnownArtifacts.java

@@ -22,9 +22,9 @@ public class KnownArtifacts extends BaseCache<Artifact, AnalysisProtosV2.Artifac
 
   private final KnownPathFragments knownPathFragments;
 
-  KnownArtifacts(StreamedOutputHandler streamedOutputHandler) {
-    super(streamedOutputHandler);
-    knownPathFragments = new KnownPathFragments(streamedOutputHandler);
+  KnownArtifacts(AqueryOutputHandler aqueryOutputHandler) {
+    super(aqueryOutputHandler);
+    knownPathFragments = new KnownPathFragments(aqueryOutputHandler);
   }
 
   @Override
@@ -39,7 +39,7 @@ AnalysisProtosV2.Artifact createProto(Artifact artifact, int id) throws IOExcept
   }
 
   @Override
-  void streamToOutput(AnalysisProtosV2.Artifact artifactProto) throws IOException {
-    streamedOutputHandler.printArtifact(artifactProto);
+  void toOutput(AnalysisProtosV2.Artifact artifactProto) throws IOException {
+    aqueryOutputHandler.outputArtifact(artifactProto);
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/KnownAspectDescriptors.java

@@ -23,8 +23,8 @@
 public class KnownAspectDescriptors
     extends BaseCache<AspectDescriptor, AnalysisProtosV2.AspectDescriptor> {
 
-  KnownAspectDescriptors(StreamedOutputHandler streamedOutputHandler) {
-    super(streamedOutputHandler);
+  KnownAspectDescriptors(AqueryOutputHandler aqueryOutputHandler) {
+    super(aqueryOutputHandler);
   }
 
   @Override
@@ -44,7 +44,7 @@ AnalysisProtosV2.AspectDescriptor createProto(AspectDescriptor aspectDescriptor,
   }
 
   @Override
-  void streamToOutput(AnalysisProtosV2.AspectDescriptor aspectDescriptorProto) throws IOException {
-    streamedOutputHandler.printAspectDescriptor(aspectDescriptorProto);
+  void toOutput(AnalysisProtosV2.AspectDescriptor aspectDescriptorProto) throws IOException {
+    aqueryOutputHandler.outputAspectDescriptor(aspectDescriptorProto);
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/KnownConfigurations.java

@@ -21,8 +21,8 @@
 /** Cache for BuildConfigurations in the action graph. */
 public class KnownConfigurations extends BaseCache<BuildEvent, Configuration> {
 
-  KnownConfigurations(StreamedOutputHandler streamedOutputHandler) {
-    super(streamedOutputHandler);
+  KnownConfigurations(AqueryOutputHandler aqueryOutputHandler) {
+    super(aqueryOutputHandler);
   }
 
   @Override
@@ -38,7 +38,7 @@ Configuration createProto(BuildEvent config, int id) throws IOException {
   }
 
   @Override
-  void streamToOutput(Configuration configurationProto) throws IOException {
-    streamedOutputHandler.printConfiguration(configurationProto);
+  void toOutput(Configuration configurationProto) throws IOException {
+    aqueryOutputHandler.outputConfiguration(configurationProto);
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/KnownNestedSets.java

@@ -22,8 +22,8 @@
 public class KnownNestedSets extends BaseCache<Object, DepSetOfFiles> {
   private final KnownArtifacts knownArtifacts;
 
-  KnownNestedSets(StreamedOutputHandler streamedOutputHandler, KnownArtifacts knownArtifacts) {
-    super(streamedOutputHandler);
+  KnownNestedSets(AqueryOutputHandler aqueryOutputHandler, KnownArtifacts knownArtifacts) {
+    super(aqueryOutputHandler);
     this.knownArtifacts = knownArtifacts;
   }
 
@@ -50,7 +50,7 @@ DepSetOfFiles createProto(Object nestedSetViewObject, int id) throws IOException
   }
 
   @Override
-  void streamToOutput(DepSetOfFiles depSetOfFilesProto) throws IOException {
-    streamedOutputHandler.printDepSetOfFiles(depSetOfFilesProto);
+  void toOutput(DepSetOfFiles depSetOfFilesProto) throws IOException {
+    aqueryOutputHandler.outputDepSetOfFiles(depSetOfFilesProto);
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/KnownPathFragments.java

@@ -19,8 +19,8 @@
 
 /** Cache for {@link PathFragment} in the action graph. */
 public class KnownPathFragments extends BaseCache<PathFragment, AnalysisProtosV2.PathFragment> {
-  KnownPathFragments(StreamedOutputHandler streamedOutputHandler) {
-    super(streamedOutputHandler);
+  KnownPathFragments(AqueryOutputHandler aqueryOutputHandler) {
+    super(aqueryOutputHandler);
   }
 
   @Override
@@ -40,8 +40,8 @@ AnalysisProtosV2.PathFragment createProto(PathFragment pathFragment, int id) thr
   }
 
   @Override
-  void streamToOutput(AnalysisProtosV2.PathFragment pathFragmentProto) throws IOException {
-    streamedOutputHandler.printPathFragment(pathFragmentProto);
+  void toOutput(AnalysisProtosV2.PathFragment pathFragmentProto) throws IOException {
+    aqueryOutputHandler.outputPathFragment(pathFragmentProto);
   }
 
   private static boolean hasParent(PathFragment pathFragment) {

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/KnownRuleClassStrings.java

@@ -19,8 +19,8 @@
 /** Cache for RuleClassStrings in the action graph. */
 public class KnownRuleClassStrings extends BaseCache<String, RuleClass> {
 
-  KnownRuleClassStrings(StreamedOutputHandler streamedOutputHandler) {
-    super(streamedOutputHandler);
+  KnownRuleClassStrings(AqueryOutputHandler aqueryOutputHandler) {
+    super(aqueryOutputHandler);
   }
 
   @Override
@@ -29,7 +29,7 @@ RuleClass createProto(String ruleClassString, int id) throws IOException {
   }
 
   @Override
-  void streamToOutput(RuleClass ruleClassProto) throws IOException {
-    streamedOutputHandler.printRuleClass(ruleClassProto);
+  void toOutput(RuleClass ruleClassProto) throws IOException {
+    aqueryOutputHandler.outputRuleClass(ruleClassProto);
   }
 }

src/main/java/com/google/devtools/build/lib/skyframe/actiongraph/v2/KnownRuleConfiguredTargets.java

@@ -24,8 +24,8 @@ public class KnownRuleConfiguredTargets extends BaseCache<RuleConfiguredTarget,
   private final KnownRuleClassStrings knownRuleClassStrings;
 
   KnownRuleConfiguredTargets(
-      StreamedOutputHandler streamedOutputHandler, KnownRuleClassStrings knownRuleClassStrings) {
-    super(streamedOutputHandler);
+      AqueryOutputHandler aqueryOutputHandler, KnownRuleClassStrings knownRuleClassStrings) {
+    super(aqueryOutputHandler);
     this.knownRuleClassStrings = knownRuleClassStrings;
   }
 
@@ -42,7 +42,7 @@ Target createProto(RuleConfiguredTarget ruleConfiguredTarget, int id) throws IOE
   }
 
   @Override
-  void streamToOutput(Target targetProto) throws IOException {
-    streamedOutputHandler.printTarget(targetProto);
+  void toOutput(Target targetProto) throws IOException {
+    aqueryOutputHandler.outputTarget(targetProto);
   }
 }