Add telemetry for the RUM injector

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

@@ -51,7 +51,7 @@ 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))) {
       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,41 @@ public class InjectingPipeOutputStream extends OutputStream {
   private final Runnable onContentInjected;
   private final int bulkWriteThreshold;
   private final OutputStream downstream;
+  private final LongConsumer onBytesWritten;
+  private final LongConsumer onInjectionTime;
+  private long bytesWritten = 0;
 
   /**
+   * This constructor is typically used for testing where we care about the logic and not the
+   * telemetry.
+   *
+   * @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.
+   */
+  public InjectingPipeOutputStream(
+      final OutputStream downstream, final byte[] marker, final byte[] contentToInject) {
+    this(downstream, marker, contentToInject, null, null, null);
+  }
+
+  /**
+   * This constructor contains the full set of parameters.
+   *
    * @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.
+   * @param onInjectionTime callback called with the time in milliseconds taken to write the
+   *     injection content.
    */
   public InjectingPipeOutputStream(
       final OutputStream downstream,
       final byte[] marker,
       final byte[] contentToInject,
-      final Runnable onContentInjected) {
+      final Runnable onContentInjected,
+      final LongConsumer onBytesWritten,
+      final LongConsumer onInjectionTime) {
     this.downstream = downstream;
     this.marker = marker;
     this.lookbehind = new byte[marker.length];
@@ -46,6 +70,8 @@ public InjectingPipeOutputStream(
     this.filter = true;
     this.contentToInject = contentToInject;
     this.onContentInjected = onContentInjected;
+    this.onBytesWritten = onBytesWritten;
+    this.onInjectionTime = onInjectionTime;
     this.bulkWriteThreshold = marker.length * 2 - 2;
   }
 
@@ -57,11 +83,13 @@ public void write(int b) throws IOException {
         drain();
       }
       downstream.write(b);
+      bytesWritten++;
       return;
     }
 
     if (count == lookbehind.length) {
       downstream.write(lookbehind[pos]);
+      bytesWritten++;
     } else {
       count++;
     }
@@ -72,7 +100,12 @@ public void write(int b) throws IOException {
     if (marker[matchingPos++] == b) {
       if (matchingPos == marker.length) {
         filter = false;
+        long injectionStart = System.nanoTime();
         downstream.write(contentToInject);
+        long injectionEnd = System.nanoTime();
+        if (onInjectionTime != null) {
+          onInjectionTime.accept((injectionEnd - injectionStart) / 1_000_000L);
+        }
         if (onContentInjected != null) {
           onContentInjected.run();
         }
@@ -91,6 +124,7 @@ public void write(byte[] array, int off, int len) throws IOException {
         drain();
       }
       downstream.write(array, off, len);
+      bytesWritten += len;
       return;
     }
 
@@ -103,12 +137,21 @@ public void write(byte[] array, int off, int len) throws IOException {
         // we have a full match. just write everything
         filter = false;
         drain();
-        downstream.write(array, off, idx);
+        int bytesToWrite = idx;
+        downstream.write(array, off, bytesToWrite);
+        bytesWritten += bytesToWrite;
+        long injectionStart = System.nanoTime();
         downstream.write(contentToInject);
+        long injectionEnd = System.nanoTime();
+        if (onInjectionTime != null) {
+          onInjectionTime.accept((injectionEnd - injectionStart) / 1_000_000L);
+        }
         if (onContentInjected != null) {
           onContentInjected.run();
         }
-        downstream.write(array, off + idx, len - idx);
+        bytesToWrite = len - idx;
+        downstream.write(array, off + idx, bytesToWrite);
+        bytesWritten += bytesToWrite;
       } else {
         // we don't have a full match. write everything in a bulk except the lookbehind buffer
         // sequentially
@@ -120,7 +163,9 @@ public void write(byte[] array, int off, int len) throws IOException {
 
         // will be reset if no errors after the following write
         filter = false;
-        downstream.write(array, off + marker.length - 1, len - bulkWriteThreshold);
+        int bytesToWrite = len - bulkWriteThreshold;
+        downstream.write(array, off + marker.length - 1, bytesToWrite);
+        bytesWritten += bytesToWrite;
         filter = wasFiltering;
         for (int i = len - marker.length + 1; i < len; i++) {
           write(array[i]);
@@ -163,6 +208,7 @@ private void drain() throws IOException {
       int cnt = count;
       for (int i = 0; i < cnt; i++) {
         downstream.write(lookbehind[(start + i) % lookbehind.length]);
+        bytesWritten++;
         count--;
       }
       filter = wasFiltering;
@@ -185,6 +231,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,41 @@ public class InjectingPipeWriter extends Writer {
   private final Runnable onContentInjected;
   private final int bulkWriteThreshold;
   private final Writer downstream;
+  private final LongConsumer onBytesWritten;
+  private final LongConsumer onInjectionTime;
+  private long bytesWritten = 0;
 
   /**
+   * This constructor is typically used for testing where we care about the logic and not the
+   * telemetry.
+   *
+   * @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.
+   */
+  public InjectingPipeWriter(
+      final Writer downstream, final char[] marker, final char[] contentToInject) {
+    this(downstream, marker, contentToInject, null, null, null);
+  }
+
+  /**
+   * This constructor contains the full set of parameters.
+   *
    * @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.
+   * @param onInjectionTime callback called with the time in milliseconds taken to write the
+   *     injection content.
    */
   public InjectingPipeWriter(
       final Writer downstream,
       final char[] marker,
       final char[] contentToInject,
-      final Runnable onContentInjected) {
+      final Runnable onContentInjected,
+      final LongConsumer onBytesWritten,
+      final LongConsumer onInjectionTime) {
     this.downstream = downstream;
     this.marker = marker;
     this.lookbehind = new char[marker.length];
@@ -46,6 +70,8 @@ public InjectingPipeWriter(
     this.filter = true;
     this.contentToInject = contentToInject;
     this.onContentInjected = onContentInjected;
+    this.onBytesWritten = onBytesWritten;
+    this.onInjectionTime = onInjectionTime;
     this.bulkWriteThreshold = marker.length * 2 - 2;
   }
 
@@ -57,11 +83,13 @@ public void write(int c) throws IOException {
         drain();
       }
       downstream.write(c);
+      bytesWritten++;
       return;
     }
 
     if (count == lookbehind.length) {
       downstream.write(lookbehind[pos]);
+      bytesWritten++;
     } else {
       count++;
     }
@@ -72,7 +100,12 @@ public void write(int c) throws IOException {
     if (marker[matchingPos++] == c) {
       if (matchingPos == marker.length) {
         filter = false;
+        long injectionStart = System.nanoTime();
         downstream.write(contentToInject);
+        long injectionEnd = System.nanoTime();
+        if (onInjectionTime != null) {
+          onInjectionTime.accept((injectionEnd - injectionStart) / 1_000_000L);
+        }
         if (onContentInjected != null) {
           onContentInjected.run();
         }
@@ -91,6 +124,7 @@ public void write(char[] array, int off, int len) throws IOException {
         drain();
       }
       downstream.write(array, off, len);
+      bytesWritten += len;
       return;
     }
 
@@ -103,12 +137,21 @@ public void write(char[] array, int off, int len) throws IOException {
         // we have a full match. just write everything
         filter = false;
         drain();
-        downstream.write(array, off, idx);
+        int bytesToWrite = idx;
+        downstream.write(array, off, bytesToWrite);
+        bytesWritten += bytesToWrite;
+        long injectionStart = System.nanoTime();
         downstream.write(contentToInject);
+        long injectionEnd = System.nanoTime();
+        if (onInjectionTime != null) {
+          onInjectionTime.accept((injectionEnd - injectionStart) / 1_000_000L);
+        }
         if (onContentInjected != null) {
           onContentInjected.run();
         }
-        downstream.write(array, off + idx, len - idx);
+        bytesToWrite = len - idx;
+        downstream.write(array, off + idx, bytesToWrite);
+        bytesWritten += bytesToWrite;
       } else {
         // we don't have a full match. write everything in a bulk except the lookbehind buffer
         // sequentially
@@ -120,7 +163,9 @@ public void write(char[] array, int off, int len) throws IOException {
 
         // will be reset if no errors after the following write
         filter = false;
-        downstream.write(array, off + marker.length - 1, len - bulkWriteThreshold);
+        int bytesToWrite = len - bulkWriteThreshold;
+        downstream.write(array, off + marker.length - 1, bytesToWrite);
+        bytesWritten += bytesToWrite;
         filter = wasFiltering;
 
         for (int i = len - marker.length + 1; i < len; i++) {
@@ -164,6 +209,7 @@ private void drain() throws IOException {
       int cnt = count;
       for (int i = 0; i < cnt; i++) {
         downstream.write(lookbehind[(start + i) % lookbehind.length]);
+        bytesWritten++;
         count--;
       }
       filter = wasFiltering;
@@ -188,6 +234,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;
     }
   }