Add new output format for cquery --output=streamed_proto.

* The current state of the output formats of cquery has a few forms now:
  - (UNCHANGED) `cquery --output=proto|jsonproto|textproto --proto:include_configurations`: A single CqueryResult of the specified `--output` format.
  - (NEW) `cquery --output=streamed_proto --proto:include_configurations`: Multiple length-delimited `CqueryResult` protos each containing a single `ConfiguredTarget` or `Configuration`.
  - (UNCHANGED) `cquery --output=proto|jsonproto|textproto --noproto:include_configurations`: A single `QueryResult` containing the corresponding `Target`(s) in the specified format.
  - (NEW) `cquery --output=streamed_proto --noproto:include_configurations`: Multiple length-delimited `Target` protos.
* Reworked the callback to write directly to a stream (instead of a `CqueryResult.Builder`) when processing a set of `ConfiguredTarget` results (instead of writing only at the end on `close()`).

Fixes #17743.

RELNOTES: Added a new output format for cquery --output=streamed_proto that writes multiple length-delimited CqueryResult protos, each containing a single ConfiguredTarget or Configuration. This allows us to "bypass" the hard limit of 2GB on the size of protocol buffers by splitting it up into multiple.
PiperOrigin-RevId: 540814164
Change-Id: Ia880218e67fdd86876ee56581634e5eabbc899c6

src/main/java/com/google/devtools/build/lib/query2/cquery/ConfiguredTargetQueryEnvironment.java

@@ -53,6 +53,7 @@
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor;
 import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.WalkableGraph;
+import com.google.protobuf.CodedOutputStream;
 import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -77,6 +78,12 @@ public class ConfiguredTargetQueryEnvironment
   /** Cquery specific functions. */
   public static final ImmutableList<QueryFunction> CQUERY_FUNCTIONS = getCqueryFunctions();
 
+  /**
+   * Pseudo-arbitrarily chosen buffer size for output. Chosen to be large enough to fit a handful of
+   * messages without needing to flush to the underlying output, which may not be buffered.
+   */
+  private static final int OUTPUT_BUFFER_SIZE = 16384;
+
   private CqueryOptions cqueryOptions;
 
   private final TopLevelArtifactContext topLevelArtifactContext;
@@ -224,6 +231,7 @@ private static ImmutableMap<String, BuildConfigurationValue> getTransitiveConfig
             eventHandler,
             cqueryOptions,
             out,
+            CodedOutputStream.newInstance(out, OUTPUT_BUFFER_SIZE),
             skyframeExecutor,
             accessor,
             aspectResolver,
@@ -233,6 +241,17 @@ private static ImmutableMap<String, BuildConfigurationValue> getTransitiveConfig
             eventHandler,
             cqueryOptions,
             out,
