Remove global registry and use instrumentation context (#318)

Signed-off-by: Pavol Loffay <p.loffay@gmail.com>

Author: pavolloffay

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

@@ -23,6 +23,8 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import io.opentelemetry.javaagent.instrumentation.hypertrace.apachehttpclient.v4_0.ApacheHttpClientObjectRegistry.SpanAndAttributeKey;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.io.InputStream;
@@ -37,8 +39,7 @@
 import net.bytebuddy.matcher.ElementMatcher;
 import org.apache.http.Header;
 import org.apache.http.HttpEntity;
-import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
-import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry.SpanAndBuffer;
+import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
 import org.hypertrace.agent.core.instrumentation.utils.ContentLengthUtils;
@@ -103,7 +104,8 @@ public static void exit(@Advice.This HttpEntity thizz, @Advice.Return InputStrea
               BoundedBuffersFactory.createStream((int) contentSize, charset),
               clientSpan.attributeKey,
               charset);
-      GlobalObjectRegistry.inputStreamToSpanAndBufferMap.put(inputStream, spanAndBuffer);
+      InstrumentationContext.get(InputStream.class, SpanAndBuffer.class)
+          .put(inputStream, spanAndBuffer);
     }
   }
 
@@ -131,7 +133,8 @@ public static void enter(
       BoundedByteArrayOutputStream byteArrayOutputStream =
           BoundedBuffersFactory.createStream((int) contentSize, charset);
 
-      GlobalObjectRegistry.outputStreamToBufferMap.put(outputStream, byteArrayOutputStream);
+      InstrumentationContext.get(OutputStream.class, BoundedByteArrayOutputStream.class)
+          .put(outputStream, byteArrayOutputStream);
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -144,9 +147,10 @@ public static void exit(
         return;
       }
 
-      BoundedByteArrayOutputStream bufferedOutStream =
-          GlobalObjectRegistry.outputStreamToBufferMap.get(outputStream);
-      GlobalObjectRegistry.outputStreamToBufferMap.remove(outputStream);
+      ContextStore<OutputStream, BoundedByteArrayOutputStream> contextStore =
+          InstrumentationContext.get(OutputStream.class, BoundedByteArrayOutputStream.class);
+      BoundedByteArrayOutputStream bufferedOutStream = contextStore.get(outputStream);
+      contextStore.put(outputStream, null);
       try {
         String requestBody = bufferedOutStream.toStringWithSuppliedCharset();
         spanAndAttributeKey.span.setAttribute(spanAndAttributeKey.attributeKey, requestBody);

instrumentation/java-streams/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/java/inputstream/InputStreamInstrumentationModule.java

@@ -27,6 +27,8 @@
 
 import com.google.auto.service.AutoService;
 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.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.io.IOException;
@@ -39,18 +41,17 @@
 import net.bytebuddy.description.method.MethodDescription;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
-import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry.SpanAndBuffer;
+import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
 
 /**
  * {@link InputStream} instrumentation. The type matcher applies to all implementations. However
- * only streams that are in the {@link GlobalObjectRegistry#inputStreamToSpanAndBufferMap} are
- * instrumented, otherwise the instrumentation is noop.
+ * only streams that are in the {@link ContextStore} are instrumented, otherwise the instrumentation
+ * is noop.
  *
- * <p>If the stream is in the {@link GlobalObjectRegistry#inputStreamToSpanAndBufferMap} then result
- * of read methods is also passed to the buffered stream (value) from the map. The instrumentation
- * adds buffer to span from the map when read is finished (return -1), creates new span with buffer
- * when the original span is not recording.
+ * <p>If the stream is in the {@link ContextStore} then result of read methods is also passed to the
+ * buffered stream (value) from the map. The instrumentation adds buffer to span from the map when
+ * read is finished (return -1), creates new span with buffer when the original span is not
+ * recording.
  *
  * <p>Maybe we could add optimization to instrument the input streams only when certain classes are
  * present in classloader e.g. classes from frameworks that we instrument.
@@ -117,7 +118,8 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
   public static class InputStream_ReadNoArgsAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
-      return InputStreamUtils.check(thizz);
+      return InputStreamUtils.check(
+          thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -133,14 +135,19 @@ public static void exit(
         return;
       }
 
-      InputStreamUtils.read(thizz, spanAndBuffer, read);
+      InputStreamUtils.read(
+          thizz,
+          spanAndBuffer,
+          InstrumentationContext.get(InputStream.class, SpanAndBuffer.class),
+          read);
     }
   }
 
   public static class InputStream_ReadByteArrayAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
-      return InputStreamUtils.check(thizz);
+      return InputStreamUtils.check(
+          thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -164,7 +171,8 @@ public static void exit(
   public static class InputStream_ReadByteArrayOffsetAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
-      return InputStreamUtils.check(thizz);
+      return InputStreamUtils.check(
+          thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -183,14 +191,22 @@ public static void exit(
         return;
       }
 
-      InputStreamUtils.read(thizz, spanAndBuffer, read, b, off, len);
+      InputStreamUtils.read(
+          thizz,
+          spanAndBuffer,
+          InstrumentationContext.get(InputStream.class, SpanAndBuffer.class),
+          read,
+          b,
+          off,
+          len);
     }
   }
 
   public static class InputStream_ReadAllBytes {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
-      return InputStreamUtils.check(thizz);
+      return InputStreamUtils.check(
+          thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -207,14 +223,19 @@ public static void exit(
         return;
       }
 
-      InputStreamUtils.readAll(thizz, spanAndBuffer, b);
+      InputStreamUtils.readAll(
+          thizz,
+          spanAndBuffer,
+          InstrumentationContext.get(InputStream.class, SpanAndBuffer.class),
+          b);
     }
   }
 
   public static class InputStream_ReadNBytes {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
-      return InputStreamUtils.check(thizz);
+      return InputStreamUtils.check(
+          thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -232,14 +253,35 @@ public static void exit(
       if (callDepth > 0) {
         return;
       }
-      InputStreamUtils.readNBytes(thizz, spanAndBuffer, read, b, off, len);
+      InputStreamUtils.readNBytes(
+          thizz,
+          spanAndBuffer,
+          InstrumentationContext.get(InputStream.class, SpanAndBuffer.class),
+          read,
+          b,
+          off,
+          len);
     }
   }
 
   public static class InputStream_Available {
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void exit(@Advice.This InputStream thizz, @Advice.Return int available) {
-      InputStreamUtils.available(thizz, available);
+      if (available != 0) {
+        return;
+      }
+      ContextStore<InputStream, SpanAndBuffer> contextStore =
+          InstrumentationContext.get(InputStream.class, SpanAndBuffer.class);
+
+      SpanAndBuffer spanAndBuffer = contextStore.get(thizz);
+      if (spanAndBuffer != null) {
+        InputStreamUtils.addBody(
+            spanAndBuffer.span,
+            spanAndBuffer.attributeKey,
+            spanAndBuffer.byteArrayBuffer,
+            spanAndBuffer.charset);
+        contextStore.put(thizz, null);
+      }
     }
   }
 }

instrumentation/java-streams/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/java/inputstream/InputStreamUtils.java

@@ -22,14 +22,15 @@
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
+import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
+import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.UnsupportedEncodingException;
 import java.nio.charset.Charset;
-import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
-import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry.SpanAndBuffer;
 import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
+import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -69,9 +70,9 @@ public static void addBody(
     }
   }
 
-  public static SpanAndBuffer check(InputStream inputStream) {
-    SpanAndBuffer spanAndBuffer =
-        GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
+  public static SpanAndBuffer check(
+      InputStream inputStream, ContextStore<InputStream, SpanAndBuffer> contextStore) {
+    SpanAndBuffer spanAndBuffer = contextStore.get(inputStream);
     if (spanAndBuffer == null) {
       return null;
     }
@@ -80,7 +81,11 @@ public static SpanAndBuffer check(InputStream inputStream) {
     return spanAndBuffer;
   }
 
-  public static void read(InputStream inputStream, SpanAndBuffer spanAndBuffer, int read) {
+  public static void read(
+      InputStream inputStream,
+      SpanAndBuffer spanAndBuffer,
+      ContextStore<InputStream, SpanAndBuffer> contextStore,
+      int read) {
     if (read != -1) {
       spanAndBuffer.byteArrayBuffer.write((byte) read);
     } else if (read == -1) {
@@ -89,7 +94,7 @@ public static void read(InputStream inputStream, SpanAndBuffer spanAndBuffer, in
           spanAndBuffer.attributeKey,
           spanAndBuffer.byteArrayBuffer,
           spanAndBuffer.charset);
-      GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
+      contextStore.put(inputStream, null);
     }
   }
 
@@ -103,12 +108,18 @@ public static void read(
           spanAndBuffer.attributeKey,
           spanAndBuffer.byteArrayBuffer,
           spanAndBuffer.charset);
-      GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
+      InstrumentationContext.get(InputStream.class, SpanAndBuffer.class).put(inputStream, null);
     }
   }
 
   public static void read(
-      InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b, int off, int len) {
+      InputStream inputStream,
+      SpanAndBuffer spanAndBuffer,
+      ContextStore<InputStream, SpanAndBuffer> contextStore,
+      int read,
+      byte[] b,
+      int off,
+      int len) {
     if (read > 0) {
       spanAndBuffer.byteArrayBuffer.write(b, off, read);
     } else if (read == -1) {
@@ -117,43 +128,37 @@ public static void read(
           spanAndBuffer.attributeKey,
           spanAndBuffer.byteArrayBuffer,
           spanAndBuffer.charset);
-      GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
+      contextStore.put(inputStream, null);
     }
   }
 
-  public static void readAll(InputStream inputStream, SpanAndBuffer spanAndBuffer, byte[] b)
+  public static void readAll(
+      InputStream inputStream,
+      SpanAndBuffer spanAndBuffer,
+      ContextStore<InputStream, SpanAndBuffer> contextStore,
+      byte[] b)
       throws IOException {
     spanAndBuffer.byteArrayBuffer.write(b);
-    GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
+    contextStore.put(inputStream, null);
   }
 
   public static void readNBytes(
-      InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b, int off, int len) {
+      InputStream inputStream,
+      SpanAndBuffer spanAndBuffer,
+      ContextStore<InputStream, SpanAndBuffer> contextStore,
+      int read,
+      byte[] b,
+      int off,
+      int len) {
     if (read == 0) {
       InputStreamUtils.addBody(
           spanAndBuffer.span,
           spanAndBuffer.attributeKey,
           spanAndBuffer.byteArrayBuffer,
           spanAndBuffer.charset);
-      GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
+      contextStore.put(inputStream, null);
     } else {
       spanAndBuffer.byteArrayBuffer.write(b, off, read);
     }
   }
-
-  public static void available(InputStream inputStream, int available) {
-    if (available != 0) {
-      return;
-    }
-    SpanAndBuffer spanAndBuffer =
-        GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
-    if (spanAndBuffer != null) {
-      InputStreamUtils.addBody(
-          spanAndBuffer.span,
-          spanAndBuffer.attributeKey,
-          spanAndBuffer.byteArrayBuffer,
-          spanAndBuffer.charset);
-      GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
-    }
-  }
 }