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
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 c4e8808..f84611e 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 @@
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 5a6a487..d8cd9f4 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
similarity index 71%
rename from src/main/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMultimapCodec.java
rename to src/main/java/com/google/devtools/build/lib/skyframe/serialization/MultimapCodec.java
index f5f407f..f990bb6 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 787b6c8..7564899 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 @@
* 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 @@
* 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);
}
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMultimapCodecTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/MultimapCodecTest.java
similarity index 86%
rename from src/test/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMultimapCodecTest.java
rename to src/test/java/com/google/devtools/build/lib/skyframe/serialization/MultimapCodecTest.java
index dbcf4aa..9209c93 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/ImmutableMultimapCodecTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/MultimapCodecTest.java
@@ -19,21 +19,25 @@
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.devtools.build.lib.skyframe.serialization.testutils.SerializationTester;
import com.google.devtools.build.lib.skyframe.serialization.testutils.SerializationTester.VerificationFunction;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
-/** Tests for {@link ImmutableMultimapCodec}. */
+/** Tests for {@link MultimapCodec}. */
@RunWith(JUnit4.class)
-public class ImmutableMultimapCodecTest {
+public class MultimapCodecTest {
@Test
- public void testImmutableMultimap() throws Exception {
+ public void testMultimap() throws Exception {
+ LinkedHashMultimap<String, String> linkedHashMultimap = LinkedHashMultimap.create();
+ linkedHashMultimap.put("A", "//foo:B");
new SerializationTester(
ImmutableMultimap.of(),
ImmutableMultimap.of("A", "//foo:A"),
- ImmutableMultimap.builder().putAll("B", "//foo:B1", "//foo:B2").build())
+ ImmutableMultimap.builder().putAll("B", "//foo:B1", "//foo:B2").build(),
+ linkedHashMultimap)
.runTests();
}
diff --git a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java
index 36f93a1..4976a28 100644
--- a/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skyframe/serialization/SerializationContextTest.java
@@ -141,17 +141,17 @@
}
@Test
- public void explicitlyAllowedClassCheck() {
+ public void explicitlyAllowedClassCheck() throws SerializationException {
SerializationContext underTest =
new SerializationContext(ObjectCodecRegistry.newBuilder().build(), ImmutableMap.of())
.getMemoizingContext();
underTest.addExplicitlyAllowedClass(String.class);
underTest.checkClassExplicitlyAllowed(String.class);
assertThrows(
- RuntimeException.class, () -> underTest.checkClassExplicitlyAllowed(Integer.class));
+ SerializationException.class, () -> underTest.checkClassExplicitlyAllowed(Integer.class));
// Explicitly registered classes do not carry over to a new context.
assertThrows(
- RuntimeException.class,
+ SerializationException.class,
() -> underTest.getNewMemoizingContext().checkClassExplicitlyAllowed(String.class));
}
@@ -159,7 +159,8 @@
public void explicitlyAllowedClassCheckFailsIfNotMemoizing() {
SerializationContext underTest =
new SerializationContext(ObjectCodecRegistry.newBuilder().build(), ImmutableMap.of());
- assertThrows(RuntimeException.class, () -> underTest.addExplicitlyAllowedClass(String.class));
+ assertThrows(
+ SerializationException.class, () -> underTest.addExplicitlyAllowedClass(String.class));
}
private static class CodecMemoizing implements ObjectCodec<ImmutableList<Object>> {
