Support async servlet for RUM injection

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStream.java

@@ -172,7 +172,6 @@ private void drain() throws IOException {
 
   public void commit() throws IOException {
     if (filter || wasDraining) {
-      filter = false;
       drain();
     }
   }
@@ -190,4 +189,8 @@ public void close() throws IOException {
       downstream.close();
     }
   }
+
+  public void setFilter(boolean filter) {
+    this.filter = filter;
+  }
 }

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriter.java

@@ -173,7 +173,6 @@ private void drain() throws IOException {
 
   public void commit() throws IOException {
     if (filter || wasDraining) {
-      filter = false;
       drain();
     }
   }
@@ -191,4 +190,8 @@ public void close() throws IOException {
       downstream.close();
     }
   }
+
+  public void setFilter(boolean filter) {
+    this.filter = filter;
+  }
 }

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumAsyncContextInstrumentation.java

@@ -0,0 +1,64 @@
+package datadog.trace.instrumentation.servlet3;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
+import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import javax.servlet.AsyncContext;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class RumAsyncContextInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  public RumAsyncContextInstrumentation() {
+    super("servlet", "servlet-3", "servlet-3-async-context");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "javax.servlet.AsyncContext";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".RumHttpServletResponseWrapper", packageName + ".WrappedServletOutputStream",
+    };
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return super.isEnabled() && InstrumenterConfig.get().isRumEnabled();
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(namedOneOf("complete", "dispatch")), getClass().getName() + "$CommitAdvice");
+  }
+
+  public static class CommitAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void commitRumBuffer(@Advice.This final AsyncContext asyncContext) {
+      final Object maybeRumWrappedResponse =
+          asyncContext.getRequest().getAttribute(DD_RUM_INJECTED);
+      if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
+        ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletRequestWrapper.java

@@ -0,0 +1,45 @@
+package datadog.trace.instrumentation.servlet3;
+
+import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
+
+import javax.servlet.AsyncContext;
+import javax.servlet.ServletRequest;
+import javax.servlet.ServletResponse;
+import javax.servlet.http.HttpServletRequest;
+import javax.servlet.http.HttpServletRequestWrapper;
+import javax.servlet.http.HttpServletResponse;
+
+public class RumHttpServletRequestWrapper extends HttpServletRequestWrapper {
+
+  private final HttpServletResponse response;
+
+  public RumHttpServletRequestWrapper(
+      final HttpServletRequest request, final HttpServletResponse response) {
+    super(request);
+    this.response = response;
+  }
+
+  @Override
+  public AsyncContext startAsync() throws IllegalStateException {
+    // need to hide this method otherwise we cannot control the wrapped response used asynchronously
+    return startAsync(getRequest(), response);
+  }
+
+  @Override
+  public AsyncContext startAsync(ServletRequest servletRequest, ServletResponse servletResponse)
+      throws IllegalStateException {
+    // deactivate the previous wrapper
+    final Object maybeRumWrappedResponse = (servletRequest.getAttribute(DD_RUM_INJECTED));
+    if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
+      ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
+      ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).stopFiltering();
+    }
+    ServletResponse actualResponse = servletResponse;
+    // rewrap it
+    if (servletResponse instanceof HttpServletResponse) {
+      actualResponse = new RumHttpServletResponseWrapper((HttpServletResponse) servletResponse);
+      servletRequest.setAttribute(DD_RUM_INJECTED, actualResponse);
+    }
+    return super.startAsync(servletRequest, actualResponse);
+  }
+}

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/RumHttpServletResponseWrapper.java

@@ -131,6 +131,7 @@ public void setContentType(String type) {
     }
     if (!shouldInject) {
       commit();
+      stopFiltering();
     }
     super.setContentType(type);
   }
@@ -149,4 +150,14 @@ public void commit() {
       }
     }
   }
+
+  public void stopFiltering() {
+    shouldInject = false;
+    if (wrappedPipeWriter != null) {
+      wrappedPipeWriter.setFilter(false);
+    }
+    if (outputStream != null) {
+      outputStream.setFilter(false);
+    }
+  }
 }

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Advice.java

