From 6d63f8a1ad33490b97e60955277c4ebca40a4f48 Mon Sep 17 00:00:00 2001
From: michajlo <michajlo@google.com>
Date: Wed, 17 Jul 2019 07:55:43 -0700
Subject: [PATCH] Add (inactive) code to compress large GenQuery results in
 memory on the fly

Dynamically switches to gzipping GenQuery results >1M, trading CPU for memory.
Hard coded to off for now while I sort out how to connect this to a flag.

PiperOrigin-RevId: 258568242
---
 .../build/lib/rules/genquery/GenQuery.java    |  56 +++--
 .../rules/genquery/GenQueryOutputStream.java  | 210 ++++++++++++++++++
 .../genquery/GenQueryOutputStreamTest.java    | 141 ++++++++++++
 3 files changed, 392 insertions(+), 15 deletions(-)
 create mode 100644 src/main/java/com/google/devtools/build/lib/rules/genquery/GenQueryOutputStream.java
 create mode 100644 src/test/java/com/google/devtools/build/lib/rules/genquery/GenQueryOutputStreamTest.java

diff --git a/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java b/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java
index 118e783acf7444..190aa3930c61a0 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQuery.java
@@ -36,7 +36,7 @@
 import com.google.devtools.build.lib.analysis.Runfiles;
 import com.google.devtools.build.lib.analysis.RunfilesProvider;
 import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction;
-import com.google.devtools.build.lib.analysis.actions.ByteStringDeterministicWriter;
+import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction.DeterministicWriter;
 import com.google.devtools.build.lib.analysis.config.CoreOptions;
 import com.google.devtools.build.lib.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
@@ -73,6 +73,7 @@
 import com.google.devtools.build.lib.query2.query.output.QueryOptions;
 import com.google.devtools.build.lib.query2.query.output.QueryOptions.OrderOutput;
 import com.google.devtools.build.lib.query2.query.output.QueryOutputUtils;
+import com.google.devtools.build.lib.rules.genquery.GenQueryOutputStream.GenQueryResult;
 import com.google.devtools.build.lib.runtime.KeepGoingOption;
 import com.google.devtools.build.lib.skyframe.PackageValue;
 import com.google.devtools.build.lib.skyframe.TargetPatternValue;
@@ -92,6 +93,7 @@
 import com.google.devtools.common.options.OptionsParsingException;
 import com.google.protobuf.ByteString;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.nio.channels.ClosedByInterruptException;
 import java.util.Collection;
 import java.util.HashSet;
@@ -117,6 +119,7 @@ public ConfiguredTarget create(RuleContext ruleContext)
     // The query string
     final String query = ruleContext.attributes().get("expression", Type.STRING);
 
+    @SuppressWarnings("unchecked")
     OptionsParser optionsParser =
         OptionsParser.builder()
             .optionsClasses(QueryOptions.class, KeepGoingOption.class)
@@ -169,7 +172,7 @@ public ConfiguredTarget create(RuleContext ruleContext)
     // force relative_locations to true so it has a deterministic output across machines.
     queryOptions.relativeLocations = true;
 
-    ByteString result;
+    GenQueryResult result;
     try (SilentCloseable c =
         Profiler.instance().profile("GenQuery.executeQuery/" + ruleContext.getLabel())) {
       result =
@@ -193,11 +196,14 @@ public ConfiguredTarget create(RuleContext ruleContext)
     NestedSet<Artifact> filesToBuild = NestedSetBuilder.create(Order.STABLE_ORDER, outputArtifact);
     return new RuleConfiguredTargetBuilder(ruleContext)
         .setFilesToBuild(filesToBuild)
-        .add(RunfilesProvider.class, RunfilesProvider.simple(
-            new Runfiles.Builder(
-                ruleContext.getWorkspaceName(),
-                ruleContext.getConfiguration().legacyExternalRunfiles())
-                .addTransitiveArtifacts(filesToBuild).build()))
+        .addProvider(
+            RunfilesProvider.class,
+            RunfilesProvider.simple(
+                new Runfiles.Builder(
+                        ruleContext.getWorkspaceName(),
+                        ruleContext.getConfiguration().legacyExternalRunfiles())
+                    .addTransitiveArtifacts(filesToBuild)
+                    .build()))
         .build();
   }
 
@@ -273,7 +279,7 @@ private ExtendedEventHandler getEventHandler(RuleContext ruleContext) {
   }
 
   @Nullable
