Allow class prefix names to be blacklisted from DynamicCodec use.

PiperOrigin-RevId: 194487570

3 files changed, 33 insertions, 23 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 6c81ccfe92..c2713459a1 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
@@ -41,8 +41,8 @@ public class AutoRegistry {
    * where Tristate lives. */
   private static final String PACKAGE_PREFIX = "com.google.devtools";
 
-  /** Package prefixes to blacklist for {@link DynamicCodec}. */
-  private static final ImmutableList<String> PACKAGE_PREFIX_BLACKLIST =
+  /** Class name prefixes to blacklist for {@link DynamicCodec}. */
+  private static final ImmutableList<String> CLASS_NAME_PREFIX_BLACKLIST =
       ImmutableList.of("com.google.devtools.build.lib.vfs");
 
   /** Classes outside {@link AutoRegistry#PACKAGE_PREFIX} that need to be serialized. */
@@ -74,8 +74,7 @@ public class AutoRegistry {
 
   private static ObjectCodecRegistry create() {
     try {
-      ObjectCodecRegistry.Builder registry =
-          CodecScanner.initializeCodecRegistry(PACKAGE_PREFIX, PACKAGE_PREFIX_BLACKLIST);
+      ObjectCodecRegistry.Builder registry = CodecScanner.initializeCodecRegistry(PACKAGE_PREFIX);
       for (String className : EXTERNAL_CLASS_NAMES_TO_REGISTER) {
         registry.addClassName(className);
       }
@@ -85,6 +84,9 @@ public class AutoRegistry {
       for (Object constant : VALUE_CONSTANTS_TO_REGISTER) {
         registry.addValueConstant(constant);
       }
+      for (String classNamePrefix : CLASS_NAME_PREFIX_BLACKLIST) {
+        registry.blacklistClassNamePrefix(classNamePrefix);
+      }
       return registry.build();
     } catch (IOException | ReflectiveOperationException e) {
       throw new IllegalStateException(e);
diff --git a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java
index dc93c5817c..61b4706731 100644
--- a/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java
+++ b/src/main/java/com/google/devtools/build/lib/skyframe/serialization/CodecScanner.java
@@ -15,7 +15,6 @@
 package com.google.devtools.build.lib.skyframe.serialization;
 
 import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableList;
 import com.google.common.reflect.ClassPath;
 import com.google.common.reflect.ClassPath.ClassInfo;
 import com.google.devtools.build.lib.skyframe.serialization.autocodec.RegisteredSingletonDoNotUse;
@@ -45,21 +44,14 @@ public class CodecScanner {
 
   private static final Logger log = Logger.getLogger(CodecScanner.class.getName());
 
-  public static ObjectCodecRegistry.Builder initializeCodecRegistry(String packagePrefix)
-      throws IOException, ReflectiveOperationException {
-    return initializeCodecRegistry(packagePrefix, ImmutableList.of());
-  }
-
   /**
    * Initializes an {@link ObjectCodecRegistry} builder by scanning a given package prefix.
    *
    * @param packagePrefix processes only classes in packages having this prefix
-   * @param packagePrefixBlacklist avoids processing classes in packages having this prefix
    * @see CodecRegisterer
    */
   @SuppressWarnings("unchecked")
-  public static ObjectCodecRegistry.Builder initializeCodecRegistry(
-      String packagePrefix, ImmutableList<String> packagePrefixBlacklist)
+  static ObjectCodecRegistry.Builder initializeCodecRegistry(String packagePrefix)
       throws IOException, ReflectiveOperationException {
     log.info("Building ObjectCodecRegistry");
     ArrayList<Class<? extends ObjectCodec<?>>> codecs = new ArrayList<>();
@@ -77,13 +69,6 @@ public class CodecScanner {
                   .endsWith(CodecScanningConstants.REGISTERED_SINGLETON_SUFFIX)) {
                 processLikelyConstant(classInfo.load(), builder);
               } else {
-                // Assumes that anything with a class name matching the above won't need to be
-                // serialized.
-                for (String prefix : packagePrefixBlacklist) {
-                  if (classInfo.getPackageName().startsWith(prefix)) {
-                    return;
-                  }
-                }
                 builder.addClassName(classInfo.getName().intern());
               }
             });
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 5c144112af..211a67e53d 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
@@ -66,6 +66,7 @@ public class ObjectCodecRegistry {
       ImmutableList<Object> referenceConstants,
       ImmutableList<Object> valueConstants,
       ImmutableSortedSet<String> classNames,
+      ImmutableList<String> blacklistedClassNamePrefixes,
       boolean allowDefaultCodec) {
     this.allowDefaultCodec = allowDefaultCodec;
 
@@ -104,8 +105,12 @@ public class ObjectCodecRegistry {
                 ImmutableMap.toImmutableMap(
                     Map.Entry::getKey, e -> ImmutableMap.copyOf(e.getValue())));
     this.valueConstants = valueConstants;
-    this.classNames = classNames.asList();
-    this.dynamicCodecs = createDynamicCodecs(classNames, nextTag);
+    this.classNames =
+        classNames
+            .stream()
+            .filter((str) -> isAllowed(str, blacklistedClassNamePrefixes))
+            .collect(ImmutableList.toImmutableList());
+    this.dynamicCodecs = createDynamicCodecs(this.classNames, nextTag);
   }
 
   public CodecDescriptor getCodecDescriptorForObject(Object obj)
@@ -286,6 +291,8 @@ public class ObjectCodecRegistry {
     private final ImmutableList.Builder<Object> referenceConstantsBuilder = ImmutableList.builder();
     private final ImmutableList.Builder<Object> valueConstantsBuilder = ImmutableList.builder();
     private final ImmutableSortedSet.Builder<String> classNames = ImmutableSortedSet.naturalOrder();
+    private final ImmutableList.Builder<String> blacklistedClassNamePrefixes =
+        ImmutableList.builder();
     private boolean allowDefaultCodec = true;
 
     /**
@@ -350,12 +357,18 @@ public class ObjectCodecRegistry {
       return this;
     }
 
+    public Builder blacklistClassNamePrefix(String classNamePrefix) {
+      blacklistedClassNamePrefixes.add(classNamePrefix);
+      return this;
+    }
+
     public ObjectCodecRegistry build() {
       return new ObjectCodecRegistry(
           ImmutableSet.copyOf(codecs.values()),
           referenceConstantsBuilder.build(),
           valueConstantsBuilder.build(),
           classNames.build(),
+          blacklistedClassNamePrefixes.build(),
           allowDefaultCodec);
     }
   }
@@ -379,7 +392,7 @@ public class ObjectCodecRegistry {
   }
 
   private static IdentityHashMap<String, Supplier<CodecDescriptor>> createDynamicCodecs(
-      ImmutableSortedSet<String> classNames, int nextTag) {
+      ImmutableList<String> classNames, int nextTag) {
     IdentityHashMap<String, Supplier<CodecDescriptor>> dynamicCodecs =
         new IdentityHashMap<>(classNames.size());
     for (String className : classNames) {
@@ -390,6 +403,16 @@ public class ObjectCodecRegistry {
     return dynamicCodecs;
   }
 
+  private static boolean isAllowed(
+      String className, ImmutableList<String> blacklistedClassNamePefixes) {
+    for (String blacklistedClassNamePrefix : blacklistedClassNamePefixes) {
+      if (className.startsWith(blacklistedClassNamePrefix)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   /** For enums, this method must only be called for the declaring class. */
   private static CodecDescriptor createDynamicCodecDescriptor(int tag, String className) {
     try {