Add Servlet no wrapping instrumentation

instrumentation/build.gradle.kts

@@ -33,6 +33,8 @@ dependencies{
     implementation(project(":instrumentation:servlet:servlet-2.3"))
     implementation(project(":instrumentation:servlet:servlet-3.0"))
     implementation(project(":instrumentation:servlet:servlet-3.1"))
+    implementation(project(":instrumentation:servlet:servlet-rw"))
+    implementation(project(":instrumentation:servlet:servlet-3.0-no-wrapping"))
     implementation(project(":instrumentation:spark-2.3"))
     implementation(project(":instrumentation:grpc-1.5"))
     implementation(project(":instrumentation:okhttp:okhttp-3.0"))

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

@@ -17,13 +17,16 @@
 package io.opentelemetry.javaagent.instrumentation.hypertrace.java.inputstream;
 
 import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.extendsClass;
+import static io.opentelemetry.javaagent.tooling.bytebuddy.matcher.AgentElementMatchers.safeHasSuperType;
 import static net.bytebuddy.matcher.ElementMatchers.is;
 import static net.bytebuddy.matcher.ElementMatchers.isPublic;
 import static net.bytebuddy.matcher.ElementMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.not;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.io.IOException;
@@ -68,7 +71,8 @@ static class InputStreamInstrumentation implements TypeInstrumentation {
 
     @Override
     public ElementMatcher<? super TypeDescription> typeMatcher() {
-      return extendsClass(named(InputStream.class.getName()));
+      return extendsClass(named(InputStream.class.getName()))
+          .and(not(safeHasSuperType(named("javax.servlet.ServletInputStream"))));
     }
 
     @Override
@@ -116,14 +120,20 @@ public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
       return InputStreamUtils.check(thizz);
     }
 
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
     public static void exit(
         @Advice.This InputStream thizz,
         @Advice.Return int read,
         @Advice.Enter SpanAndBuffer spanAndBuffer) {
-      if (spanAndBuffer != null) {
-        InputStreamUtils.read(thizz, spanAndBuffer, read);
+      if (spanAndBuffer == null) {
+        return;
+      }
+      int callDepth = CallDepthThreadLocalMap.decrementCallDepth(InputStream.class);
+      if (callDepth > 0) {
+        return;
       }
+
+      InputStreamUtils.read(thizz, spanAndBuffer, read);
     }
   }
 
@@ -133,15 +143,21 @@ public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
       return InputStreamUtils.check(thizz);
     }
 
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
     public static void exit(
         @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);
+      if (spanAndBuffer == null) {
+        return;
       }
+      int callDepth = CallDepthThreadLocalMap.decrementCallDepth(InputStream.class);
+      if (callDepth > 0) {
+        return;
+      }
+
+      InputStreamUtils.read(thizz, spanAndBuffer, read, b);
     }
   }
 
@@ -151,17 +167,23 @@ public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
       return InputStreamUtils.check(thizz);
     }
 
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = 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,
         @Advice.Enter SpanAndBuffer spanAndBuffer) {
-      if (spanAndBuffer != null) {
-        InputStreamUtils.read(thizz, spanAndBuffer, read, b, off, len);
+      if (spanAndBuffer == null) {
+        return;
+      }
+      int callDepth = CallDepthThreadLocalMap.decrementCallDepth(InputStream.class);
+      if (callDepth > 0) {
+        return;
       }
+
+      InputStreamUtils.read(thizz, spanAndBuffer, read, b, off, len);
     }
   }
 
@@ -171,15 +193,21 @@ public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
       return InputStreamUtils.check(thizz);
     }
 
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
     public static void exit(
         @Advice.This InputStream thizz,
         @Advice.Return byte[] b,
         @Advice.Enter SpanAndBuffer spanAndBuffer)
         throws IOException {
-      if (spanAndBuffer != null) {
-        InputStreamUtils.readAll(thizz, spanAndBuffer, b);
+      if (spanAndBuffer == null) {
+        return;
       }
+      int callDepth = CallDepthThreadLocalMap.decrementCallDepth(InputStream.class);
+      if (callDepth > 0) {
+        return;
+      }
+
+      InputStreamUtils.readAll(thizz, spanAndBuffer, b);
     }
   }
 
@@ -189,14 +217,21 @@ public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
       return InputStreamUtils.check(thizz);
     }
 
-    @Advice.OnMethodExit(suppress = Throwable.class)
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = 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,
         @Advice.Enter SpanAndBuffer spanAndBuffer) {
+      if (spanAndBuffer == null) {
+        return;
+      }
+      int callDepth = CallDepthThreadLocalMap.decrementCallDepth(InputStream.class);
+      if (callDepth > 0) {
+        return;
+      }
       InputStreamUtils.readNBytes(thizz, spanAndBuffer, read, b, off, len);
     }
   }

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

@@ -76,10 +76,7 @@ public static SpanAndBuffer check(InputStream inputStream) {
       return null;
     }
 
-    int callDepth = CallDepthThreadLocalMap.incrementCallDepth(InputStream.class);
-    if (callDepth > 0) {
-      return null;
-    }
+    CallDepthThreadLocalMap.incrementCallDepth(InputStream.class);
     return spanAndBuffer;
   }
 