@@ -46,11 +46,16 @@ public static boolean onEnter(
     final HttpServletRequest httpServletRequest = (HttpServletRequest) request;
     HttpServletResponse httpServletResponse = (HttpServletResponse) response;
 
-    if (RumInjector.get().isEnabled() && httpServletRequest.getAttribute(DD_RUM_INJECTED) == null) {
-      httpServletRequest.setAttribute(DD_RUM_INJECTED, Boolean.TRUE);
-      rumServletWrapper = new RumHttpServletResponseWrapper(httpServletResponse);
-      httpServletResponse = rumServletWrapper;
-      response = httpServletResponse;
+    if (RumInjector.get().isEnabled()) {
+      final Object maybeRumWrapper = httpServletRequest.getAttribute(DD_RUM_INJECTED);
+      if (maybeRumWrapper instanceof RumHttpServletResponseWrapper) {
+        rumServletWrapper = (RumHttpServletResponseWrapper) maybeRumWrapper;
+      } else {
+        rumServletWrapper = new RumHttpServletResponseWrapper((HttpServletResponse) response);
+        httpServletRequest.setAttribute(DD_RUM_INJECTED, rumServletWrapper);
+        response = rumServletWrapper;
+        request = new RumHttpServletRequestWrapper(httpServletRequest, rumServletWrapper);
+      }
     }
 
     Object dispatchSpan = request.getAttribute(DD_DISPATCH_SPAN_ATTRIBUTE);

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/Servlet3Instrumentation.java

@@ -52,6 +52,7 @@ public String[] helperClassNames() {
       packageName + ".Servlet3Decorator",
       packageName + ".ServletRequestURIAdapter",
       packageName + ".FinishAsyncDispatchListener",
+      packageName + ".RumHttpServletRequestWrapper",
       packageName + ".RumHttpServletResponseWrapper",
       packageName + ".WrappedServletOutputStream",
       "datadog.trace.instrumentation.servlet.ServletBlockingHelper",

dd-java-agent/instrumentation/servlet/request-3/src/main/java/datadog/trace/instrumentation/servlet3/WrappedServletOutputStream.java

@@ -86,4 +86,8 @@ public void setWriteListener(WriteListener writeListener) {
   public void commit() throws IOException {
     filtered.commit();
   }
+
+  public void setFilter(boolean filter) {
+    filtered.setFilter(filter);
+  }
 }

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/JakartaServletInstrumentation.java

@@ -42,7 +42,9 @@ public String hierarchyMarkerType() {
   @Override
   public String[] helperClassNames() {
     return new String[] {
-      packageName + ".RumHttpServletResponseWrapper", packageName + ".WrappedServletOutputStream",
+      packageName + ".RumHttpServletRequestWrapper",
+      packageName + ".RumHttpServletResponseWrapper",
+      packageName + ".WrappedServletOutputStream",
     };
   }
 
@@ -67,7 +69,7 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class JakartaServletAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static AgentSpan before(
-        @Advice.Argument(0) final ServletRequest request,
+        @Advice.Argument(value = 0, readOnly = false) ServletRequest request,
         @Advice.Argument(value = 1, readOnly = false) ServletResponse response,
         @Advice.Local("rumServletWrapper") RumHttpServletResponseWrapper rumServletWrapper) {
       if (!(request instanceof HttpServletRequest)) {
@@ -77,11 +79,16 @@ public static AgentSpan before(
       if (response instanceof HttpServletResponse) {
         final HttpServletRequest httpServletRequest = (HttpServletRequest) request;
 
-        if (RumInjector.get().isEnabled()
-            && httpServletRequest.getAttribute(DD_RUM_INJECTED) == null) {
-          httpServletRequest.setAttribute(DD_RUM_INJECTED, Boolean.TRUE);
-          rumServletWrapper = new RumHttpServletResponseWrapper((HttpServletResponse) response);
-          response = rumServletWrapper;
+        if (RumInjector.get().isEnabled()) {
+          final Object maybeRumWrapper = httpServletRequest.getAttribute(DD_RUM_INJECTED);
+          if (maybeRumWrapper instanceof RumHttpServletResponseWrapper) {
+            rumServletWrapper = (RumHttpServletResponseWrapper) maybeRumWrapper;
+          } else {
+            rumServletWrapper = new RumHttpServletResponseWrapper((HttpServletResponse) response);
+            httpServletRequest.setAttribute(DD_RUM_INJECTED, rumServletWrapper);
+            response = rumServletWrapper;
+            request = new RumHttpServletRequestWrapper(httpServletRequest, rumServletWrapper);
+          }
         }
       }
 

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumAsyncContextInstrumentation.java

@@ -0,0 +1,64 @@
+package datadog.trace.instrumentation.servlet5;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
+import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.InstrumenterConfig;
+import jakarta.servlet.AsyncContext;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+@AutoService(InstrumenterModule.class)
+public class RumAsyncContextInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
+
+  public RumAsyncContextInstrumentation() {
+    super("servlet", "servlet-5", "servlet-5-async-context");
+  }
+
+  @Override
+  public String hierarchyMarkerType() {
+    return "jakarta.servlet.AsyncContext";
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {
+      packageName + ".RumHttpServletResponseWrapper", packageName + ".WrappedServletOutputStream",
+    };
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return implementsInterface(named(hierarchyMarkerType()));
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return super.isEnabled() && InstrumenterConfig.get().isRumEnabled();
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod().and(namedOneOf("complete", "dispatch")), getClass().getName() + "$CommitAdvice");
+  }
+
+  public static class CommitAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void commitRumBuffer(@Advice.This final AsyncContext asyncContext) {
+      final Object maybeRumWrappedResponse =
+          asyncContext.getRequest().getAttribute(DD_RUM_INJECTED);
+      if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
+        ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
+      }
+    }
+  }
+}

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletRequestWrapper.java

@@ -0,0 +1,45 @@
+package datadog.trace.instrumentation.servlet5;
+
+import static datadog.trace.bootstrap.instrumentation.decorator.HttpServerDecorator.DD_RUM_INJECTED;
+
+import jakarta.servlet.AsyncContext;
+import jakarta.servlet.ServletRequest;
+import jakarta.servlet.ServletResponse;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.servlet.http.HttpServletRequestWrapper;
+import jakarta.servlet.http.HttpServletResponse;
+
+public class RumHttpServletRequestWrapper extends HttpServletRequestWrapper {
+
+  private final HttpServletResponse response;
+
+  public RumHttpServletRequestWrapper(
+      final HttpServletRequest request, final HttpServletResponse response) {
+    super(request);
+    this.response = response;
+  }
+
+  @Override
+  public AsyncContext startAsync() throws IllegalStateException {
+    // need to hide this method otherwise we cannot control the wrapped response used asynchronously
+    return startAsync(getRequest(), response);
+  }
+
+  @Override
+  public AsyncContext startAsync(ServletRequest servletRequest, ServletResponse servletResponse)
+      throws IllegalStateException {
+    // deactivate the previous wrapper
+    final Object maybeRumWrappedResponse = (servletRequest.getAttribute(DD_RUM_INJECTED));
+    if (maybeRumWrappedResponse instanceof RumHttpServletResponseWrapper) {
+      ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).commit();
+      ((RumHttpServletResponseWrapper) maybeRumWrappedResponse).stopFiltering();
+    }
+    ServletResponse actualResponse = servletResponse;
+    // rewrap it
+    if (servletResponse instanceof HttpServletResponse) {
+      actualResponse = new RumHttpServletResponseWrapper((HttpServletResponse) servletResponse);
+      servletRequest.setAttribute(DD_RUM_INJECTED, actualResponse);
+    }
+    return super.startAsync(servletRequest, actualResponse);
+  }
+}

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/RumHttpServletResponseWrapper.java

@@ -108,6 +108,7 @@ public void setContentType(String type) {
     }
     if (!shouldInject) {
       commit();
+      stopFiltering();
     }
     super.setContentType(type);
   }
@@ -126,4 +127,14 @@ public void commit() {
       }
     }
   }
+
+  public void stopFiltering() {
+    shouldInject = false;
+    if (wrappedPipeWriter != null) {
+      wrappedPipeWriter.setFilter(false);
+    }
+    if (outputStream != null) {
+      outputStream.setFilter(false);
+    }
+  }
 }

dd-java-agent/instrumentation/servlet/request-5/src/main/java/datadog/trace/instrumentation/servlet5/WrappedServletOutputStream.java

@@ -53,4 +53,8 @@ public boolean isReady() {
   public void setWriteListener(WriteListener writeListener) {
     delegate.setWriteListener(writeListener);
   }
+
+  public void setFilter(boolean filter) {
+    filtered.setFilter(filter);
+  }
 }