-  private ByteString executeQuery(
+  private GenQueryResult executeQuery(
       RuleContext ruleContext, QueryOptions queryOptions, Collection<Label> scope, String query)
       throws InterruptedException {
     SkyFunction.Environment env = ruleContext.getAnalysisEnvironment().getSkyframeEnv();
@@ -300,7 +306,7 @@ private ByteString executeQuery(
 
   @SuppressWarnings("unchecked")
   @Nullable
-  private ByteString doQuery(
+  private GenQueryResult doQuery(
       QueryOptions queryOptions,
       PreloadedMapPackageProvider packageProvider,
       Predicate<Label> labelFilter,
@@ -372,37 +378,39 @@ private ByteString doQuery(
       throw new RuntimeException(e);
     }
 
-    ByteString.Output outputStream = ByteString.newOutput();
+    // TODO(b/137379942): Enable compression.
+    GenQueryOutputStream outputStream = new GenQueryOutputStream(/*compressionEnabled=*/ false);
     try {
       QueryOutputUtils
           .output(queryOptions, queryResult, targets.getResult(), formatter, outputStream,
           queryOptions.aspectDeps.createResolver(packageProvider, getEventHandler(ruleContext)));
+      outputStream.close();
     } catch (ClosedByInterruptException e) {
       throw new InterruptedException(e.getMessage());
     } catch (IOException e) {
       throw new RuntimeException(e);
     }
 
-    return outputStream.toByteString();
+    return outputStream.getResult();
   }
 
   @Immutable // assuming no other reference to result
   private static final class QueryResultAction extends AbstractFileWriteAction {
-    private final ByteString result;
+    private final GenQueryResult result;
 
-    private QueryResultAction(ActionOwner owner, Artifact output, ByteString result) {
+    private QueryResultAction(ActionOwner owner, Artifact output, GenQueryResult result) {
       super(owner, ImmutableList.<Artifact>of(), output, /*makeExecutable=*/false);
       this.result = result;
     }
 
     @Override
     public DeterministicWriter newDeterministicWriter(ActionExecutionContext ctx) {
-      return new ByteStringDeterministicWriter(result);
+      return new GenQueryResultWriter(result);
     }
 
     @Override
     protected void computeKey(ActionKeyContext actionKeyContext, Fingerprint fp) {
-      fp.addBytes(result);
+      result.fingerprint(fp);
     }
   }
 
@@ -575,4 +583,22 @@ public BrokenQueryScopeException(String message) {
       super(message);
     }
   }
+
+  private static class GenQueryResultWriter implements DeterministicWriter {
+    private final GenQueryResult genQueryResult;
+
+    GenQueryResultWriter(GenQueryResult genQueryResult) {
+      this.genQueryResult = genQueryResult;
+    }
+
+    @Override
+    public void writeOutputFile(OutputStream out) throws IOException {
+      genQueryResult.writeTo(out);
+    }
+
+    @Override
+    public ByteString getBytes() throws IOException {
+      return genQueryResult.getBytes();
+    }
+  }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQueryOutputStream.java b/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQueryOutputStream.java
new file mode 100644
index 00000000000000..0e1b36ac847058
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/rules/genquery/GenQueryOutputStream.java
@@ -0,0 +1,210 @@
+// Copyright 2019 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.rules.genquery;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import com.google.common.io.ByteStreams;
+import com.google.devtools.build.lib.util.Fingerprint;
+import com.google.protobuf.ByteString;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.zip.GZIPInputStream;
+import java.util.zip.GZIPOutputStream;
+
+/**
+ * {@link OutputStream} implementation optimized for {@link GenQuery} by (optionally) compressing
+ * query results on the fly. Produces {@link GenQueryResult}s which are preferred for storing the
+ * output of {@link GenQuery}'s underlying queries.
+ */
+class GenQueryOutputStream extends OutputStream {
+
+  /**
+   * When compression is enabled, the threshold at which the stream will switch to compressing
+   * output. The value of this constant is arbitrary but effective.
+   */
+  private static final int COMPRESSION_THRESHOLD = 1 << 20;
+
+  /**
+   * Encapsulates the output of a {@link GenQuery}'s query. CPU and memory overhead of individual
+   * methods depends on the underlying content and settings.
+   */
+  interface GenQueryResult {
+    /** Returns the query output as a {@link ByteString}. */
+    ByteString getBytes() throws IOException;
+
+    /**
+     * Adds the query output to the supplied {@link Fingerprint}. Equivalent to {@code
+     * fingerprint.addBytes(genQueryResult.getBytes())}, but potentially more efficient.
+     */
+    void fingerprint(Fingerprint fingerprint);
+
+    /**
+     * Returns the size of the output. This must be a constant-time operation for all
+     * implementations.
+     */
+    int size();
+
+    /**
+     * Writes the query output to the provided {@link OutputStream}. Equivalent to {@code
+     * genQueryResult.getBytes().writeTo(out)}, but potentially more efficient.
+     */
+    void writeTo(OutputStream out) throws IOException;
+  }
+
+  private final boolean compressionEnabled;
+  private int bytesWritten = 0;
+  private boolean compressed = false;
+  private boolean closed = false;
+  private ByteString.Output bytesOut = ByteString.newOutput();
+  private OutputStream out = bytesOut;
+
+  GenQueryOutputStream(boolean compressionEnabled) {
+    this.compressionEnabled = compressionEnabled;
+  }
+
+  @Override
+  public void write(int b) throws IOException {
+    maybeStartCompression(1);
+    out.write(b);
+    bytesWritten += 1;
+  }
+
+  @Override
+  public void write(byte[] bytes) throws IOException {
+    write(bytes, 0, bytes.length);
+  }
+
+  @Override
+  public void write(byte[] bytes, int off, int len) throws IOException {
+    maybeStartCompression(len);
+    out.write(bytes, off, len);
+    bytesWritten += len;
+  }
+
+  @Override
+  public void flush() throws IOException {
+    out.flush();
+  }
+
+  @Override
+  public void close() throws IOException {
+    out.close();
+    closed = true;
+  }
+
+  GenQueryResult getResult() {
+    Preconditions.checkState(closed, "Must be closed");
+    return compressed
+        ? new CompressedResult(bytesOut.toByteString(), bytesWritten)
+        : new RegularResult(bytesOut.toByteString());
+  }
+
+  private void maybeStartCompression(int additionalBytes) throws IOException {
+    if (!compressionEnabled) {
+      return;
+    }
+
+    if (compressed) {
+      return;
+    }
+
+    if (bytesWritten + additionalBytes < COMPRESSION_THRESHOLD) {
+      return;
+    }
+
+    ByteString.Output compressedBytesOut = ByteString.newOutput();
+    GZIPOutputStream gzipOut = new GZIPOutputStream(compressedBytesOut);
+    bytesOut.writeTo(gzipOut);
+    bytesOut = compressedBytesOut;
+    out = gzipOut;
+    compressed = true;
+  }
+
+  @VisibleForTesting
+  static class RegularResult implements GenQueryResult {
+    private final ByteString data;
+
+    RegularResult(ByteString data) {
+      this.data = data;
+    }
+
+    @Override
+    public ByteString getBytes() {
+      return data;
+    }
+
+    @Override
+    public int size() {
+      return data.size();
+    }
+
+    @Override
+    public void fingerprint(Fingerprint fingerprint) {
+      fingerprint.addBytes(data);
+    }
+
+    @Override
+    public void writeTo(OutputStream out) throws IOException {
+      data.writeTo(out);
+    }
+  }
+
+  @VisibleForTesting
+  static class CompressedResult implements GenQueryResult {
+    private final ByteString compressedData;
+    private final int size;
+
+    CompressedResult(ByteString compressedData, int size) {
+      this.compressedData = compressedData;
+      this.size = size;
+    }
+
+    @Override
+    public ByteString getBytes() throws IOException {
+      ByteString.Output out = ByteString.newOutput(size);
+      try (GZIPInputStream gzipIn = new GZIPInputStream(compressedData.newInput())) {
+        ByteStreams.copy(gzipIn, out);
+      }
+      return out.toByteString();
+    }
+
+    @Override
+    public int size() {
+      return size;
+    }
+
+    @Override
+    public void writeTo(OutputStream out) throws IOException {
+      try (GZIPInputStream gzipIn = new GZIPInputStream(compressedData.newInput())) {
+        ByteStreams.copy(gzipIn, out);
+      }
+    }
+
+    @Override
+    public void fingerprint(Fingerprint fingerprint) {
+      try (GZIPInputStream gzipIn = new GZIPInputStream(compressedData.newInput())) {
+        byte[] chunk = new byte[4092];
+        int bytesRead;
+        while ((bytesRead = gzipIn.read(chunk)) > 0) {
+          fingerprint.addBytes(chunk, 0, bytesRead);
+        }
+      } catch (IOException e) {
+        // Unexpected, everything should be in memory!
+        throw new IllegalStateException("Unexpected IOException", e);
+      }
+    }
+  }
+}
diff --git a/src/test/java/com/google/devtools/build/lib/rules/genquery/GenQueryOutputStreamTest.java b/src/test/java/com/google/devtools/build/lib/rules/genquery/GenQueryOutputStreamTest.java
new file mode 100644
index 00000000000000..e69d5393e36f9c
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/rules/genquery/GenQueryOutputStreamTest.java
@@ -0,0 +1,141 @@
+// Copyright 2019 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.rules.genquery;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.common.base.Strings;
+import com.google.devtools.build.lib.rules.genquery.GenQueryOutputStream.GenQueryResult;
+import com.google.devtools.build.lib.util.Fingerprint;
+import com.google.protobuf.ByteString;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link GenQueryOutputStream}. */
+@RunWith(JUnit4.class)
+public class GenQueryOutputStreamTest {
+
+  @Test
+  public void testSmallOutputMultibyteWriteWithCompressionEnabled() throws IOException {
+    runMultibyteWriteTest(
+        Strings.repeat("xyz", 10_000),
+        /*compressionEnabled=*/ true,
+        GenQueryOutputStream.RegularResult.class);
+  }
+
+  @Test
+  public void testSmallOutputMultibyteWriteWithCompressionDisabled() throws IOException {
+    runMultibyteWriteTest(
+        Strings.repeat("xyz", 10_000),
+        /*compressionEnabled=*/ false,
+        GenQueryOutputStream.RegularResult.class);
+  }
+
+  @Test
+  public void testBigOutputMultibyteWriteWithCompressionEnabled() throws IOException {
+    runMultibyteWriteTest(
+        Strings.repeat("xyz", 1_000_000),
+        /*compressionEnabled=*/ true,
+        GenQueryOutputStream.CompressedResult.class);
+  }
+
+  @Test
+  public void testBigOutputMultibyteWriteWithCompressionDisabled() throws IOException {
+    runMultibyteWriteTest(
+        Strings.repeat("xyz", 1_000_000),
+        /*compressionEnabled=*/ false,
+        GenQueryOutputStream.RegularResult.class);
+  }
+
+  @Test
+  public void testSmallOutputSingleByteWritesWithCompressionEnabled() throws IOException {
+    runSingleByteWriteTest(
+        Strings.repeat("xyz", 10_000),
+        /*compressionEnabled=*/ true,
+        GenQueryOutputStream.RegularResult.class);
+  }
+
+  @Test
+  public void testSmallOutputSingleByteWritesWithCompressionDisabled() throws IOException {
+    runSingleByteWriteTest(
+        Strings.repeat("xyz", 10_000),
+        /*compressionEnabled=*/ false,
+        GenQueryOutputStream.RegularResult.class);
+  }
+
+  @Test
+  public void testBigOutputSingleByteWritesWithCompressionEnabled() throws IOException {
+    runSingleByteWriteTest(
+        Strings.repeat("xyz", 1_000_000),
+        /*compressionEnabled=*/ true,
+        GenQueryOutputStream.CompressedResult.class);
+  }
+
+  @Test
+  public void testBigOutputSingleByteWritesWithCompressionDisabled() throws IOException {
+    runSingleByteWriteTest(
+        Strings.repeat("xyz", 1_000_000),
+        /*compressionEnabled=*/ false,
+        GenQueryOutputStream.RegularResult.class);
+  }
+
+  private static void runMultibyteWriteTest(
+      String data, boolean compressionEnabled, Class<? extends GenQueryResult> resultClass)
+      throws IOException {
+    GenQueryOutputStream underTest = new GenQueryOutputStream(compressionEnabled);
+    underTest.write(data.getBytes(StandardCharsets.UTF_8));
+    underTest.close();
+
+    GenQueryOutputStream.GenQueryResult result = underTest.getResult();
+    assertThat(result).isInstanceOf(resultClass);
+    assertThat(result.getBytes()).isEqualTo(ByteString.copyFromUtf8(data));
+    assertThat(result.size()).isEqualTo(data.length());
+
+    Fingerprint fingerprint = new Fingerprint();
+    result.fingerprint(fingerprint);
+    assertThat(fingerprint.hexDigestAndReset()).isEqualTo(Fingerprint.getHexDigest(data));
+
+    ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
+    result.writeTo(bytesOut);
+    assertThat(new String(bytesOut.toByteArray(), StandardCharsets.UTF_8)).isEqualTo(data);
+  }
+
+  private static void runSingleByteWriteTest(
+      String data, boolean compressionEnabled, Class<? extends GenQueryResult> resultClass)
+      throws IOException {
+    GenQueryOutputStream underTest = new GenQueryOutputStream(compressionEnabled);
+    for (byte b : data.getBytes(StandardCharsets.UTF_8)) {
+      underTest.write(b);
+    }
+    underTest.close();
+
+    GenQueryOutputStream.GenQueryResult result = underTest.getResult();
+    assertThat(result).isInstanceOf(resultClass);
+    assertThat(result.getBytes()).isEqualTo(ByteString.copyFromUtf8(data));
+    assertThat(result.size()).isEqualTo(data.length());
+
+    Fingerprint fingerprint = new Fingerprint();
+    result.fingerprint(fingerprint);
+    assertThat(fingerprint.hexDigestAndReset()).isEqualTo(Fingerprint.getHexDigest(data));
+
+    ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
+    result.writeTo(bytesOut);
+    assertThat(new String(bytesOut.toByteArray(), StandardCharsets.UTF_8)).isEqualTo(data);
+  }
+}