@@ -94,7 +91,6 @@ public static void read(InputStream inputStream, SpanAndBuffer spanAndBuffer, in
           spanAndBuffer.charset);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
     }
-    CallDepthThreadLocalMap.reset(InputStream.class);
   }
 
   public static void read(
@@ -109,7 +105,6 @@ public static void read(
           spanAndBuffer.charset);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
     }
-    CallDepthThreadLocalMap.reset(InputStream.class);
   }
 
   public static void read(
@@ -124,14 +119,12 @@ public static void read(
           spanAndBuffer.charset);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
     }
-    CallDepthThreadLocalMap.reset(InputStream.class);
   }
 
   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(
@@ -146,7 +139,6 @@ public static void readNBytes(
     } else {
       spanAndBuffer.byteArrayBuffer.write(b, off, read);
     }
-    CallDepthThreadLocalMap.reset(InputStream.class);
   }
 
   public static void available(InputStream inputStream, int available) {

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

@@ -24,6 +24,7 @@
 import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import com.google.auto.service.AutoService;
+import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
 import io.opentelemetry.javaagent.tooling.InstrumentationModule;
 import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
 import java.io.IOException;
@@ -37,6 +38,7 @@
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
 
 /**
  * {@link OutputStream} instrumentation. The type matcher applies to all implementations. However
@@ -96,51 +98,79 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
 
   static class OutputStream_WriteIntAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static boolean enter(@Advice.This OutputStream thizz) {
-      return OutputStreamUtils.check(thizz);
+    public static BoundedByteArrayOutputStream enter(
+        @Advice.This OutputStream thizz, @Advice.Argument(0) int b) {
+      BoundedByteArrayOutputStream buffer = GlobalObjectRegistry.outputStreamToBufferMap.get(thizz);
+      if (buffer == null) {
+        return null;
+      }
+      int callDepth = CallDepthThreadLocalMap.incrementCallDepth(OutputStream.class);
+      if (callDepth > 0) {
+        return buffer;
+      }
+
+      buffer.write(b);
+      return buffer;
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
-    public static void exit(
-        @Advice.This OutputStream thizz,
-        @Advice.Argument(0) int b,
-        @Advice.Enter boolean retEnterAdvice)
-        throws IOException {
-      OutputStreamUtils.write(thizz, retEnterAdvice, b);
+    public static void exit(@Advice.Enter BoundedByteArrayOutputStream buffer) {
+      if (buffer != null) {
+        CallDepthThreadLocalMap.decrementCallDepth(OutputStream.class);
+      }
     }
   }
 
   static class OutputStream_WriteByteArrAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static boolean enter(@Advice.This OutputStream thizz) {
-      return OutputStreamUtils.check(thizz);
+    public static BoundedByteArrayOutputStream enter(
+        @Advice.This OutputStream thizz, @Advice.Argument(0) byte b[]) throws IOException {
+      BoundedByteArrayOutputStream buffer = GlobalObjectRegistry.outputStreamToBufferMap.get(thizz);
+      if (buffer == null) {
+        return null;
+      }
+      int callDepth = CallDepthThreadLocalMap.incrementCallDepth(OutputStream.class);
+      if (callDepth > 0) {
+        return buffer;
+      }
+
+      buffer.write(b);
+      return buffer;
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
-    public static void exit(
-        @Advice.This OutputStream thizz,
-        @Advice.Argument(0) byte b[],
-        @Advice.Enter boolean retEnterAdvice)
-        throws IOException {
-      OutputStreamUtils.write(thizz, retEnterAdvice, b);
+    public static void exit(@Advice.Enter BoundedByteArrayOutputStream buffer) {
+      if (buffer != null) {
+        CallDepthThreadLocalMap.decrementCallDepth(OutputStream.class);
+      }
     }
   }
 
   static class OutputStream_WriteByteArrOffsetAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static boolean enter(@Advice.This OutputStream thizz) {
-      return OutputStreamUtils.check(thizz);
-    }
-
-    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
-    public static void exit(
+    public static BoundedByteArrayOutputStream enter(
         @Advice.This OutputStream thizz,
         @Advice.Argument(0) byte b[],
         @Advice.Argument(1) int off,
-        @Advice.Argument(2) int len,
-        @Advice.Enter boolean retEnterAdvice)
-        throws IOException {
-      OutputStreamUtils.write(thizz, retEnterAdvice, b, off, len);
+        @Advice.Argument(2) int len) {
+      BoundedByteArrayOutputStream buffer = GlobalObjectRegistry.outputStreamToBufferMap.get(thizz);
+      if (buffer == null) {
+        return null;
+      }
+      int callDepth = CallDepthThreadLocalMap.incrementCallDepth(OutputStream.class);
+      if (callDepth > 0) {
+        return buffer;
+      }
+
+      buffer.write(b, off, len);
+      return buffer;
+    }
+
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(@Advice.Enter BoundedByteArrayOutputStream buffer) {
+      if (buffer != null) {
+        CallDepthThreadLocalMap.decrementCallDepth(OutputStream.class);
+      }
     }
   }
 }