revisit peer tags aggregation rules according to the rfc

Author: amarziali

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

@@ -1,10 +1,12 @@
 package datadog.trace.common.metrics;
 
 import static datadog.communication.ddagent.DDAgentFeaturesDiscovery.V6_METRICS_ENDPOINT;
+import static datadog.trace.api.DDTags.BASE_SERVICE;
 import static datadog.trace.api.Functions.UTF8_ENCODE;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CLIENT;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_CONSUMER;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_INTERNAL;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_PRODUCER;
 import static datadog.trace.bootstrap.instrumentation.api.Tags.SPAN_KIND_SERVER;
 import static datadog.trace.common.metrics.AggregateMetric.ERROR_TAG;
@@ -75,12 +77,15 @@ public final class ConflatingMetricsAggregator implements MetricsAggregator, Eve
                   value -> UTF8BytesString.create(key + ":" + value));
   private static final CharSequence SYNTHETICS_ORIGIN = "synthetics";
 
-  private static final Set<String> ELIGIBLE_SPAN_KINDS =
+  private static final Set<String> ELIGIBLE_SPAN_KINDS_FOR_METRICS =
       unmodifiableSet(
           new HashSet<>(
               Arrays.asList(
                   SPAN_KIND_SERVER, SPAN_KIND_CLIENT, SPAN_KIND_CONSUMER, SPAN_KIND_PRODUCER)));
 
+  private static final Set<String> ELIGIBLE_SPAN_KINDS_FOR_PEER_AGGREGATION =
+      unmodifiableSet(new HashSet<>(Arrays.asList(SPAN_KIND_CLIENT, SPAN_KIND_PRODUCER)));
+
   private final Set<String> ignoredResources;
   private final Queue<Batch> batchPool;
   private final NonBlockingHashMap<MetricKey, Batch> pending;
