Defer discovery as much as possible

Author: amarziali

dd-trace-core/src/main/java/datadog/trace/common/metrics/ConflatingMetricsAggregator.java

@@ -199,9 +199,20 @@ public ConflatingMetricsAggregator(
     this.reportingIntervalTimeUnit = timeUnit;
   }
 
+  private DDAgentFeaturesDiscovery getFeatures() {
+    DDAgentFeaturesDiscovery ret = features;
+    if (ret == null) {
+      return ret;
+    }
+    ret = sharedCommunicationObjects.featuresDiscovery(Config.get());
+    features = ret;
+    return ret;
+  }
+
   @Override
   public void start() {
-    features = sharedCommunicationObjects.featuresDiscovery(Config.get());
+    AgentTaskScheduler.get()
+        .execute(() -> features = sharedCommunicationObjects.featuresDiscovery(Config.get()));
     sink.register(this);
     thread.start();
     cancellation =
@@ -216,13 +227,7 @@ public void start() {
   }
 
   private boolean isMetricsEnabled() {
-    if (features != null) {
-      if (features.getMetricsEndpoint() == null) {
-        features.discoverIfOutdated();
-      }
-      return features.supportsMetrics();
-    }
-    return false;
+    return getFeatures().supportsMetrics();
   }
 
   @Override
@@ -278,6 +283,7 @@ public Future<Boolean> forceReport() {
   public boolean publish(List<? extends CoreSpan<?>> trace) {
     boolean forceKeep = false;
     int counted = 0;
+    final DDAgentFeaturesDiscovery features = getFeatures();
     if (features != null && features.supportsMetrics()) {
       for (CoreSpan<?> span : trace) {
         boolean isTopLevel = span.isTopLevel();
@@ -288,7 +294,7 @@ public boolean publish(List<? extends CoreSpan<?>> trace) {
             break;
           }
           counted++;
-          forceKeep |= publish(span, isTopLevel);
+          forceKeep |= publish(span, isTopLevel, features);
         }
       }
       healthMetrics.onClientStatTraceComputed(
@@ -310,7 +316,7 @@ private boolean spanKindEligible(CoreSpan<?> span) {
     return spanKind != null && ELIGIBLE_SPAN_KINDS_FOR_METRICS.contains(spanKind.toString());
   }
 
-  private boolean publish(CoreSpan<?> span, boolean isTopLevel) {
+  private boolean publish(CoreSpan<?> span, boolean isTopLevel, DDAgentFeaturesDiscovery features) {
     final CharSequence spanKind = span.getTag(SPAN_KIND, "");
     MetricKey newKey =
         new MetricKey(
@@ -323,7 +329,7 @@ private boolean publish(CoreSpan<?> span, boolean isTopLevel) {
             span.getParentId() == 0,
             SPAN_KINDS.computeIfAbsent(
                 spanKind, UTF8BytesString::create), // save repeated utf8 conversions
-            getPeerTags(span, spanKind.toString()));
+            getPeerTags(span, spanKind.toString(), features));
     boolean isNewKey = false;
     MetricKey key = keys.putIfAbsent(newKey, newKey);
     if (null == key) {
@@ -358,7 +364,8 @@ private boolean publish(CoreSpan<?> span, boolean isTopLevel) {
     return isNewKey || span.getError() > 0;
   }
 
-  private List<UTF8BytesString> getPeerTags(CoreSpan<?> span, String spanKind) {
+  private List<UTF8BytesString> getPeerTags(
+      CoreSpan<?> span, String spanKind, DDAgentFeaturesDiscovery features) {
     if (features != null) {
       if (ELIGIBLE_SPAN_KINDS_FOR_PEER_AGGREGATION.contains(spanKind)) {
         List<UTF8BytesString> peerTags = new ArrayList<>();
@@ -425,8 +432,7 @@ public void onEvent(EventType eventType, String message) {
     switch (eventType) {
       case DOWNGRADED:
         log.debug("Agent downgrade was detected");
-        disable();
-        healthMetrics.onClientStatDowngraded();
+        AgentTaskScheduler.get().execute(this::disable);
         break;
       case BAD_PAYLOAD:
         log.debug("bad metrics payload sent to trace agent: {}", message);
@@ -442,10 +448,11 @@ public void onEvent(EventType eventType, String message) {
   }
 
   private void disable() {
-    // note: disable is called only if started so we're not nullchecking before accessing features
+    final DDAgentFeaturesDiscovery features = getFeatures();
     features.discover();
     if (!features.supportsMetrics()) {
       log.debug("Disabling metric reporting because an agent downgrade was detected");
+      healthMetrics.onClientStatDowngraded();
       this.pending.clear();
       this.batchPool.clear();
       this.inbox.clear();

dd-trace-core/src/test/groovy/datadog/trace/common/metrics/MetricsReliabilityTest.groovy

@@ -12,6 +12,8 @@ import java.util.concurrent.CountDownLatch
 
 import static datadog.trace.agent.test.server.http.TestHttpServer.httpServer
 
+import java.util.concurrent.TimeUnit
+
 class MetricsReliabilityTest extends DDCoreSpecification {
 
   static class State {
@@ -87,13 +89,19 @@ class MetricsReliabilityTest extends DDCoreSpecification {
 
 
     when: "simulate an agent downgrade"
+    def discoveryState = featuresDiscovery.discoveryState
     state.reset(false, 404)
     tracer.startSpan("test", "test").finish()
     tracer.flush()
     tracer.flushMetrics()
 
     then: "a discovery should have done - we do not support anymore stats calculation"
     state.latch.await()
+    // wait at least 5 seconds. the discovery is done asynchronously so we should wait for the internal state flip
+    def start = System.nanoTime()
+    while (discoveryState == featuresDiscovery.discoveryState && TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() - start) < 5) {
+      Thread.yield()
+    }
     assert !featuresDiscovery.supportsMetrics()
     // 2 traces processed. 1 p0 dropped. 2 requests and 1 downgrade no errors
     assertMetrics(healthMetrics, 2, 1, 2, 0, 1)