Instrumentation Context outline for HttpURLConnection as an example
MODIFY
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Instrumenter.java
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);
}

ADD
dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/context/InstrumentationContext.java
+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");
+ }
+}

MODIFY
dd-java-agent/instrumentation/http-url-connection/src/main/java/datadog/trace/instrumentation/http_url_connection/HttpUrlConnectionInstrumentation.java
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() {}
}
}

Instrumentation Context outline for HttpURLConnection as an example