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/main/java/datadog/trace/bootstrap/instrumentation/buffer/InjectingPipeWriter.java

@@ -2,6 +2,7 @@
 
 import java.io.IOException;
 import java.io.Writer;
+import java.util.function.LongConsumer;
 import javax.annotation.concurrent.NotThreadSafe;
 
 /**
@@ -23,18 +24,22 @@ public class InjectingPipeWriter extends Writer {
   private final Runnable onContentInjected;
   private final int bulkWriteThreshold;
   private final Writer downstream;
+  private final LongConsumer onBytesWritten;
+  private long bytesWritten = 0;
 
   /**
    * @param downstream the delegate writer
    * @param marker the marker to find in the stream. Must at least be one char.
    * @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 writer is closed to report total bytes written.
    */
   public InjectingPipeWriter(
       final Writer downstream,
       final char[] marker,
       final char[] contentToInject,
-      final Runnable onContentInjected) {
+      final Runnable onContentInjected,
+      final LongConsumer onBytesWritten) {
     this.downstream = downstream;
     this.marker = marker;
     this.lookbehind = new char[marker.length];
@@ -46,11 +51,13 @@ public InjectingPipeWriter(
     this.filter = true;
     this.contentToInject = contentToInject;
     this.onContentInjected = onContentInjected;
+    this.onBytesWritten = onBytesWritten;
     this.bulkWriteThreshold = marker.length * 2 - 2;
   }
 
   @Override
   public void write(int c) throws IOException {
+    bytesWritten++;
     if (!filter) {
       if (wasDraining) {
         // continue draining
@@ -85,6 +92,7 @@ public void write(int c) throws IOException {
 
   @Override
   public void write(char[] array, int off, int len) throws IOException {
+    bytesWritten += len;
     if (!filter) {
       if (wasDraining) {
         // needs drain
@@ -113,6 +121,7 @@ public void write(char[] 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();
@@ -124,12 +133,14 @@ public void write(char[] array, int off, int len) throws IOException {
         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]);
       }
     }
@@ -188,6 +199,11 @@ public void close() throws IOException {
       commit();
     } finally {
       downstream.close();
+      // report the size of the original HTTP response before injecting via callback
+      if (onBytesWritten != null) {
+        onBytesWritten.accept(bytesWritten);
+      }
+      bytesWritten = 0;
     }
   }
 

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) {
@@ -87,4 +87,94 @@ class InjectingPipeOutputStreamTest extends DDSpecification {
     // expected broken since the real write happens at close (drain) being the content smaller than the buffer. And retry on close is not a common practice. Hence, we suppose loosing content
     ["<foo/>"]                                                      | "<longerThanFoo>" | "<nothing>"             | 3         | "<f"
   }
+
+  def 'should count bytes correctly when writing byte arrays'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, onBytesWritten)
+
+    when:
+    piped.write("test content".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    bytesWritten.size() == 1
+    bytesWritten[0] == 12
+    downstream.toByteArray() == "test content".getBytes("UTF-8")
+  }
+
+  def 'should count bytes correctly when writing bytes individually'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, onBytesWritten)
+
+    when:
+    def bytes = "test".getBytes("UTF-8")
+    for (int i = 0; i < bytes.length; i++) {
+      piped.write((int) bytes[i])
+    }
+    piped.close()
+
+    then:
+    bytesWritten.size() == 1
+    bytesWritten[0] == 4
+    downstream.toByteArray() == "test".getBytes("UTF-8")
+  }
+
+  def 'should count bytes correctly with multiple writes'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, onBytesWritten)
+
+    when:
+    piped.write("test".getBytes("UTF-8"))
+    piped.write(" ".getBytes("UTF-8"))
+    piped.write("content".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    bytesWritten.size() == 1
+    bytesWritten[0] == 12
+    downstream.toByteArray() == "test content".getBytes("UTF-8")
+  }
+
+  def 'should be resilient to exceptions when onBytesWritten callback is null'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, null)
+
+    when:
+    piped.write("test content".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    noExceptionThrown()
+    downstream.toByteArray() == "test content".getBytes("UTF-8")
+  }
+
+  def 'should reset byte count after close'() {
+    setup:
+    def downstream = new ByteArrayOutputStream()
+    def bytesWritten = []
+    def onBytesWritten = { long bytes -> bytesWritten.add(bytes) }
+    def piped = new InjectingPipeOutputStream(downstream, "</head>".getBytes("UTF-8"), "<script></script>".getBytes("UTF-8"), null, onBytesWritten)
+
+    when:
+    piped.write("test".getBytes("UTF-8"))
+    piped.close()
+
+    piped.write("content".getBytes("UTF-8"))
+    piped.close()
+
+    then:
+    bytesWritten.size() == 2
+    bytesWritten[0] == 4
+    bytesWritten[1] == 7
+  }
 }