From de6f33c035a982f854529c06ed54a2fb35558c97 Mon Sep 17 00:00:00 2001
From: Luca Abbati <luca.abbati@datadoghq.com>
Date: Mon, 8 Jul 2019 18:41:48 -0400
Subject: [PATCH] Draft of per-classloader netty AttributeKey definition

---
 .../tooling/ClassLoaderScopedWeakMap.java     | 33 +++++++++
 .../netty40/AttributeKeys.java                | 71 +++++++++++--------
 2 files changed, 73 insertions(+), 31 deletions(-)
 create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderScopedWeakMap.java

diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderScopedWeakMap.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderScopedWeakMap.java
new file mode 100644
index 0000000000..e996f76651
--- /dev/null
+++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ClassLoaderScopedWeakMap.java
@@ -0,0 +1,33 @@
+package datadog.trace.agent.tooling;
+
+import datadog.trace.bootstrap.WeakMap;
+import java.util.HashMap;
+import java.util.Map;
+
+public class ClassLoaderScopedWeakMap {
+
+  public static final ClassLoaderScopedWeakMap INSTANCE = new ClassLoaderScopedWeakMap();
+
+  private final WeakMap<ClassLoader, Map<Object, Object>> map = WeakMap.Supplier.DEFAULT.get();
+
+  public synchronized Object getOrCreate(
+      ClassLoader classLoader, Object key, Supplier valueSupplier) {
+    Map<Object, Object> classLoaderMap = map.get(classLoader);
+    if (classLoaderMap == null) {
+      classLoaderMap = new HashMap<>();
+      map.put(classLoader, classLoaderMap);
+    }
+
+    if (classLoaderMap.containsKey(key)) {
+      return classLoaderMap.get(key);
+    }
+
+    Object value = valueSupplier.get();
+    classLoaderMap.put(key, value);
+    return value;
+  }
+
+  public interface Supplier {
+    Object get();
+  }
+}
diff --git a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java
index 3aae4b30c8..7b21590ce8 100644
--- a/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java
+++ b/dd-java-agent/instrumentation/netty-4.0/src/main/java/datadog/trace/instrumentation/netty40/AttributeKeys.java
@@ -1,5 +1,6 @@
 package datadog.trace.instrumentation.netty40;
 
+import datadog.trace.agent.tooling.ClassLoaderScopedWeakMap;
 import datadog.trace.context.TraceScope;
 import datadog.trace.instrumentation.netty40.client.HttpClientTracingHandler;
 import datadog.trace.instrumentation.netty40.server.HttpServerTracingHandler;
@@ -8,41 +9,49 @@ import io.opentracing.Span;
 
 public class AttributeKeys {
 
+  private static final String PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY_NAME =
+      "datadog.trace.instrumentation.netty40.parent.connect.continuation";
+
   public static final AttributeKey<TraceScope.Continuation>
       PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY =
-          new AttributeKey<>(
-              buildContextSpecificKey(
-                  "datadog.trace.instrumentation.netty40.parent.connect.continuation"));
+          (AttributeKey<TraceScope.Continuation>)
+              ClassLoaderScopedWeakMap.INSTANCE.getOrCreate(
+                  AttributeKey.class.getClassLoader(),
+                  PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY_NAME,
+                  new ClassLoaderScopedWeakMap.Supplier() {
+                    @Override
+                    public Object get() {
+                      return new AttributeKey<>(PARENT_CONNECT_CONTINUATION_ATTRIBUTE_KEY_NAME);
+                    }
+                  });
+
+  private static final String SERVER_ATTRIBUTE_KEY_NAME =
+      HttpServerTracingHandler.class.getName() + ".span";
 
   public static final AttributeKey<Span> SERVER_ATTRIBUTE_KEY =
-      new AttributeKey<>(
-          buildContextSpecificKey(HttpServerTracingHandler.class.getName() + ".span"));
+      (AttributeKey<Span>)
+          ClassLoaderScopedWeakMap.INSTANCE.getOrCreate(
+              AttributeKey.class.getClassLoader(),
+              SERVER_ATTRIBUTE_KEY_NAME,
+              new ClassLoaderScopedWeakMap.Supplier() {
+                @Override
+                public Object get() {
+                  return new AttributeKey<>(SERVER_ATTRIBUTE_KEY_NAME);
+                }
+              });
+
+  private static final String CLIENT_ATTRIBUTE_KEY_NAME =
+      HttpClientTracingHandler.class.getName() + ".span";
 
   public static final AttributeKey<Span> CLIENT_ATTRIBUTE_KEY =
-      new AttributeKey<>(
-          buildContextSpecificKey(HttpClientTracingHandler.class.getName() + ".span"));
-
-  /**
-   * Netty 4.0 before 4.0.26 handled differently how unique attributes where handled, with 4.0.26+
-   * being more lenient with duplicates. We found a use case in Apache Atlas 1.1.0 where for some
-   * reason, this class gets loaded by multiple class loaders generating an error in 4.0.25- before
-   * an exception was thrown if that attribute was already defined.
-   *
-   * @param simpleKey The logical key assigned.
-   * @return A key scoped to the current class loader in use, if not null.
-   */
-  private static String buildContextSpecificKey(final String simpleKey) {
-    ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader();
-    String key = simpleKey;
-    if (contextClassLoader != null) {
-      key =
-          "ClassLoader."
-              + contextClassLoader.getClass().getName()
-              + "."
-              + contextClassLoader.hashCode()
-              + "."
-              + key;
-    }
-    return key;
-  }
+      (AttributeKey<Span>)
+          ClassLoaderScopedWeakMap.INSTANCE.getOrCreate(
+              AttributeKey.class.getClassLoader(),
+              CLIENT_ATTRIBUTE_KEY_NAME,
+              new ClassLoaderScopedWeakMap.Supplier() {
+                @Override
+                public Object get() {
+                  return new AttributeKey<>(CLIENT_ATTRIBUTE_KEY_NAME);
+                }
+              });
 }