Add telemetry for the RUM injector

dd-java-agent/agent-bootstrap/src/jmh/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStreamBenchmark.java

@@ -51,7 +51,8 @@ public class InjectingPipeOutputStreamBenchmark {
   public void withPipe() throws Exception {
     try (final PrintWriter out =
         new PrintWriter(
-            new InjectingPipeOutputStream(new ByteArrayOutputStream(), marker, content, null))) {
+            new InjectingPipeOutputStream(
+                new ByteArrayOutputStream(), marker, content, null, null))) {
       htmlContent.forEach(out::println);
     }
   }

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

@@ -2,6 +2,7 @@
 
 import java.io.IOException;
 import java.io.OutputStream;
+import java.util.function.LongConsumer;
 import javax.annotation.concurrent.NotThreadSafe;
 
 /**
@@ -23,18 +24,22 @@ public class InjectingPipeOutputStream extends OutputStream {
   private final Runnable onContentInjected;
   private final int bulkWriteThreshold;
   private final OutputStream downstream;
+  private final LongConsumer onBytesWritten;
+  private long bytesWritten = 0;
 
   /**
    * @param downstream the delegate output stream
    * @param marker the marker to find in the stream. Must at least be one byte.
    * @param contentToInject the content to inject once before the marker if found.
    * @param onContentInjected callback called when and if the content is injected.
+   * @param onBytesWritten callback called when stream is closed to report total bytes written.
    */
   public InjectingPipeOutputStream(
       final OutputStream downstream,
       final byte[] marker,
       final byte[] contentToInject,
-      final Runnable onContentInjected) {
+      final Runnable onContentInjected,
+      final LongConsumer onBytesWritten) {
     this.downstream = downstream;
     this.marker = marker;
     this.lookbehind = new byte[marker.length];
@@ -46,11 +51,13 @@ public InjectingPipeOutputStream(
     this.filter = true;
     this.contentToInject = contentToInject;
     this.onContentInjected = onContentInjected;
+    this.onBytesWritten = onBytesWritten;
     this.bulkWriteThreshold = marker.length * 2 - 2;
   }
 
   @Override
   public void write(int b) throws IOException {
+    bytesWritten++;
     if (!filter) {
       if (wasDraining) {
         // continue draining
@@ -85,6 +92,7 @@ public void write(int b) throws IOException {
 
   @Override
   public void write(byte[] array, int off, int len) throws IOException {
+    bytesWritten += len;
     if (!filter) {
       if (wasDraining) {
         // needs drain
@@ -113,6 +121,7 @@ public void write(byte[] array, int off, int len) throws IOException {
         // we don't have a full match. write everything in a bulk except the lookbehind buffer
         // sequentially
         for (int i = off; i < off + marker.length - 1; i++) {
+          bytesWritten--; // avoid double counting
           write(array[i]);
         }
         drain();
@@ -123,12 +132,14 @@ public void write(byte[] array, int off, int len) throws IOException {
         downstream.write(array, off + marker.length - 1, len - bulkWriteThreshold);
         filter = wasFiltering;
         for (int i = len - marker.length + 1; i < len; i++) {
+          bytesWritten--; // avoid double counting
           write(array[i]);
         }
       }
     } else {
       // use slow path because the length to write is small and within the lookbehind buffer size
       for (int i = off; i < off + len; i++) {
+        bytesWritten--; // avoid double counting
         write(array[i]);
       }
     }
@@ -185,6 +196,11 @@ public void flush() throws IOException {
   public void close() throws IOException {
     try {
       commit();
+      // report the size of the original HTTP response before injecting via callback
+      if (onBytesWritten != null) {
+        onBytesWritten.accept(bytesWritten);
+      }
+      bytesWritten = 0;
     } finally {
       downstream.close();
     }

dd-java-agent/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeOutputStreamTest.groovy

@@ -36,7 +36,7 @@ class InjectingPipeOutputStreamTest extends DDSpecification {
   def 'should filter a buffer and inject if found #found'() {
     setup:
     def downstream = new ByteArrayOutputStream()
-    def piped = new OutputStreamWriter(new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8"), null),
+    def piped = new OutputStreamWriter(new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8"), null, null),
     "UTF-8")
     when:
     try (def closeme = piped) {
@@ -55,7 +55,7 @@ class InjectingPipeOutputStreamTest extends DDSpecification {
     setup:
     def baos = new ByteArrayOutputStream()
     def downstream = new GlitchedOutputStream(baos, glichesAt)
-    def piped = new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8"), null)
+    def piped = new InjectingPipeOutputStream(downstream, marker.getBytes("UTF-8"), contentToInject.getBytes("UTF-8"), null, null)
     when:
     try {
       for (String line : body) {

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

@@ -18,6 +18,8 @@ public class RumHttpServletResponseWrapper extends HttpServletResponseWrapper {
   private PrintWriter printWriter;
   private InjectingPipeWriter wrappedPipeWriter;
   private boolean shouldInject = true;
+  private long injectionStartTime = -1;
+  private String contentEncoding = "none";
 
   private static final MethodHandle SET_CONTENT_LENGTH_LONG = getMh("setContentLengthLong");
 
@@ -45,18 +47,30 @@ public ServletOutputStream getOutputStream() throws IOException {
       return outputStream;
     }
     if (!shouldInject) {
+      RumInjector.getTelemetryCollector().onInjectionSkipped("3");
       return super.getOutputStream();
     }
-    String encoding = getCharacterEncoding();
-    if (encoding == null) {
-      encoding = Charset.defaultCharset().name();
+    // start timing injection
+    if (injectionStartTime == -1) {
+      injectionStartTime = System.nanoTime();
+    }
+    try {
+      String encoding = getCharacterEncoding();
+      if (encoding == null) {
+        encoding = Charset.defaultCharset().name();
+      }
+      outputStream =
+          new WrappedServletOutputStream(
+              super.getOutputStream(),
+              rumInjector.getMarkerBytes(encoding),
+              rumInjector.getSnippetBytes(encoding),
+              this::onInjected,
+              bytes -> RumInjector.getTelemetryCollector().onInjectionResponseSize("3", bytes));
+    } catch (Exception e) {
+      injectionStartTime = -1;
+      RumInjector.getTelemetryCollector().onInjectionFailed("3", contentEncoding);
+      throw e;
     }
-    outputStream =
-        new WrappedServletOutputStream(
-            super.getOutputStream(),
-            rumInjector.getMarkerBytes(encoding),
-            rumInjector.getSnippetBytes(encoding),
-            this::onInjected);
 
     return outputStream;
   }
@@ -67,19 +81,56 @@ public PrintWriter getWriter() throws IOException {
       return printWriter;
     }
     if (!shouldInject) {
+      RumInjector.getTelemetryCollector().onInjectionSkipped("3");
       return super.getWriter();
     }
-    wrappedPipeWriter =
-        new InjectingPipeWriter(
-            super.getWriter(),
-            rumInjector.getMarkerChars(),
-            rumInjector.getSnippetChars(),
-            this::onInjected);
-    printWriter = new PrintWriter(wrappedPipeWriter);
+    // start timing injection
+    if (injectionStartTime == -1) {
+      injectionStartTime = System.nanoTime();
+    }
+    try {
+      wrappedPipeWriter =
+          new InjectingPipeWriter(
+              super.getWriter(),
+              rumInjector.getMarkerChars(),
+              rumInjector.getSnippetChars(),
+              this::onInjected);
+      printWriter = new PrintWriter(wrappedPipeWriter);
+    } catch (Exception e) {
+      injectionStartTime = -1;
+      RumInjector.getTelemetryCollector().onInjectionFailed("3", contentEncoding);
+      throw e;
+    }
 
     return printWriter;
   }
 
+  @Override
+  public void setHeader(String name, String value) {
+    if (name != null) {
+      String lowerName = name.toLowerCase();
+      if (lowerName.startsWith("content-security-policy")) {
+        RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected("3");
+      } else if (lowerName.contains("content-encoding")) {
+        this.contentEncoding = value;
+      }
+    }
+    super.setHeader(name, value);
+  }
+
+  @Override
+  public void addHeader(String name, String value) {
+    if (name != null) {
+      String lowerName = name.toLowerCase();
+      if (lowerName.startsWith("content-security-policy")) {
+        RumInjector.getTelemetryCollector().onContentSecurityPolicyDetected("3");
+      } else if (lowerName.contains("content-encoding")) {
+        this.contentEncoding = value;
+      }
+    }
+    super.addHeader(name, value);
+  }
+
   @Override
   public void setContentLength(int len) {
     // don't set it since we don't know if we will inject
@@ -105,6 +156,7 @@ public void reset() {
     this.wrappedPipeWriter = null;
     this.printWriter = null;
     this.shouldInject = false;
+    this.injectionStartTime = -1;
     super.reset();
   }
 
@@ -117,6 +169,16 @@ public void resetBuffer() {
   }
 
   public void onInjected() {
+    RumInjector.getTelemetryCollector().onInjectionSucceed("3");
+
+    // calculate total injection time
+    if (injectionStartTime != -1) {
+      long nanoseconds = System.nanoTime() - injectionStartTime;
+      long milliseconds = nanoseconds / 1_000_000L;
+      RumInjector.getTelemetryCollector().onInjectionTime("3", milliseconds);
+      injectionStartTime = -1;
+    }
+
     try {
       setHeader("x-datadog-rum-injected", "1");
     } catch (Throwable ignored) {

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

@@ -4,6 +4,7 @@
 import datadog.trace.util.MethodHandles;
 import java.io.IOException;
 import java.lang.invoke.MethodHandle;
+import java.util.function.LongConsumer;
 import javax.servlet.ServletOutputStream;
 import javax.servlet.WriteListener;
 
@@ -29,8 +30,14 @@ private static <E extends Throwable> void sneakyThrow(Throwable e) throws E {
   }
 
   public WrappedServletOutputStream(
-      ServletOutputStream delegate, byte[] marker, byte[] contentToInject, Runnable onInjected) {
-    this.filtered = new InjectingPipeOutputStream(delegate, marker, contentToInject, onInjected);
+      ServletOutputStream delegate,
+      byte[] marker,
+      byte[] contentToInject,
+      Runnable onInjected,
+      LongConsumer onBytesWritten) {
+    this.filtered =
+        new InjectingPipeOutputStream(
+            delegate, marker, contentToInject, onInjected, onBytesWritten);
     this.delegate = delegate;
   }