Use WeakMap provider in global registry or context store in instrumentations (#187)

* Use dynamic context store

Signed-off-by: Pavol Loffay <[email protected]>

* use map provider

Signed-off-by: Pavol Loffay <[email protected]>

Author: pavolloffay

instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/ApacheClientInstrumentationModule.java

@@ -30,6 +30,8 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.api.Java8BytecodeBridge;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
@@ -72,6 +74,13 @@ public int getOrder() {
     return 1;
   }
 
+  @Override
+  protected Map<String, String> contextStore() {
+    Map<String, String> context = new HashMap<>();
+    context.put("org.apache.http.HttpEntity", Span.class.getName());
+    return context;
+  }
+
   @Override
   public List<TypeInstrumentation> typeInstrumentations() {
     return Arrays.asList(new HttpEntityInstrumentation(), new ApacheClientInstrumentation());
@@ -118,7 +127,9 @@ public static boolean enter(@Advice.Argument(0) HttpMessage request) {
       if (callDepth > 0) {
         return false;
       }
-      ApacheHttpClientUtils.traceRequest(request);
+      ContextStore<HttpEntity, Span> contextStore =
+          InstrumentationContext.get(HttpEntity.class, Span.class);
+      ApacheHttpClientUtils.traceRequest(contextStore, request);
       return true;
     }
 
@@ -138,7 +149,9 @@ public static boolean enter(@Advice.Argument(1) HttpMessage request) {
       if (callDepth > 0) {
         return false;
       }
-      ApacheHttpClientUtils.traceRequest(request);
+      ContextStore<HttpEntity, Span> contextStore =
+          InstrumentationContext.get(HttpEntity.class, Span.class);
+      ApacheHttpClientUtils.traceRequest(contextStore, request);
       return true;
     }
 
@@ -178,8 +191,13 @@ public static void exit(@Advice.Return Object response, @Advice.Enter boolean re
 
         if (agentConfig.getDataCapture().getHttpBody().getResponse().getValue()) {
           HttpEntity entity = httpResponse.getEntity();
+          ContextStore<HttpEntity, Span> contextStore =
+              InstrumentationContext.get(HttpEntity.class, Span.class);
           ApacheHttpClientUtils.traceEntity(
-              currentSpan, HypertraceSemanticAttributes.HTTP_RESPONSE_BODY.getKey(), entity);
+              contextStore,
+              currentSpan,
+              HypertraceSemanticAttributes.HTTP_RESPONSE_BODY.getKey(),
+              entity);
         }
       }
     }
