Merge pull request #203 from DataDog/ark/the_weak_must_die

Use a thread-safe synchronized weak hash map

Author: realark

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/ConnectionInstrumentation.java

@@ -14,6 +14,7 @@
 import datadog.trace.agent.tooling.Instrumenter;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
+import java.util.Collections;
 import java.util.Map;
 import java.util.WeakHashMap;
 import lombok.Data;
@@ -22,8 +23,10 @@
 
 @AutoService(Instrumenter.class)
 public final class ConnectionInstrumentation implements Instrumenter {
-  public static final Map<Connection, DBInfo> connectionInfo = new WeakHashMap<>();
-  public static final Map<PreparedStatement, String> preparedStatements = new WeakHashMap<>();
+  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>());
 
   @Override
   public AgentBuilder instrument(final AgentBuilder agentBuilder) {

dd-java-agent/instrumentation/spring-web/src/main/java/datadog/trace/instrumentation/springweb/SpringWebInstrumentation.java

@@ -16,17 +16,13 @@
 import datadog.trace.api.DDTags;
 import io.opentracing.Scope;
 import io.opentracing.util.GlobalTracer;
-import java.sql.PreparedStatement;
-import java.util.Map;
-import java.util.WeakHashMap;
 import javax.servlet.http.HttpServletRequest;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 import org.springframework.web.servlet.HandlerMapping;
 
 @AutoService(Instrumenter.class)
 public final class SpringWebInstrumentation implements Instrumenter {
-  public static final Map<PreparedStatement, String> preparedStatements = new WeakHashMap<>();
 
   @Override
   public AgentBuilder instrument(final AgentBuilder agentBuilder) {

dd-java-agent/instrumentation/trace-annotation/src/main/java/datadog/trace/instrumentation/trace_annotation/TraceAnnotationInstrumentation.java

@@ -13,16 +13,12 @@
 import io.opentracing.tag.Tags;
 import io.opentracing.util.GlobalTracer;
 import java.lang.reflect.Method;
-import java.sql.PreparedStatement;
 import java.util.Collections;
-import java.util.Map;
-import java.util.WeakHashMap;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.asm.Advice;
 
 @AutoService(Instrumenter.class)
 public final class TraceAnnotationInstrumentation implements Instrumenter {
-  public static final Map<PreparedStatement, String> preparedStatements = new WeakHashMap<>();
 
   @Override
   public AgentBuilder instrument(final AgentBuilder agentBuilder) {

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

@@ -1,5 +1,7 @@
 package datadog.trace.agent.tooling;
 
+import java.util.Collections;
+import java.util.Map;
 import java.util.WeakHashMap;
 import net.bytebuddy.matcher.ElementMatcher;
 
@@ -46,7 +48,8 @@ public boolean matches(final ClassLoader target) {
   public static class ClassLoaderHasClassMatcher
       extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
 
-    private final WeakHashMap<ClassLoader, Boolean> cache = new WeakHashMap<>();
+    private final Map<ClassLoader, Boolean> cache =
+        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
 
     private final String[] names;
 
@@ -80,7 +83,8 @@ public boolean matches(final ClassLoader target) {
   public static class ClassLoaderHasClassWithFieldMatcher
       extends ElementMatcher.Junction.AbstractBase<ClassLoader> {
 
-    private final WeakHashMap<ClassLoader, Boolean> cache = new WeakHashMap<>();
+    private final Map<ClassLoader, Boolean> cache =
+        Collections.synchronizedMap(new WeakHashMap<ClassLoader, Boolean>());
 
     private final String className;
     private final String fieldName;