Small tweaks to tags map in DDSpanContext

dd-trace-core/jfr-openjdk/src/test/groovy/datadog/trace/core/jfr/openjdk/ScopeEventTest.groovy

@@ -40,7 +40,7 @@ class ScopeEventTest extends DDSpecification {
       [:],
       false,
       "fakeType",
-      null,
+      0,
       new PendingTrace(tracer, DDId.from(123)),
       tracer,
       [:])

dd-trace-core/src/main/java/datadog/trace/common/writer/ddagent/TraceMapper.java

@@ -18,14 +18,15 @@
 
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.core.DDSpan;
+import datadog.trace.core.DDSpanContext;
 import datadog.trace.core.serialization.msgpack.Mapper;
 import datadog.trace.core.serialization.msgpack.Writable;
 import java.util.List;
 import java.util.Map;
 
 public final class TraceMapper implements Mapper<List<DDSpan>> {
   @Override
-  public void map(List<DDSpan> trace, Writable writable) {
+  public void map(List<DDSpan> trace, final Writable writable) {
     writable.startArray(trace.size());
     for (DDSpan span : trace) {
       writable.startMap(12);
@@ -64,42 +65,48 @@ public void map(List<DDSpan> trace, Writable writable) {
       writable.writeMap(span.getMetrics(), CONSTANT_KEYS);
       /* 12 */
       writable.writeUTF8(META);
-      Map<String, String> baggage = span.context().getBaggageItems();
-      Map<String, Object> tags = span.context().getTags();
-      // since tags can "override" baggage, we need to count the non overlapping ones
-      int size = tags.size();
-      boolean overlap = false;
-      if (baggage.size() > 0) {
-        for (String key : baggage.keySet()) {
-          if (!tags.containsKey(key)) {
-            size++;
-          } else {
-            overlap = true;
-          }
-        }
-      }
-      writable.startMap(size);
-      for (Map.Entry<String, String> entry : baggage.entrySet()) {
-        // tags and baggage may intersect, but tags take priority
-        if (!overlap || !tags.containsKey(entry.getKey())) {
-          writable.writeString(entry.getKey(), CONSTANT_KEYS);
-          writable.writeObject(entry.getValue(), NO_CACHING);
-        }
-      }
-      for (Map.Entry<String, Object> entry : tags.entrySet()) {
-        writable.writeString(entry.getKey(), CONSTANT_KEYS);
-        if (entry.getValue() instanceof Long || entry.getValue() instanceof Integer) {
-          // TODO it would be nice not to need to do this, either because
-          //  the agent would accept variably typed tag values, or numeric
-          //  tags get moved to the metrics
-          writeLongAsString(((Number) entry.getValue()).longValue(), writable);
-        } else if (entry.getValue() instanceof UTF8BytesString) {
-          // TODO assess whether this is still worth it
-          writable.writeObject(entry.getValue(), NO_CACHING);
-        } else {
-          writable.writeString(String.valueOf(entry.getValue()), NO_CACHING);
-        }
-      }
+      span.context()
+          .processTagsAndBaggage(
+              new DDSpanContext.TagsAndBaggageConsumer() {
+                @Override
+                public void accept(Map<String, Object> tags, Map<String, String> baggage) {
+                  // since tags can "override" baggage, we need to count the non overlapping ones
+                  int size = tags.size();
+                  boolean overlap = false;
+                  if (baggage.size() > 0) {
+                    for (String key : baggage.keySet()) {
+                      if (!tags.containsKey(key)) {
+                        size++;
+                      } else {
+                        overlap = true;
+                      }
+                    }
+                  }
+                  writable.startMap(size);
+                  for (Map.Entry<String, String> entry : baggage.entrySet()) {
+                    // tags and baggage may intersect, but tags take priority
+                    if (!overlap || !tags.containsKey(entry.getKey())) {
+                      writable.writeString(entry.getKey(), CONSTANT_KEYS);
+                      writable.writeObject(entry.getValue(), NO_CACHING);
+                    }
+                  }
+                  for (Map.Entry<String, Object> entry : tags.entrySet()) {
+                    writable.writeString(entry.getKey(), CONSTANT_KEYS);
+                    if (entry.getValue() instanceof Long || entry.getValue() instanceof Integer) {
+                      // TODO it would be nice not to need to do this, either because
+                      //  the agent would accept variably typed tag values, or numeric
+                      //  tags get moved to the metrics
+                      TraceMapper.this.writeLongAsString(
+                          ((Number) entry.getValue()).longValue(), writable);
+                    } else if (entry.getValue() instanceof UTF8BytesString) {
+                      // TODO assess whether this is still worth it
+                      writable.writeObject(entry.getValue(), NO_CACHING);
+                    } else {
+                      writable.writeString(String.valueOf(entry.getValue()), NO_CACHING);
+                    }
+                  }
+                }
+              });
     }
   }
 

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

@@ -42,14 +42,14 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
 import java.util.ServiceConfigurationError;
 import java.util.ServiceLoader;
 import java.util.SortedSet;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.TimeUnit;
 import lombok.Builder;
@@ -100,9 +100,11 @@ public class CoreTracer implements AgentTracer.TracerAPI {
    */
   private final Thread shutdownCallback;
 
-  /** Span tag interceptors */
-  private final Map<String, List<AbstractTagInterceptor>> spanTagInterceptors =
-      new ConcurrentHashMap<>();
+  /**
+   * Span tag interceptors. This Map is only ever added to during initialization, so it doesn't need
+   * to be concurrent.
+   */
+  private final Map<String, List<AbstractTagInterceptor>> spanTagInterceptors = new HashMap<>();
 
   private final SortedSet<TraceInterceptor> interceptors =
       new ConcurrentSkipListSet<>(
@@ -582,7 +584,7 @@ public class CoreSpanBuilder implements AgentTracer.SpanBuilder {
     private final String operationName;
 
     // Builder attributes
-    private final Map<String, Object> tags = new LinkedHashMap<String, Object>(defaultSpanTags);
+    private Map<String, Object> tags;
     private long timestampMicro;
     private Object parent;
     private String serviceName;
@@ -669,10 +671,14 @@ public CoreSpanBuilder asChildOf(final AgentSpan agentSpan) {
 
     @Override
     public CoreSpanBuilder withTag(final String tag, final Object value) {
+      Map<String, Object> tagMap = this.tags;
+      if (tagMap == null) {
+        tags = tagMap = new LinkedHashMap<>(); // Insertion order is important
+      }
       if (value == null || (value instanceof String && ((String) value).isEmpty())) {
-        tags.remove(tag);
+        tagMap.remove(tag);
       } else {
-        tags.put(tag, value);
+        tagMap.put(tag, value);
       }
       return this;
     }
@@ -691,6 +697,8 @@ private DDSpanContext buildSpanContext() {
       final PendingTrace parentTrace;
       final int samplingPriority;
       final String origin;
+      final Map<String, String> coreTags;
+      final Map<String, String> rootSpanTags;
 
       final DDSpanContext context;
 
@@ -716,6 +724,8 @@ private DDSpanContext buildSpanContext() {
         parentTrace = ddsc.getTrace();
         samplingPriority = PrioritySampling.UNSET;
         origin = null;
+        coreTags = null;
+        rootSpanTags = null;
         if (serviceName == null) {
           serviceName = ddsc.getServiceName();
         }
@@ -738,13 +748,14 @@ private DDSpanContext buildSpanContext() {
 
         // Get header tags and set origin whether propagating or not.
         if (parentContext instanceof TagContext) {
-          tags.putAll(((TagContext) parentContext).getTags());
+          coreTags = ((TagContext) parentContext).getTags();
           origin = ((TagContext) parentContext).getOrigin();
         } else {
+          coreTags = null;
           origin = null;
         }
 
-        tags.putAll(localRootSpanTags);
+        rootSpanTags = localRootSpanTags;
 
         parentTrace = PendingTrace.create(CoreTracer.this, traceId);
       }
@@ -755,6 +766,11 @@ private DDSpanContext buildSpanContext() {
 
       final String operationName = this.operationName != null ? this.operationName : resourceName;
 
+      int tagsSize =
+          (null == tags ? 0 : tags.size())
+              + defaultSpanTags.size()
+              + (null == coreTags ? 0 : coreTags.size())
+              + (null == rootSpanTags ? 0 : rootSpanTags.size());
       // some attributes are inherited from the parent
       context =
           new DDSpanContext(
@@ -769,40 +785,18 @@ private DDSpanContext buildSpanContext() {
               baggage,
               errorFlag,
               spanType,
-              tags,
+              tagsSize,
               parentTrace,
               CoreTracer.this,
               serviceNameMappings);
 
-      // Apply Decorators to handle any tags that may have been set via the builder.
-      for (final Map.Entry<String, Object> tag : tags.entrySet()) {
-        if (tag.getValue() == null) {
-          context.setTag(tag.getKey(), null);
-          continue;
-        }
-
-        boolean addTag = true;
-
-        // Call interceptors
-        final List<AbstractTagInterceptor> interceptors = getSpanTagInterceptors(tag.getKey());
-        if (interceptors != null) {
-          for (final AbstractTagInterceptor interceptor : interceptors) {
-            try {
-              addTag &= interceptor.shouldSetTag(context, tag.getKey(), tag.getValue());
-            } catch (final Throwable ex) {
-              log.debug(
-                  "Could not intercept the span interceptor={}: {}",
-                  interceptor.getClass().getSimpleName(),
-                  ex.getMessage());
-            }
-          }
-        }
-
-        if (!addTag) {
-          context.setTag(tag.getKey(), null);
-        }
-      }
-
+      // By setting the tags on the context we apply decorators to any tags that have been set via
+      // the builder. This is the order that the tags were added previously, but maybe the `tags`
+      // set in the builder should come last, so that they override other tags.
+      context.setAllTags(defaultSpanTags);
+      context.setAllTags(tags);
+      context.setAllTags(coreTags);
+      context.setAllTags(rootSpanTags);
       return context;
     }
   }

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

@@ -9,7 +9,6 @@
 import java.io.PrintWriter;
 import java.io.StringWriter;
 import java.lang.ref.WeakReference;
-import java.util.Collections;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
@@ -212,12 +211,12 @@ public AgentSpan removeTag(final String tag) {
   }
 
   public Object getAndRemoveTag(final String tag) {
-    return context.getTags().remove(tag);
+    return context.getAndRemoveTag(tag);
   }
 
   @Override
   public Object getTag(final String tag) {
-    return context.getTags().get(tag);
+    return context.getTag(tag);
   }
 
   @Override
@@ -352,7 +351,8 @@ public String getSpanType() {
 
   @Override
   public Map<String, Object> getTags() {
-    return Collections.unmodifiableMap(context.getTags());
+    // This is an imutable copy of the tags
+    return context.getTags();
   }
 
   public String getType() {