Package serialization operates on Coded{Input,Output}Streams

--
MOS_MIGRATED_REVID=108056940

2 files changed, 42 insertions, 33 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/packages/PackageDeserializer.java b/src/main/java/com/google/devtools/build/lib/packages/PackageDeserializer.java
index de36382811..2a57d5e567 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/PackageDeserializer.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/PackageDeserializer.java
@@ -44,9 +44,10 @@ import com.google.devtools.build.lib.syntax.GlobList;
 import com.google.devtools.build.lib.syntax.Type;
 import com.google.devtools.build.lib.vfs.Path;
 import com.google.devtools.build.lib.vfs.PathFragment;
+import com.google.protobuf.CodedInputStream;
+import com.google.protobuf.ExtensionRegistryLite;
 
 import java.io.IOException;
-import java.io.InputStream;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -371,7 +372,7 @@ public class PackageDeserializer {
    * @throws InterruptedException
    */
   private void deserializeInternal(Build.Package packagePb, StoredEventHandler eventHandler,
-      Package.Builder builder, InputStream in)
+      Package.Builder builder, CodedInputStream codedIn)
       throws PackageDeserializationException, IOException, InterruptedException {
     Path buildFile = packageDeserializationEnvironment.getPath(packagePb.getBuildFilePath());
     Preconditions.checkNotNull(buildFile);
@@ -433,13 +434,13 @@ public class PackageDeserializer {
 
     builder.setWorkspaceName(packagePb.getWorkspaceName());
 
-    deserializeTargets(in, context);
+    deserializeTargets(codedIn, context);
   }
 
-  private void deserializeTargets(InputStream in, DeserializationContext context)
+  private void deserializeTargets(CodedInputStream codedIn, DeserializationContext context)
       throws IOException, PackageDeserializationException, InterruptedException {
     Build.TargetOrTerminator tot;
-    while (!(tot = Build.TargetOrTerminator.parseDelimitedFrom(in)).getIsTerminator()) {
+    while (!(tot = readTargetOrTerminator(codedIn)).getIsTerminator()) {
       Build.Target target = tot.getTarget();
       switch (target.getType()) {
         case SOURCE_FILE:
@@ -457,37 +458,44 @@ public class PackageDeserializer {
     }
   }
 
+  private static Build.TargetOrTerminator readTargetOrTerminator(CodedInputStream codedIn)
+      throws IOException {
+    Build.TargetOrTerminator.Builder builder = Build.TargetOrTerminator.newBuilder();
+    codedIn.readMessage(builder, ExtensionRegistryLite.getEmptyRegistry());
+    return builder.build();
+  }
+
   /**
    * Deserializes a {@link Package} from {@code in}. The inverse of
    * {@link PackageSerializer#serialize}.
    *
-   * <p>Expects {@code in} to contain a single
+   * <p>Expects {@code codedIn} to contain a single
    * {@link com.google.devtools.build.lib.query2.proto.proto2api.Build.Package} message followed
    * by a series of
    * {@link com.google.devtools.build.lib.query2.proto.proto2api.Build.TargetOrTerminator}
    * messages encoding the associated targets.
    *
-   * @param in stream to read from
+   * @param codedIn stream to read from
    * @return a new {@link Package} as read from {@code in}
    * @throws PackageDeserializationException on failures deserializing the input
    * @throws IOException on failures reading from {@code in}
    * @throws InterruptedException
    */
-  public Package deserialize(InputStream in)
+  public Package deserialize(CodedInputStream codedIn)
       throws PackageDeserializationException, IOException, InterruptedException {
     try {
-      return deserializeInternal(in);
+      return deserializeInternal(codedIn);
     } catch (PackageDeserializationException | RuntimeException e) {
       LOG.log(Level.WARNING, "Failed to deserialize Package object", e);
       throw e;
     }
   }
 
-  private Package deserializeInternal(InputStream in)
+  private Package deserializeInternal(CodedInputStream codedIn)
       throws PackageDeserializationException, IOException, InterruptedException {
     // Read the initial Package message so we have the data to initialize the builder. We will read
     // the Targets in individually later.
-    Build.Package packagePb = Build.Package.parseDelimitedFrom(in);
+    Build.Package packagePb = readPackageProto(codedIn);
     Package.Builder builder;
     try {
       builder = new Package.Builder(
@@ -498,11 +506,17 @@ public class PackageDeserializer {
       throw new PackageDeserializationException(e);
     }
     StoredEventHandler eventHandler = new StoredEventHandler();
-    deserializeInternal(packagePb, eventHandler, builder, in);
+    deserializeInternal(packagePb, eventHandler, builder, codedIn);
     builder.addEvents(eventHandler.getEvents());
     return builder.build();
   }
 
+  private static Build.Package readPackageProto(CodedInputStream codedIn) throws IOException {
+    Build.Package.Builder builder = Build.Package.newBuilder();
+    codedIn.readMessage(builder, ExtensionRegistryLite.getEmptyRegistry());
+    return builder.build();
+  }
+
   private static void deserializeEvent(StoredEventHandler eventHandler, Build.Event event) {
     String message = event.getMessage();
     switch (event.getKind()) {
diff --git a/src/main/java/com/google/devtools/build/lib/packages/PackageSerializer.java b/src/main/java/com/google/devtools/build/lib/packages/PackageSerializer.java
index 2462423c7f..899651277a 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/PackageSerializer.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/PackageSerializer.java
@@ -45,9 +45,9 @@ import com.google.devtools.build.lib.query2.proto.proto2api.Build;
 import com.google.devtools.build.lib.query2.proto.proto2api.Build.Rule.Builder;
 import com.google.devtools.build.lib.syntax.GlobCriteria;
 import com.google.devtools.build.lib.syntax.GlobList;
+import com.google.protobuf.CodedOutputStream;
 
 import java.io.IOException;
-import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.LinkedList;
@@ -131,15 +131,10 @@ public class PackageSerializer {
    * encoding the targets.
    *
    * @param pkg the {@link Package} to be serialized
-   * @param out the stream to pkg's serialized representation to
+   * @param codedOut the stream to pkg's serialized representation to
    * @throws IOException on failure writing to {@code out}
    */
-  public void serialize(Package pkg, OutputStream out) throws IOException {
-    serializePackageInternal(pkg, out);
-  }
-
-  /** Serializes pkg to out as a series of protocol buffers */
-  private void serializePackageInternal(Package pkg, OutputStream out) throws IOException {
+  public void serialize(Package pkg, CodedOutputStream codedOut) throws IOException {
     Build.Package.Builder builder = Build.Package.newBuilder();
     builder.setName(pkg.getName());
     builder.setRepository(pkg.getPackageIdentifier().getRepository().toString());
@@ -197,11 +192,11 @@ public class PackageSerializer {
 
     builder.setWorkspaceName(pkg.getWorkspaceName());
 
-    builder.build().writeDelimitedTo(out);
+    codedOut.writeMessageNoTag(builder.build());
 
     // Targets are emitted separately as individual protocol buffers as to prevent overwhelming
     // protocol buffer deserialization size limits.
-    emitTargets(pkg.getTargets(), out);
+    emitTargets(pkg.getTargets(), codedOut);
   }
 
   /**
@@ -541,31 +536,31 @@ public class PackageSerializer {
   }
 
   /** Writes targets as a series of separate TargetOrTerminator messages to out. */
-  private void emitTargets(Collection<Target> targets, OutputStream out) throws IOException {
+  private void emitTargets(Collection<Target> targets, CodedOutputStream codedOut)
+      throws IOException {
     for (Target target : targets) {
       if (target instanceof InputFile) {
-        emitTarget(serializeInputFile((InputFile) target), out);
+        emitTarget(serializeInputFile((InputFile) target), codedOut);
       } else if (target instanceof OutputFile) {
         // Output files are not serialized; they are recreated by the RuleClass on deserialization.
       } else if (target instanceof PackageGroup) {
-        emitTarget(serializePackageGroup((PackageGroup) target), out);
+        emitTarget(serializePackageGroup((PackageGroup) target), codedOut);
       } else if (target instanceof Rule) {
-        emitTarget(serializeRule((Rule) target), out);
+        emitTarget(serializeRule((Rule) target), codedOut);
       }
     }
 
     // Terminate stream with isTerminator = true.
-    Build.TargetOrTerminator.newBuilder()
+    codedOut.writeMessageNoTag(Build.TargetOrTerminator.newBuilder()
         .setIsTerminator(true)
-        .build()
-        .writeDelimitedTo(out);
+        .build());
   }
 
-  private static void emitTarget(Build.Target target, OutputStream out) throws IOException {
-    Build.TargetOrTerminator.newBuilder()
+  private static void emitTarget(Build.Target target, CodedOutputStream codedOut)
+      throws IOException {
+    codedOut.writeMessageNoTag(Build.TargetOrTerminator.newBuilder()
         .setTarget(target)
-        .build()
-        .writeDelimitedTo(out);
+        .build());
   }
 
   // This is needed because I do not want to use the SymlinkBehavior from the