Merge pull request #428 from DataDog/tyler/use-WeakConcurrentMap

Use WeakConcurrentMap

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/JDBCMaps.java

@@ -1,10 +1,9 @@
 package datadog.trace.bootstrap;
 
+import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
+
 import java.sql.Connection;
 import java.sql.PreparedStatement;
-import java.util.Collections;
-import java.util.Map;
-import java.util.WeakHashMap;
 import lombok.Data;
 
 /**
@@ -13,10 +12,8 @@ import lombok.Data;
  * <p>In the bootstrap project to ensure visibility by all classes.
  */
 public class JDBCMaps {
-  public static final Map<Connection, DBInfo> connectionInfo =
-      Collections.synchronizedMap(new WeakHashMap<Connection, DBInfo>());
-  public static final Map<PreparedStatement, String> preparedStatements =
-      Collections.synchronizedMap(new WeakHashMap<PreparedStatement, String>());
+  public static final WeakMap<Connection, DBInfo> connectionInfo = newWeakMap();
+  public static final WeakMap<PreparedStatement, String> preparedStatements = newWeakMap();
 
   public static final String DB_QUERY = "DB Query";
 

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/WeakMap.java

@@ -0,0 +1,82 @@
+package datadog.trace.bootstrap;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
+import java.util.concurrent.atomic.AtomicReference;
+import lombok.extern.slf4j.Slf4j;
+
+public interface WeakMap<K, V> {
+
+  int size();
+
+  boolean containsKey(K target);
+
+  V get(K key);
+
+  void put(K key, V value);
+
+  class Provider {
+    private static final AtomicReference<Supplier> provider =
+        new AtomicReference<>(Supplier.DEFAULT);
+
+    public static void registerIfAbsent(final Supplier provider) {
+      if (provider != null && provider != Supplier.DEFAULT) {
+        Provider.provider.compareAndSet(Supplier.DEFAULT, provider);
+      }
+    }
+
+    public static <K, V> WeakMap<K, V> newWeakMap() {
+      return provider.get().get();
+    }
+  }
+
+  interface Supplier {
+    <K, V> WeakMap<K, V> get();
+
+    Supplier DEFAULT = new Default();
+
+    @Slf4j
+    class Default implements Supplier {
+
+      @Override
+      public <K, V> WeakMap<K, V> get() {
+        log.warn("WeakMap.Supplier not registered. Returning a synchronized WeakHashMap.");
+        return new MapAdapter<>(Collections.synchronizedMap(new WeakHashMap<K, V>()));
+      }
+    }
+  }
+
+  class MapAdapter<K, V> implements WeakMap<K, V> {
+    private final Map<K, V> map;
+
+    public MapAdapter(final Map<K, V> map) {
+      this.map = map;
+    }
+
+    @Override
+    public int size() {
+      return map.size();
+    }
+
+    @Override
+    public boolean containsKey(final K key) {
+      return map.containsKey(key);
+    }
+
+    @Override
+    public V get(final K key) {
+      return map.get(key);
+    }
+
+    @Override
+    public void put(final K key, final V value) {
+      map.put(key, value);
+    }
+
+    @Override
+    public String toString() {
+      return map.toString();
+    }
+  }
+}

dd-java-agent/agent-tooling/agent-tooling.gradle

@@ -9,6 +9,7 @@ configurations {
 
 dependencies {
   compile project(':dd-java-agent:agent-bootstrap')
+  compile group: 'com.blogspot.mydailyjava', name: 'weak-lock-free', version: '0.13'
   compile deps.bytebuddy
   compile deps.bytebuddyagent
   annotationProcessor deps.autoservice
@@ -17,6 +18,7 @@ dependencies {
   compileOnly project(':dd-trace-ot')
 
   testCompile deps.opentracing
+  testCompile project(':dd-java-agent:testing')
 
   instrumentationMuzzle sourceSets.main.output
   instrumentationMuzzle configurations.compile

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java

@@ -8,6 +8,7 @@ import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 
+import datadog.trace.bootstrap.WeakMap;
 import java.lang.instrument.Instrumentation;
 import java.util.ServiceLoader;
 import lombok.extern.slf4j.Slf4j;
@@ -38,6 +39,8 @@ public class AgentInstaller {
   public static ResettableClassFileTransformer installBytebuddyAgent(
       final Instrumentation inst, final AgentBuilder.Listener... listeners) {
     INSTRUMENTATION = inst;
+    registerWeakMapProvider();
+
     AgentBuilder agentBuilder =
         new AgentBuilder.Default()
             .disableClassFormatChanges()
@@ -86,6 +89,12 @@ public class AgentInstaller {
     return agentBuilder.installOn(inst);
   }
 
+  private static void registerWeakMapProvider() {
+    WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent());
+    //    WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.WeakConcurrent.Inline());
+    //    WeakMap.Provider.registerIfAbsent(new WeakMapSuppliers.Guava());
+  }
+
   @Slf4j
   static class LoggingListener implements AgentBuilder.Listener {
 

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderMatcher.java

@@ -1,13 +1,14 @@
 package datadog.trace.agent.tooling;
 
+import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
+
 import datadog.trace.bootstrap.DatadogClassLoader;
 import datadog.trace.bootstrap.PatchLogger;
+import datadog.trace.bootstrap.WeakMap;
 import io.opentracing.util.GlobalTracer;
 import java.util.Collections;
 import java.util.HashSet;
-import java.util.Map;
 import java.util.Set;
-import java.util.WeakHashMap;
 import lombok.extern.slf4j.Slf4j;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -43,8 +44,7 @@ public class ClassLoaderMatcher {
       extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
     public static final SkipClassLoaderMatcher INSTANCE = new SkipClassLoaderMatcher();
     /* Cache of classloader-instance -> (true|false). True = skip instrumentation. False = safe to instrument. */
-    private static final Map<ClassLoader, Boolean> SKIP_CACHE =
-        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
+    private static final WeakMap<ClassLoader, Boolean> SKIP_CACHE = newWeakMap();
     private static final Set<String> CLASSLOADER_CLASSES_TO_SKIP;
 
     static {
@@ -131,8 +131,7 @@ public class ClassLoaderMatcher {
   public static class ClassLoaderHasClassMatcher
       extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
 
-    private final Map<ClassLoader, Boolean> cache =
-        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
+    private final WeakMap<ClassLoader, Boolean> cache = newWeakMap();
 
     private final String[] names;
 
@@ -164,8 +163,7 @@ public class ClassLoaderMatcher {
   public static class ClassLoaderHasClassWithFieldMatcher
       extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
 
-    private final Map<ClassLoader, Boolean> cache =
-        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
+    private final WeakMap<ClassLoader, Boolean> cache = newWeakMap();
 
     private final String className;
     private final String fieldName;
@@ -203,8 +201,7 @@ public class ClassLoaderMatcher {
   public static class ClassLoaderHasClassWithMethodMatcher
       extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
 
-    private final Map<ClassLoader, Boolean> cache =
-        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
+    private final WeakMap<ClassLoader, Boolean> cache = newWeakMap();
 
     private final String className;
     private final String methodName;

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Utils.java

@@ -22,6 +22,8 @@ public class Utils {
     "com.google.auto",
     "com.google.common",
     "com.google.thirdparty.publicsuffix",
+    // WeakConcurrentMap
+    "com.blogspot.mydailyjava.weaklockfree",
     // bytebuddy
     "net.bytebuddy",
     // OT contribs for dd trace resolver
@@ -34,13 +36,13 @@ public class Utils {
 
   private static Method findLoadedClassMethod = null;
 
-  private static BootstrapClassLoaderProxy unitTestBootstrapProxy =
+  private static final BootstrapClassLoaderProxy unitTestBootstrapProxy =
       new BootstrapClassLoaderProxy(new URL[0], null);
 
   static {
     try {
       findLoadedClassMethod = ClassLoader.class.getDeclaredMethod("findLoadedClass", String.class);
-    } catch (NoSuchMethodException | SecurityException e) {
+    } catch (final NoSuchMethodException | SecurityException e) {
       throw new IllegalStateException(e);
     }
   }

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMapSuppliers.java

@@ -0,0 +1,158 @@
+package datadog.trace.agent.tooling;
+
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.MapMaker;
+import datadog.trace.bootstrap.WeakMap;
+import java.lang.ref.WeakReference;
+import java.util.Iterator;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+class WeakMapSuppliers {
+  // Comparison with using WeakConcurrentMap vs Guava's implementation:
+  // Cleaning:
+  // * `WeakConcurrentMap`: centralized but we have to maintain out own code and thread for it
+  // * `Guava`: inline on application's thread, with constant max delay
+  // Jar Size:
+  // * `WeakConcurrentMap`: small
+  // * `Guava`: large, but we may use other features, like immutable collections - and we already
+  //          ship Guava as part of distribution now, so using Guava for this doesn’t increase size.
+  // Must go on bootstrap classpath:
+  // * `WeakConcurrentMap`: version conflict is unlikely, so we can directly inject for now
+  // * `Guava`: need to implement shadow copy (might eventually be necessary for other dependencies)
+  // Used by other javaagents for similar purposes:
+  // * `WeakConcurrentMap`: anecdotally used by other agents
+  // * `Guava`: specifically agent use is unknown at the moment, but Guava is a well known library
+  //            backed by big company with many-many users
+
+  /**
+   * Provides instances of {@link WeakConcurrentMap} and retains weak reference to them to allow a
+   * single thread to clean void weak references out for all instances. Cleaning is done every
+   * second.
+   */
+  static class WeakConcurrent implements WeakMap.Supplier {
+
+    private static final long SHUTDOWN_WAIT_SECONDS = 5;
+
+    @VisibleForTesting static final long CLEAN_FREQUENCY_SECONDS = 1;
+
+    private static final ThreadFactory THREAD_FACTORY =
+        new ThreadFactory() {
+          @Override
+          public Thread newThread(final Runnable r) {
+            final Thread thread = new Thread(r, "dd-weak-ref-cleaner");
+            thread.setDaemon(true);
+            thread.setPriority(Thread.MIN_PRIORITY);
+            return thread;
+          }
+        };
+
+    private final ScheduledExecutorService cleanerExecutorService;
+
+    private final Queue<WeakReference<WeakConcurrentMap>> suppliedMaps =
+        new ConcurrentLinkedQueue<>();
+
+    private final Runnable runnable =
+        new Runnable() {
+          @Override
+          public void run() {
+            for (final Iterator<WeakReference<WeakConcurrentMap>> iterator =
+                    suppliedMaps.iterator();
+                iterator.hasNext(); ) {
+              final WeakConcurrentMap map = iterator.next().get();
+              if (map == null) {
+                iterator.remove();
+              } else {
+                map.expungeStaleEntries();
+              }
+            }
+          }
+        };
+
+    WeakConcurrent() {
+      cleanerExecutorService = Executors.newScheduledThreadPool(1, THREAD_FACTORY);
+      cleanerExecutorService.scheduleAtFixedRate(
+          runnable, CLEAN_FREQUENCY_SECONDS, CLEAN_FREQUENCY_SECONDS, TimeUnit.SECONDS);
+
+      try {
+        Runtime.getRuntime()
+            .addShutdownHook(
+                new Thread() {
+                  @Override
+                  public void run() {
+                    try {
+                      cleanerExecutorService.shutdownNow();
+                      cleanerExecutorService.awaitTermination(
+                          SHUTDOWN_WAIT_SECONDS, TimeUnit.SECONDS);
+                    } catch (final InterruptedException e) {
+                      // Don't bother waiting then...
+                    }
+                  }
+                });
+      } catch (final IllegalStateException ex) {
+        // The JVM is already shutting down.
+      }
+    }
+
+    @Override
+    public <K, V> WeakMap<K, V> get() {
+      final WeakConcurrentMap<K, V> map = new WeakConcurrentMap<>(false);
+      suppliedMaps.add(new WeakReference<WeakConcurrentMap>(map));
+      return new Adapter<>(map);
+    }
+
+    private static class Adapter<K, V> implements WeakMap<K, V> {
+      private final WeakConcurrentMap<K, V> map;
+
+      private Adapter(final WeakConcurrentMap<K, V> map) {
+        this.map = map;
+      }
+
+      @Override
+      public int size() {
+        return map.approximateSize();
+      }
+
+      @Override
+      public boolean containsKey(final K key) {
+        return map.containsKey(key);
+      }
+
+      @Override
+      public V get(final K key) {
+        return map.get(key);
+      }
+
+      @Override
+      public void put(final K key, final V value) {
+        map.put(key, value);
+      }
+    }
+
+    static class Inline implements WeakMap.Supplier {
+
+      @Override
+      public <K, V> WeakMap<K, V> get() {
+        return new Adapter<>(new WeakConcurrentMap.WithInlinedExpunction<K, V>());
+      }
+    }
+  }
+
+  static class Guava implements WeakMap.Supplier {
+
+    @Override
+    public <K, V> WeakMap<K, V> get() {
+      return new WeakMap.MapAdapter<>(new MapMaker().weakKeys().<K, V>makeMap());
+    }
+
+    public <K, V> WeakMap<K, V> get(int concurrencyLevel) {
+      return new WeakMap.MapAdapter<>(
+          new MapMaker().concurrencyLevel(concurrencyLevel).weakKeys().<K, V>makeMap());
+    }
+  }
+}

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/ReferenceMatcher.java

@@ -1,23 +1,21 @@
 package datadog.trace.agent.tooling.muzzle;
 
+import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
 import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
 
 import datadog.trace.agent.tooling.Utils;
+import datadog.trace.bootstrap.WeakMap;
 import java.util.ArrayList;
 import java.util.Arrays;
-import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
-import java.util.Map;
 import java.util.Set;
-import java.util.WeakHashMap;
 import lombok.extern.slf4j.Slf4j;
 
 /** Matches a set of references against a classloader. */
 @Slf4j
 public class ReferenceMatcher {
-  private final Map<ClassLoader, List<Reference.Mismatch>> mismatchCache =
-      Collections.synchronizedMap(new WeakHashMap<ClassLoader, List<Reference.Mismatch>>());
+  private final WeakMap<ClassLoader, List<Reference.Mismatch>> mismatchCache = newWeakMap();
   private final Reference[] references;
   private final Set<String> helperClassNames;
 

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/WeakConcurrentSupplierTest.groovy

@@ -0,0 +1,106 @@
+package datadog.trace.agent.tooling
+
+import datadog.trace.agent.test.TestUtils
+import datadog.trace.bootstrap.WeakMap
+import spock.lang.Shared
+import spock.lang.Specification
+
+import java.lang.ref.WeakReference
+import java.util.concurrent.TimeUnit
+
+class WeakConcurrentSupplierTest extends Specification {
+  @Shared
+  def weakConcurrentSupplier = new WeakMapSuppliers.WeakConcurrent()
+  @Shared
+  def weakInlineSupplier = new WeakMapSuppliers.WeakConcurrent.Inline()
+  @Shared
+  def guavaSupplier = new WeakMapSuppliers.Guava()
+
+  def "Calling newWeakMap on #name creates independent maps"() {
+    setup:
+    WeakMap.Provider.provider.set(supplier)
+    def key = new Object()
+    def map1 = WeakMap.Provider.newWeakMap()
+    def map2 = WeakMap.Provider.newWeakMap()
+
+    when:
+    map1.put(key, "value1")
+    map2.put(key, "value2")
+
+    then:
+    map1.get(key) == "value1"
+    map2.get(key) == "value2"
+
+    where:
+    name             | supplier
+    "WeakConcurrent" | weakConcurrentSupplier
+    "WeakInline"     | weakInlineSupplier
+    "Guava"          | guavaSupplier
+  }
+
+  def "Unreferenced map get cleaned up on #name"() {
+    setup:
+    WeakMap.Provider.provider.set(supplier)
+    def map = WeakMap.Provider.newWeakMap()
+    def ref = new WeakReference(map)
+
+    when:
+    map = null
+    TestUtils.awaitGC()
+
+    then:
+    ref.get() == null
+
+    where:
+    name             | supplier
+    "WeakConcurrent" | weakConcurrentSupplier
+    "WeakInline"     | weakInlineSupplier
+    "Guava"          | guavaSupplier
+  }
+
+  def "Unreferenced keys get cleaned up on #name"() {
+    setup:
+    def key = new Object()
+    map.put(key, "value")
+    TestUtils.awaitGC()
+
+    expect:
+    map.size() == 1
+
+    when:
+    key = null
+    TestUtils.awaitGC()
+
+    if (name == "WeakConcurrent") {
+      // Sleep enough time for cleanup thread to get scheduled.
+      // But on a very slow box (or high load) scheduling may not be exactly predictable
+      // so we try a few times.
+      int count = 0
+      while (map.size() != 0 && count < 10) {
+        Thread.sleep(TimeUnit.SECONDS.toMillis(WeakMapSuppliers.WeakConcurrent.CLEAN_FREQUENCY_SECONDS))
+        count++
+      }
+    }
+
+    // Hit map a few times to trigger unreferenced entries cleanup.
+    // Exact number of times that we need to hit map is implementation dependent.
+    // For Guava it i specified in
+    // com.google.common.collect.MapMakerInternalMap.DRAIN_THRESHOLD = 0x3F
+    if (name == "Guava" || name == "WeakInline") {
+      for (int i = 0; i <= 0x3F; i++) {
+        map.get("test")
+      }
+    }
+
+    then:
+    map.size() == 0
+
+    where:
+    name             | map
+    "WeakConcurrent" | weakConcurrentSupplier.get()
+    "WeakInline"     | weakInlineSupplier.get()
+    // Guava's cleanup process depends on concurrency level,
+    // and in order to be able to test it we need to set concurrency to 1
+    "Guava"          | guavaSupplier.get(1)
+  }
+}

dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java

@@ -1,6 +1,7 @@
 package datadog.trace.instrumentation.http_url_connection;
 
 import static datadog.trace.agent.tooling.ByteBuddyElementMatchers.safeHasSuperType;
+import static datadog.trace.bootstrap.WeakMap.Provider.newWeakMap;
 import static io.opentracing.log.Fields.ERROR_OBJECT;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -12,6 +13,7 @@ import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.api.DDSpanTypes;
 import datadog.trace.api.DDTags;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.WeakMap;
 import io.opentracing.Scope;
 import io.opentracing.Span;
 import io.opentracing.Tracer;
@@ -21,7 +23,6 @@ import java.net.HttpURLConnection;
 import java.net.URL;
 import java.util.Collections;
 import java.util.Map;
-import java.util.WeakHashMap;
 import javax.net.ssl.HttpsURLConnection;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -175,8 +176,7 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default {
   }
 
   public static class HttpURLState {
-    private static final Map<HttpURLConnection, HttpURLState> STATE_MAP =
-        Collections.synchronizedMap(new WeakHashMap<HttpURLConnection, HttpURLState>());
+    private static final WeakMap<HttpURLConnection, HttpURLState> STATE_MAP = newWeakMap();
 
     public static HttpURLState get(final HttpURLConnection connection) {
       HttpURLState state = STATE_MAP.get(connection);

dd-java-agent/testing/src/main/java/datadog/trace/agent/test/TestUtils.java

@@ -13,6 +13,7 @@ import java.io.File;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
+import java.lang.ref.WeakReference;
 import java.lang.reflect.Field;
 import java.net.ServerSocket;
 import java.net.URL;
@@ -204,4 +205,13 @@ public class TestUtils {
       return -1;
     }
   }
+
+  public static void awaitGC() {
+    Object obj = new Object();
+    final WeakReference ref = new WeakReference<>(obj);
+    obj = null;
+    while (ref.get() != null) {
+      System.gc();
+    }
+  }
 }

dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java

@@ -29,6 +29,7 @@ import io.opentracing.ScopeManager;
 import io.opentracing.Span;
 import io.opentracing.SpanContext;
 import io.opentracing.propagation.Format;
+import java.io.Closeable;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -48,7 +49,7 @@ import lombok.extern.slf4j.Slf4j;
 
 /** DDTracer makes it easy to send traces and span to DD using the OpenTracing API. */
 @Slf4j
-public class DDTracer implements io.opentracing.Tracer {
+public class DDTracer implements io.opentracing.Tracer, Closeable {
 
   public static final String UNASSIGNED_DEFAULT_SERVICE_NAME = "unnamed-java-app";
 
@@ -302,6 +303,7 @@ public class DDTracer implements io.opentracing.Tracer {
     }
   }
 
+  @Override
   public void close() {
     PendingTrace.close();
     writer.close();