@@ -213,7 +231,9 @@ static class HttpEntity_GetContentAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void exit(@Advice.This HttpEntity thizz, @Advice.Return InputStream inputStream) {
       // here the Span.current() is finished for response entities
-      Span clientSpan = ApacheHttpClientObjectRegistry.httpEntityToSpanMap.remove(thizz);
+      ContextStore<HttpEntity, Span> contextStore =
+          InstrumentationContext.get(HttpEntity.class, Span.class);
+      Span clientSpan = contextStore.get(thizz);
       // HttpEntity might be wrapped multiple times
       // this ensures that the advice runs only for the most outer one
       // the returned inputStream is put into globally accessible map
@@ -250,7 +270,9 @@ static class HttpEntity_WriteToAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(
         @Advice.This HttpEntity thizz, @Advice.Argument(0) OutputStream outputStream) {
-      if (!ApacheHttpClientObjectRegistry.httpEntityToSpanMap.containsKey(thizz)) {
+      ContextStore<HttpEntity, Span> contextStore =
+          InstrumentationContext.get(HttpEntity.class, Span.class);
+      if (contextStore.get(thizz) == null) {
         return;
       }
 
@@ -267,7 +289,9 @@ public static void enter(
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
     public static void exit(
         @Advice.This HttpEntity thizz, @Advice.Argument(0) OutputStream outputStream) {
-      Span clientSpan = ApacheHttpClientObjectRegistry.httpEntityToSpanMap.remove(thizz);
+      ContextStore<HttpEntity, Span> contextStore =
+          InstrumentationContext.get(HttpEntity.class, Span.class);
+      Span clientSpan = contextStore.get(thizz);
       if (clientSpan == null) {
         return;
       }

instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/ApacheHttpClientObjectRegistry.java

@@ -18,6 +18,7 @@
 
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UnsupportedEncodingException;
@@ -59,7 +60,8 @@ private static void addHeaders(
     }
   }
 
-  public static void traceRequest(HttpMessage request) {
+  public static void traceRequest(
+      ContextStore<HttpEntity, Span> contextStore, HttpMessage request) {
     Span currentSpan = Span.current();
     AgentConfig agentConfig = HypertraceConfig.get();
     if (agentConfig.getDataCapture().getHttpHeaders().getRequest().getValue()) {
@@ -71,11 +73,18 @@ public static void traceRequest(HttpMessage request) {
       HttpEntityEnclosingRequest entityRequest = (HttpEntityEnclosingRequest) request;
       HttpEntity entity = entityRequest.getEntity();
       ApacheHttpClientUtils.traceEntity(
-          currentSpan, HypertraceSemanticAttributes.HTTP_REQUEST_BODY.getKey(), entity);
+          contextStore,
+          currentSpan,
+          HypertraceSemanticAttributes.HTTP_REQUEST_BODY.getKey(),
+          entity);
     }
   }
 
-  public static void traceEntity(Span span, String bodyAttributeKey, HttpEntity entity) {
+  public static void traceEntity(
+      ContextStore<HttpEntity, Span> contextStore,
+      Span span,
+      String bodyAttributeKey,
+      HttpEntity entity) {
     if (entity == null) {
       return;
     }
@@ -105,6 +114,6 @@ public static void traceEntity(Span span, String bodyAttributeKey, HttpEntity en
     // request body is traced via HttpEntity.writeTo(OutputStream) and OutputStream instrumentation
     // response body is traced via InputStream HttpEntity.getContent() and InputStream
     // instrumentation
-    ApacheHttpClientObjectRegistry.httpEntityToSpanMap.put(entity, span);
+    contextStore.put(entity, span);
   }
 }

instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/ApacheHttpClientUtils.java

@@ -25,6 +25,7 @@ val versions: Map<String, String> by extra
 
 dependencies {
     api("io.opentelemetry:opentelemetry-api:${versions["opentelemetry"]}")
+    api("io.opentelemetry.javaagent:opentelemetry-javaagent-api:${versions["opentelemetry_java_agent"]}")
     implementation("org.slf4j:slf4j-api:1.7.30")
 
     api("com.google.protobuf:protobuf-java:3.11.4")

javaagent-core/build.gradle.kts

@@ -16,9 +16,9 @@
 
 package org.hypertrace.agent.core;
 
-import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.javaagent.instrumentation.api.WeakMap;
 import java.io.ByteArrayOutputStream;
 import java.io.InputStream;
 import java.io.OutputStream;
@@ -27,17 +27,17 @@
 public class GlobalObjectRegistry {
 
   // original input stream to span and byte buffer
-  public static final WeakConcurrentMap<InputStream, SpanAndBuffer> inputStreamToSpanAndBufferMap =
-      new WeakConcurrentMap<>(false);
+  public static final WeakMap<InputStream, SpanAndBuffer> inputStreamToSpanAndBufferMap =
+      WeakMap.Provider.newWeakMap();
 
   // original output stream to byte buffer
-  public static final WeakConcurrentMap<OutputStream, ByteArrayOutputStream>
-      outputStreamToBufferMap = new WeakConcurrentMap<>(false);
+  public static final WeakMap<OutputStream, ByteArrayOutputStream> outputStreamToBufferMap =
+      WeakMap.Provider.newWeakMap();
 
   // original input stream to buffered one
-  public static final WeakConcurrentMap<InputStream, InputStream> inputStreamMap =
-      new WeakConcurrentMap<>(false);
-  public static final WeakConcurrentMap<Object, Object> objectMap = new WeakConcurrentMap<>(false);
+  public static final WeakMap<InputStream, InputStream> inputStreamMap =
+      WeakMap.Provider.newWeakMap();
+  public static final WeakMap<Object, Object> objectMap = WeakMap.Provider.newWeakMap();
 
   public static class SpanAndBuffer {
     public final Span span;