Cache superclass traversal to find proper codec in ObjectCodecRegistry.

PiperOrigin-RevId: 197508159

1 files changed, 11 insertions, 10 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 1ad3e1f270..2fb0ac39e7 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
@@ -18,7 +18,6 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedSet;
 import com.google.protobuf.CodedInputStream;
@@ -29,6 +28,8 @@ import java.util.Comparator;
 import java.util.HashMap;
 import java.util.IdentityHashMap;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentMap;
 import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
@@ -45,7 +46,7 @@ public class ObjectCodecRegistry {
 
   private final boolean allowDefaultCodec;
 
-  private final ImmutableMap<Class<?>, CodecDescriptor> classMappedCodecs;
+  private final ConcurrentMap<Class<?>, CodecDescriptor> classMappedCodecs;
   private final ImmutableList<CodecDescriptor> tagMappedCodecs;
 
   private final int referenceConstantsStartTag;
@@ -66,15 +67,13 @@ public class ObjectCodecRegistry {
     this.allowDefaultCodec = allowDefaultCodec;
 
     int nextTag = 1; // 0 is reserved for null.
-    ImmutableMap.Builder<Class<?>, CodecDescriptor> memoizingCodecsBuilder =
-        ImmutableMap.builderWithExpectedSize(memoizingCodecs.size());
+    this.classMappedCodecs =
+        new ConcurrentHashMap<>(
+            memoizingCodecs.size(), 0.75f, Runtime.getRuntime().availableProcessors());
     ImmutableList.Builder<CodecDescriptor> tagMappedMemoizingCodecsBuilder =
         ImmutableList.builderWithExpectedSize(memoizingCodecs.size());
     nextTag =
-        processCodecs(
-            memoizingCodecs, nextTag, tagMappedMemoizingCodecsBuilder, memoizingCodecsBuilder);
-
-    this.classMappedCodecs = memoizingCodecsBuilder.build();
+        processCodecs(memoizingCodecs, nextTag, tagMappedMemoizingCodecsBuilder, classMappedCodecs);
     this.tagMappedCodecs = tagMappedMemoizingCodecsBuilder.build();
 
     referenceConstantsStartTag = nextTag;
@@ -116,10 +115,12 @@ public class ObjectCodecRegistry {
    * <p>Also checks if there are codecs for a superclass of the given type.
    */
   private @Nullable CodecDescriptor getCodecDescriptor(Class<?> type) {
-    // TODO(blaze-team): consider caching this traversal.
     for (Class<?> nextType = type; nextType != null; nextType = nextType.getSuperclass()) {
       CodecDescriptor result = classMappedCodecs.get(nextType);
       if (result != null) {
+        if (nextType != type) {
+          classMappedCodecs.put(type, result);
+        }
         return result;
       }
     }
@@ -325,7 +326,7 @@ public class ObjectCodecRegistry {
       Iterable<? extends ObjectCodec<?>> memoizingCodecs,
       int nextTag,
       ImmutableList.Builder<CodecDescriptor> tagMappedCodecsBuilder,
-      ImmutableMap.Builder<Class<?>, CodecDescriptor> codecsBuilder) {
+      ConcurrentMap<Class<?>, CodecDescriptor> codecsBuilder) {
     for (ObjectCodec<?> codec :
         ImmutableList.sortedCopyOf(
             Comparator.comparing(o -> o.getEncodedClass().getName()), memoizingCodecs)) {