Use WeakConcurrentMap

Instead of `Collections.synchronizedMap(new WeakHashMap<>())`.

dd-java-agent/agent-bootstrap/agent-bootstrap.gradle

@@ -9,10 +9,13 @@ excludedClassesConverage += ['datadog.trace.bootstrap.*']
 
 dependencies {
   compile project(':dd-trace-api')
+  compile group: 'com.blogspot.mydailyjava', name: 'weak-lock-free', version: '0.13'
   compile deps.opentracing
   compile deps.slf4j
   compile group: 'org.slf4j', name: 'slf4j-simple', version: versions.slf4j
   // ^ Generally a bad idea for libraries, but we're shadowing.
+
+  testCompile project(':dd-java-agent:testing')
 }
 
 jar {

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

@@ -1,10 +1,10 @@
 package datadog.trace.bootstrap;
 
+import static datadog.trace.bootstrap.WeakMapManager.newWeakMap;
+
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 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 +13,9 @@ import lombok.Data;
  * <p>In the bootstrap project to ensure visibility by all classes.
  */
 public class JDBCMaps {
-  public static final Map<Connection, DBInfo> connectionInfo =
+  public static final WeakConcurrentMap<Connection, DBInfo> connectionInfo = newWeakMap();
-      Collections.synchronizedMap(new WeakHashMap<Connection, DBInfo>());
+  public static final WeakConcurrentMap<PreparedStatement, String> preparedStatements =
-  public static final Map<PreparedStatement, String> preparedStatements =
+      newWeakMap();
-      Collections.synchronizedMap(new WeakHashMap<PreparedStatement, String>());
 
   public static final String DB_QUERY = "DB Query";
 

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

@@ -0,0 +1,73 @@
+package datadog.trace.bootstrap;
+
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
+import java.util.List;
+import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+public class WeakMapManager {
+  private static final long CLEAN_FREQUENCY_SECONDS = 1;
+  private static final List<WeakConcurrentMap> maps = new CopyOnWriteArrayList<>();
+
+  private static final ThreadFactory weakMapCleanerThreadFactory =
+      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 static final ScheduledExecutorService cleaner =
+      Executors.newScheduledThreadPool(1, weakMapCleanerThreadFactory);
+
+  private static final Runnable runnable = new Cleaner();
+
+  static {
+    cleaner.scheduleAtFixedRate(
+        runnable, CLEAN_FREQUENCY_SECONDS, CLEAN_FREQUENCY_SECONDS, TimeUnit.SECONDS);
+
+    try {
+      Runtime.getRuntime()
+          .addShutdownHook(
+              new Thread() {
+                @Override
+                public void run() {
+                  try {
+                    cleaner.shutdownNow();
+                    cleaner.awaitTermination(5, TimeUnit.SECONDS);
+                  } catch (final InterruptedException e) {
+                    // Don't bother waiting then...
+                  }
+                }
+              });
+    } catch (final IllegalStateException ex) {
+      // The JVM is already shutting down.
+    }
+  }
+
+  public static <K, V> WeakConcurrentMap<K, V> newWeakMap() {
+    final WeakConcurrentMap<K, V> map = new WeakConcurrentMap<>(false);
+    maps.add(map);
+    return map;
+  }
+
+  public static void cleanMaps() {
+    for (final WeakConcurrentMap map : maps) {
+      map.expungeStaleEntries();
+    }
+  }
+
+  private static class Cleaner implements Runnable {
+
+    @Override
+    public void run() {
+      cleanMaps();
+    }
+  }
+}

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/WeakMapManagerTest.groovy

@@ -0,0 +1,41 @@
+package datadog.trace.bootstrap
+
+import datadog.trace.agent.test.TestUtils
+import spock.lang.Specification
+
+class WeakMapManagerTest extends Specification {
+
+  def setup() {
+    WeakMapManager.maps.clear()
+  }
+
+  def "calling new adds to the list"() {
+    when:
+    def map1 = WeakMapManager.newWeakMap()
+
+    then:
+    WeakMapManager.maps == [map1]
+
+    when:
+    def map2 = WeakMapManager.newWeakMap()
+
+    then:
+    WeakMapManager.maps == [map1, map2]
+  }
+
+  def "calling cleanMaps does cleanup"() {
+    setup:
+    def map = WeakMapManager.newWeakMap()
+    map.put(new Object(), "value")
+    TestUtils.awaitGC()
+
+    expect:
+    map.approximateSize() == 1
+
+    when:
+    WeakMapManager.cleanMaps()
+
+    then:
+    map.approximateSize() == 0
+  }
+}

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.WeakMapManager.newWeakMap;
+
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 import datadog.trace.bootstrap.DatadogClassLoader;
 import datadog.trace.bootstrap.PatchLogger;
 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 =
+    private static final WeakConcurrentMap<ClassLoader, Boolean> SKIP_CACHE = newWeakMap();
-        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
     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 =
+    private final WeakConcurrentMap<ClassLoader, Boolean> cache = newWeakMap();
-        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
 
     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 =
+    private final WeakConcurrentMap<ClassLoader, Boolean> cache = newWeakMap();
-        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
 
     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 =
+    private final WeakConcurrentMap<ClassLoader, Boolean> cache = newWeakMap();
-        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
 
     private final String className;
     private final String methodName;

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

@@ -1,23 +1,22 @@
 package datadog.trace.agent.tooling.muzzle;
 
+import static datadog.trace.bootstrap.WeakMapManager.newWeakMap;
 import static net.bytebuddy.dynamic.loading.ClassLoadingStrategy.BOOTSTRAP_LOADER;
 
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 import datadog.trace.agent.tooling.Utils;
 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 =
+  private final WeakConcurrentMap<ClassLoader, List<Reference.Mismatch>> mismatchCache =
-      Collections.synchronizedMap(new WeakHashMap<ClassLoader, List<Reference.Mismatch>>());
+      newWeakMap();
   private final Reference[] references;
   private final Set<String> helperClassNames;
 

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.WeakMapManager.newWeakMap;
 import static io.opentracing.log.Fields.ERROR_OBJECT;
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
@@ -8,6 +9,7 @@ import static net.bytebuddy.matcher.ElementMatchers.nameStartsWith;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.not;
 
+import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.api.DDSpanTypes;
@@ -24,7 +26,6 @@ import java.net.URL;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.WeakHashMap;
 import javax.net.ssl.HttpsURLConnection;
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.matcher.ElementMatcher;
@@ -174,8 +175,8 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default {
   }
 
   public static class HttpURLState {
-    private static final Map<HttpURLConnection, HttpURLState> STATE_MAP =
+    private static final WeakConcurrentMap<HttpURLConnection, HttpURLState> STATE_MAP =
-        Collections.synchronizedMap(new WeakHashMap<HttpURLConnection, HttpURLState>());
+        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();

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

@@ -239,6 +239,7 @@ public class PendingTrace extends ConcurrentLinkedDeque<DDSpan> {
           public Thread newThread(final Runnable r) {
             final Thread thread = new Thread(r, "dd-span-cleaner");
             thread.setDaemon(true);
+            thread.setPriority(Thread.MIN_PRIORITY);
             return thread;
           }
         };