@@ -269,10 +274,11 @@ private boolean shouldComputeMetric(CoreSpan<?> span) {
   private boolean spanKindEligible(CoreSpan<?> span) {
     final Object spanKind = span.getTag(SPAN_KIND);
     // use toString since it could be a CharSequence...
-    return spanKind != null && ELIGIBLE_SPAN_KINDS.contains(spanKind.toString());
+    return spanKind != null && ELIGIBLE_SPAN_KINDS_FOR_METRICS.contains(spanKind.toString());
   }
 
   private boolean publish(CoreSpan<?> span, boolean isTopLevel) {
+    final CharSequence spanKind = span.getTag(SPAN_KIND, "");
     MetricKey newKey =
         new MetricKey(
             span.getResourceName(),
@@ -282,8 +288,9 @@ private boolean publish(CoreSpan<?> span, boolean isTopLevel) {
             span.getHttpStatusCode(),
             isSynthetic(span),
             span.isTopLevel(),
-            SPAN_KINDS.computeIfAbsent(span.getTag(SPAN_KIND, ""), UTF8BytesString::create),
-            getPeerTags(span));
+            SPAN_KINDS.computeIfAbsent(
+                spanKind, UTF8BytesString::create), // save repeated utf8 conversions
+            getPeerTags(span, spanKind.toString()));
     boolean isNewKey = false;
     MetricKey key = keys.putIfAbsent(newKey, newKey);
     if (null == key) {
@@ -318,17 +325,32 @@ private boolean publish(CoreSpan<?> span, boolean isTopLevel) {
     return isNewKey || span.getError() > 0;
   }
 
-  private List<UTF8BytesString> getPeerTags(CoreSpan<?> span) {
-    List<UTF8BytesString> peerTags = new ArrayList<>();
-    for (String peerTag : features.peerTags()) {
-      Object value = span.getTag(peerTag);
-      if (value != null) {
-        final Pair<DDCache<String, UTF8BytesString>, Function<String, UTF8BytesString>> pair =
-            PEER_TAGS_CACHE.computeIfAbsent(peerTag, PEER_TAGS_CACHE_ADDER);
-        peerTags.add(pair.getLeft().computeIfAbsent(value.toString(), pair.getRight()));
+  private List<UTF8BytesString> getPeerTags(CoreSpan<?> span, String spanKind) {
+    if (ELIGIBLE_SPAN_KINDS_FOR_PEER_AGGREGATION.contains(spanKind)) {
+      List<UTF8BytesString> peerTags = new ArrayList<>();
+      for (String peerTag : features.peerTags()) {
+        Object value = span.getTag(peerTag);
+        if (value != null) {
+          final Pair<DDCache<String, UTF8BytesString>, Function<String, UTF8BytesString>>
+              cacheAndCreator = PEER_TAGS_CACHE.computeIfAbsent(peerTag, PEER_TAGS_CACHE_ADDER);
+          peerTags.add(
+              cacheAndCreator
+                  .getLeft()
+                  .computeIfAbsent(value.toString(), cacheAndCreator.getRight()));
+        }
+      }
+      return peerTags;
+    } else if (SPAN_KIND_INTERNAL.equals(spanKind)) {
+      // in this case only the base service should be aggregated if present
+      final String baseService = span.getTag(BASE_SERVICE);
+      if (baseService != null) {
+        final Pair<DDCache<String, UTF8BytesString>, Function<String, UTF8BytesString>>
+            cacheAndCreator = PEER_TAGS_CACHE.computeIfAbsent(BASE_SERVICE, PEER_TAGS_CACHE_ADDER);
+        return Collections.singletonList(
+            cacheAndCreator.getLeft().computeIfAbsent(baseService, cacheAndCreator.getRight()));
       }
     }
-    return peerTags;
+    return Collections.emptyList();
   }
 
   private static boolean isSynthetic(CoreSpan<?> span) {

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

@@ -201,9 +201,9 @@ class ConflatingMetricAggregatorTest extends DDSpecification {
     CountDownLatch latch = new CountDownLatch(1)
     aggregator.publish([
       new SimpleSpan("service", "operation", "resource", "type", true, false, false, 0, 100, HTTP_OK)
-      .setTag(SPAN_KIND, "grault").setTag("country", "france").setTag("georegion", "europe"),
+      .setTag(SPAN_KIND, "client").setTag("country", "france").setTag("georegion", "europe"),
       new SimpleSpan("service", "operation", "resource", "type", true, false, false, 0, 100, HTTP_OK)
-      .setTag(SPAN_KIND, "grault").setTag("country", "france").setTag("georegion", "europe")
+      .setTag(SPAN_KIND, "client").setTag("country", "france").setTag("georegion", "europe")
     ])
     aggregator.report()
     def latchTriggered = latch.await(2, SECONDS)
@@ -220,7 +220,7 @@ class ConflatingMetricAggregatorTest extends DDSpecification {
       HTTP_OK,
       false,
       false,
-      "grault",
+      "client",
       [UTF8BytesString.create("country:france")]
       ), { AggregateMetric aggregateMetric ->
         aggregateMetric.getHitCount() == 1 && aggregateMetric.getTopLevelCount() == 0 && aggregateMetric.getDuration() == 100
@@ -234,7 +234,7 @@ class ConflatingMetricAggregatorTest extends DDSpecification {
       HTTP_OK,
       false,
       false,
-      "grault",
+      "client",
       [UTF8BytesString.create("country:france"), UTF8BytesString.create("georegion:europe")]
       ), { AggregateMetric aggregateMetric ->
         aggregateMetric.getHitCount() == 1 && aggregateMetric.getTopLevelCount() == 0 && aggregateMetric.getDuration() == 100
@@ -245,6 +245,55 @@ class ConflatingMetricAggregatorTest extends DDSpecification {
     aggregator.close()
   }
 
+  def "should aggregate the right peer tags for kind #kind"() {
+    setup:
+    MetricWriter writer = Mock(MetricWriter)
+    Sink sink = Stub(Sink)
+    DDAgentFeaturesDiscovery features = Mock(DDAgentFeaturesDiscovery)
+    features.supportsMetrics() >> true
+    features.peerTags() >> ["peer.hostname", "_dd.base_service"]
+    ConflatingMetricsAggregator aggregator = new ConflatingMetricsAggregator(empty,
+      features, sink, writer, 10, queueSize, reportingInterval, SECONDS)
+    aggregator.start()
+
+    when:
+    CountDownLatch latch = new CountDownLatch(1)
+    aggregator.publish([
+      new SimpleSpan("service", "operation", "resource", "type", true, false, false, 0, 100, HTTP_OK)
+      .setTag(SPAN_KIND, kind).setTag("peer.hostname", "localhost").setTag("_dd.base_service", "test")
+    ])
+    aggregator.report()
+    def latchTriggered = latch.await(2, SECONDS)
+
+    then:
+    latchTriggered
+    1 * writer.startBucket(1, _, _)
+    1 * writer.add(
+      new MetricKey(
+      "resource",
+      "service",
+      "operation",
+      "type",
+      HTTP_OK,
+      false,
+      false,
+      kind,
+      expectedPeerTags
+      ), { AggregateMetric aggregateMetric ->
+        aggregateMetric.getHitCount() == 1 && aggregateMetric.getTopLevelCount() == 0 && aggregateMetric.getDuration() == 100
+      })
+    1 * writer.finishBucket() >> { latch.countDown() }
+
+    cleanup:
+    aggregator.close()
+
+    where:
+    kind       | expectedPeerTags
+    "client"   | [UTF8BytesString.create("peer.hostname:localhost"), UTF8BytesString.create("_dd.base_service:test")]
+    "internal" | [UTF8BytesString.create("_dd.base_service:test")]
+    "server"   | []
+  }
+
   def "measured spans do not contribute to top level count"() {
     setup:
     MetricWriter writer = Mock(MetricWriter)