Get rid of non-class-based lookups for ObjectCodecs. There were basically no remaining non-test users (ErrorInfoEncoder was it). Should help to simplify our codec registration.

I want to converge MemoizingCodec and ObjectCodec, and the unnecessary complexity around ObjectCodec registration is annoying.

PiperOrigin-RevId: 188620988

4 files changed, 92 insertions, 332 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
index 920a7e2a64..80ec3281ae 100644
--- 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
@@ -20,14 +20,13 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.lib.skyframe.serialization.Memoizer.MemoizingCodec;
-import com.google.protobuf.ByteString;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
+import java.util.Collection;
 import java.util.Comparator;
 import java.util.IdentityHashMap;
 import java.util.Map;
-import java.util.Map.Entry;
 import javax.annotation.Nullable;
 
 /**
@@ -41,8 +40,7 @@ public class ObjectCodecRegistry {
     return new Builder();
   }
 
-  private final ImmutableMap<String, CodecDescriptor> stringMappedCodecs;
-  private final ImmutableMap<ByteString, CodecDescriptor> byteStringMappedCodecs;
+  private final ImmutableMap<Class<?>, CodecDescriptor> classMappedCodecs;
   private final ImmutableList<CodecDescriptor> tagMappedCodecs;
   @Nullable
   private final CodecDescriptor defaultCodecDescriptor;
@@ -56,25 +54,24 @@ public class ObjectCodecRegistry {
   private final int constantsStartTag;
 
   private ObjectCodecRegistry(
-      Map<String, CodecHolder> codecs,
+      Map<Class<?>, CodecHolder> codecs,
       ImmutableSet<MemoizingCodec<?>> memoizingCodecs,
       ImmutableList<Object> constants,
       boolean allowDefaultCodec) {
-    ImmutableMap.Builder<String, CodecDescriptor> codecMappingsBuilder = ImmutableMap.builder();
+    ImmutableMap.Builder<Class<?>, CodecDescriptor> codecMappingsBuilder = ImmutableMap.builder();
     int nextTag = 1; // 0 is reserved for null.
-    for (String classifier : ImmutableList.sortedCopyOf(codecs.keySet())) {
-      codecMappingsBuilder.put(classifier, codecs.get(classifier).createDescriptor(nextTag));
+    for (Class<?> clazz :
+        ImmutableList.sortedCopyOf(Comparator.comparing(Class::getName), codecs.keySet())) {
+      codecMappingsBuilder.put(clazz, codecs.get(clazz).createDescriptor(nextTag));
       nextTag++;
     }
-    this.stringMappedCodecs = codecMappingsBuilder.build();
-
-    this.byteStringMappedCodecs = makeByteStringMappedCodecs(stringMappedCodecs);
+    this.classMappedCodecs = codecMappingsBuilder.build();
 
     this.defaultCodecDescriptor =
         allowDefaultCodec
             ? new TypedCodecDescriptor<>(nextTag++, new JavaSerializableCodec())
             : null;
-    this.tagMappedCodecs = makeTagMappedCodecs(stringMappedCodecs, defaultCodecDescriptor);
+    this.tagMappedCodecs = makeTagMappedCodecs(classMappedCodecs.values(), defaultCodecDescriptor);
 
     this.memoizingCodecsStartTag = nextTag;
     ImmutableMap.Builder<Class<?>, MemoizingCodecDescriptor<?>> memoizingCodecsBuilder =
@@ -95,34 +92,6 @@ public class ObjectCodecRegistry {
     this.constants = constants;
   }
 
-  /** 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 a {@link CodecDescriptor} for the given type.
    *
@@ -133,7 +102,7 @@ public class ObjectCodecRegistry {
       throws SerializationException.NoCodecException {
     // TODO(blaze-team): consider caching this traversal.
     for (Class<?> nextType = type; nextType != null; nextType = nextType.getSuperclass()) {
-      CodecDescriptor result = stringMappedCodecs.get(nextType.getName());
+      CodecDescriptor result = classMappedCodecs.get(nextType);
       if (result != null) {
         return result;
       }
@@ -205,8 +174,9 @@ public class ObjectCodecRegistry {
   public Builder getBuilder() {
     Builder builder = newBuilder();
     builder.setAllowDefaultCodec(defaultCodecDescriptor != null);
-    for (Map.Entry<String, CodecDescriptor> entry : stringMappedCodecs.entrySet()) {
-      builder.add(entry.getKey(), entry.getValue().getCodec());
+    for (Map.Entry<Class<?>, CodecDescriptor> entry : classMappedCodecs.entrySet()) {
+      // Cast here is safe because the original #add in the Builder was type-checked.
+      builder.add(entry.getKey(), getObjectCodec(entry.getValue()));
     }
 
     for (MemoizingCodecDescriptor<?> descriptor : tagMappedMemoizingCodecs) {
@@ -219,6 +189,11 @@ public class ObjectCodecRegistry {
     return builder;
   }
 
+  @SuppressWarnings("unchecked")
+  private static ObjectCodec<Object> getObjectCodec(CodecDescriptor descriptor) {
+    return (ObjectCodec<Object>) descriptor.getCodec();
+  }
+
   /** Describes encoding logic. */
   interface CodecDescriptor {
     void serialize(SerializationContext context, Object obj, CodedOutputStream codedOut)
@@ -326,25 +301,15 @@ public class ObjectCodecRegistry {
 
   /** Builder for {@link ObjectCodecRegistry}. */
   public static class Builder {
-    private final ImmutableMap.Builder<String, CodecHolder> codecsBuilder = ImmutableMap.builder();
+    private final ImmutableMap.Builder<Class<?>, CodecHolder> codecsBuilder =
+        ImmutableMap.builder();
     private final ImmutableSet.Builder<MemoizingCodec<?>> memoizingCodecBuilder =
         ImmutableSet.builder();
     private final ImmutableList.Builder<Object> constantsBuilder = ImmutableList.builder();
     private boolean allowDefaultCodec = true;
 
-    /**
-     * 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.
-     */
-    <T> Builder add(String classifier, ObjectCodec<T> codec) {
-      codecsBuilder.put(classifier, new TypedCodecHolder<>(codec));
-      return this;
-    }
-
     public <T> Builder add(Class<? extends T> type, ObjectCodec<T> codec) {
-      add(type.getName(), codec);
+      codecsBuilder.put(type, new TypedCodecHolder<>(codec));
       return this;
     }
 
@@ -366,11 +331,6 @@ public class ObjectCodecRegistry {
       return this;
     }
 
-    /** Wrap this builder with a {@link ClassKeyedBuilder}. */
-    public ClassKeyedBuilder asClassKeyedBuilder() {
-      return new ClassKeyedBuilder(this);
-    }
-
     public ObjectCodecRegistry build() {
       return new ObjectCodecRegistry(
           codecsBuilder.build(),
@@ -380,40 +340,12 @@ public class ObjectCodecRegistry {
     }
   }
 
-  /** 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, 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) {
+      Collection<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() - 1] = entry.getValue();
+    for (CodecDescriptor codecDescriptor : codecs) {
+      codecTable[codecDescriptor.getTag() - 1] = codecDescriptor;
     }
 
     if (defaultCodecDescriptor != null) {
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 d2703e2de4..1d0fbe66b6 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
@@ -26,8 +26,6 @@ import java.io.IOException;
  * serving as a layer between the streaming-oriented {@link ObjectCodec} interface and users.
  */
 public class ObjectCodecs {
-
-  private final ObjectCodecRegistry codecRegistry;
   // TODO(shahan): when per-invocation state is needed, for example, memoization, these may
   // need to be constructed each time.
   private final SerializationContext serializationContext;
@@ -39,7 +37,6 @@ public class ObjectCodecs {
    */
   public ObjectCodecs(
       ObjectCodecRegistry codecRegistry, ImmutableMap<Class<?>, Object> dependencies) {
-    this.codecRegistry = codecRegistry;
     serializationContext = new SerializationContext(codecRegistry, dependencies);
     deserializationContext = new DeserializationContext(codecRegistry, dependencies);
   }
@@ -64,42 +61,6 @@ public class ObjectCodecs {
     }
   }
 
-  /**
-   * Serialize {@code subject}, using the serialization strategy determined by {@code classifier},
-   * returning a {@link ByteString} containing the serialized representation.
-   */
-  public ByteString serialize(String classifier, Object subject) throws SerializationException {
-    ByteString.Output resultOut = ByteString.newOutput();
-    CodedOutputStream codedOut = CodedOutputStream.newInstance(resultOut);
-    ObjectCodec<?> codec = codecRegistry.getCodecDescriptor(classifier).getCodec();
-    try {
-      doSerialize(classifier, codec, subject, codedOut);
-      codedOut.flush();
-      return resultOut.toByteString();
-    } catch (IOException e) {
-      throw new SerializationException(
-          "Failed to serialize " + subject + " using " + codec + " for " + classifier, e);
-    }
-  }
-
-  /**
-   * Similar to {@link #serialize(String, Object)}, except allows the caller to specify a {@link
-   * CodedOutputStream} to serialize {@code subject} to. Has less object overhead than {@link
-   * #serialize(String, Object)} and as such is preferrable when serializing objects in bulk.
-   *
-   * <p>{@code codedOut} is not flushed by this method.
-   */
-  public void serialize(String classifier, Object subject, CodedOutputStream codedOut)
-      throws SerializationException {
-    ObjectCodec<?> codec = codecRegistry.getCodecDescriptor(classifier).getCodec();
-    try {
-      doSerialize(classifier, codec, subject, codedOut);
-    } catch (IOException e) {
-      throw new SerializationException(
-          "Failed to serialize " + subject + " using " + codec + " for " + classifier, e);
-    }
-  }
-
   public Object deserialize(ByteString data) throws SerializationException {
     return deserialize(data.newCodedInput());
   }
@@ -112,61 +73,6 @@ public class ObjectCodecs {
     }
   }
 
-  /**
-   * Deserialize {@code data} using the serialization strategy determined by {@code classifier}.
-   * {@code classifier} should be the utf-8 encoded {@link ByteString} representation of the {@link
-   * String} classifier used to serialize {@code data}. This is preferred since callers typically
-   * have parsed {@code classifier} from a protocol buffer, for which {@link ByteString}s are
-   * cheaper to use.
-   */
-  public Object deserialize(ByteString classifier, ByteString data) throws SerializationException {
-    return deserialize(classifier, data.newCodedInput());
-  }
-
-  /**
-   * Similar to {@link #deserialize(ByteString, ByteString)}, except allows the caller to specify a
-   * {@link CodedInputStream} to deserialize data from. This is useful for decoding objects
-   * serialized in bulk by {@link #serialize(String, Object, CodedOutputStream)}.
-   */
-  public Object deserialize(ByteString classifier, CodedInputStream codedIn)
-      throws SerializationException {
-    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);
-    try {
-      Object result = codec.deserialize(deserializationContext, codedIn);
-      if (result == null) {
-        throw new NullPointerException(
-            "ObjectCodec " + codec + " for " + classifier.toStringUtf8() + " returned null");
-      }
-      return result;
-    } catch (IOException e) {
-      throw new SerializationException(
-          "Failed to deserialize data using " + codec + " for " + classifier.toStringUtf8(), e);
-    }
-  }
-
-  private <T> void doSerialize(
-      String classifier, ObjectCodec<T> codec, Object subject, CodedOutputStream codedOut)
-      throws SerializationException, IOException {
-    try {
-      codec.serialize(serializationContext, codec.getEncodedClass().cast(subject), codedOut);
-    } catch (ClassCastException e) {
-      throw new SerializationException(
-          "Codec "
-              + codec
-              + " for "
-              + classifier
-              + " is incompatible with "
-              + subject
-              + " (of type "
-              + subject.getClass().getName()
-              + ")",
-          e);
-    }
-  }
-
   @VisibleForTesting
   public SerializationContext getSerializationContextForTesting() {
     return serializationContext;
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
index 6c6cc99604..efcb760d44 100644
--- 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
@@ -19,7 +19,6 @@ 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;
@@ -31,32 +30,26 @@ public class ObjectCodecRegistryTest {
   @Test
   public void testDescriptorLookups() throws NoCodecException {
     SingletonCodec<String> codec1 = SingletonCodec.of("value1", "mnemonic1");
-    SingletonCodec<String> codec2 = SingletonCodec.of("value2", "mnemonic2");
+    SingletonCodec<Integer> codec2 = SingletonCodec.of(1, "mnemonic2");
 
-    ObjectCodecRegistry underTest = ObjectCodecRegistry.newBuilder()
-        .setAllowDefaultCodec(false)
-        .add("foo", codec1)
-        .add("bar", codec2)
-        .build();
+    ObjectCodecRegistry underTest =
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(false)
+            .add(String.class, codec1)
+            .add(Integer.class, codec2)
+            .build();
 
-    CodecDescriptor fooDescriptor = underTest.getCodecDescriptor("foo");
+    CodecDescriptor fooDescriptor = underTest.getCodecDescriptor(String.class);
     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");
+    CodecDescriptor barDescriptor = underTest.getCodecDescriptor(Integer.class);
     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.getCodecDescriptor(Byte.class));
     assertThrows(NoCodecException.class, () -> underTest.getCodecDescriptorByTag(42));
   }
 
@@ -64,21 +57,22 @@ public class ObjectCodecRegistryTest {
   public void testDefaultCodecFallback() throws NoCodecException {
     SingletonCodec<String> codec = SingletonCodec.of("value1", "mnemonic1");
 
-    ObjectCodecRegistry underTest = ObjectCodecRegistry.newBuilder()
-        .setAllowDefaultCodec(true)
-        .add("foo", codec)
-        .build();
+    ObjectCodecRegistry underTest =
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(true)
+            .add(String.class, codec)
+            .build();
 
-    CodecDescriptor fooDescriptor = underTest.getCodecDescriptor("foo");
+    CodecDescriptor fooDescriptor = underTest.getCodecDescriptor(String.class);
     assertThat(fooDescriptor.getCodec()).isSameAs(codec);
 
-    CodecDescriptor barDefaultDescriptor = underTest.getCodecDescriptor("bar");
+    CodecDescriptor barDefaultDescriptor = underTest.getCodecDescriptor(Integer.class);
     assertThat(barDefaultDescriptor.getCodec()).isNotSameAs(codec);
     assertThat(barDefaultDescriptor.getTag()).isNotEqualTo(fooDescriptor.getTag());
     assertThat(underTest.getCodecDescriptorByTag(barDefaultDescriptor.getTag()))
         .isSameAs(barDefaultDescriptor);
 
-    assertThat(underTest.getCodecDescriptor("baz")).isSameAs(barDefaultDescriptor);
+    assertThat(underTest.getCodecDescriptor(Byte.class)).isSameAs(barDefaultDescriptor);
 
     // Bogus tags still throw.
     assertThrows(NoCodecException.class, () -> underTest.getCodecDescriptorByTag(42));
@@ -87,27 +81,29 @@ public class ObjectCodecRegistryTest {
   @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());
+    SingletonCodec<Integer> codec2 = SingletonCodec.of(1, "mnemonic2");
+
+    ObjectCodecRegistry underTest1 =
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(true)
+            .add(String.class, codec1)
+            .add(Integer.class, codec2)
+            .build();
+
+    ObjectCodecRegistry underTest2 =
+        ObjectCodecRegistry.newBuilder()
+            .setAllowDefaultCodec(true)
+            .add(Integer.class, codec2)
+            .add(String.class, codec1)
+            .build();
+
+    assertThat(underTest1.getCodecDescriptor(String.class).getTag())
+        .isEqualTo(underTest2.getCodecDescriptor(String.class).getTag());
+    assertThat(underTest1.getCodecDescriptor(Integer.class).getTag())
+        .isEqualTo(underTest2.getCodecDescriptor(Integer.class).getTag());
     // Default codec.
-    assertThat(underTest1.getCodecDescriptor("baz").getTag())
-        .isEqualTo(underTest2.getCodecDescriptor("baz").getTag());
+    assertThat(underTest1.getCodecDescriptor(Byte.class).getTag())
+        .isEqualTo(underTest2.getCodecDescriptor(Byte.class).getTag());
   }
 
   @Test
@@ -157,25 +153,25 @@ public class ObjectCodecRegistryTest {
   @Test
   public void testGetBuilder() throws NoCodecException {
     SingletonCodec<String> codec1 = SingletonCodec.of("value1", "mnemonic1");
-    SingletonCodec<String> codec2 = SingletonCodec.of("value2", "mnemonic2");
+    SingletonCodec<Integer> codec2 = SingletonCodec.of(1, "mnemonic2");
     Object constant = new Object();
     Memoizer.MemoizingCodec<String> memoizingCodec = new Memoizer.ObjectCodecAdaptor<>(codec1);
 
     ObjectCodecRegistry underTest =
         ObjectCodecRegistry.newBuilder()
             .setAllowDefaultCodec(false)
-            .add("foo", codec1)
-            .add("bar", codec2)
+            .add(String.class, codec1)
+            .add(Integer.class, codec2)
             .addConstant(constant)
             .addMemoizing(memoizingCodec)
             .build();
 
     ObjectCodecRegistry copy = underTest.getBuilder().build();
-    assertThat(copy.getCodecDescriptor("bar").getTag()).isEqualTo(1);
-    assertThat(copy.getCodecDescriptor("foo").getTag()).isEqualTo(2);
+    assertThat(copy.getCodecDescriptor(Integer.class).getTag()).isEqualTo(1);
+    assertThat(copy.getCodecDescriptor(String.class).getTag()).isEqualTo(2);
     assertThat(copy.getMemoizingCodecDescriptor(String.class).getMemoizingCodec())
         .isEqualTo(memoizingCodec);
     assertThat(copy.maybeGetTagForConstant(constant)).isNotNull();
-    assertThrows(NoCodecException.class, () -> copy.getCodecDescriptor("baz"));
+    assertThrows(NoCodecException.class, () -> copy.getCodecDescriptor(Byte.class));
   }
 }
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 8abb4c8320..5577087b37 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
@@ -22,15 +22,12 @@ import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.doThrow;
-import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.verify;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
-import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.NotSerializableException;
 import org.junit.Before;
@@ -68,14 +65,6 @@ public class ObjectCodecsTest {
     private static class IntegerCodecRegisterer implements CodecRegisterer<IntegerCodec> {}
   }
 
-  private static final String KNOWN_CLASSIFIER = "KNOWN_CLASSIFIER";
-  private static final ByteString KNOWN_CLASSIFIER_BYTES =
-      ByteString.copyFromUtf8("KNOWN_CLASSIFIER");
-
-  private static final String UNKNOWN_CLASSIFIER = "UNKNOWN_CLASSIFIER";
-  private static final ByteString UNKNOWN_CLASSIFIER_BYTES =
-      ByteString.copyFromUtf8("UNKNOWN_CLASSIFIER");
-
   private ObjectCodec<Integer> spyObjectCodec;
 
   private ObjectCodecs underTest;
@@ -85,7 +74,7 @@ public class ObjectCodecsTest {
     spyObjectCodec = spy(new IntegerCodec());
     this.underTest =
         new ObjectCodecs(
-            ObjectCodecRegistry.newBuilder().add(KNOWN_CLASSIFIER, spyObjectCodec).build(),
+            ObjectCodecRegistry.newBuilder().add(Integer.class, spyObjectCodec).build(),
             ImmutableMap.of());
   }
 
@@ -109,39 +98,14 @@ public class ObjectCodecsTest {
         .when(spyObjectCodec)
         .deserialize(any(DeserializationContext.class), captor.capture());
 
-    ByteString serialized = underTest.serialize(KNOWN_CLASSIFIER, original);
-    Object deserialized = underTest.deserialize(KNOWN_CLASSIFIER_BYTES, serialized);
+    ByteString serialized = underTest.serialize(original);
+    Object deserialized = underTest.deserialize(serialized);
     assertThat(deserialized).isEqualTo(original);
 
     assertThat(captor.getValue().readInt32()).isEqualTo(42);
   }
 
   @Test
-  public void testMismatchedArgRaisesException() throws Exception {
-    Long notRight = Long.valueOf(123456789);
-    try {
-      underTest.serialize(KNOWN_CLASSIFIER, notRight);
-      fail("Expected exception");
-    } catch (SerializationException expected) {
-    }
-  }
-
-  @Test
-  public void testSerializeDeserializeWhenNocustomLogicAvailable() throws Exception {
-    Integer original = Integer.valueOf(12345);
-
-    ByteString serialized = underTest.serialize(UNKNOWN_CLASSIFIER, original);
-    Object deserialized = underTest.deserialize(UNKNOWN_CLASSIFIER_BYTES, serialized);
-    assertThat(deserialized).isEqualTo(original);
-
-    verify(spyObjectCodec, never())
-        .serialize(
-            any(SerializationContext.class), any(Integer.class), any(CodedOutputStream.class));
-    verify(spyObjectCodec, never())
-        .deserialize(any(DeserializationContext.class), any(CodedInputStream.class));
-  }
-
-  @Test
   public void testSerializePropagatesSerializationExceptionFromCustomCodec() throws Exception {
     Integer original = Integer.valueOf(12345);
 
@@ -150,7 +114,7 @@ public class ObjectCodecsTest {
         .when(spyObjectCodec)
         .serialize(any(SerializationContext.class), eq(original), any(CodedOutputStream.class));
     try {
-      underTest.serialize(KNOWN_CLASSIFIER, original);
+      underTest.serialize(original);
       fail("Expected exception");
     } catch (SerializationException e) {
       assertThat(e).isSameAs(staged);
@@ -167,7 +131,7 @@ public class ObjectCodecsTest {
         .when(spyObjectCodec)
         .serialize(any(SerializationContext.class), eq(original), any(CodedOutputStream.class));
     try {
-      underTest.serialize(KNOWN_CLASSIFIER, original);
+      underTest.serialize(original);
       fail("Expected exception");
     } catch (SerializationException e) {
       assertThat(e).hasCauseThat().isSameAs(staged);
@@ -180,12 +144,10 @@ public class ObjectCodecsTest {
     doThrow(staged)
         .when(spyObjectCodec)
         .deserialize(any(DeserializationContext.class), any(CodedInputStream.class));
-    try {
-      underTest.deserialize(KNOWN_CLASSIFIER_BYTES, ByteString.EMPTY);
-      fail("Expected exception");
-    } catch (SerializationException e) {
-      assertThat(e).isSameAs(staged);
-    }
+    SerializationException thrown =
+        assertThrows(
+            SerializationException.class, () -> underTest.deserialize(underTest.serialize(1)));
+    assertThat(thrown).isSameAs(staged);
   }
 
   @Test
@@ -196,7 +158,7 @@ public class ObjectCodecsTest {
         .when(spyObjectCodec)
         .deserialize(any(DeserializationContext.class), any(CodedInputStream.class));
     try {
-      underTest.deserialize(KNOWN_CLASSIFIER_BYTES, ByteString.EMPTY);
+      underTest.deserialize(underTest.serialize(1));
       fail("Expected exception");
     } catch (SerializationException e) {
       assertThat(e).hasCauseThat().isSameAs(staged);
@@ -204,76 +166,40 @@ public class ObjectCodecsTest {
   }
 
   @Test
-  public void testSerializePropagatesSerializationExceptionFromDefaultCodec() throws Exception {
-    Object nonSerializable = new Object();
-
-    try {
-      underTest.serialize(UNKNOWN_CLASSIFIER, nonSerializable);
-      fail("Expected exception");
-    } catch (SerializationException e) {
-      assertThat(e).hasCauseThat().isInstanceOf(NotSerializableException.class);
-    }
+  public void testSerializePropagatesSerializationExceptionFromDefaultCodec() {
+    SerializationException exception =
+        assertThrows(SerializationException.class, () -> underTest.serialize(new Object()));
+    assertThat(exception).hasCauseThat().isInstanceOf(NotSerializableException.class);
   }
 
   @Test
-  public void testDeserializePropagatesSerializationExceptionFromDefaultCodec() throws Exception {
+  public void testDeserializePropagatesSerializationExceptionFromDefaultCodec() {
     ByteString serialized = ByteString.copyFromUtf8("probably not serialized anything");
 
-    try {
-      underTest.deserialize(UNKNOWN_CLASSIFIER_BYTES, serialized);
-      fail("Expected exception");
-    } catch (SerializationException expected) {
-    }
+    assertThrows(SerializationException.class, () -> underTest.deserialize(serialized));
   }
 
   @Test
-  public void testSerializeFailsWhenNoCustomCodecAndFallbackDisabled() throws Exception {
+  public void testSerializeFailsWhenNoCustomCodecAndFallbackDisabled() {
     ObjectCodecs underTest =
         new ObjectCodecs(
             ObjectCodecRegistry.newBuilder().setAllowDefaultCodec(false).build(),
             ImmutableMap.of());
     SerializationException.NoCodecException expected =
-        assertThrows(
-            SerializationException.NoCodecException.class, () -> underTest.serialize("X", "Y"));
+        assertThrows(SerializationException.NoCodecException.class, () -> underTest.serialize("Y"));
     assertThat(expected)
         .hasMessageThat()
-        .isEqualTo("No codec available for X and default fallback disabled");
+        .isEqualTo("No codec available for class java.lang.String and default fallback disabled");
   }
 
   @Test
-  public void testDeserializeFailsWhenNoCustomCodecAndFallbackDisabled() throws Exception {
-    ByteString serialized = ByteString.copyFromUtf8("doesn't matter");
+  public void testDeserializeFailsWithNoCodecs() throws Exception {
+    ByteString serialized = underTest.serialize(1);
     ObjectCodecs underTest =
         new ObjectCodecs(
             ObjectCodecRegistry.newBuilder().setAllowDefaultCodec(false).build(),
             ImmutableMap.of());
-    SerializationException.NoCodecException expected =
-        assertThrows(
-            SerializationException.NoCodecException.class,
-            () -> underTest.deserialize(ByteString.copyFromUtf8("X"), serialized));
-
-    assertThat(expected)
-        .hasMessageThat()
-        .isEqualTo("No codec available for X and default fallback disabled");
-  }
-
-  @Test
-  public void testSerializeDeserializeInBulk() throws Exception {
-    Integer value1 = 12345;
-    Integer value2 = 67890;
-    Integer value3 = 42;
-
-    ByteArrayOutputStream bytesOut = new ByteArrayOutputStream();
-    CodedOutputStream codedOut = CodedOutputStream.newInstance(bytesOut);
-    underTest.serialize(KNOWN_CLASSIFIER, value1, codedOut);
-    underTest.serialize(KNOWN_CLASSIFIER, value2, codedOut);
-    underTest.serialize(KNOWN_CLASSIFIER, value3, codedOut);
-    codedOut.flush();
-
-    CodedInputStream codedIn = CodedInputStream.newInstance(bytesOut.toByteArray());
-    assertThat(underTest.deserialize(KNOWN_CLASSIFIER_BYTES, codedIn)).isEqualTo(value1);
-    assertThat(underTest.deserialize(KNOWN_CLASSIFIER_BYTES, codedIn)).isEqualTo(value2);
-    assertThat(underTest.deserialize(KNOWN_CLASSIFIER_BYTES, codedIn)).isEqualTo(value3);
+    assertThrows(RuntimeException.class, () -> underTest.deserialize(serialized));
   }
 
   @Test