Connect rum injector, telemetry collector, and statsdclient

Author: sarahchen6

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

@@ -80,7 +80,7 @@ public void resetBuffer() {
   }
 
   public void onInjected() {
-    RumInjector.reportInjectionSucceed();
+    RumInjector.getTelemetryCollector().onInjectionSucceed();
     try {
       setHeader("x-datadog-rum-injected", "1");
     } catch (Throwable ignored) {

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

@@ -80,7 +80,7 @@ public void resetBuffer() {
   }
 
   public void onInjected() {
-    RumInjector.reportInjectionSucceed();
+    RumInjector.getTelemetryCollector().onInjectionSucceed();
     try {
       setHeader("x-datadog-rum-injected", "1");
     } catch (Throwable ignored) {

dd-trace-core/src/main/java/datadog/trace/core/CoreTracer.java

@@ -50,8 +50,6 @@
 import datadog.trace.api.naming.SpanNaming;
 import datadog.trace.api.remoteconfig.ServiceNameCollector;
 import datadog.trace.api.rum.RumInjector;
-import datadog.trace.api.rum.RumInjectorMetrics;
-import datadog.trace.api.rum.RumTelemetryCollector;
 import datadog.trace.api.sampling.PrioritySampling;
 import datadog.trace.api.scopemanager.ScopeListener;
 import datadog.trace.api.time.SystemTimeSource;
@@ -207,7 +205,6 @@ public static CoreTracerBuilder builder() {
   private final Monitoring performanceMonitoring;
 
   private final HealthMetrics healthMetrics;
-  private final RumTelemetryCollector rumInjectorHealthMetrics;
   private final Recording traceWriteTimer;
   private final IdGenerationStrategy idGenerationStrategy;
   private final TraceCollector.Factory traceCollectorFactory;
@@ -707,14 +704,10 @@ private CoreTracer(
             : HealthMetrics.NO_OP;
     healthMetrics.start();
 
-    // Start RUM injector metrics
-    rumInjectorHealthMetrics =
-        config.isHealthMetricsEnabled() && config.isRumEnabled()
-            ? new RumInjectorMetrics(this.statsDClient)
-            : RumTelemetryCollector.NO_OP;
-    rumInjectorHealthMetrics.start();
-    // Register rumInjectorHealthMetrics as the RumInjector's telemetry collector
-    RumInjector.setTelemetryCollector(rumInjectorHealthMetrics);
+    // Start RUM injector telemetry
+    if (config.isRumEnabled()) {
+      RumInjector.enableTelemetry(this.statsDClient);
+    }
 
     performanceMonitoring =
         config.isPerfMetricsEnabled()
@@ -1261,7 +1254,7 @@ public void close() {
     tracingConfigPoller.stop();
     pendingTraceBuffer.close();
     writer.close();
-    rumInjectorHealthMetrics.close();
+    RumInjector.shutdownTelemetry();
     statsDClient.close();
     metricsAggregator.close();
     dataStreamsMonitoring.close();

internal-api/src/main/java/datadog/trace/api/rum/RumInjector.java

@@ -27,7 +27,7 @@ public final class RumInjector {
   private final DDCache<String, byte[]> markerCache;
   private final Function<String, byte[]> snippetBytes;
 
-  // Health metrics telemetry collector (set by CoreTracer)
+  // telemetry collector defaults to NO_OP
   private static volatile RumTelemetryCollector telemetryCollector = RumTelemetryCollector.NO_OP;
 
   RumInjector(Config config) {
@@ -124,26 +124,30 @@ public byte[] getMarkerBytes(String encoding) {
     return this.markerCache.computeIfAbsent(encoding, MARKER_BYTES);
   }
 
-  // set the telemetry collector for the RumInjector
-  public static void setTelemetryCollector(RumTelemetryCollector collector) {
-    telemetryCollector = collector != null ? collector : RumTelemetryCollector.NO_OP;
+  // start telemetry collection and report metrics via the given StatsDClient
+  public static void enableTelemetry(datadog.trace.api.StatsDClient statsDClient) {
+    if (statsDClient != null) {
+      RumInjectorMetrics metrics = new RumInjectorMetrics(statsDClient);
+      metrics.start();
+      telemetryCollector = metrics;
+    } else {
+      telemetryCollector = RumTelemetryCollector.NO_OP;
+    }
   }
 
-  // report to the telemetry collector that the RUM injector succeeded in injecting the SDK in an
-  // HTTP response
-  public static void reportInjectionSucceed() {
-    telemetryCollector.onInjectionSucceed();
+  // shutdown telemetry and reset to NO_OP
+  public static void shutdownTelemetry() {
+    telemetryCollector.close();
+    telemetryCollector = RumTelemetryCollector.NO_OP;
   }
 
-  // report to the telemetry collector that the RUM injector failed to inject the SDK in an HTTP
-  // response
-  public static void reportInjectionFailed() {
-    telemetryCollector.onInjectionFailed();
+  // set the telemetry collector
+  public static void setTelemetryCollector(RumTelemetryCollector collector) {
+    telemetryCollector = collector != null ? collector : RumTelemetryCollector.NO_OP;
   }
 
-  // report to the telemetry collector that the RUM injector skipped injecting the SDK in an HTTP
-  // response
-  public static void reportInjectionSkipped() {
-    telemetryCollector.onInjectionSkipped();
+  // get the telemetry collector. this is used to directly report telemetry
+  public static RumTelemetryCollector getTelemetryCollector() {
+    return telemetryCollector;
   }
 }

internal-api/src/main/java/datadog/trace/api/rum/RumInjectorMetrics.java

@@ -10,9 +10,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-// Default implementation of RumInjectorHealthMetrics that reports metrics via StatsDClient
 // This class implements the RumTelemetryCollector interface, which is used to collect telemetry
-// from the RumInjector in the internal-api module
+// from the RumInjector. Metrics are then reported via StatsDClient.
 public class RumInjectorMetrics implements RumTelemetryCollector {
   private static final Logger log = LoggerFactory.getLogger(RumInjectorMetrics.class);
 

internal-api/src/main/java/datadog/trace/api/rum/RumTelemetryCollector.java

@@ -1,11 +1,14 @@
 package datadog.trace.api.rum;
 
 // Collect RUM injection telemetry from the RumInjector
-// This is implemented by the DefaultRumInjectorHealthMetrics class in the dd-trace-core module
+// This is implemented by the RumInjectorMetrics class
 public interface RumTelemetryCollector {
 
   RumTelemetryCollector NO_OP =
       new RumTelemetryCollector() {
+        @Override
+        public void start() {}
+
         @Override
         public void onInjectionSucceed() {}
 
@@ -14,8 +17,18 @@ public void onInjectionFailed() {}
 
         @Override
         public void onInjectionSkipped() {}
+
+        @Override
+        public void close() {}
+
+        @Override
+        public String summary() {
+          return "";
+        }
       };
 
+  default void start() {}
+
   // call when RUM injection succeeds
   void onInjectionSucceed();
 
@@ -24,4 +37,11 @@ public void onInjectionSkipped() {}
 
   // call when RUM injection is skipped
   void onInjectionSkipped();
+
+  default void close() {}
+
+  // human-readable summary of the current health metrics
+  default String summary() {
+    return "";
+  }
 }

internal-api/src/test/groovy/datadog/trace/api/rum/RumInjectorTest.groovy

@@ -61,4 +61,59 @@ class RumInjectorTest extends DDSpecification {
     injector.getSnippetChars() != null
     injector.getMarkerChars() != null
   }
+
+  void 'set telemetry collector'() {
+    setup:
+    def mockTelemetryCollector = mock(RumTelemetryCollector)
+
+    when:
+    RumInjector.setTelemetryCollector(mockTelemetryCollector)
+    def telemetryCollector = RumInjector.getTelemetryCollector()
+
+    then:
+    telemetryCollector == mockTelemetryCollector
+  }
+
+  void 'return NO_OP when telemetry collector is not set'() {
+    when:
+    RumInjector.setTelemetryCollector(null)
+    def telemetryCollector = RumInjector.getTelemetryCollector()
+
+    then:
+    telemetryCollector == RumTelemetryCollector.NO_OP
+  }
+
+  void 'enable telemetry with StatsDClient'() {
+    setup:
+    def mockStatsDClient = mock(datadog.trace.api.StatsDClient)
+    RumInjector.enableTelemetry(mockStatsDClient)
+
+    when:
+    def telemetryCollector = RumInjector.getTelemetryCollector()
+
+    then:
+    telemetryCollector instanceof datadog.trace.api.rum.RumInjectorMetrics
+  }
+
+  void 'enabling telemetry with a null StatsDClient sets the telemetry collector to NO_OP'() {
+    when:
+    RumInjector.enableTelemetry(null)
+    def telemetryCollector = RumInjector.getTelemetryCollector()
+
+    then:
+    telemetryCollector == RumTelemetryCollector.NO_OP
+  }
+
+  void 'shutdown telemetry'() {
+    setup:
+    def mockStatsDClient = mock(datadog.trace.api.StatsDClient)
+    RumInjector.enableTelemetry(mockStatsDClient)
+
+    when:
+    RumInjector.shutdownTelemetry()
+    def telemetryCollector = RumInjector.getTelemetryCollector()
+
+    then:
+    telemetryCollector == RumTelemetryCollector.NO_OP
+  }
 }