Extract registry functionality from ObjectCodecs

This partially solves the age old problem of how to find a codec for a value we
don't know the type of at compile time, and allows us to represent such values
on the wire more compactly. @AutoCodec's injecting codec should be able to make
use of this right away - we'll need to make an API change to the ObjectCodec
interface to allow the existing system to make use.

PiperOrigin-RevId: 184918173

6 files changed, 351 insertions, 185 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
new file mode 100644
index 0000000000..3fa0e63b5f
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistry.java
@@ -0,0 +1,201 @@
+// 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;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.protobuf.ByteString;
+import java.util.Map;
+import java.util.Map.Entry;
+import javax.annotation.Nullable;
+
+/**
+ * Registry class for handling {@link ObjectCodec} mappings. Codecs are indexed by {@link String}
+ * classifiers and assigned deterministic numeric identifiers for more compact on-the-wire
+ * representation if desired.
+ */
+class ObjectCodecRegistry {
+
+  static Builder newBuilder() {
+    return new Builder();
+  }
+
+  private final ImmutableMap<String, CodecDescriptor> stringMappedCodecs;
+  private final ImmutableMap<ByteString, CodecDescriptor> byteStringMappedCodecs;
+  private final ImmutableList<CodecDescriptor> tagMappedCodecs;
+  @Nullable
+  private final CodecDescriptor defaultCodecDescriptor;
+
+  private ObjectCodecRegistry(Map<String, ObjectCodec<?>> codecs, boolean allowDefaultCodec) {
+    ImmutableMap.Builder<String, CodecDescriptor> codecMappingsBuilder = ImmutableMap.builder();
+    int nextTag = 0;
+    for (String classifier : ImmutableList.sortedCopyOf(codecs.keySet())) {
+      codecMappingsBuilder.put(classifier, new CodecDescriptor(nextTag, codecs.get(classifier)));
+      nextTag++;
+    }
+    this.stringMappedCodecs = codecMappingsBuilder.build();
+    this.byteStringMappedCodecs = makeByteStringMappedCodecs(stringMappedCodecs);
+
+    this.defaultCodecDescriptor = allowDefaultCodec
+        ? new CodecDescriptor(nextTag, new JavaSerializableCodec())
+        : null;
+    this.tagMappedCodecs = makeTagMappedCodecs(stringMappedCodecs, defaultCodecDescriptor);
+  }
+
+  /** Returns the {@link CodecDescriptor} associated with the supplied classifier. */
+  public CodecDescriptor getCodecDescriptor(String classifier)
+      throws SerializationException.NoCodecException {
+    CodecDescriptor result = stringMappedCodecs.getOrDefault(classifier, defaultCodecDescriptor);
+    if (result != null) {
+      return result;
+    } else {
+      throw new SerializationException.NoCodecException(
+          "No codec available for " + classifier + " and default fallback disabled");
+    }
+  }
+
+  /**
+   * Returns the {@link CodecDescriptor} associated with the supplied classifier. This method is a
+   * specialization of {@link #getCodecDescriptor(String)} for performance purposes.
+   */
+  public CodecDescriptor getCodecDescriptor(ByteString classifier)
+      throws SerializationException.NoCodecException {
+    CodecDescriptor result =
+        byteStringMappedCodecs.getOrDefault(classifier, defaultCodecDescriptor);
+    if (result != null) {
+      return result;
+    } else {
+      throw new SerializationException.NoCodecException(
+          "No codec available for " + classifier.toStringUtf8() + " and default fallback disabled");
+    }
+  }
+
+  /** Returns the {@link CodecDescriptor} associated with the supplied tag. */
+  public CodecDescriptor getCodecDescriptorByTag(int tag)
+      throws SerializationException.NoCodecException {
+    if (tag < 0 || tag > tagMappedCodecs.size()) {
+      throw new SerializationException.NoCodecException("No codec available for tag " + tag);
+    }
+
+    CodecDescriptor result = tagMappedCodecs.get(tag);
+    if (result != null) {
+      return result;
+    } else {
+      throw new SerializationException.NoCodecException("No codec available for tag " + tag);
+    }
+  }
+
+  /** Describes encoding logic. */
+  static class CodecDescriptor {
+    private final int tag;
+    private final ObjectCodec<?> codec;
+
+    private CodecDescriptor(int tag, ObjectCodec<?> codec) {
+      this.tag = tag;
+      this.codec = codec;
+    }
+
+    /**
+     * Unique identifier identifying the associated codec. Intended to be used as a compact
+     * on-the-wire representation of an encoded object's type.
+     */
+    int getTag() {
+      return tag;
+    }
+
+    ObjectCodec<?> getCodec() {
+      return codec;
+    }
+  }
+
+  /** Builder for {@link ObjectCodecRegistry}. */
+  static class Builder {
+    private final ImmutableMap.Builder<String, ObjectCodec<?>> codecsBuilder =
+        ImmutableMap.builder();
+    private boolean allowDefaultCodec = true;
+
+    private Builder() {}
+
+    /**
+     * Add custom serialization strategy ({@code codec}) for {@code classifier}.
+     *
+     * <p>Intended for package-internal usage only. Consider using the specialized build types
+     * returned by {@link #asClassKeyedBuilder()} before using this method.
+     */
+    Builder add(String classifier, ObjectCodec<?> codec) {
+      codecsBuilder.put(classifier, codec);
+      return this;
+    }
+
+    /**
+     * Set whether or not we allow fallback to java serialization when no matching codec is found.
+     */
+    public Builder setAllowDefaultCodec(boolean allowDefaultCodec) {
+      this.allowDefaultCodec = allowDefaultCodec;
+      return this;
+    }
+
+    /** Wrap this builder with a {@link ClassKeyedBuilder}. */
+    public ClassKeyedBuilder asClassKeyedBuilder() {
+      return new ClassKeyedBuilder(this);
+    }
+
+    public ObjectCodecRegistry build() {
+      return new ObjectCodecRegistry(codecsBuilder.build(), allowDefaultCodec);
+    }
+  }
+
+  /** Convenience builder for adding codecs classified by class name. */
+  static class ClassKeyedBuilder {
+    private final Builder underlying;
+
+    private ClassKeyedBuilder(Builder underlying) {
+      this.underlying = underlying;
+    }
+
+    public <T> ClassKeyedBuilder add(Class<? extends T> clazz, ObjectCodec<T> codec) {
+      underlying.add(clazz.getName(), codec);
+      return this;
+    }
+
+    public ObjectCodecRegistry build() {
+      return underlying.build();
+    }
+  }
+
+  private static ImmutableMap<ByteString, CodecDescriptor> makeByteStringMappedCodecs(
+      Map<String, CodecDescriptor> stringMappedCodecs) {
+    ImmutableMap.Builder<ByteString, CodecDescriptor> result = ImmutableMap.builder();
+    for (Entry<String, CodecDescriptor> entry : stringMappedCodecs.entrySet()) {
+      result.put(ByteString.copyFromUtf8(entry.getKey()), entry.getValue());
+    }
+    return result.build();
+  }
+
+  private static ImmutableList<CodecDescriptor> makeTagMappedCodecs(
+      Map<String, CodecDescriptor> codecs,
+      @Nullable CodecDescriptor defaultCodecDescriptor) {
+    CodecDescriptor[] codecTable =
+        new CodecDescriptor[codecs.size() + (defaultCodecDescriptor != null ? 1 : 0)];
+    for (Entry<String, CodecDescriptor> entry : codecs.entrySet()) {
+      codecTable[entry.getValue().getTag()] = entry.getValue();
+    }
+
+    if (defaultCodecDescriptor != null) {
+      codecTable[defaultCodecDescriptor.getTag()] = defaultCodecDescriptor;
+    }
+    return ImmutableList.copyOf(codecTable);
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
index c20b734264..c1a0d599c5 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecs.java
@@ -14,65 +14,25 @@
 
 package com.google.devtools.build.lib.skyframe.serialization;
 
-import com.google.common.collect.ImmutableMap;
-import com.google.devtools.build.skyframe.SkyFunctionName;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
-import java.util.Map;
-import java.util.Map.Entry;
 
 /**
  * Wrapper for the minutiae of serializing and deserializing objects using {@link ObjectCodec}s,
  * serving as a layer between the streaming-oriented {@link ObjectCodec} interface and users.
- * Handles the mapping and selection of custom serialization implementations, falling back on less
- * performant java serialization by default when no better option is available and it is allowed by
- * the configuration.
- *
- * <p>To use, create a {@link ObjectCodecs.Builder} and add custom classifier to {@link ObjectCodec}
- * mappings using {@link ObjectCodecs.Builder#add} directly or by using one of the convenience
- * builders returned by {@link ObjectCodecs.Builder#asSkyFunctionNameKeyedBuilder()} or
- * {@link ObjectCodecs.Builder#asClassKeyedBuilder()}. The provided mappings are then used to
- * determine serialization/deserialization logic. For example:
- *
- * <pre>{@code
- * // Create an instance for which anything identified as "foo" will use FooCodec.
- * ObjectCodecs objectCodecs = ObjectCodecs.newBuilder()
- *     .add("foo", new FooCodec())
- *     .build();
- *
- * // This will use the custom supplied FooCodec to serialize obj:
- * ByteString serialized = objectCodecs.serialize("foo", obj);
- * Object deserialized = objectCodecs.deserialize(ByteString.copyFromUtf8("foo"), serialized);
- *
- * // This will use default java object serialization to serialize obj:
- * ByteString serialized = objectCodecs.serialize("bar", obj);
- * Object deserialized = objectCodecs.deserialize(ByteString.copyFromUtf8("bar"), serialized);
- * }</pre>
- *
- * <p>Classifiers will typically be class names or SkyFunction names.
  */
 public class ObjectCodecs {
 
-  private static final ObjectCodec<Object> DEFAULT_CODEC = new JavaSerializableCodec();
+  private final ObjectCodecRegistry codecRegistry;
 
-  /** Create new ObjectCodecs.Builder, the preferred instantiation method. */
-  // TODO(janakr,michajlo): Specialize builders into ones keyed by class (even if the class isn't
-  // the one specified by the codec) and ones keyed by string, and expose a getClassifier() method
-  // for ObjectCodecs keyed by class.
-  public static ObjectCodecs.Builder newBuilder() {
-    return new Builder();
-  }
-
-  private final Map<String, ObjectCodec<?>> stringMappedCodecs;
-  private final Map<ByteString, ObjectCodec<?>> byteStringMappedCodecs;
-  private final boolean allowDefaultCodec;
-
-  private ObjectCodecs(Map<String, ObjectCodec<?>> codecs, boolean allowDefaultCodec) {
-    this.stringMappedCodecs = codecs;
-    this.byteStringMappedCodecs = makeByteStringMappedCodecs(codecs);
-    this.allowDefaultCodec = allowDefaultCodec;
+  /**
+   * Creates an instance using the supplied {@link ObjectCodecRegistry} for looking up
+   * {@link ObjectCodec}s.
+   */
+  ObjectCodecs(ObjectCodecRegistry codecRegistry) {
+    this.codecRegistry = codecRegistry;
   }
 
   /**
@@ -82,7 +42,7 @@ public class ObjectCodecs {
   public ByteString serialize(String classifier, Object subject) throws SerializationException {
     ByteString.Output resultOut = ByteString.newOutput();
     CodedOutputStream codedOut = CodedOutputStream.newInstance(resultOut);
-    ObjectCodec<?> codec = getCodec(classifier);
+    ObjectCodec<?> codec = codecRegistry.getCodecDescriptor(classifier).getCodec();
     try {
       doSerialize(classifier, codec, subject, codedOut);
       codedOut.flush();
@@ -102,7 +62,7 @@ public class ObjectCodecs {
    */
   public void serialize(String classifier, Object subject, CodedOutputStream codedOut)
       throws SerializationException {
-    ObjectCodec<?> codec = getCodec(classifier);
+    ObjectCodec<?> codec = codecRegistry.getCodecDescriptor(classifier).getCodec();
     try {
       doSerialize(classifier, codec, subject, codedOut);
     } catch (IOException e) {
@@ -129,7 +89,7 @@ public class ObjectCodecs {
    */
   public Object deserialize(ByteString classifier, CodedInputStream codedIn)
       throws SerializationException {
-    ObjectCodec<?> codec = getCodec(classifier);
+    ObjectCodec<?> codec = codecRegistry.getCodecDescriptor(classifier).getCodec();
     // If safe, this will allow CodedInputStream to return a direct view of the underlying bytes
     // in some situations, bypassing a copy.
     codedIn.enableAliasing(true);
@@ -146,31 +106,6 @@ public class ObjectCodecs {
     }
   }
 
-  private ObjectCodec<?> getCodec(String classifier)
-      throws SerializationException.NoCodecException {
-    ObjectCodec<?> result = stringMappedCodecs.get(classifier);
-    if (result != null) {
-      return result;
-    } else if (allowDefaultCodec) {
-      return DEFAULT_CODEC;
-    } else {
-      throw new SerializationException.NoCodecException(
-          "No codec available for " + classifier + " and default fallback disabled");
-    }
-  }
-
-  private ObjectCodec<?> getCodec(ByteString classifier) throws SerializationException {
-    ObjectCodec<?> result = byteStringMappedCodecs.get(classifier);
-    if (result != null) {
-      return result;
-    } else if (allowDefaultCodec) {
-      return DEFAULT_CODEC;
-    } else {
-      throw new SerializationException.NoCodecException(
-          "No codec available for " + classifier.toStringUtf8() + " and default fallback disabled");
-    }
-  }
-
   private static <T> void doSerialize(
       String classifier, ObjectCodec<T> codec, Object subject, CodedOutputStream codedOut)
       throws SerializationException, IOException {
@@ -190,91 +125,4 @@ public class ObjectCodecs {
           e);
     }
   }
-
-  /** Builder for {@link ObjectCodecs}. */
-  static class Builder {
-    private final ImmutableMap.Builder<String, ObjectCodec<?>> codecsBuilder =
-        ImmutableMap.builder();
-    private boolean allowDefaultCodec = true;
-
-    private Builder() {}
-
-    /**
-     * Add custom serialization strategy ({@code codec}) for {@code classifier}.
-     *
-     * <p>Intended for package-internal usage only. Consider using the specialized build types
-     * returned by {@link #asClassKeyedBuilder()} or {@link #asSkyFunctionNameKeyedBuilder()}
-     * before using this method.
-     */
-    Builder add(String classifier, ObjectCodec<?> codec) {
-      codecsBuilder.put(classifier, codec);
-      return this;
-    }
-
-    /** Set whether or not we allow fallback to the default codec, java serialization. */
-    public Builder setAllowDefaultCodec(boolean allowDefaultCodec) {
-      this.allowDefaultCodec = allowDefaultCodec;
-      return this;
-    }
-
-    /** Wrap this builder with a {@link ClassKeyedBuilder}. */
-    public ClassKeyedBuilder asClassKeyedBuilder() {
-      return new ClassKeyedBuilder(this);
-    }
-
-    /** Wrap this builder with a {@link SkyFunctionNameKeyedBuilder}. */
-    public SkyFunctionNameKeyedBuilder asSkyFunctionNameKeyedBuilder() {
-      return new SkyFunctionNameKeyedBuilder(this);
-    }
-
-    public ObjectCodecs build() {
-      return new ObjectCodecs(codecsBuilder.build(), allowDefaultCodec);
-    }
-  }
-
-  /** Convenience builder for adding codecs classified by class name. */
-  static class ClassKeyedBuilder {
-    private final Builder underlying;
-
-    private ClassKeyedBuilder(Builder underlying) {
-      this.underlying = underlying;
-    }
-
-    public <T> ClassKeyedBuilder add(Class<? extends T> clazz, ObjectCodec<T> codec) {
-      underlying.add(clazz.getName(), codec);
-      return this;
-    }
-
-    public ObjectCodecs build() {
-      return underlying.build();
-    }
-  }
-
-  /** Convenience builder for adding codecs classified by SkyFunctionName. */
-  static class SkyFunctionNameKeyedBuilder {
-    private final Builder underlying;
-
-    private SkyFunctionNameKeyedBuilder(Builder underlying) {
-      this.underlying = underlying;
-    }
-
-    public SkyFunctionNameKeyedBuilder add(SkyFunctionName skyFuncName, ObjectCodec<?> codec) {
-      underlying.add(skyFuncName.getName(), codec);
-      return this;
-    }
-
-    public ObjectCodecs build() {
-      return underlying.build();
-    }
-  }
-
-  private static Map<ByteString, ObjectCodec<?>> makeByteStringMappedCodecs(
-      Map<String, ObjectCodec<?>> stringMappedCodecs) {
-    ImmutableMap.Builder<ByteString, ObjectCodec<?>> result = ImmutableMap.builder();
-    for (Entry<String, ObjectCodec<?>> entry : stringMappedCodecs.entrySet()) {
-      result.put(ByteString.copyFromUtf8(entry.getKey()), entry.getValue());
-    }
-    return result.build();
-  }
-
 }
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/BUILD b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/BUILD
index bd3b55bfc7..b2fd6f5154 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/BUILD
@@ -21,6 +21,7 @@ java_library(
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization",
         "//src/main/java/com/google/devtools/build/lib/skyframe/serialization/testutils",
         "//src/main/java/com/google/devtools/build/lib/vfs",
+        "//src/test/java/com/google/devtools/build/lib:testutil",
         "//third_party:guava",
         "//third_party:guava-testlib",
         "//third_party:jsr305",
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java
new file mode 100644
index 0000000000..c4391ecbd6
--- /dev/null
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecRegistryTest.java
@@ -0,0 +1,114 @@
+// 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;
+
+import static com.google.common.truth.Truth.assertThat;
+import static com.google.devtools.build.lib.testutil.MoreAsserts.assertThrows;
+
+import com.google.devtools.build.lib.skyframe.serialization.ObjectCodecRegistry.CodecDescriptor;
+import com.google.devtools.build.lib.skyframe.serialization.SerializationException.NoCodecException;
+import com.google.protobuf.ByteString;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests for {@link ObjectCodecRegistry}. */
+@RunWith(JUnit4.class)
+public class ObjectCodecRegistryTest {
+
+  @Test
+  public void testDescriptorLookups() throws NoCodecException {
+    SingletonCodec<String> codec1 = SingletonCodec.of("value1", "mnemonic1");
+    SingletonCodec<String> codec2 = SingletonCodec.of("value2", "mnemonic2");
+
+    ObjectCodecRegistry underTest = ObjectCodecRegistry.newBuilder()
+        .setAllowDefaultCodec(false)
+        .add("foo", codec1)
+        .add("bar", codec2)
+        .build();
+
+    CodecDescriptor fooDescriptor = underTest.getCodecDescriptor("foo");
+    assertThat(fooDescriptor.getCodec()).isSameAs(codec1);
+    assertThat(underTest.getCodecDescriptor(ByteString.copyFromUtf8("foo")))
+        .isSameAs(fooDescriptor);
+    assertThat(underTest.getCodecDescriptorByTag(fooDescriptor.getTag())).isSameAs(fooDescriptor);
+
+    CodecDescriptor barDescriptor = underTest.getCodecDescriptor("bar");
+    assertThat(barDescriptor.getCodec()).isSameAs(codec2);
+    assertThat(underTest.getCodecDescriptor(ByteString.copyFromUtf8("bar")))
+        .isSameAs(barDescriptor);
+    assertThat(underTest.getCodecDescriptorByTag(barDescriptor.getTag())).isSameAs(barDescriptor);
+
+    assertThat(barDescriptor.getTag()).isNotEqualTo(fooDescriptor.getTag());
+
+    assertThrows(NoCodecException.class, () -> underTest.getCodecDescriptor("baz"));
+    assertThrows(
+        NoCodecException.class,
+        () -> underTest.getCodecDescriptor(ByteString.copyFromUtf8("baz")));
+    assertThrows(NoCodecException.class, () -> underTest.getCodecDescriptorByTag(42));
+  }
+
+  @Test
+  public void testDefaultCodecFallback() throws NoCodecException {
+    SingletonCodec<String> codec = SingletonCodec.of("value1", "mnemonic1");
+
+    ObjectCodecRegistry underTest = ObjectCodecRegistry.newBuilder()
+        .setAllowDefaultCodec(true)
+        .add("foo", codec)
+        .build();
+
+    CodecDescriptor fooDescriptor = underTest.getCodecDescriptor("foo");
+    assertThat(fooDescriptor.getCodec()).isSameAs(codec);
+
+    CodecDescriptor barDefaultDescriptor = underTest.getCodecDescriptor("bar");
+    assertThat(barDefaultDescriptor.getCodec()).isNotSameAs(codec);
+    assertThat(barDefaultDescriptor.getTag()).isNotEqualTo(fooDescriptor.getTag());
+    assertThat(underTest.getCodecDescriptorByTag(barDefaultDescriptor.getTag()))
+        .isSameAs(barDefaultDescriptor);
+
+    assertThat(underTest.getCodecDescriptor("baz")).isSameAs(barDefaultDescriptor);
+
+    // Bogus tags still throw.
+    assertThrows(NoCodecException.class, () -> underTest.getCodecDescriptorByTag(42));
+  }
+
+  @Test
+  public void testStableTagOrdering() throws NoCodecException {
+    SingletonCodec<String> codec1 = SingletonCodec.of("value1", "mnemonic1");
+    SingletonCodec<String> codec2 = SingletonCodec.of("value2", "mnemonic2");
+
+    ObjectCodecRegistry underTest1 = ObjectCodecRegistry.newBuilder()
+        .setAllowDefaultCodec(true)
+        .add("foo", codec1)
+        .add("bar", codec2)
+        .build();
+
+    ObjectCodecRegistry underTest2 = ObjectCodecRegistry.newBuilder()
+        .setAllowDefaultCodec(true)
+        .add("bar", codec2)
+        .add("foo", codec1)
+        .build();
+
+    assertThat(underTest1.getCodecDescriptor("foo").getTag())
+        .isEqualTo(underTest2.getCodecDescriptor("foo").getTag());
+    assertThat(underTest1.getCodecDescriptor("bar").getTag())
+        .isEqualTo(underTest2.getCodecDescriptor("bar").getTag());
+    // Default codec.
+    assertThat(underTest1.getCodecDescriptor("baz").getTag())
+        .isEqualTo(underTest2.getCodecDescriptor("baz").getTag());
+
+
+  }
+}
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecsTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecsTest.java
index 3ad0d3c7b1..7e98f045b5 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ObjectCodecsTest.java
@@ -15,6 +15,7 @@
 package com.google.devtools.build.lib.skyframe.serialization;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.devtools.build.lib.testutil.MoreAsserts.expectThrows;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.eq;
@@ -78,7 +79,10 @@ public class ObjectCodecsTest {
   @Before
   public final void setup() {
     spyObjectCodec = spy(new IntegerCodec());
-    this.underTest = ObjectCodecs.newBuilder().add(KNOWN_CLASSIFIER, spyObjectCodec).build();
+    this.underTest = new ObjectCodecs(
+        ObjectCodecRegistry.newBuilder()
+            .add(KNOWN_CLASSIFIER, spyObjectCodec)
+            .build());
   }
 
   @Test
@@ -207,30 +211,28 @@ public class ObjectCodecsTest {
 
   @Test
   public void testSerializeFailsWhenNoCustomCodecAndFallbackDisabled() throws Exception {
-    try {
-      ObjectCodecs.newBuilder().setAllowDefaultCodec(false).build().serialize("X", "Y");
-      fail("Expected exception");
-    } catch (SerializationException e) {
-      assertThat(e)
-          .hasMessageThat()
-          .isEqualTo("No codec available for X and default fallback disabled");
-    }
+    ObjectCodecs underTest = new ObjectCodecs(
+        ObjectCodecRegistry.newBuilder().setAllowDefaultCodec(false).build());
+    SerializationException.NoCodecException expected = expectThrows(
+        SerializationException.NoCodecException.class,
+        () -> underTest.serialize("X", "Y"));
+    assertThat(expected)
+        .hasMessageThat()
+        .isEqualTo("No codec available for X and default fallback disabled");
   }
 
   @Test
   public void testDeserializeFailsWhenNoCustomCodecAndFallbackDisabled() throws Exception {
     ByteString serialized = ByteString.copyFromUtf8("doesn't matter");
-    try {
-      ObjectCodecs.newBuilder()
-          .setAllowDefaultCodec(false)
-          .build()
-          .deserialize(ByteString.copyFromUtf8("X"), serialized);
-      fail("Expected exception");
-    } catch (SerializationException e) {
-      assertThat(e)
-          .hasMessageThat()
-          .isEqualTo("No codec available for X and default fallback disabled");
-    }
+    ObjectCodecs underTest = new ObjectCodecs(
+        ObjectCodecRegistry.newBuilder().setAllowDefaultCodec(false).build());
+    SerializationException.NoCodecException expected = expectThrows(
+        SerializationException.NoCodecException.class,
+        () -> underTest.deserialize(ByteString.copyFromUtf8("X"), serialized));
+
+    assertThat(expected)
+        .hasMessageThat()
+        .isEqualTo("No codec available for X and default fallback disabled");
   }
 
   @Test
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/PrecomputedValueCodecTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/PrecomputedValueCodecTest.java
index 692f602bc2..f6532fcb61 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/PrecomputedValueCodecTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/PrecomputedValueCodecTest.java
@@ -30,13 +30,13 @@ public class PrecomputedValueCodecTest extends AbstractObjectCodecTest<Precomput
   public PrecomputedValueCodecTest() {
     super(
         new PrecomputedValueCodec(
-            () ->
-                ObjectCodecs.newBuilder()
+            () -> new ObjectCodecs(
+                ObjectCodecRegistry.newBuilder()
                     .asClassKeyedBuilder()
                     // Note no PathFragmentCodec.
                     .add(String.class, StringCodecs.asciiOptimized())
                     .add(Label.class, LabelCodec.INSTANCE)
-                    .build()),
+                    .build())),
         new PrecomputedValue(PathFragment.create("java serializable 1")),
         new PrecomputedValue(PathFragment.create("java serializable 2")),
         new PrecomputedValue("first string"),