wip: add poolable map of tag

Author: amarziali

dd-trace-core/src/main/java/datadog/trace/civisibility/writer/ddintake/CiTestCycleMapperV1.java

@@ -9,7 +9,6 @@
 import datadog.communication.serialization.msgpack.MsgPackWriter;
 import datadog.trace.api.DDTags;
 import datadog.trace.api.DDTraceId;
-import datadog.trace.api.TagMap;
 import datadog.trace.api.civisibility.CiVisibilityWellKnownTags;
 import datadog.trace.api.civisibility.InstrumentationBridge;
 import datadog.trace.api.civisibility.telemetry.CiVisibilityDistributionMetric;
@@ -317,7 +316,7 @@ MetaWriter withWritable(Writable writable) {
 
     @Override
     public void accept(Metadata metadata) {
-      TagMap tags = metadata.getTags().copy();
+      Map<String, Object> tags = new HashMap<>(metadata.getTags());
 
       for (String ignoredTag : DEFAULT_TOP_LEVEL_TAGS) {
         tags.remove(ignoredTag);

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

@@ -1256,6 +1256,9 @@ void write(final List<DDSpan> trace) {
     if (null != rootSpan) {
       onRootSpanFinished(rootSpan, rootSpan.getEndpointTracker());
     }
+    for (DDSpan span : writtenTrace) {
+      span.context().releaseTags();
+    }
   }
 
   private List<DDSpan> interceptCompleteTrace(List<DDSpan> trace) {

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

@@ -11,6 +11,7 @@
 import datadog.trace.api.DDTraceId;
 import datadog.trace.api.Functions;
 import datadog.trace.api.ProcessTags;
+import datadog.trace.api.SimplePooledMap;
 import datadog.trace.api.TagMap;
 import datadog.trace.api.cache.DDCache;
 import datadog.trace.api.cache.DDCaches;
@@ -93,6 +94,7 @@ public class DDSpanContext
 
   private volatile short httpStatusCode;
   private CharSequence integrationName;
+  private final String poolKey;
 
   /**
    * Tags are associated to the current span, they will not propagate to the children span.
@@ -103,7 +105,7 @@ public class DDSpanContext
    * rather read and accessed in a serial fashion on thread after thread. The synchronization can
    * then be wrapped around bulk operations to minimize the costly atomic operations.
    */
-  private final TagMap unsafeTags;
+  private final SimplePooledMap unsafeTags;
 
   /** The service name is required, otherwise the span are dropped by the agent */
   private volatile String serviceName;
@@ -112,10 +114,13 @@ public class DDSpanContext
   private volatile CharSequence resourceName;
 
   private volatile byte resourceNamePriority = ResourceNamePriorities.DEFAULT;
+
   /** Each span have an operation name describing the current span */
   private volatile CharSequence operationName;
+
   /** The type of the span. If null, the Datadog Agent will report as a custom */
   private volatile CharSequence spanType;
+
   /** True indicates that the span reports an error */
   private volatile boolean errorFlag;
 
@@ -351,7 +356,8 @@ public DDSpanContext(
     // The +1 is the magic number from the tags below that we set at the end,
     // and "* 4 / 3" is to make sure that we don't resize immediately
     final int capacity = Math.max((tagsSize <= 0 ? 3 : (tagsSize + 1)) * 4 / 3, 8);
-    this.unsafeTags = TagMap.create(capacity);
+    this.poolKey = operationName != null ? operationName.toString() : "";
+    this.unsafeTags = SimplePooledMap.acquire(poolKey, capacity);
 
     // must set this before setting the service and resource names below
     this.profilingContextIntegration = profilingContextIntegration;
@@ -790,7 +796,7 @@ void setAllTags(final TagMap map, boolean needsIntercept) {
               Object value = tagEntry.objectValue();
 
               if (!tagInterceptor.interceptTag(ctx, tag, value)) {
-                ctx.unsafeTags.set(tagEntry);
+                ctx.unsafeTags.put(tag, value);
               }
             });
       } else {
@@ -816,7 +822,7 @@ void setAllTags(final TagMap.Ledger ledger) {
           Object value = entry.objectValue();
 
           if (!tagInterceptor.interceptTag(this, tag, value)) {
-            unsafeTags.set(entry);
+            unsafeTags.put(tag, value);
           }
         }
       }
@@ -863,6 +869,10 @@ Object getTag(final String key) {
     }
   }
 
+  void releaseTags() {
+    SimplePooledMap.release(poolKey, unsafeTags);
+  }
+
   /**
    * This is not thread-safe and must only be used when it can be guaranteed that the context will
    * not be mutated. This is internal API and must not be exposed to users.
@@ -871,13 +881,13 @@ Object getTag(final String key) {
    * @return the value associated with the tag
    */
   public Object unsafeGetTag(final String tag) {
-    return unsafeTags.getObject(tag);
+    return unsafeTags.get(tag);
   }
 
   @Deprecated
   public TagMap getTags() {
     synchronized (unsafeTags) {
-      TagMap tags = unsafeTags.copy();
+      TagMap tags = TagMap.fromMap(unsafeTags);
 
       tags.put(DDTags.THREAD_ID, threadId);
       // maintain previously observable type of the thread name :|

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

@@ -2,15 +2,14 @@
 
 import static datadog.trace.api.sampling.PrioritySampling.UNSET;
 
-import datadog.trace.api.TagMap;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import java.util.Map;
 
 public final class Metadata {
   private final long threadId;
   private final UTF8BytesString threadName;
   private final UTF8BytesString httpStatusCode;
-  private final TagMap tags;
+  private final Map<String, Object> tags;
   private final Map<String, String> baggage;
 
   private final int samplingPriority;
@@ -23,7 +22,7 @@ public final class Metadata {
   public Metadata(
       long threadId,
       UTF8BytesString threadName,
-      TagMap tags,
+      Map<String, Object> tags,
       Map<String, String> baggage,
       int samplingPriority,
       boolean measured,
@@ -61,7 +60,7 @@ public UTF8BytesString getThreadName() {
     return threadName;
   }
 
-  public TagMap getTags() {
+  public Map<String, Object> getTags() {
     return this.tags;
   }
 

dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/BaseServiceAdder.java

@@ -1,28 +1,29 @@
 package datadog.trace.core.tagprocessor;
 
 import datadog.trace.api.DDTags;
-import datadog.trace.api.TagMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.core.DDSpanContext;
 import java.util.List;
+import java.util.Map;
 import javax.annotation.Nullable;
 
-public final class BaseServiceAdder extends TagsPostProcessor {
+public final class BaseServiceAdder implements TagsPostProcessor {
   private final UTF8BytesString ddService;
 
   public BaseServiceAdder(@Nullable final String ddService) {
     this.ddService = ddService != null ? UTF8BytesString.create(ddService) : null;
   }
 
   @Override
-  public void processTags(
-      TagMap unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
+  public Map<String, Object> processTags(
+      Map<String, Object> unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
     if (ddService != null
         && spanContext != null
         && !ddService.toString().equalsIgnoreCase(spanContext.getServiceName())) {
       unsafeTags.put(DDTags.BASE_SERVICE, ddService);
       unsafeTags.remove("version");
     }
+    return unsafeTags;
   }
 }

dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/IntegrationAdder.java

@@ -2,20 +2,21 @@
 
 import static datadog.trace.api.DDTags.DD_INTEGRATION;
 
-import datadog.trace.api.TagMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
 import datadog.trace.core.DDSpanContext;
 import java.util.List;
+import java.util.Map;
 
-public class IntegrationAdder extends TagsPostProcessor {
+public class IntegrationAdder implements TagsPostProcessor {
   @Override
-  public void processTags(
-      TagMap unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
+  public Map<String, Object> processTags(
+      Map<String, Object> unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
     final CharSequence instrumentationName = spanContext.getIntegrationName();
     if (instrumentationName != null) {
-      unsafeTags.set(DD_INTEGRATION, instrumentationName);
+      unsafeTags.put(DD_INTEGRATION, instrumentationName);
     } else {
       unsafeTags.remove(DD_INTEGRATION);
     }
+    return unsafeTags;
   }
 }

dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/PayloadTagsProcessor.java

@@ -4,7 +4,6 @@
 
 import datadog.trace.api.Config;
 import datadog.trace.api.ConfigDefaults;
-import datadog.trace.api.TagMap;
 import datadog.trace.api.telemetry.LogCollector;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
 import datadog.trace.core.DDSpanContext;
@@ -22,7 +21,7 @@
 import org.slf4j.LoggerFactory;
 
 /** Post-processor that extracts tags from payload data injected as tags by instrumentations. */
-public final class PayloadTagsProcessor extends TagsPostProcessor {
+public final class PayloadTagsProcessor implements TagsPostProcessor {
   private static final Logger log = LoggerFactory.getLogger(PayloadTagsProcessor.class);
 
   private static final String REDACTED = "redacted";
@@ -70,19 +69,17 @@ public static PayloadTagsProcessor create(Config config) {
   }
 
   @Override
-  public void processTags(
-      TagMap unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
+  public Map<String, Object> processTags(
+      Map<String, Object> unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
     int spanMaxTags = maxTags + unsafeTags.size();
     for (Map.Entry<String, RedactionRules> tagPrefixRedactionRules :
         redactionRulesByTagPrefix.entrySet()) {
       String tagPrefix = tagPrefixRedactionRules.getKey();
       RedactionRules redactionRules = tagPrefixRedactionRules.getValue();
-      Object tagValue = unsafeTags.getObject(tagPrefix);
+      Object tagValue = unsafeTags.get(tagPrefix);
       if (tagValue instanceof PayloadTagsData) {
-        if (unsafeTags.remove(tagPrefix)) {
-          spanMaxTags -= 1;
-        }
-
+        unsafeTags.remove(tagPrefix);
+        spanMaxTags -= 1;
         PayloadTagsData payloadTagsData = (PayloadTagsData) tagValue;
         PayloadTagsCollector payloadTagsCollector =
             new PayloadTagsCollector(maxDepth, spanMaxTags, redactionRules, tagPrefix, unsafeTags);
@@ -95,6 +92,7 @@ public void processTags(
             tagValue);
       }
     }
+    return unsafeTags;
   }
 
   private void collectPayloadTags(
@@ -181,14 +179,14 @@ private static final class PayloadTagsCollector implements JsonStreamParser.Visi
     private final RedactionRules redactionRules;
     private final String tagPrefix;
 
-    private final TagMap collectedTags;
+    private final Map<String, Object> collectedTags;
 
     public PayloadTagsCollector(
         int maxDepth,
         int maxTags,
         RedactionRules redactionRules,
         String tagPrefix,
-        TagMap collectedTags) {
+        Map<String, Object> collectedTags) {
       this.maxDepth = maxDepth;
       this.maxTags = maxTags;
       this.redactionRules = redactionRules;

dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/PeerServiceCalculator.java

@@ -2,7 +2,6 @@
 
 import datadog.trace.api.Config;
 import datadog.trace.api.DDTags;
-import datadog.trace.api.TagMap;
 import datadog.trace.api.naming.NamingSchema;
 import datadog.trace.api.naming.SpanNaming;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
@@ -12,7 +11,7 @@
 import java.util.Map;
 import javax.annotation.Nonnull;
 
-public final class PeerServiceCalculator extends TagsPostProcessor {
+public final class PeerServiceCalculator implements TagsPostProcessor {
   private final NamingSchema.ForPeerService peerServiceNaming;
 
   private final Map<String, String> peerServiceMapping;
@@ -33,26 +32,27 @@ public PeerServiceCalculator() {
   }
 
   @Override
-  public void processTags(
-      TagMap unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
-    Object peerService = unsafeTags.getObject(Tags.PEER_SERVICE);
+  public Map<String, Object> processTags(
+      Map<String, Object> unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
+    Object peerService = unsafeTags.get(Tags.PEER_SERVICE);
     // the user set it
     if (peerService != null) {
       if (canRemap) {
         remapPeerService(unsafeTags, peerService);
-        return;
+        return unsafeTags;
       }
     } else if (peerServiceNaming.supports()) {
       // calculate the defaults (if any)
       peerServiceNaming.tags(unsafeTags);
       // only remap if the mapping is not empty (saves one get)
-      remapPeerService(unsafeTags, canRemap ? unsafeTags.getObject(Tags.PEER_SERVICE) : null);
-      return;
+      remapPeerService(unsafeTags, canRemap ? unsafeTags.get(Tags.PEER_SERVICE) : null);
+      return unsafeTags;
     }
     // we have no peer.service and we do not compute defaults. Leave the map untouched
+    return unsafeTags;
   }
 
-  private void remapPeerService(TagMap unsafeTags, Object value) {
+  private void remapPeerService(Map<String, Object> unsafeTags, Object value) {
     if (value != null) {
       String mapped = peerServiceMapping.get(value);
       if (mapped != null) {

dd-trace-core/src/main/java/datadog/trace/core/tagprocessor/PostProcessorChain.java

@@ -1,24 +1,25 @@
 package datadog.trace.core.tagprocessor;
 
-import datadog.trace.api.TagMap;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpanLink;
 import datadog.trace.core.DDSpanContext;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import javax.annotation.Nonnull;
 
-public final class PostProcessorChain extends TagsPostProcessor {
+public final class PostProcessorChain implements TagsPostProcessor {
   private final TagsPostProcessor[] chain;
 
   public PostProcessorChain(@Nonnull final TagsPostProcessor... processors) {
     chain = Objects.requireNonNull(processors);
   }
 
   @Override
-  public void processTags(
-      TagMap unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
+  public Map<String, Object> processTags(
+      Map<String, Object> unsafeTags, DDSpanContext spanContext, List<AgentSpanLink> spanLinks) {
     for (final TagsPostProcessor tagsPostProcessor : chain) {
       tagsPostProcessor.processTags(unsafeTags, spanContext, spanLinks);
     }
+    return unsafeTags;
   }
 }