From 7330ceaf4ee73f517b41a5bdba94441fe632c609 Mon Sep 17 00:00:00 2001
From: Andrew Kent <andrew.kent@datadoghq.com>
Date: Mon, 1 Oct 2018 14:42:54 -0400
Subject: [PATCH] Instrumentation Context outline for HttpURLConnection as an
 example

---
 .../trace/agent/tooling/Instrumenter.java     |  8 ++++++
 .../context/InstrumentationContext.java       | 12 +++++++++
 .../HttpUrlConnectionInstrumentation.java     | 27 +++++--------------
 3 files changed, 27 insertions(+), 20 deletions(-)
 create mode 100644 dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContext.java

diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
index de370a0529..96446a426e 100644
--- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
+++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
@@ -8,6 +8,7 @@ import datadog.trace.agent.tooling.muzzle.Reference;
 import datadog.trace.agent.tooling.muzzle.ReferenceMatcher;
 import java.security.ProtectionDomain;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -187,6 +188,13 @@ public interface Instrumenter {
     @Override
     public abstract Map<? extends ElementMatcher, String> transformers();
 
+    /**
+     * A map of {class-name -> context-class-name}. Keys (and their subclasses) will be associated with a context of the value.
+     */
+    public Map<String, String> contextStore() {
+      return Collections.EMPTY_MAP;
+    }
+
     protected boolean defaultEnabled() {
       return getConfigEnabled("dd.integrations.enabled", true);
     }
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContext.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContext.java
new file mode 100644
index 0000000000..f9c8d91ee8
--- /dev/null
+++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContext.java
@@ -0,0 +1,12 @@
+package datadog.trace.agent.tooling.context;
+
+/**
+ * Instrumentation Context API
+ */
+public class InstrumentationContext {
+  private InstrumentationContext() {}
+
+  public static <T> T get(Object contextInstance, Class<T> contextClass) {
+    throw new RuntimeException("calls to this method should be rewritten");
+  }
+}
diff --git a/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java b/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java
index 5565dfe17d..9d793e1269 100644
--- a/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java
+++ b/dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java
@@ -10,6 +10,7 @@ import static net.bytebuddy.matcher.ElementMatchers.not;
 
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.context.InstrumentationContext;
 import datadog.trace.api.DDSpanTypes;
 import datadog.trace.api.DDTags;
 import datadog.trace.bootstrap.CallDepthThreadLocalMap;
@@ -52,6 +53,10 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default {
     return new String[] {HttpUrlConnectionInstrumentation.class.getName() + "$HttpURLState"};
   }
 
+  public Map<String, String> contextStore() {
+    return Collections.singletonMap("java.net.HttpURLConnection", getClass().getName() + "$HttpURLState");
+  }
+
   @Override
   public Map<ElementMatcher, String> transformers() {
     return Collections.<ElementMatcher, String>singletonMap(
@@ -69,7 +74,8 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default {
         @Advice.FieldValue("connected") final boolean connected,
         @Advice.Origin("#m") final String methodName) {
 
-      final HttpURLState state = HttpURLState.get(thiz);
+      final HttpURLState state = InstrumentationContext.get(thiz, HttpURLState.class);
+
       String operationName = "http.request";
 
       switch (methodName) {
@@ -177,26 +183,7 @@ public class HttpUrlConnectionInstrumentation extends Instrumenter.Default {
   }
 
   public static class HttpURLState {
-    private static final WeakMap<HttpURLConnection, HttpURLState> STATE_MAP = newWeakMap();
-
-    public static HttpURLState get(final HttpURLConnection connection) {
-      HttpURLState state = STATE_MAP.get(connection);
-      if (state == null) {
-        synchronized (connection) {
-          // might not be a good idea to synchronize on a method parameter...
-          state = STATE_MAP.get(connection);
-          if (state == null) {
-            state = new HttpURLState();
-            STATE_MAP.put(connection, state);
-          }
-        }
-      }
-      return state;
-    }
-
     public boolean calledOutputStream = false;
     public boolean calledInputStream = false;
-
-    private HttpURLState() {}
   }
 }