Serializer implementation for NestedSet

Adds some logging to test helpers for size of serialized data.

Jan 25, 2018 7:16:25 AM com.google.devtools.build.lib.skyframe.serialization.testutils.SerializerTester testSerializeDeserialize
INFO: total serialized bytes = 70
Jan 25, 2018 7:16:25 AM com.google.devtools.build.lib.skyframe.serialization.testutils.ObjectCodecTester testSerializeDeserialize
INFO: total serialized bytes = 208

Kryo output is significantly smaller.

PiperOrigin-RevId: 183300353

5 files changed, 122 insertions, 24 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD b/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
index 8e3b861b9d..92e058bb1f 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/BUILD
@@ -26,10 +26,14 @@ java_library(
 
 java_library(
     name = "serialization",
-    srcs = ["NestedSetCodec.java"],
+    srcs = [
+        "NestedSetCodec.java",
+        "NestedSetSerializer.java",
+    ],
     deps = [
         ":nestedset",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
+        "//src/main/java/com/google/devtools/build/lib/skyframe/serialization:kryo",
         "//third_party:guava",
         "//third_party/protobuf:protobuf_java",
     ],
diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetSerializer.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetSerializer.java
new file mode 100644
index 0000000000..3bb2552915
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetSerializer.java
@@ -0,0 +1,56 @@
+// Copyright 2018 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.collect.nestedset;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.Serializer;
+import com.esotericsoftware.kryo.io.Input;
+import com.esotericsoftware.kryo.io.Output;
+
+/**
+ * {@link Serializer} for {@link NestedSet}.
+ *
+ * <p>Needed to handle {@link NestedSet}'s sentinel values correctly.
+ */
+public class NestedSetSerializer extends Serializer<NestedSet<Object>> {
+
+  @Override
+  public void write(Kryo kryo, Output output, NestedSet<Object> nestedSet) {
+    kryo.writeObject(output, nestedSet.getOrder());
+    Object children = nestedSet.rawChildren();
+    if (children == NestedSet.EMPTY_CHILDREN) {
+      output.writeBoolean(false);
+    } else {
+      output.writeBoolean(true);
+      kryo.writeClassAndObject(output, children);
+    }
+  }
+
+  @Override
+  public NestedSet<Object> read(Kryo kryo, Input input, Class<NestedSet<Object>> unusedType) {
+    Order order = kryo.readObject(input, Order.class);
+    if (input.readBoolean()) {
+      return new NestedSet<>(order, kryo.readClassAndObject(input));
+    } else {
+      return new NestedSet<>(order, NestedSet.EMPTY_CHILDREN);
+    }
+  }
+
+  public static void registerSerializers(Kryo kryo) {
+    kryo.register(NestedSet.class, new NestedSetSerializer());
+    kryo.register(Order.class);
+    kryo.register(Object[].class);
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/ObjectCodecTester.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/ObjectCodecTester.java
index 3251364461..6d537597ae 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/ObjectCodecTester.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/ObjectCodecTester.java
@@ -24,9 +24,12 @@ import com.google.devtools.build.lib.skyframe.serialization.SerializationExcepti
 import com.google.protobuf.CodedInputStream;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 
 /** Utility for testing {@link ObjectCodec} instances. */
 public class ObjectCodecTester<T> {
+  private static final Logger logger = Logger.getLogger(SerializerTester.class.getName());
 
   /** Interface for testing successful deserialization of an object. */
   @FunctionalInterface
@@ -75,11 +78,14 @@ public class ObjectCodecTester<T> {
 
   /** Runs serialization/deserialization tests. */
   void testSerializeDeserialize() throws Exception {
+    int totalBytes = 0;
     for (T subject : subjects) {
       byte[] serialized = toBytes(subject);
+      totalBytes += serialized.length;
       T deserialized = fromBytes(serialized);
       verificationFunction.verifyDeserialized(subject, deserialized);
     }
+    logger.log(Level.INFO, "total serialized bytes = " + totalBytes);
   }
 
   /** Runs serialized bytes stability tests. */
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializerTester.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializerTester.java
index 77319a3253..6780b80107 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializerTester.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils/SerializerTester.java
@@ -26,6 +26,9 @@ import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.skyframe.serialization.serializers.RegistrationUtil;
 import java.io.ByteArrayOutputStream;
 import java.util.Random;
+import java.util.function.Consumer;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import org.objenesis.instantiator.ObjectInstantiator;
 
 /** Utility for testing {@link Serializer} instances. */
@@ -33,6 +36,8 @@ public class SerializerTester<SubjectT, SerializerT extends SubjectT> {
   public static final int DEFAULT_JUNK_INPUTS = 20;
   public static final int JUNK_LENGTH_UPPER_BOUND = 20;
 
+  private static final Logger logger = Logger.getLogger(SerializerTester.class.getName());
+
   /** Interface for testing successful deserialization of an object. */
   @FunctionalInterface
   public interface VerificationFunction<T> {
@@ -78,11 +83,14 @@ public class SerializerTester<SubjectT, SerializerT extends SubjectT> {
 
   /** Runs serialization/deserialization tests. */
   void testSerializeDeserialize() throws Exception {
+    int totalBytes = 0;
     for (SubjectT subject : subjects) {
       byte[] serialized = toBytes(subject);
+      totalBytes += serialized.length;
       SubjectT deserialized = fromBytes(serialized);
       verificationFunction.verifyDeserialized(subject, deserialized);
     }
+    logger.log(Level.INFO, "total serialized bytes = " + totalBytes);
   }
 
   /** Runs serialized bytes stability tests. */
@@ -174,6 +182,16 @@ public class SerializerTester<SubjectT, SerializerT extends SubjectT> {
       return this;
     }
 
+    /**
+     * Hands a {@link Kryo} instance to the visitor.
+     *
+     * @param visitor usually a reference to a {@code registerSerializers} method
+     */
+    public Builder<SubjectT, SerializerT> visitKryo(Consumer<Kryo> visitor) {
+      visitor.accept(kryo);
+      return this;
+    }
+
     /** Adds subjects to be tested for serialization/deserialization. */
     @SafeVarargs
     public final Builder<SubjectT, SerializerT> addSubjects(
diff --git a/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java b/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java
index cc71bee36a..9adb1d62a0 100644
--- a/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java
+++ b/src/test/java/com/google/devtools/build/lib/collect/nestedset/NestedSetCodecTest.java
@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.skyframe.serialization.strings.StringCodecs;
 import com.google.devtools.build.lib.skyframe.serialization.testutils.ObjectCodecTester;
+import com.google.devtools.build.lib.skyframe.serialization.testutils.SerializerTester;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -29,36 +30,49 @@ public class NestedSetCodecTest {
   private static final NestedSet<String> SHARED_NESTED_SET =
       NestedSetBuilder.<String>stableOrder().add("e").build();
 
+  private static final ImmutableList<NestedSet<String>> SUBJECTS =
+      ImmutableList.of(
+          NestedSetBuilder.emptySet(Order.STABLE_ORDER),
+          NestedSetBuilder.emptySet(Order.NAIVE_LINK_ORDER),
+          NestedSetBuilder.create(Order.STABLE_ORDER, "a"),
+          NestedSetBuilder.create(Order.STABLE_ORDER, "a", "b", "c"),
+          NestedSetBuilder.<String>stableOrder()
+              .add("a")
+              .add("b")
+              .addTransitive(
+                  NestedSetBuilder.<String>stableOrder()
+                      .add("c")
+                      .addTransitive(SHARED_NESTED_SET)
+                      .build())
+              .addTransitive(
+                  NestedSetBuilder.<String>stableOrder()
+                      .add("d")
+                      .addTransitive(SHARED_NESTED_SET)
+                      .build())
+              .addTransitive(NestedSetBuilder.emptySet(Order.STABLE_ORDER))
+              .build());
+
   @Test
   public void testCodec() throws Exception {
-    ImmutableList<NestedSet<String>> subjects =
-        ImmutableList.of(
-            NestedSetBuilder.emptySet(Order.STABLE_ORDER),
-            NestedSetBuilder.emptySet(Order.NAIVE_LINK_ORDER),
-            NestedSetBuilder.create(Order.STABLE_ORDER, "a"),
-            NestedSetBuilder.create(Order.STABLE_ORDER, "a", "b", "c"),
-            NestedSetBuilder.<String>stableOrder()
-                .add("a")
-                .add("b")
-                .addTransitive(
-                    NestedSetBuilder.<String>stableOrder()
-                        .add("c")
-                        .addTransitive(SHARED_NESTED_SET)
-                        .build())
-                .addTransitive(
-                    NestedSetBuilder.<String>stableOrder()
-                        .add("d")
-                        .addTransitive(SHARED_NESTED_SET)
-                        .build())
-                .addTransitive(NestedSetBuilder.emptySet(Order.STABLE_ORDER))
-                .build());
-
     ObjectCodecTester.newBuilder(new NestedSetCodec<>(StringCodecs.simple()))
-        .addSubjects(subjects)
+        .addSubjects(SUBJECTS)
         .verificationFunction(NestedSetCodecTest::verifyDeserialization)
         .buildAndRunTests();
   }
 
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  @Test
+  public void testSerializer() throws Exception {
+    SerializerTester.Builder<NestedSet, NestedSet> builder =
+        SerializerTester.newBuilder(NestedSet.class)
+            .visitKryo(NestedSetSerializer::registerSerializers)
+            .setVerificationFunction(NestedSetCodecTest::verifyDeserialization);
+    for (NestedSet<String> subject : SUBJECTS) {
+      builder.addSubjects(subject);
+    }
+    builder.buildAndRunTests();
+  }
+
   private static void verifyDeserialization(
       NestedSet<String> subject, NestedSet<String> deserialized) {
     assertThat(subject.getOrder()).isEqualTo(deserialized.getOrder());