+            CodedOutputStream.newInstance(out, OUTPUT_BUFFER_SIZE),
+            skyframeExecutor,
+            accessor,
+            aspectResolver,
+            OutputType.DELIMITED_BINARY,
+            ruleClassProvider),
+        new ProtoOutputFormatterCallback(
+            eventHandler,
+            cqueryOptions,
+            out,
+            CodedOutputStream.newInstance(out, OUTPUT_BUFFER_SIZE),
             skyframeExecutor,
             accessor,
             aspectResolver,
@@ -242,6 +261,7 @@ private static ImmutableMap<String, BuildConfigurationValue> getTransitiveConfig
             eventHandler,
             cqueryOptions,
             out,
+            CodedOutputStream.newInstance(out, OUTPUT_BUFFER_SIZE),
             skyframeExecutor,
             accessor,
             aspectResolver,

src/main/java/com/google/devtools/build/lib/query2/cquery/CqueryOptions.java

@@ -45,8 +45,9 @@ public enum Transitions {
       effectTags = {OptionEffectTag.TERMINAL_OUTPUT},
       help =
           "The format in which the cquery results should be printed. Allowed values for cquery "
-              + "are: label, label_kind, textproto, transitions, proto, jsonproto. If you select "
-              + "'transitions', you also have to specify the --transitions=(lite|full) option.")
+              + "are: label, label_kind, textproto, transitions, proto, streamed_proto, jsonproto. "
+              + "If you select 'transitions', you also have to specify the "
+              + "--transitions=(lite|full) option.")
   public String outputFormat;
 
   @Option(

src/main/java/com/google/devtools/build/lib/query2/cquery/ProtoOutputFormatterCallback.java

@@ -15,14 +15,15 @@
 
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
-import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Ordering;
 import com.google.devtools.build.lib.actions.BuildConfigurationEvent;
 import com.google.devtools.build.lib.analysis.AnalysisProtosV2;
 import com.google.devtools.build.lib.analysis.AnalysisProtosV2.Configuration;
+import com.google.devtools.build.lib.analysis.AnalysisProtosV2.CqueryResult;
 import com.google.devtools.build.lib.analysis.ConfiguredTarget;
 import com.google.devtools.build.lib.analysis.config.BuildOptions;
 import com.google.devtools.build.lib.analysis.config.ConfigMatchingProvider;
@@ -39,11 +40,13 @@
 import com.google.devtools.build.lib.query2.cquery.CqueryTransitionResolver.EvaluateException;
 import com.google.devtools.build.lib.query2.engine.QueryEnvironment.TargetAccessor;
 import com.google.devtools.build.lib.query2.proto.proto2api.Build;
+import com.google.devtools.build.lib.query2.proto.proto2api.Build.QueryResult;
 import com.google.devtools.build.lib.query2.query.aspectresolvers.AspectResolver;
 import com.google.devtools.build.lib.query2.query.output.ProtoOutputFormatter;
 import com.google.devtools.build.lib.skyframe.BuildConfigurationKey;
 import com.google.devtools.build.lib.skyframe.ConfiguredTargetKey;
 import com.google.devtools.build.lib.skyframe.SkyframeExecutor;
+import com.google.protobuf.CodedOutputStream;
 import com.google.protobuf.Message;
 import com.google.protobuf.TextFormat;
 import com.google.protobuf.util.JsonFormat;
@@ -60,6 +63,7 @@ class ProtoOutputFormatterCallback extends CqueryThreadsafeCallback {
   /** Defines the types of proto output this class can handle. */
   public enum OutputType {
     BINARY("proto"),
+    DELIMITED_BINARY("streamed_proto"),
     TEXT("textproto"),
     JSON("jsonproto");
 
@@ -102,61 +106,71 @@ public ImmutableList<Configuration> getConfigurations() {
     }
   }
 
+  private final CodedOutputStream codedOut;
   private final OutputType outputType;
   private final AspectResolver resolver;
   private final SkyframeExecutor skyframeExecutor;
   private final ConfigurationCache configurationCache = new ConfigurationCache();
   private final JsonFormat.Printer jsonPrinter = JsonFormat.printer();
   private final RuleClassProvider ruleClassProvider;
 
-  private AnalysisProtosV2.CqueryResult.Builder protoResult;
-
   private final Map<Label, Target> partialResultMap;
   private ConfiguredTarget currentTarget;
 
   ProtoOutputFormatterCallback(
       ExtendedEventHandler eventHandler,
       CqueryOptions options,
       OutputStream out,
+      CodedOutputStream codedOut,
       SkyframeExecutor skyframeExecutor,
       TargetAccessor<ConfiguredTarget> accessor,
       AspectResolver resolver,
       OutputType outputType,
       RuleClassProvider ruleClassProvider) {
     super(eventHandler, options, out, skyframeExecutor, accessor, /*uniquifyResults=*/ false);
+    this.codedOut = codedOut;
     this.outputType = outputType;
     this.skyframeExecutor = skyframeExecutor;
     this.resolver = resolver;
     this.ruleClassProvider = ruleClassProvider;
     this.partialResultMap = Maps.newHashMap();
   }
 
-  @Override
-  public void start() {
-    protoResult = AnalysisProtosV2.CqueryResult.newBuilder();
-  }
-
   @Override
   public void close(boolean failFast) throws IOException {
-    if (!failFast && printStream != null) {
-      if (options.protoIncludeConfigurations) {
-        writeData(getProtoResult());
-      } else {
-        // Documentation promises that setting this flag to false means we convert directly
-        // to the build.proto format. This is hard to test in integration testing due to the way
-        // proto output is turned readable (codex). So change the following code with caution.
-        Build.QueryResult.Builder queryResult = Build.QueryResult.newBuilder();
-        protoResult.getResultsList().forEach(ct -> queryResult.addTarget(ct.getTarget()));
-        writeData(queryResult.build());
+    if (failFast || printStream == null) {
+      return;
+    }
+    if (options.protoIncludeConfigurations) {
+      for (Configuration configuration : configurationCache.getConfigurations()) {
+        if (outputType == OutputType.DELIMITED_BINARY) {
+          // For streamed protos, we wrap each Configuration in its own CqueryResult that will be
+          // written length delimited to the stream.
+          writeData(
+              AnalysisProtosV2.CqueryResult.newBuilder().addConfigurations(configuration).build());
+        } else {
+          writeData(configuration, CqueryResult.CONFIGURATIONS_FIELD_NUMBER);
+        }
       }
-      printStream.flush();
     }
+    codedOut.flush();
+    outputStream.flush();
+    printStream.flush();
   }
 
   private void writeData(Message message) throws IOException {
+    writeData(message, 0);
+  }
+
+  private void writeData(Message message, int fieldNumber) throws IOException {
     switch (outputType) {
       case BINARY:
-        message.writeTo(outputStream);
+        Preconditions.checkState(
+            fieldNumber != 0, "Cannot have fieldNumber of 0 when outputType is BINARY");
+        codedOut.writeMessage(fieldNumber, message);
+        break;
+      case DELIMITED_BINARY:
+        message.writeDelimitedTo(outputStream);
         break;
       case TEXT:
         TextFormat.printer().print(message, printStream);
@@ -175,14 +189,9 @@ public String getName() {
     return outputType.formatName();
   }
 
-  @VisibleForTesting
-  public AnalysisProtosV2.CqueryResult getProtoResult() {
-    protoResult.addAllConfigurations(configurationCache.getConfigurations());
-    return protoResult.build();
-  }
-
   @Override
-  public void processOutput(Iterable<ConfiguredTarget> partialResult) throws InterruptedException {
+  public void processOutput(Iterable<ConfiguredTarget> partialResult)
+      throws InterruptedException, IOException {
     partialResult.forEach(
         kct -> partialResultMap.put(kct.getOriginalLabel(), accessor.getTarget(kct)));
 
@@ -239,6 +248,7 @@ public void processOutput(Iterable<ConfiguredTarget> partialResult) throws Inter
           throw new InterruptedException(e.getMessage());
         }
       }
+
       builder.setTarget(targetBuilder);
 
       if (options.protoIncludeConfigurations) {
@@ -259,7 +269,33 @@ public void processOutput(Iterable<ConfiguredTarget> partialResult) throws Inter
         }
       }
 
-      protoResult.addResults(builder.build());
+      // There are a few cases that affect the shape of the output:
+      //   1. --output=proto|textproto|jsonproto --proto:include_configurations =>
+      //        Writes a single CqueryResult containing all the ConfiguredTarget(s) and
+      //        Configuration(s) in the specified output format.
+      //   2. --output=streamed_proto --proto:include_configurations =>
+      //        Writes multiple length delimited CqueryResult protos, each containing a single
+      //        ConfiguredTarget or Configuration.
+      //   3. --output=proto|textproto|jsonproto --noproto:include_configurations =>
+      //        Writes a single QueryResult containing all the corresponding Target(s) in the
+      //        specified output format.
+      //   4.--output=streamed_proto --noproto:include_configurations =>
+      //        Writes multiple length delimited Target protos.
+      if (options.protoIncludeConfigurations) {
+        if (outputType == OutputType.DELIMITED_BINARY) {
+          // Case 2.
+          writeData(AnalysisProtosV2.CqueryResult.newBuilder().addResults(builder).build());
+        } else {
+          // Case 1.
+          writeData(builder.build(), CqueryResult.RESULTS_FIELD_NUMBER);
+        }
+      } else {
+        // Case 3 & 4.
+        // Documentation promises that setting this flag to false means we convert directly
+        // to the build.proto format. This is hard to test in integration testing due to the way
+        // proto output is turned readable (codex). So change the following code with caution.
+        writeData(builder.build().getTarget(), QueryResult.TARGET_FIELD_NUMBER);
+      }
     }
   }
 

src/test/java/com/google/devtools/build/lib/query2/cquery/BUILD

@@ -137,12 +137,14 @@ java_test(
         "//src/main/java/com/google/devtools/build/lib/query2/engine",
         "//src/main/java/com/google/devtools/build/lib/query2/query/aspectresolvers",
         "//src/main/java/com/google/devtools/build/lib/util:filetype",
+        "//src/main/java/com/google/devtools/build/lib/util:pair",
         "//src/main/protobuf:analysis_v2_java_proto",
         "//src/main/protobuf:build_java_proto",
         "//src/test/java/com/google/devtools/build/lib/analysis/util",
         "//third_party:guava",
         "//third_party:junit4",
         "//third_party:truth",
+        "//third_party/protobuf:protobuf_java",
     ],
 )