Add a #getBytes() method to DeterministicWriter that returns a ByteString. By default it just delegates to the existing #writeOutputFile, but implementations may choose to override if they have easy access to a ByteString.

Also change some DeterministicWriter implementations that do have easy access to the ByteString.

PiperOrigin-RevId: 160550028

10 files changed, 143 insertions, 67 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/actions/cache/VirtualActionInput.java b/src/main/java/com/google/devtools/build/lib/actions/cache/VirtualActionInput.java
index cdf9bcfaee..87432f2c6e 100644
--- a/src/main/java/com/google/devtools/build/lib/actions/cache/VirtualActionInput.java
+++ b/src/main/java/com/google/devtools/build/lib/actions/cache/VirtualActionInput.java
@@ -14,7 +14,7 @@
 package com.google.devtools.build.lib.actions.cache;
 
 import com.google.devtools.build.lib.actions.ActionInput;
-
+import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.io.OutputStream;
 
@@ -24,8 +24,14 @@ import java.io.OutputStream;
  */
 public interface VirtualActionInput extends ActionInput {
   /**
-   * Writes the the fake file to an OutputStream. MUST be deterministic, in that multiple calls
-   * to write the same VirtualActionInput must write identical bytes.
+   * Writes the fake file to an OutputStream. MUST be deterministic, in that multiple calls to write
+   * the same VirtualActionInput must write identical bytes.
    */
   void writeTo(OutputStream out) throws IOException;
+
+  /**
+   * Gets a {@link ByteString} representation of the fake file. Used to avoid copying if the fake
+   * file is internally represented as a {@link ByteString}.
+   */
+  ByteString getBytes() throws IOException;
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/AbstractFileWriteAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/AbstractFileWriteAction.java
index d0050ba7fd..53a03d6308 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/AbstractFileWriteAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/AbstractFileWriteAction.java
@@ -23,6 +23,7 @@ import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ExecException;
 import com.google.devtools.build.lib.cmdline.Label;
+import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.io.OutputStream;
 
@@ -111,6 +112,16 @@ public abstract class AbstractFileWriteAction extends AbstractAction {
    * on every invocation of writeOutputFile().
    */
   public interface DeterministicWriter {
-    public void writeOutputFile(OutputStream out) throws IOException;
+    void writeOutputFile(OutputStream out) throws IOException;
+
+    /**
+     * Returns the contents that would be written, as a {@link ByteString}. Used when the caller
+     * wants a {@link ByteString} in the end, to avoid making unnecessary copies.
+     */
+    default ByteString getBytes() throws IOException {
+      ByteString.Output out = ByteString.newOutput();
+      writeOutputFile(out);
+      return out.toByteString();
+    }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/ByteStringDeterministicWriter.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/ByteStringDeterministicWriter.java
new file mode 100644
index 0000000000..d46bb797ad
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/ByteStringDeterministicWriter.java
@@ -0,0 +1,41 @@
+// Copyright 2017 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.analysis.actions;
+
+import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction.DeterministicWriter;
+import com.google.protobuf.ByteString;
+import java.io.IOException;
+import java.io.OutputStream;
+
+/**
+ * A {@link DeterministicWriter} that wraps a {@link ByteString}. Use to avoid {@link ByteString}
+ * copies.
+ */
+public class ByteStringDeterministicWriter implements DeterministicWriter {
+  private final ByteString byteString;
+
+  public ByteStringDeterministicWriter(ByteString byteString) {
+    this.byteString = byteString;
+  }
+
+  @Override
+  public void writeOutputFile(OutputStream out) throws IOException {
+    byteString.writeTo(out);
+  }
+
+  @Override
+  public ByteString getBytes() throws IOException {
+    return byteString;
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/ProtoDeterministicWriter.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/ProtoDeterministicWriter.java
new file mode 100644
index 0000000000..239a7ea0ae
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/ProtoDeterministicWriter.java
@@ -0,0 +1,39 @@
+// Copyright 2017 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.analysis.actions;
+
+import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction.DeterministicWriter;
+import com.google.protobuf.AbstractMessageLite;
+import com.google.protobuf.ByteString;
+import java.io.IOException;
+import java.io.OutputStream;
+
+/** A {@link DeterministicWriter} wrapping an {@link AbstractMessageLite} object. */
+public class ProtoDeterministicWriter implements DeterministicWriter {
+  private final AbstractMessageLite message;
+
+  public ProtoDeterministicWriter(AbstractMessageLite message) {
+    this.message = message;
+  }
+
+  @Override
+  public void writeOutputFile(OutputStream out) throws IOException {
+    message.writeTo(out);
+  }
+
+  @Override
+  public ByteString getBytes() throws IOException {
+    return message.toByteString();
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionAction.java b/src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionAction.java
index 6fac2d2228..930cf9b9dc 100644
--- a/src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionAction.java
+++ b/src/main/java/com/google/devtools/build/lib/analysis/actions/TemplateExpansionAction.java
@@ -30,8 +30,8 @@ import com.google.devtools.build.lib.util.ResourceFileLoader;
 import com.google.devtools.build.lib.util.StringUtilities;
 import com.google.devtools.build.lib.vfs.FileSystemUtils;
 import com.google.devtools.build.lib.vfs.Path;
+import com.google.protobuf.ByteString;
 import java.io.IOException;
-import java.io.OutputStream;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.Collection;
@@ -353,13 +353,8 @@ public final class TemplateExpansionAction extends AbstractFileWriteAction {
 
   @Override
   public DeterministicWriter newDeterministicWriter(ActionExecutionContext ctx) throws IOException {
-    final byte[] bytes = getFileContents().getBytes(Template.DEFAULT_CHARSET);
-    return new DeterministicWriter() {
-      @Override
-      public void writeOutputFile(OutputStream out) throws IOException {
-        out.write(bytes);
-      }
-    };
+    return new ByteStringDeterministicWriter(
+        ByteString.copyFrom(getFileContents().getBytes(Template.DEFAULT_CHARSET)));
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/remote/EmptyActionInput.java b/src/main/java/com/google/devtools/build/lib/remote/EmptyActionInput.java
index 8d25a0b6be..5d19070b59 100644
--- a/src/main/java/com/google/devtools/build/lib/remote/EmptyActionInput.java
+++ b/src/main/java/com/google/devtools/build/lib/remote/EmptyActionInput.java
@@ -16,6 +16,7 @@ package com.google.devtools.build.lib.remote;
 import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
 import com.google.devtools.build.lib.util.Preconditions;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.protobuf.ByteString;
 import java.io.IOException;
 import java.io.OutputStream;
 
@@ -46,7 +47,12 @@ final class EmptyActionInput implements VirtualActionInput {
   }
 
   @Override
+  public ByteString getBytes() throws IOException {
+    return ByteString.EMPTY;
+  }
+
+  @Override
   public String toString() {
     return "EmptyActionInput: " + execPath;
   }
-}
\ No newline at end of file
+}
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidDeployInfoAction.java b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidDeployInfoAction.java
index 1f38303611..ba4b09dcaf 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/AndroidDeployInfoAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/AndroidDeployInfoAction.java
@@ -16,23 +16,21 @@ package com.google.devtools.build.lib.rules.android;
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
-import com.google.common.io.ByteStreams;
 import com.google.devtools.build.lib.actions.Action;
 import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.RuleContext;
 import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction;
+import com.google.devtools.build.lib.analysis.actions.ByteStringDeterministicWriter;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.rules.android.deployinfo.AndroidDeployInfoOuterClass;
 import com.google.devtools.build.lib.rules.android.deployinfo.AndroidDeployInfoOuterClass.AndroidDeployInfo;
 import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.InvalidProtocolBufferException;
-
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.OutputStream;
 
 /**
  * Writes AndroidDeployInfo proto message. This proto describes how
@@ -96,16 +94,7 @@ public final class AndroidDeployInfoAction extends AbstractFileWriteAction {
 
   @Override
   public DeterministicWriter newDeterministicWriter(ActionExecutionContext ctx) throws IOException {
-
-    return new DeterministicWriter() {
-      @Override
-      public void writeOutputFile(OutputStream out) throws IOException {
-        try (InputStream in = byteString.newInput()) {
-          ByteStreams.copy(in, out);
-        }
-        out.flush();
-      }
-    };
+    return new ByteStringDeterministicWriter(byteString);
   }
 
   @VisibleForTesting
diff --git a/src/main/java/com/google/devtools/build/lib/rules/android/ApkManifestAction.java b/src/main/java/com/google/devtools/build/lib/rules/android/ApkManifestAction.java
index 6461a03db0..15c85583a7 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/android/ApkManifestAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/android/ApkManifestAction.java
@@ -22,6 +22,7 @@ import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.analysis.FilesToRunProvider;
 import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction;
+import com.google.devtools.build.lib.analysis.actions.ProtoDeterministicWriter;
 import com.google.devtools.build.lib.collect.CollectionUtils;
 import com.google.devtools.build.lib.collect.nestedset.NestedSet;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
@@ -118,16 +119,16 @@ public final class ApkManifestAction extends AbstractFileWriteAction {
 
     final ApkManifest manifest = manifestCreator.createManifest();
 
-    return new DeterministicWriter() {
-      @Override
-      public void writeOutputFile(OutputStream out) throws IOException {
-        if (textOutput) {
+    if (textOutput) {
+      return new DeterministicWriter() {
+        @Override
+        public void writeOutputFile(OutputStream out) throws IOException {
           TextFormat.print(manifest, new PrintStream(out));
-        } else {
-          manifest.writeTo(out);
         }
-      }
-    };
+      };
+    } else {
+      return new ProtoDeterministicWriter(manifest);
+    }
   }
 
   @Override
diff --git a/src/main/java/com/google/devtools/build/lib/rules/extra/ExtraActionInfoFileWriteAction.java b/src/main/java/com/google/devtools/build/lib/rules/extra/ExtraActionInfoFileWriteAction.java
index f1db755e8f..36fee0385f 100644
--- a/src/main/java/com/google/devtools/build/lib/rules/extra/ExtraActionInfoFileWriteAction.java
+++ b/src/main/java/com/google/devtools/build/lib/rules/extra/ExtraActionInfoFileWriteAction.java
@@ -19,13 +19,12 @@ import com.google.devtools.build.lib.actions.ActionExecutionContext;
 import com.google.devtools.build.lib.actions.ActionOwner;
 import com.google.devtools.build.lib.actions.Artifact;
 import com.google.devtools.build.lib.actions.ExecException;
-import com.google.devtools.build.lib.actions.extra.ExtraActionInfo;
 import com.google.devtools.build.lib.analysis.actions.AbstractFileWriteAction;
+import com.google.devtools.build.lib.analysis.actions.ProtoDeterministicWriter;
 import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
 import com.google.devtools.build.lib.util.Fingerprint;
 import com.google.devtools.build.lib.util.Preconditions;
 import java.io.IOException;
-import java.io.OutputStream;
 
 /**
  * Requests extra action info from shadowed action and writes it, in protocol buffer format, to an
@@ -48,18 +47,7 @@ public final class ExtraActionInfoFileWriteAction extends AbstractFileWriteActio
   @Override
   public DeterministicWriter newDeterministicWriter(ActionExecutionContext ctx)
       throws IOException, InterruptedException, ExecException {
-    return new DeterministicWriter() {
-      // Instantiate the extra action info only on execution, so it is computed freshly each
-      // execution, but is constant for the lifetime of this action's execution. These are not
-      // large objects, so keeping them in memory for the duration of a single action's execution
-      // is acceptable.
-      private final ExtraActionInfo extraActionInfo = shadowedAction.getExtraActionInfo().build();
-
-      @Override
-      public void writeOutputFile(OutputStream out) throws IOException {
-        extraActionInfo.writeTo(out);
-      }
-    };
+    return new ProtoDeterministicWriter(shadowedAction.getExtraActionInfo().build());
   }
 
   @Override
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 fbd0943744..d967df0956 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
@@ -30,6 +30,7 @@ import com.google.devtools.build.lib.analysis.RuleContext;
 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.cmdline.Label;
 import com.google.devtools.build.lib.cmdline.PackageIdentifier;
 import com.google.devtools.build.lib.cmdline.ResolvedTargets;
@@ -80,9 +81,8 @@ import com.google.devtools.build.skyframe.SkyKey;
 import com.google.devtools.build.skyframe.ValueOrException;
 import com.google.devtools.common.options.OptionsParser;
 import com.google.devtools.common.options.OptionsParsingException;
-import java.io.ByteArrayOutputStream;
+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;
@@ -146,7 +146,7 @@ public class GenQuery implements RuleConfiguredTargetFactory {
     // force relative_locations to true so it has a deterministic output across machines.
     queryOptions.relativeLocations = true;
 
-    final byte[] result = executeQuery(ruleContext, queryOptions, getScope(ruleContext), query);
+    ByteString result = executeQuery(ruleContext, queryOptions, getScope(ruleContext), query);
     if (result == null || ruleContext.hasErrors()) {
       return null;
     }
@@ -252,8 +252,9 @@ public class GenQuery implements RuleConfiguredTargetFactory {
   }
 
   @Nullable
-  private byte[] executeQuery(RuleContext ruleContext, QueryOptions queryOptions,
-      Set<Target> scope, String query) throws InterruptedException {
+  private ByteString executeQuery(
+      RuleContext ruleContext, QueryOptions queryOptions, Set<Target> scope, String query)
+      throws InterruptedException {
     SkyFunction.Environment env = ruleContext.getAnalysisEnvironment().getSkyframeEnv();
     Pair<ImmutableMap<PackageIdentifier, Package>, ImmutableMap<Label, Target>> closureInfo;
     try {
@@ -277,9 +278,13 @@ public class GenQuery implements RuleConfiguredTargetFactory {
 
   @SuppressWarnings("unchecked")
   @Nullable
-  private byte[] doQuery(QueryOptions queryOptions, PackageProvider packageProvider,
-                         Predicate<Label> labelFilter, TargetPatternEvaluator evaluator,
-                         String query, RuleContext ruleContext)
+  private ByteString doQuery(
+      QueryOptions queryOptions,
+      PackageProvider packageProvider,
+      Predicate<Label> labelFilter,
+      TargetPatternEvaluator evaluator,
+      String query,
+      RuleContext ruleContext)
       throws InterruptedException {
 
     DigraphQueryEvalResult<Target> queryResult;
@@ -339,7 +344,7 @@ public class GenQuery implements RuleConfiguredTargetFactory {
       throw new RuntimeException(e);
     }
 
-    ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
+    ByteString.Output outputStream = ByteString.newOutput();
     try {
       QueryOutputUtils
           .output(queryOptions, queryResult, targets.getResult(), formatter, outputStream,
@@ -350,32 +355,27 @@ public class GenQuery implements RuleConfiguredTargetFactory {
       throw new RuntimeException(e);
     }
 
-    return outputStream.toByteArray();
+    return outputStream.toByteString();
   }
 
   @Immutable // assuming no other reference to result
   private static final class QueryResultAction extends AbstractFileWriteAction {
-    private final byte[] result;
+    private final ByteString result;
 
-    private QueryResultAction(ActionOwner owner, Artifact output, byte[] result) {
+    private QueryResultAction(ActionOwner owner, Artifact output, ByteString result) {
       super(owner, ImmutableList.<Artifact>of(), output, /*makeExecutable=*/false);
       this.result = result;
     }
 
     @Override
     public DeterministicWriter newDeterministicWriter(ActionExecutionContext ctx) {
-      return new DeterministicWriter() {
-        @Override
-        public void writeOutputFile(OutputStream out) throws IOException {
-          out.write(result);
-        }
-      };
+      return new ByteStringDeterministicWriter(result);
     }
 
     @Override
     protected String computeKey() {
       Fingerprint f = new Fingerprint();
-      f.addBytes(result);
+      f.addBytes(result.toByteArray());
       return f.hexDigestAndReset();
     }
   }