Expand the abilities of codecs. Let ImmutableMultimapCodec actually handle LinkedHashMultimap as well. Let ImmutableSetRuntimeCodec handle some other sets. We squash these extra types into the same codec because we want them to become immutable anyway, and they can be dealt with easily by the existing codecs.

Add some more non-Bazel classes that should be DynamicCodec'ed. Also, when a disallowed value is serialized, throw a SerializationException instead of crashing. This makes it easier to recover in tests.

PiperOrigin-RevId: 201237631

4 files changed, 68 insertions, 19 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java
index c4e880881e..f84611e96b 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/AutoRegistry.java
@@ -49,7 +49,12 @@ public class AutoRegistry {
       ImmutableList.of(
           "java.io.FileNotFoundException",
           "java.io.IOException",
-          "java.lang.invoke.SerializedLambda");
+          "java.lang.invoke.SerializedLambda",
+          "com.google.common.base.Predicates$InPredicate",
+          // Sadly, these builders are serialized as part of SkylarkCustomCommandLine$Builder, which
+          // apparently can be preserved through analysis. We may investigate if this actually has
+          // performance/correctness implications.
+          "com.google.common.collect.ImmutableList$Builder");
 
   private static final ImmutableList<Object> REFERENCE_CONSTANTS_TO_REGISTER =
       ImmutableList.of(
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableSetRuntimeCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableSetRuntimeCodec.java
index 5a6a487909..d8cd9f4eb0 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableSetRuntimeCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableSetRuntimeCodec.java
@@ -14,22 +14,44 @@
 
 package com.google.devtools.build.lib.skyframe.serialization;
 
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.LinkedHashMultimap;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
 
-/** {@link ObjectCodec} for {@link ImmutableSet}. */
+/** {@link ObjectCodec} for {@link ImmutableSet} and other sets that should be immutable. */
 @SuppressWarnings("rawtypes") // Intentional erasure of ImmutableSet.
-class ImmutableSetRuntimeCodec implements ObjectCodec<ImmutableSet> {
+class ImmutableSetRuntimeCodec implements ObjectCodec<Set> {
+  @SuppressWarnings("unchecked")
+  private static final Class<Set> LINKED_HASH_MULTIMAP_CLASS =
+      (Class<Set>) LinkedHashMultimap.create(ImmutableMultimap.of("a", "b")).get("a").getClass();
+
+  @SuppressWarnings("unchecked")
+  private static final Class<Set> SINGLETON_SET_CLASS =
+      (Class<Set>) Collections.singleton("a").getClass();
+
+  @SuppressWarnings("unchecked")
+  private static final Class<Set> EMPTY_SET_CLASS = (Class<Set>) Collections.emptySet().getClass();
+
   @Override
   public Class<ImmutableSet> getEncodedClass() {
     return ImmutableSet.class;
   }
 
   @Override
-  public void serialize(
-      SerializationContext context, ImmutableSet object, CodedOutputStream codedOut)
+  public ImmutableList<Class<? extends Set>> additionalEncodedClasses() {
+    return ImmutableList.of(
+        LINKED_HASH_MULTIMAP_CLASS, SINGLETON_SET_CLASS, EMPTY_SET_CLASS, HashSet.class);
+  }
+
+  @Override
+  public void serialize(SerializationContext context, Set object, CodedOutputStream codedOut)
       throws SerializationException, IOException {
     codedOut.writeInt32NoTag(object.size());
     for (Object obj : object) {
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMultimapCodec.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/MultimapCodec.java
index f5f407f062..f990bb65fd 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMultimapCodec.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/MultimapCodec.java
@@ -13,34 +13,46 @@
 // limitations under the License.
 package com.google.devtools.build.lib.skyframe.serialization;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSetMultimap;
+import com.google.common.collect.LinkedHashMultimap;
+import com.google.common.collect.ListMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.SetMultimap;
 import com.google.protobuf.CodedInputStream;
 import com.google.protobuf.CodedOutputStream;
 import java.io.IOException;
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 
 /**
- * A codec for {@link ImmutableMultimap}. Handles both {@link ImmutableListMultimap} and {@link
- * ImmutableSetMultimap}.
+ * A codec for {@link Multimap}. Handles {@link ImmutableListMultimap}, {@link ImmutableSetMultimap}
+ * and {@link LinkedHashMultimap}. Makes all multimaps immutable, since serialized state shouldn't
+ * be mutable.
  */
-public class ImmutableMultimapCodec<K, V> implements ObjectCodec<ImmutableMultimap<K, V>> {
-
+public class MultimapCodec<K, V> implements ObjectCodec<Multimap<K, V>> {
   @SuppressWarnings("unchecked")
   @Override
   public Class<? extends ImmutableMultimap<K, V>> getEncodedClass() {
     return (Class<ImmutableMultimap<K, V>>) ((Class<?>) ImmutableMultimap.class);
   }
 
+  @SuppressWarnings("unchecked")
+  @Override
+  public List<Class<? extends Multimap<K, V>>> additionalEncodedClasses() {
+    return ImmutableList.of((Class<? extends Multimap<K, V>>) (Class<?>) LinkedHashMultimap.class);
+  }
+
   @Override
   public void serialize(
-      SerializationContext context, ImmutableMultimap<K, V> obj, CodedOutputStream codedOut)
+      SerializationContext context, Multimap<K, V> obj, CodedOutputStream codedOut)
       throws SerializationException, IOException {
-    if (obj instanceof ImmutableListMultimap) {
+    if (obj instanceof ListMultimap) {
       codedOut.writeBoolNoTag(true);
-    } else if (obj instanceof ImmutableSetMultimap) {
+    } else if (obj instanceof SetMultimap) {
       codedOut.writeBoolNoTag(false);
     } else {
       throw new SerializationException("Unexpected multimap type: " + obj.getClass());
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java
index 787b6c8e2b..75648991df 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContext.java
@@ -190,10 +190,18 @@ public class SerializationContext {
    * com.google.devtools.build.lib.packages.Package} inside {@link
    * com.google.devtools.build.lib.skyframe.PackageValue}.
    */
-  public void checkClassExplicitlyAllowed(Class<?> allowedClass) {
-    Preconditions.checkNotNull(
-        serializer, "Cannot check explicitly allowed class %s without memoization", allowedClass);
-    Preconditions.checkState(explicitlyAllowedClasses.contains(allowedClass), allowedClass);
+  public void checkClassExplicitlyAllowed(Class<?> allowedClass) throws SerializationException {
+    if (serializer == null) {
+      throw new SerializationException(
+          "Cannot check explicitly allowed class " + allowedClass + " without memoization");
+    }
+    if (!explicitlyAllowedClasses.contains(allowedClass)) {
+      throw new SerializationException(
+          allowedClass
+              + " not explicitly allowed (allowed classes were: "
+              + explicitlyAllowedClasses
+              + ")");
+    }
   }
 
   /**
@@ -212,9 +220,11 @@ public class SerializationContext {
    * be determined if the inclusion of the expensive object is legitimate, before it is whitelisted
    * using this method.
    */
-  public void addExplicitlyAllowedClass(Class<?> allowedClass) {
-    Preconditions.checkNotNull(
-        serializer, "Cannot add explicitly allowed class %s without memoization", allowedClass);
+  public void addExplicitlyAllowedClass(Class<?> allowedClass) throws SerializationException {
+    if (serializer == null) {
+      throw new SerializationException(
+          "Cannot add explicitly allowed class %s without memoization: " + allowedClass);
+    }
     explicitlyAllowedClasses.add(allowedClass);
   }