CanonicalReferenceResolver for Kryo

Uses type and object equality to determine references. This results in a canonical serialized representation.

ObjectCodecTester and SerializerTester now log some timing information.

PiperOrigin-RevId: 183403658

4 files changed, 141 insertions, 8 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/CanonicalReferenceResolver.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/CanonicalReferenceResolver.java
new file mode 100644
index 0000000000..daa338984b
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/CanonicalReferenceResolver.java
@@ -0,0 +1,113 @@
+// 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.skyframe.serialization.serializers;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.ReferenceResolver;
+import com.esotericsoftware.kryo.util.Util;
+import java.util.ArrayList;
+import java.util.HashMap;
+
+/**
+ * {@link ReferenceResolver} implementation that uses object equality.
+ *
+ * <p>Provided that underlying objects implement equals correctly, this produces canonical, stable,
+ * serialized representations.
+ *
+ * <p>References must match also on class for stability of serialized representation. For example,
+ * {@code ArrayList} and {@code LinkedList} objects might evaluate equals, but Kryo has different
+ * serialized representations because each distinct type has a different registration index.
+ *
+ * <p>TODO(shahan): consider changing ClassResolver to allow multiple classes to share a
+ * registration in Kryo so this can be avoided.
+ */
+public class CanonicalReferenceResolver implements ReferenceResolver {
+
+  private final HashMap<Element, Integer> written;
+  private final ArrayList<Object> read;
+
+  CanonicalReferenceResolver() {
+    this.written = new HashMap<>();
+    this.read = new ArrayList<>();
+  }
+
+  @Override
+  public void setKryo(Kryo unusedKryo) {}
+
+  @Override
+  public int getWrittenId(Object object) {
+    return written.getOrDefault(new Element(object), -1);
+  }
+
+  @Override
+  public int addWrittenObject(Object object) {
+    int id = written.size();
+    written.put(new Element(object), id);
+    return id;
+  }
+
+  @Override
+  public int nextReadId(Class unusedType) {
+    int id = read.size();
+    read.add(null);
+    return id;
+  }
+
+  @Override
+  public void setReadObject(int id, Object object) {
+    read.set(id, object);
+  }
+
+  @Override
+  public Object getReadObject(Class unusedType, int id) {
+    return read.get(id);
+  }
+
+  @Override
+  public void reset() {
+    written.clear();
+    read.clear();
+  }
+
+  @Override
+  public boolean useReferences(Class type) {
+    return !Util.isWrapperClass(type);
+  }
+
+  /**
+   * Wrapper to ensure types are equal in addition to underlying objects.
+   *
+   * <p>Despite the overhead introduced by churn, this is empirically more efficient than {@link
+   * com.google.common.collect.Table}.
+   */
+  private static class Element {
+    public final Object contents;
+
+    private Element(Object contents) {
+      this.contents = contents;
+    }
+
+    @Override
+    public int hashCode() {
+      return contents.hashCode();
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+      Element that = (Element) obj;
+      return contents.getClass().equals(that.contents.getClass()) && contents.equals(that.contents);
+    }
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/RegistrationUtil.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/KryoConfigUtil.java
index ce2668c1b8..d0adf345fb 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/RegistrationUtil.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/serializers/KryoConfigUtil.java
@@ -18,10 +18,19 @@ import com.esotericsoftware.kryo.Kryo;
 import com.google.common.collect.Ordering;
 import java.util.Collections;
 
-/** Utility for registering Serializers defined in this package. */
-public interface RegistrationUtil {
+/** Utility for Kryo configuration. */
+public class KryoConfigUtil {
+  private KryoConfigUtil() {}
 
-  static void registerSerializers(Kryo kryo) {
+  /** Creates a {@link Kryo} instance with extensions defined in this package. */
+  public static Kryo create() {
+    Kryo kryo = new Kryo(new CanonicalReferenceResolver());
+    registerSerializers(kryo);
+    return kryo;
+  }
+
+  /** Registers Serializers defined in this package. */
+  private static void registerSerializers(Kryo kryo) {
     kryo.register(Ordering.natural().getClass());
     kryo.register(Collections.reverseOrder().getClass());
 
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 6d537597ae..3f7071db29 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
@@ -18,6 +18,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.fail;
 
 import com.google.common.base.Preconditions;
+import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableList;
 import com.google.devtools.build.lib.skyframe.serialization.ObjectCodec;
 import com.google.devtools.build.lib.skyframe.serialization.SerializationException;
@@ -78,6 +79,7 @@ public class ObjectCodecTester<T> {
 
   /** Runs serialization/deserialization tests. */
   void testSerializeDeserialize() throws Exception {
+    Stopwatch timer = Stopwatch.createStarted();
     int totalBytes = 0;
     for (T subject : subjects) {
       byte[] serialized = toBytes(subject);
@@ -85,7 +87,13 @@ public class ObjectCodecTester<T> {
       T deserialized = fromBytes(serialized);
       verificationFunction.verifyDeserialized(subject, deserialized);
     }
-    logger.log(Level.INFO, "total serialized bytes = " + totalBytes);
+    logger.log(
+        Level.INFO,
+        underTest.getEncodedClass().getSimpleName()
+            + " total serialized bytes = "
+            + totalBytes
+            + ", "
+            + timer);
   }
 
   /** 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 52a68a9a2b..2a5d2ccee4 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
@@ -22,8 +22,9 @@ import com.esotericsoftware.kryo.Serializer;
 import com.esotericsoftware.kryo.io.UnsafeInput;
 import com.esotericsoftware.kryo.io.UnsafeOutput;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Stopwatch;
 import com.google.common.collect.ImmutableList;
-import com.google.devtools.build.lib.skyframe.serialization.serializers.RegistrationUtil;
+import com.google.devtools.build.lib.skyframe.serialization.serializers.KryoConfigUtil;
 import java.io.ByteArrayOutputStream;
 import java.util.Random;
 import java.util.function.Consumer;
@@ -83,6 +84,7 @@ public class SerializerTester<SubjectT, SerializerT extends SubjectT> {
 
   /** Runs serialization/deserialization tests. */
   void testSerializeDeserialize() throws Exception {
+    Stopwatch timer = Stopwatch.createStarted();
     int totalBytes = 0;
     for (SubjectT subject : subjects) {
       byte[] serialized = toBytes(subject);
@@ -90,7 +92,9 @@ public class SerializerTester<SubjectT, SerializerT extends SubjectT> {
       SubjectT deserialized = fromBytes(serialized);
       verificationFunction.verifyDeserialized(subject, deserialized);
     }
-    logger.log(Level.INFO, type.getSimpleName() + " total serialized bytes = " + totalBytes);
+    logger.log(
+        Level.INFO,
+        type.getSimpleName() + " total serialized bytes = " + totalBytes + ", " + timer);
   }
 
   /** Runs serialized bytes stability tests. */
@@ -152,8 +156,7 @@ public class SerializerTester<SubjectT, SerializerT extends SubjectT> {
 
     private Builder(Class<SerializerT> type) {
       this.type = type;
-      this.kryo = new Kryo();
-      RegistrationUtil.registerSerializers(kryo);
+      this.kryo = KryoConfigUtil.create();
       kryo.setRegistrationRequired(true);
     }