Improve safety of NestedSetFingerprintCache by detecting multiple instances of the same mapFn class.

This code tries to add protection against the user creating new mapFn instances per-rule. This would cause the nested set cache to be computed per-rule instead of shared across rule instances, causing memory bloat and slowdowns.

Since this can only happen in native code, we can get away with detecting this and crashing blaze. I think this is a better choice than silently allowing it / falling back to slow computations.

The user can override this behaviour by inheriting from CommandLineItem.CapturingMapFn, in which case the user is explicitly saying they assume responsibility for the number of instances of the mapFn the application will use.

PiperOrigin-RevId: 184061642

1 files changed, 49 insertions, 4 deletions

diff --git a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCache.java b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCache.java
index b48d4a2298..083ef06269 100644
--- a/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCache.java
+++ b/src/main/java/com/google/devtools/build/lib/collect/nestedset/NestedSetFingerprintCache.java
@@ -15,9 +15,14 @@
 package com.google.devtools.build.lib.collect.nestedset;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.HashMultiset;
+import com.google.common.collect.Multiset;
 import com.google.devtools.build.lib.analysis.actions.CommandLineItem;
+import com.google.devtools.build.lib.analysis.actions.CommandLineItem.MapFn;
 import com.google.devtools.build.lib.util.Fingerprint;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 
 /** Computes fingerprints for nested sets, reusing sub-computations from children. */
@@ -26,7 +31,10 @@ public class NestedSetFingerprintCache {
   private static final int EMPTY_SET_DIGEST = 104_395_303;
 
   /** Memoize the subresults. We have to have one cache per type of command item map function. */
-  private Map<CommandLineItem.MapFn<?>, DigestMap> mapFnToFingerprints = createMap();
+  private Map<CommandLineItem.MapFn<?>, DigestMap> mapFnToDigestMap = createMap();
+
+  private final Set<Class<?>> seenMapFns = new HashSet<>();
+  private final Multiset<Class<?>> seenParametrizedMapFns = HashMultiset.create();
 
   public <T> void addNestedSetToFingerprint(Fingerprint fingerprint, NestedSet<T> nestedSet) {
     addNestedSetToFingerprint(CommandLineItem.MapFn.DEFAULT, fingerprint, nestedSet);
@@ -34,20 +42,32 @@ public class NestedSetFingerprintCache {
 
   public <T> void addNestedSetToFingerprint(
       CommandLineItem.MapFn<? super T> mapFn, Fingerprint fingerprint, NestedSet<T> nestedSet) {
+    if (mapFn instanceof CommandLineItem.CapturingMapFn) {
+      addNestedSetToFingerprintSlow(mapFn, fingerprint, nestedSet);
+      return;
+    }
     // Only top-level nested sets can be empty, so we can bail here
     if (nestedSet.isEmpty()) {
       fingerprint.addInt(EMPTY_SET_DIGEST);
       return;
     }
-    DigestMap digestMap =
-        mapFnToFingerprints.computeIfAbsent(mapFn, k -> new DigestMap(DIGEST_SIZE, 1024));
+    DigestMap digestMap = mapFnToDigestMap.computeIfAbsent(mapFn, this::newDigestMap);
     fingerprint.addInt(nestedSet.getOrder().ordinal());
     Object children = nestedSet.rawChildren();
     addToFingerprint(mapFn, fingerprint, digestMap, children);
   }
 
+  private <T> void addNestedSetToFingerprintSlow(
+      MapFn<? super T> mapFn, Fingerprint fingerprint, NestedSet<T> nestedSet) {
+    for (T object : nestedSet) {
+      fingerprint.addString(mapFn.expandToCommandLine(object));
+    }
+  }
+
   public void clear() {
-    mapFnToFingerprints = createMap();
+    mapFnToDigestMap = createMap();
+    seenMapFns.clear();
+    seenParametrizedMapFns.clear();
   }
 
   @SuppressWarnings("unchecked")
@@ -78,4 +98,29 @@ public class NestedSetFingerprintCache {
   private static Map<CommandLineItem.MapFn<?>, DigestMap> createMap() {
     return new ConcurrentHashMap<>();
   }
+
+  private DigestMap newDigestMap(CommandLineItem.MapFn<?> mapFn) {
+    Class<?> mapFnClass = mapFn.getClass();
+    if (mapFn instanceof CommandLineItem.ParametrizedMapFn) {
+      int occurrences = seenParametrizedMapFns.add(mapFnClass, 1) + 1;
+      if (occurrences > ((CommandLineItem.ParametrizedMapFn) mapFn).maxInstancesAllowed()) {
+        throw new IllegalArgumentException(
+            String.format(
+                "Too many instances of CommandLineItem.ParametrizedMapFn '%s' detected. "
+                    + "Please construct fewer instances or use CommandLineItem.CapturingMapFn.",
+                mapFnClass.getName()));
+      }
+    } else {
+      if (!seenMapFns.add(mapFnClass)) {
+        throw new IllegalArgumentException(
+            String.format(
+                "Illegal mapFn implementation: '%s'. "
+                    + "CommandLineItem.MapFn instances must be singletons. "
+                    + "Please see CommandLineItem.ParametrizedMapFn or "
+                    + "CommandLineItem.CapturingMapFn for alternatives.",
+                mapFnClass.getName()));
+      }
+    }
+    return new DigestMap(DIGEST_SIZE, 1024);
+  }
 }