Java streams instrumentation tests (#188)

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

Author: pavolloffay

instrumentation/java-streams/build.gradle.kts

@@ -1,6 +1,14 @@
 plugins {
     `java-library`
     id("net.bytebuddy.byte-buddy")
+    id("io.opentelemetry.instrumentation.auto-instrumentation")
+    muzzle
+}
+
+muzzle {
+    pass {
+        coreJdk()
+    }
 }
 
 afterEvaluate{
@@ -10,3 +18,7 @@ afterEvaluate{
             project(":javaagent-tooling").configurations["instrumentationMuzzle"] + configurations.runtimeClasspath
     ).configure()
 }
+
+dependencies {
+    testImplementation(project(":testing-common"))
+}

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

@@ -37,6 +37,7 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.hypertrace.agent.core.GlobalObjectRegistry;
+import org.hypertrace.agent.core.GlobalObjectRegistry.SpanAndBuffer;
 
 /**
  * {@link InputStream} instrumentation. The type matcher applies to all implementations. However
@@ -67,7 +68,7 @@ static class InputStreamInstrumentation implements TypeInstrumentation {
 
     @Override
     public ElementMatcher<? super TypeDescription> typeMatcher() {
-      return extendsClass(named("java.io.InputStream"));
+      return extendsClass(named(InputStream.class.getName()));
     }
 
     @Override
@@ -93,7 +94,7 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
               + "$InputStream_ReadByteArrayOffsetAdvice");
       transformers.put(
           named("readAllBytes").and(takesArguments(0)).and(isPublic()),
-          InputStream_ReadAllBytes.class.getName());
+          InputStreamInstrumentationModule.class.getName() + "$InputStream_ReadAllBytes");
       transformers.put(
           named("readNBytes")
               .and(takesArguments(0))
@@ -107,50 +108,93 @@ 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);
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void exit(@Advice.This InputStream thizz, @Advice.Return int read) {
-      InputStreamUtils.read(thizz, read);
+    public static void exit(
+        @Advice.This InputStream thizz,
+        @Advice.Return int read,
+        @Advice.Enter SpanAndBuffer spanAndBuffer) {
+      if (spanAndBuffer != null) {
+        InputStreamUtils.read(thizz, spanAndBuffer, read);
+      }
     }
   }
 
   public static class InputStream_ReadByteArrayAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
+      return InputStreamUtils.check(thizz);
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void exit(
-        @Advice.This InputStream thizz, @Advice.Return int read, @Advice.Argument(0) byte b[]) {
-      InputStreamUtils.read(thizz, read, b);
+        @Advice.This InputStream thizz,
+        @Advice.Return int read,
+        @Advice.Argument(0) byte b[],
+        @Advice.Enter SpanAndBuffer spanAndBuffer) {
+      if (spanAndBuffer != null) {
+        InputStreamUtils.read(thizz, spanAndBuffer, read, b);
+      }
     }
   }
 
   public static class InputStream_ReadByteArrayOffsetAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
+      return InputStreamUtils.check(thizz);
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void exit(
         @Advice.This InputStream thizz,
         @Advice.Return int read,
         @Advice.Argument(0) byte b[],
         @Advice.Argument(1) int off,
-        @Advice.Argument(2) int len) {
-      InputStreamUtils.read(thizz, read, b, off, len);
+        @Advice.Argument(2) int len,
+        @Advice.Enter SpanAndBuffer spanAndBuffer) {
+      if (spanAndBuffer != null) {
+        InputStreamUtils.read(thizz, spanAndBuffer, 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);
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
-    public static void exit(@Advice.This InputStream thizz, @Advice.Return byte[] b)
+    public static void exit(
+        @Advice.This InputStream thizz,
+        @Advice.Return byte[] b,
+        @Advice.Enter SpanAndBuffer spanAndBuffer)
         throws IOException {
-      InputStreamUtils.readAll(thizz, b);
+      if (spanAndBuffer != null) {
+        InputStreamUtils.readAll(thizz, spanAndBuffer, b);
+      }
     }
   }
 
   public static class InputStream_ReadNBytes {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
+      return InputStreamUtils.check(thizz);
+    }
+
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void exit(
         @Advice.This InputStream thizz,
         @Advice.Return int read,
         @Advice.Argument(0) byte[] b,
         @Advice.Argument(1) int off,
-        @Advice.Argument(2) int len)
-        throws IOException {
-      InputStreamUtils.readNBytes(thizz, read, b, off, len);
+        @Advice.Argument(2) int len,
+        @Advice.Enter SpanAndBuffer spanAndBuffer) {
+      InputStreamUtils.readNBytes(thizz, spanAndBuffer, read, b, off, len);
     }
   }
 }

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

@@ -21,6 +21,7 @@
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Context;
+import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -69,12 +70,21 @@ public static void addBody(
     }
   }
 
-  public static void read(InputStream inputStream, int read) {
+  public static SpanAndBuffer check(InputStream inputStream) {
     SpanAndBuffer spanAndBuffer =
         GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
     if (spanAndBuffer == null) {
-      return;
+      return null;
     }
+
+    int callDepth = CallDepthThreadLocalMap.incrementCallDepth(InputStream.class);
+    if (callDepth > 0) {
+      return null;
+    }
+    return spanAndBuffer;
+  }
+
+  public static void read(InputStream inputStream, SpanAndBuffer spanAndBuffer, int read) {
     if (read != -1) {
       spanAndBuffer.byteArrayBuffer.write((byte) read);
     } else if (read == -1) {
@@ -85,14 +95,11 @@ public static void read(InputStream inputStream, int read) {
           spanAndBuffer.charset);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
     }
+    CallDepthThreadLocalMap.reset(InputStream.class);
   }
 
-  public static void read(InputStream inputStream, int read, byte[] b) {
-    SpanAndBuffer spanAndBuffer =
-        GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
-    if (spanAndBuffer == null) {
-      return;
-    }
+  public static void read(
+      InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b) {
     if (read > 0) {
       spanAndBuffer.byteArrayBuffer.write(b, 0, read);
     } else if (read == -1) {
@@ -103,14 +110,11 @@ public static void read(InputStream inputStream, int read, byte[] b) {
           spanAndBuffer.charset);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
     }
+    CallDepthThreadLocalMap.reset(InputStream.class);
   }
 
-  public static void read(InputStream inputStream, int read, byte[] b, int off, int len) {
-    SpanAndBuffer spanAndBuffer =
-        GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
-    if (spanAndBuffer == null) {
-      return;
-    }
+  public static void read(
+      InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b, int off, int len) {
     if (read > 0) {
       spanAndBuffer.byteArrayBuffer.write(b, off, read);
     } else if (read == -1) {
@@ -121,24 +125,18 @@ public static void read(InputStream inputStream, int read, byte[] b, int off, in
           spanAndBuffer.charset);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
     }
+    CallDepthThreadLocalMap.reset(InputStream.class);
   }
 
-  public static void readAll(InputStream inputStream, byte[] b) throws IOException {
-    SpanAndBuffer spanAndBuffer =
-        GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
-    if (spanAndBuffer == null) {
-      return;
-    }
+  public static void readAll(InputStream inputStream, SpanAndBuffer spanAndBuffer, byte[] b)
+      throws IOException {
     spanAndBuffer.byteArrayBuffer.write(b);
     GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
+    CallDepthThreadLocalMap.reset(InputStream.class);
   }
 
-  public static void readNBytes(InputStream inputStream, int read, byte[] b, int off, int len) {
-    SpanAndBuffer spanAndBuffer =
-        GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
-    if (spanAndBuffer == null) {
-      return;
-    }
+  public static void readNBytes(
+      InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b, int off, int len) {
     if (read == 0) {
       InputStreamUtils.addBody(
           spanAndBuffer.span,
@@ -149,5 +147,6 @@ public static void readNBytes(InputStream inputStream, int read, byte[] b, int o
     } else {
       spanAndBuffer.byteArrayBuffer.write(b, off, read);
     }
+    CallDepthThreadLocalMap.reset(InputStream.class);
   }
 }

instrumentation/java-streams/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/java/outputstream/OutputStreamInstrumentationModule.java

@@ -63,7 +63,7 @@ static class OutputStreamInstrumentation implements TypeInstrumentation {
 
     @Override
     public ElementMatcher<? super TypeDescription> typeMatcher() {
-      return extendsClass(named("java.io.OutputStream"));
+      return extendsClass(named(OutputStream.class.getName()));
     }
 
     @Override