Replace use of lambdas/method-refs in low-level storage implementations

field-inject/src/main/java/datadog/instrument/fieldinject/GlobalObjectStore.java

@@ -17,7 +17,6 @@
 import java.util.concurrent.atomic.AtomicIntegerArray;
 import java.util.concurrent.atomic.AtomicReferenceArray;
 import java.util.function.Function;
-import java.util.function.Supplier;
 import javax.annotation.Nullable;
 
 /**
@@ -51,9 +50,6 @@ public final class GlobalObjectStore {
   /** Randomly sample underlying map size, approximately once every 1024 writes per-thread. */
   private static final int SIZE_SAMPLE_RATE = 1 << 10;
 
-  /** Constant supplier used when nothing in a shard is considered old. */
-  private static final Supplier<Object> NO_OLD_STALE_KEYS = () -> null;
-
   /** The current generation of each shard of the global object store. */
   private static final AtomicReferenceArray<GlobalObjectStore> shards = initShards();
 
@@ -79,7 +75,7 @@ private static AtomicReferenceArray<GlobalObjectStore> initShards() {
   private final ConcurrentHashMap<StoreKey, Object> map;
 
   /** Supplies old keys where the key object is unused and eligible for collection. */
-  private final Supplier<Object> oldStaleKeys;
+  @Nullable private final ReferenceQueue<Object> oldStaleKeys;
 
   /** Map of old store keys to value objects. */
   private final Map<StoreKey, Object> oldMap;
@@ -89,14 +85,14 @@ private static AtomicReferenceArray<GlobalObjectStore> initShards() {
   private GlobalObjectStore(int shardIndex) {
     this.shardIndex = shardIndex;
     this.map = new ConcurrentHashMap<>();
-    this.oldStaleKeys = NO_OLD_STALE_KEYS;
+    this.oldStaleKeys = null;
     this.oldMap = Collections.emptyMap();
   }
 
   private GlobalObjectStore(GlobalObjectStore oldStore) {
     this.shardIndex = oldStore.shardIndex;
     this.map = new ConcurrentHashMap<>(INLINE_CLEANUP_THRESHOLD);
-    this.oldStaleKeys = oldStore.staleKeys::poll;
+    this.oldStaleKeys = oldStore.staleKeys;
     this.oldMap = oldStore.map;
   }
 
@@ -130,9 +126,11 @@ private int doRemoveStaleEntries() {
     sampledYoungSize = estimatedTotal;
 
     // next remove stale entries from the old map
-    while ((staleKey = oldStaleKeys.get()) != null) {
-      //noinspection All: we know staleKey is a store key
-      oldMap.remove(staleKey);
+    if (oldStaleKeys != null) {
+      while ((staleKey = oldStaleKeys.poll()) != null) {
+        //noinspection All: we know staleKey is a store key
+        oldMap.remove(staleKey);
+      }
     }
 
     estimatedTotal += oldMap.size();
@@ -246,10 +244,10 @@ public static Object getOrCompute(Object key, int storeId, Function valueFunctio
     }
   }
 
-  @SuppressWarnings({"rawtypes", "unchecked"})
+  @SuppressWarnings("rawtypes")
   private Object doGetOrCompute(Object key, int storeId, Function valueFunction) {
     return map.computeIfAbsent(
-        new StoreKey(staleKeys, key, storeId), unused -> valueFunction.apply(key));
+        new StoreKey(staleKeys, key, storeId), new ValueComputer(key, valueFunction));
   }
 
   /**
@@ -368,13 +366,14 @@ public boolean equals(Object o) {
   private static final class LookupKey {
 
     /** Avoid allocation by maintaining a reusable lookup key per-thread. */
-    private static final ThreadLocal<LookupKey> LOOKUP_KEY_CACHE =
-        ThreadLocal.withInitial(LookupKey::new);
+    private static final LookupKeyCache LOOKUP_KEY_CACHE = new LookupKeyCache();
 
     Object key;
     int hash;
     int storeId;
 
+    LookupKey() {} // avoids synthetic accessor
+
     /**
      * Returns a temporary lookup key for the current thread with the given object key and store-id.
      * This key must be reset by calling {@link #reset} as soon as the get/remove request completes.
@@ -415,4 +414,32 @@ public boolean equals(Object o) {
       }
     }
   }
+
+  /** Per-thread cache of reusable lookup keys. */
+  private static final class LookupKeyCache extends ThreadLocal<LookupKey> {
+    LookupKeyCache() {} // avoids synthetic accessor
+
+    @Override
+    protected LookupKey initialValue() {
+      return new LookupKey();
+    }
+  }
+
+  /** Computes the value for a missing store key by applying the value function to the key. */
+  @SuppressWarnings("rawtypes")
+  private static final class ValueComputer implements Function<Object, Object> {
+    private final Object key; // captured strongly so the key stays reachable during computation
+    private final Function valueFunction;
+
+    ValueComputer(Object key, Function valueFunction) {
+      this.key = key;
+      this.valueFunction = valueFunction;
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public Object apply(Object unused) {
+      return valueFunction.apply(key);
+    }
+  }
 }

field-inject/src/main/java/datadog/instrument/fieldinject/ObjectStoreIds.java

@@ -9,15 +9,14 @@
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
 
 /** Manages unique {@link ObjectStore} ids for each key-value type combination. */
 public final class ObjectStoreIds {
 
+  private static final TypeIdGenerator typeIdGenerator = new TypeIdGenerator();
   private static final Map<String, Short> typeIds = new ConcurrentHashMap<>();
 
-  // start at 1 to give a range of store ids we can allocate another way (0-32767)
-  private static final AtomicInteger nextTypeId = new AtomicInteger(1);
-
   private ObjectStoreIds() {}
 
   /**
@@ -39,6 +38,19 @@ public static int objectStoreId(String keyType, String valueType) {
    * @return the unique type id
    */
   private static short typeId(String type) {
-    return typeIds.computeIfAbsent(type, unused -> (short) nextTypeId.getAndIncrement());
+    return typeIds.computeIfAbsent(type, typeIdGenerator);
+  }
+
+  /** Generates the next unique type id. */
+  private static final class TypeIdGenerator implements Function<String, Short> {
+    // start at 1 to give a slice of store ids we can allocate another way (0-32767)
+    private final AtomicInteger nextTypeId = new AtomicInteger(1);
+
+    TypeIdGenerator() {} // avoids synthetic accessor
+
+    @Override
+    public Short apply(String unused) {
+      return (short) nextTypeId.getAndIncrement();
+    }
   }
 }

utils/src/main/java/datadog/instrument/utils/ClassLoaderKey.java

@@ -10,9 +10,9 @@
 import java.lang.ref.ReferenceQueue;
 import java.lang.ref.WeakReference;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.Consumer;
 
 /**
  * Reference key used to weakly associate a class-loader with a computed value.
@@ -34,21 +34,21 @@ final class ClassLoaderKey extends WeakReference<ClassLoader> {
   // stale class-loader keys that are now eligible for collection
   private static final ReferenceQueue<ClassLoader> staleKeys = new ReferenceQueue<>();
 
-  // registered cleaners of stale class-loader keys and their values
-  private static final List<Consumer<ClassLoaderKey>> cleaners = new CopyOnWriteArrayList<>();
+  // registered maps of class-loader keys to values; stale keys are removed from these maps
+  private static final List<Map<ClassLoaderKey, ?>> valueMaps = new CopyOnWriteArrayList<>();
 
-  /** Registers a cleaner of stale class-loader keys. */
-  static void registerCleaner(Consumer<ClassLoaderKey> cleaner) {
-    cleaners.add(cleaner);
+  /** Registers a map of class-loader keys to values for cleaning. */
+  static void registerValueMap(Map<ClassLoaderKey, ?> valueMap) {
+    valueMaps.add(valueMap);
   }
 
-  /** Checks for stale class-loader keys; stale keys are cleaned by the registered cleaners. */
+  /** Checks for stale class-loader keys; stale keys are removed from the registered maps. */
   static void cleanStaleKeys() {
     ClassLoaderKey key;
     while ((key = (ClassLoaderKey) staleKeys.poll()) != null) {
       //noinspection ForLoopReplaceableByForEach - indexed loop performs better
-      for (int i = 0, size = cleaners.size(); i < size; i++) {
-        cleaners.get(i).accept(key);
+      for (int i = 0, size = valueMaps.size(); i < size; i++) {
+        valueMaps.get(i).remove(key);
       }
     }
   }

utils/src/main/java/datadog/instrument/utils/ClassLoaderValue.java

@@ -11,12 +11,12 @@
 import static datadog.instrument.utils.ClassLoaderKey.SYSTEM_CLASS_LOADER;
 
 import datadog.instrument.utils.ClassLoaderKey.LookupKey;
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.function.BiConsumer;
+import java.util.function.Function;
 import javax.annotation.Nullable;
 
 /**
@@ -47,9 +47,8 @@ public abstract class ClassLoaderValue<V> {
   private final Map<ClassLoaderKey, V> otherValues = new ConcurrentHashMap<>();
 
   /** Register subclass instances for cleaning. */
-  @SuppressFBWarnings("CT_CONSTRUCTOR_THROW") // registerCleaner never sees partial instance
   protected ClassLoaderValue() {
-    ClassLoaderKey.registerCleaner(otherValues::remove);
+    ClassLoaderKey.registerValueMap(otherValues);
   }
 
   /**
@@ -180,18 +179,29 @@ private V getSystemValue() {
     return value;
   }
 
-  /** Helper to make {@code computeValue} compatible with {@code computeIfAbsent}. */
-  private V computeValueForKey(ClassLoaderKey key) {
-    return computeValue(key.get());
-  }
-
   /** Lazily associate a computed value with a custom class-loader. */
   private V getOtherValue(ClassLoader cl) {
     //noinspection All: intentionally use lookup key without reference overhead
     V value = otherValues.get(new LookupKey(cl));
     if (value == null) {
-      value = otherValues.computeIfAbsent(getClassLoaderKey(cl), this::computeValueForKey);
+      value = otherValues.computeIfAbsent(getClassLoaderKey(cl), new ValueComputer<>(cl, this));
     }
     return value;
   }
+
+  /** Computes the value for a missing class-loader key by applying computeValue to the loader. */
+  private static final class ValueComputer<V> implements Function<ClassLoaderKey, V> {
+    private final ClassLoader cl; // captured strongly so the loader stays reachable during compute
+    private final ClassLoaderValue<V> classLoaderValue;
+
+    ValueComputer(ClassLoader cl, ClassLoaderValue<V> classLoaderValue) {
+      this.cl = cl;
+      this.classLoaderValue = classLoaderValue;
+    }
+
+    @Override
+    public V apply(ClassLoaderKey unused) {
+      return classLoaderValue.computeValue(cl);
+    }
+  }
 }