Add primitive setters to DDSpanContext to take advantage of TagMap

dd-java-agent/instrumentation-testing/src/main/groovy/datadog/trace/agent/test/TrackingSpanDecorator.groovy

@@ -134,6 +134,11 @@ class TrackingSpanDecorator implements AgentSpan {
     return delegate.setTag(key, value)
   }
 
+  @Override
+  AgentSpan setTag(String key, float value) {
+    return delegate.setTag(key, value)
+  }
+
   @Override
   AgentSpan setTag(String key, double value) {
     return delegate.setTag(key, value)

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

@@ -432,6 +432,12 @@ public DDSpan setTag(final String tag, final long value) {
     return this;
   }
 
+  @Override
+  public DDSpan setTag(final String tag, final float value) {
+    context.setTag(tag, value);
+    return this;
+  }
+
   @Override
   public DDSpan setTag(final String tag, final double value) {
     context.setTag(tag, value);

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

@@ -76,6 +76,8 @@ public class DDSpanContext
   /** The collection of all span related to this one */
   private final TraceCollector traceCollector;
 
+  private final TagInterceptor tagInterceptor;
+
   /** Baggage is associated with the whole trace and shared with other spans */
   private volatile Map<String, String> baggageItems;
 
@@ -330,6 +332,7 @@ public DDSpanContext(
 
     assert traceCollector != null;
     this.traceCollector = traceCollector;
+    this.tagInterceptor = this.traceCollector.getTracer().getTagInterceptor();
 
     assert traceId != null;
     this.traceId = traceId;
@@ -766,9 +769,122 @@ public void setTag(final String tag, final Object value) {
       synchronized (unsafeTags) {
         unsafeTags.remove(tag);
       }
-    } else if (!traceCollector.getTracer().getTagInterceptor().interceptTag(this, tag, value)) {
+    } else if (!tagInterceptor.interceptTag(this, tag, value)) {
+      synchronized (unsafeTags) {
+        unsafeTags.set(tag, value);
+      }
+    }
+  }
+
+  public void setTag(final String tag, final String value) {
+    if (null == tag) {
+      return;
+    }
+    if (null == value) {
+      synchronized (unsafeTags) {
+        unsafeTags.remove(tag);
+      }
+    } else if (!tagInterceptor.interceptTag(this, tag, value)) {
+      synchronized (unsafeTags) {
+        unsafeTags.set(tag, value);
+      }
+    }
+  }
+
+  /*
+   * Uses to determine if there's an opportunity to avoid primitve boxing.
+   * If the underlying map doesn't support efficient primitives, then boxing is used.
+   * If the tag may be intercepted, then boxing is also used.
+   */
+  private boolean precheckIntercept(String tag) {
+    // Usually only a single instanceof TagMap will be loaded,
+    // so isOptimized is turned into a direct call and then inlines to a constant
+    // Since isOptimized just returns a constant - doesn't require synchronization
+    return !unsafeTags.isOptimized() || tagInterceptor.needsIntercept(tag);
+  }
+
+  /*
+   * Used when precheckIntercept determines that boxing is unavoidable
+   *
+   * Either because the tagInterceptor needs to be fully checked (which requires boxing)
+   * In that case, a box has already been created so it makes sense to pass the box
+   * onto TagMap, since optimized TagMap will cache the box
+   *
+   * -- OR --
+   *
+   * The TagMap isn't optimized and will need to box the primitive regardless of
+   * tag interception
+   */
+  private void setBox(String tag, Object box) {
+    if (!tagInterceptor.interceptTag(this, tag, box)) {
+      synchronized (unsafeTags) {
+        unsafeTags.set(tag, box);
+      }
+    }
+  }
+
+  public void setTag(final String tag, final boolean value) {
+    if (null == tag) {
+      return;
+    }
+    if (precheckIntercept(tag)) {
+      this.setBox(tag, value);
+    } else {
       synchronized (unsafeTags) {
-        unsafeSetTag(tag, value);
+        unsafeTags.set(tag, value);
+      }
+    }
+  }
+
+  public void setTag(final String tag, final int value) {
+    if (null == tag) {
+      return;
+    }
+    if (precheckIntercept(tag)) {
+      this.setBox(tag, value);
+    } else {
+      synchronized (unsafeTags) {
+        unsafeTags.set(tag, value);
+      }
+    }
+  }
+
+  public void setTag(final String tag, final long value) {
+    if (null == tag) {
+      return;
+    }
+    // check needsIntercept first to avoid unnecessary boxing
+    boolean intercepted =
+        tagInterceptor.needsIntercept(tag) && tagInterceptor.interceptTag(this, tag, value);
+    if (!intercepted) {
+      synchronized (unsafeTags) {
+        unsafeTags.set(tag, value);
+      }
+    }
+  }
+
+  public void setTag(final String tag, final float value) {
+    if (null == tag) {
+      return;
+    }
+    if (precheckIntercept(tag)) {
+      this.setBox(tag, value);
+    } else {
+      synchronized (unsafeTags) {
+        unsafeTags.set(tag, value);
+      }
+    }
+  }
+
+  public void setTag(final String tag, final double value) {
+    if (null == tag) {
+      return;
+    }
+    if (precheckIntercept(tag)) {
+      this.setBox(tag, value);
+    } else {
+      synchronized (unsafeTags) {
+        unsafeTags.set(tag, value);
       }
     }
   }
@@ -789,12 +905,11 @@ void setAllTags(final TagMap map, boolean needsIntercept) {
         // to avoid using a capturing lambda
         map.forEach(
             this,
-            traceCollector.getTracer().getTagInterceptor(),
-            (ctx, tagInterceptor, tagEntry) -> {
+            (ctx, tagEntry) -> {
               String tag = tagEntry.tag();
               Object value = tagEntry.objectValue();
 
-              if (!tagInterceptor.interceptTag(ctx, tag, value)) {
+              if (!ctx.tagInterceptor.interceptTag(ctx, tag, value)) {
                 ctx.unsafeTags.set(tagEntry);
               }
             });
@@ -809,7 +924,6 @@ void setAllTags(final TagMap.Ledger ledger) {
       return;
     }
 
-    TagInterceptor tagInterceptor = traceCollector.getTracer().getTagInterceptor();
     synchronized (unsafeTags) {
       for (final TagMap.EntryChange entryChange : ledger) {
         if (entryChange.isRemoval()) {
@@ -834,7 +948,6 @@ void setAllTags(final Map<String, ?> map) {
     } else if (map instanceof TagMap) {
       setAllTags((TagMap) map);
     } else if (!map.isEmpty()) {
-      TagInterceptor tagInterceptor = traceCollector.getTracer().getTagInterceptor();
       synchronized (unsafeTags) {
         for (final Map.Entry<String, ?> tag : map.entrySet()) {
           if (!tagInterceptor.interceptTag(this, tag.getKey(), tag.getValue())) {
@@ -846,7 +959,19 @@ void setAllTags(final Map<String, ?> map) {
   }
 
   void unsafeSetTag(final String tag, final Object value) {
-    unsafeTags.put(tag, value);
+    unsafeTags.set(tag, value);
+  }
+
+  void unsafeSetTag(final String tag, final CharSequence value) {
+    unsafeTags.set(tag, value);
+  }
+
+  void unsafeSetTag(final String tag, final byte value) {
+    unsafeTags.set(tag, value);
+  }
+
+  void unsafeSetTag(final String tag, final double value) {
+    unsafeTags.set(tag, value);
   }
 
   Object getTag(final String key) {

dd-trace-core/src/test/groovy/datadog/trace/core/PendingTraceBufferTest.groovy

@@ -104,6 +104,7 @@ class PendingTraceBufferTest extends DDSpecification {
     1 * bufferSpy.longRunningSpansEnabled()
     1 * bufferSpy.enqueue(trace)
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     1 * tracer.getTimeWithNanoTicks(_)
     1 * tracer.onRootSpanPublished(span)
     0 * _
@@ -137,6 +138,7 @@ class PendingTraceBufferTest extends DDSpecification {
     1 * bufferSpy.enqueue(trace)
     _ * bufferSpy.longRunningSpansEnabled()
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     1 * tracer.getTimeWithNanoTicks(_)
     1 * tracer.onRootSpanPublished(parent)
     0 * _
@@ -151,6 +153,7 @@ class PendingTraceBufferTest extends DDSpecification {
     1 * tracer.write({ it.size() == 2 })
     1 * tracer.writeTimer() >> Monitoring.DISABLED.newTimer("")
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     1 * tracer.getTimeWithNanoTicks(_)
     0 * _
   }
@@ -168,6 +171,7 @@ class PendingTraceBufferTest extends DDSpecification {
 
     then:
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     _ * traceConfig.getServiceMapping() >> [:]
     _ * tracer.getTimeWithNanoTicks(_)
     1 * tracer.writeTimer() >> Monitoring.DISABLED.newTimer("")
@@ -194,6 +198,7 @@ class PendingTraceBufferTest extends DDSpecification {
     buffer.queue.capacity() * bufferSpy.enqueue(_)
     _ * bufferSpy.longRunningSpansEnabled()
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     _ * traceConfig.getServiceMapping() >> [:]
     _ * tracer.getTimeWithNanoTicks(_)
     buffer.queue.capacity() * tracer.onRootSpanPublished(_)
@@ -210,6 +215,7 @@ class PendingTraceBufferTest extends DDSpecification {
     _ * bufferSpy.longRunningSpansEnabled()
     1 * tracer.write({ it.size() == 1 })
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     _ * traceConfig.getServiceMapping() >> [:]
     2 * tracer.getTimeWithNanoTicks(_)
     1 * tracer.onRootSpanPublished(_)
@@ -232,6 +238,7 @@ class PendingTraceBufferTest extends DDSpecification {
     buffer.queue.capacity() * bufferSpy.enqueue(_)
     _ * bufferSpy.longRunningSpansEnabled()
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     _ * traceConfig.getServiceMapping() >> [:]
     _ * tracer.getTimeWithNanoTicks(_)
     buffer.queue.capacity() * tracer.onRootSpanPublished(_)
@@ -250,6 +257,7 @@ class PendingTraceBufferTest extends DDSpecification {
     _ * bufferSpy.longRunningSpansEnabled()
     0 * tracer.write({ it.size() == 1 })
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     _ * traceConfig.getServiceMapping() >> [:]
     _ * tracer.getTimeWithNanoTicks(_)
     1 * tracer.onRootSpanPublished(_)
@@ -279,6 +287,7 @@ class PendingTraceBufferTest extends DDSpecification {
     _ * bufferSpy.longRunningSpansEnabled()
     1 * bufferSpy.enqueue(trace)
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     1 * tracer.getTimeWithNanoTicks(_)
     1 * tracer.onRootSpanPublished(parent)
     0 * _
@@ -307,6 +316,7 @@ class PendingTraceBufferTest extends DDSpecification {
       latch.countDown()
     }
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     0 * _
   }
 
@@ -333,6 +343,7 @@ class PendingTraceBufferTest extends DDSpecification {
       parentLatch.countDown()
     }
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     1 * tracer.getTimeWithNanoTicks(_)
     1 * tracer.onRootSpanPublished(parent)
     0 * _
@@ -349,6 +360,7 @@ class PendingTraceBufferTest extends DDSpecification {
     1 * bufferSpy.enqueue(trace)
     _ * bufferSpy.longRunningSpansEnabled()
     _ * tracer.getPartialFlushMinSpans() >> 10
+    _ * tracer.getTagInterceptor()
     1 * tracer.writeTimer() >> Monitoring.DISABLED.newTimer("")
     1 * tracer.write({ it.size() == 1 }) >> {
       childLatch.countDown()
@@ -428,6 +440,7 @@ class PendingTraceBufferTest extends DDSpecification {
     1 * tracer.writeTimer() >> Monitoring.DISABLED.newTimer("")
     1 * tracer.write({ it.size() == 1 })
     1 * tracer.getPartialFlushMinSpans() >> 10000
+    _ * tracer.getTagInterceptor()
     1 * traceConfig.getServiceMapping() >> [:]
     2 * tracer.getTimeWithNanoTicks(_)
     1 * tracer.onRootSpanPublished(_)
@@ -444,6 +457,7 @@ class PendingTraceBufferTest extends DDSpecification {
     buffer.queue.capacity() * bufferSpy.enqueue(_)
     _ * bufferSpy.longRunningSpansEnabled()
     _ * tracer.getPartialFlushMinSpans() >> 10000
+    _ * tracer.getTagInterceptor()
     _ * traceConfig.getServiceMapping() >> [:]
     _ * tracer.getTimeWithNanoTicks(_)
     0 * _

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentSpan.java

@@ -66,6 +66,8 @@ default boolean isValid() {
 
   AgentSpan setTag(String key, long value);
 
+  AgentSpan setTag(String key, float value);
+
   AgentSpan setTag(String key, double value);
 
   @Override

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/ImmutableSpan.java

@@ -26,6 +26,11 @@ public AgentSpan setTag(String key, long value) {
     return this;
   }
 
+  @Override
+  public AgentSpan setTag(String key, float value) {
+    return this;
+  }
+
   @Override
   public AgentSpan setTag(String key, double value) {
     return this;