fix(propagation): enforce tracestate limits and bound pruning work

benchmark/sirun/propagation/index.js

@@ -1,17 +1,33 @@
 'use strict'
 
-const guard = require('../startup-guard')
 // eslint-disable-next-line import/order -- The startup guard must run before every other require.
+const guard = require('../startup-guard')
 const assert = require('node:assert/strict')
 
+const { channel } = require('dc-polyfill')
+
 const { getAllBaggageItems } = require('../../../packages/dd-trace/src/baggage')
 const id = require('../../../packages/dd-trace/src/id')
+const {
+  PROPAGATED_ML_APP_KEY,
+  PROPAGATED_PARENT_AGENT_ID_KEY,
+  PROPAGATED_PARENT_AGENT_NAME_KEY,
+  PROPAGATED_PARENT_ID_KEY,
+} = require('../../../packages/dd-trace/src/llmobs/constants/tags')
 const SpanContext = require('../../../packages/dd-trace/src/opentracing/span_context')
 const TextMapPropagator = require('../../../packages/dd-trace/src/opentracing/propagation/text_map')
+const TraceState = require('../../../packages/dd-trace/src/opentracing/propagation/tracestate')
 
 const { VARIANT } = process.env
 
 const OPERATIONS = Number(process.env.OPERATIONS)
+const WITH_TRACE_TAG_REPLACEMENTS = VARIANT === 'inject-replacements'
+
+/**
+ * @typedef {object} TraceTagInjection
+ * @property {Array<string | undefined>} [traceTagReplacements]
+ * @property {number} [optionalTraceTagCount]
+ */
 
 // Duck-typed config keeps the bench out of the full `Config` singleton (telemetry
 // registration, env reads). The propagator only reads the fields below.
@@ -51,6 +67,13 @@ const EXTRACT_CARRIER_DATADOG = {
   'x-datadog-tags': '_dd.p.dm=-1,_dd.p.tid=1234567890abcdef',
 }
 
+const TRACE_TAG_REPLACEMENTS = [
+  PROPAGATED_PARENT_ID_KEY, 'abcdef1234567890',
+  PROPAGATED_ML_APP_KEY, 'benchmark-app',
+  PROPAGATED_PARENT_AGENT_ID_KEY, '0123456789abcdef',
+  PROPAGATED_PARENT_AGENT_NAME_KEY, 'benchmark-agent',
+]
+
 const injectContext = new SpanContext({
   traceId: id('1234567890abcdef'),
   spanId: id('abcdef1234567890'),
@@ -61,8 +84,19 @@ const injectContext = new SpanContext({
     started: [],
     finished: [],
   },
+  tracestate: WITH_TRACE_TAG_REPLACEMENTS ? TraceState.fromString(EXTRACT_CARRIER_ASCII.tracestate) : undefined,
 })
 
+const injectCh = channel('dd-trace:span:inject')
+
+/** @param {TraceTagInjection} injection */
+function addTraceTagReplacements (injection) {
+  injection.traceTagReplacements = TRACE_TAG_REPLACEMENTS
+  injection.optionalTraceTagCount = 2
+}
+
+if (WITH_TRACE_TAG_REPLACEMENTS) injectCh.subscribe(addTraceTagReplacements)
+
 // Pre-flight: confirm the selected variant does real work; catches a silent
 // breakage where the duck-typed config is missing a field the propagator now reads.
 if (VARIANT === 'extract') {
@@ -81,10 +115,26 @@ if (VARIANT === 'extract') {
     '_dd.p.dm': '-1',
     '_dd.p.tid': '1234567890abcdef',
   })
-} else if (VARIANT === 'inject') {
+} else if (VARIANT === 'inject' || WITH_TRACE_TAG_REPLACEMENTS) {
   const sanityInjected = {}
   propagator.inject(injectContext, sanityInjected)
-  assert.ok(sanityInjected.traceparent && sanityInjected['x-datadog-trace-id'], 'inject populated no headers')
+  if (WITH_TRACE_TAG_REPLACEMENTS) {
+    assert.strictEqual(
+      sanityInjected['x-datadog-tags'],
+      '_dd.p.dm=-1,_dd.p.tid=1234567890abcdef,' +
+      '_dd.p.llmobs_parent_id=abcdef1234567890,_dd.p.llmobs_ml_app=benchmark-app,' +
+      '_dd.p.llmobs_pagent_span_id=0123456789abcdef,_dd.p.llmobs_pagent_name=benchmark-agent'
+    )
+    assert.strictEqual(
+      sanityInjected.tracestate,
+      'dd=t.llmobs_pagent_name:benchmark-agent;t.llmobs_pagent_span_id:0123456789abcdef;' +
+      't.llmobs_ml_app:benchmark-app;t.llmobs_parent_id:abcdef1234567890;' +
+      't.tid:1234567890abcdef;t.dm:-1;s:1;p:abc,' +
+      'vendor1=k1:v1,vendor2=k1:v1;k2:v2,foo=bar'
+    )
+  } else {
+    assert.ok(sanityInjected.traceparent && sanityInjected['x-datadog-trace-id'], 'inject populated no headers')
+  }
 }
 
 guard.loopStart()
@@ -100,9 +150,10 @@ if (VARIANT === 'extract') {
   for (let iteration = 0; iteration < OPERATIONS; iteration++) {
     propagator.extract(EXTRACT_CARRIER_DATADOG)
   }
-} else if (VARIANT === 'inject') {
+} else if (VARIANT === 'inject' || WITH_TRACE_TAG_REPLACEMENTS) {
   for (let iteration = 0; iteration < OPERATIONS; iteration++) {
     propagator.inject(injectContext, {})
   }
 }
 guard.done()
+if (WITH_TRACE_TAG_REPLACEMENTS) injectCh.unsubscribe(addTraceTagReplacements)

benchmark/sirun/propagation/meta.json

@@ -21,6 +21,15 @@
         "OPERATIONS": "3000000"
       }
     },
+    "inject-replacements": {
+      "run": "node --predictable-gc-schedule index.js",
+      "run_with_affinity": "bash -c \"taskset -c $CPU_AFFINITY node --predictable-gc-schedule index.js\"",
+      "iterations": 10,
+      "env": {
+        "VARIANT": "inject-replacements",
+        "OPERATIONS": "1500000"
+      }
+    },
     "extract-baggage-percent": {
       "env": {
         "VARIANT": "extract-baggage-percent",

packages/dd-trace/src/opentracing/propagation/text_map.js

@@ -87,6 +87,13 @@ let otelSampling
  * @property {string} [traceId]
  */
 
+/**
+ * @typedef {object} TraceTagInjection
+ * @property {DatadogSpanContext} spanContext
+ * @property {Array<string | undefined>} [traceTagReplacements]
+ * @property {number} [optionalTraceTagCount]
+ */
+
 /**
  * @param {Array<string | undefined>} traceTagReplacements
  * @param {string} key
@@ -98,6 +105,28 @@ function hasTraceTagReplacement (traceTagReplacements, key) {
   return false
 }
 
+/** @param {string} key */
+function toTraceStateTagKey (key) {
+  return 't.' + key.slice(6).replaceAll(tracestateTagKeyFilter, '_')
+}
+
+/**
+ * @param {string} key
+ * @param {TraceTagInjection | undefined} injection
+ */
+function isOptionalDatadogTraceStateField (key, injection) {
+  if (!key.startsWith('t.') || key === 't.dm' || key === 't.ts') return false
+
+  const traceTagReplacements = injection?.traceTagReplacements
+  if (!traceTagReplacements) return true
+  const firstOptionalTraceTagIndex = traceTagReplacements.length - (injection.optionalTraceTagCount ?? 0) * 2
+  for (let index = 0; index < firstOptionalTraceTagIndex; index += 2) {
+    const traceTagKey = traceTagReplacements[index]
+    if (traceTagKey.startsWith('_dd.p.') && toTraceStateTagKey(traceTagKey) === key) return false
+  }
+  return true
+}
+
 /**
  * @param {string | undefined} traceId
  * @param {string | undefined} spanId
@@ -331,13 +360,15 @@ class TextMapPropagator {
       hasTraceSourcePropagationTag(spanContext._trace.tags)
     let traceTagReplacements
     let optionalTraceTagCount = 0
+    /** @type {TraceTagInjection | undefined} */
+    let traceTagInjection
     if (injectTraceContext && injectCh.hasSubscribers && (
       this.#hasPropagationStyle('inject', 'datadog') || this.#hasPropagationStyle('inject', 'tracecontext')
     )) {
-      const injection = { spanContext }
-      injectCh.publish(injection)
-      traceTagReplacements = injection.traceTagReplacements
-      optionalTraceTagCount = injection.optionalTraceTagCount ?? 0
+      traceTagInjection = { spanContext }
+      injectCh.publish(traceTagInjection)
+      traceTagReplacements = traceTagInjection.traceTagReplacements
+      optionalTraceTagCount = traceTagInjection.optionalTraceTagCount ?? 0
     }
     if (injectTraceContext) {
       injectedCarrier = this.#injectDatadog(
@@ -348,7 +379,7 @@ class TextMapPropagator {
     }
     injectedCarrier = this
       .#injectTraceparent(
-        spanContext, injectedCarrier ?? carrier, injectTraceContext, traceTagReplacements
+        spanContext, injectedCarrier ?? carrier, injectTraceContext, traceTagReplacements, traceTagInjection
       ) ?? injectedCarrier
 
     if (injectedCarrier === undefined) return
@@ -514,7 +545,7 @@ class TextMapPropagator {
     }
 
     if (traceTagReplacements) {
-      const requiredEntryCount = traceTagReplacements.length - optionalTraceTagCount * 2
+      const firstOptionalTraceTagIndex = traceTagReplacements.length - optionalTraceTagCount * 2
       for (let index = 0; index < traceTagReplacements.length; index += 2) {
         const key = traceTagReplacements[index]
         const value = traceTagReplacements[index + 1]
@@ -525,7 +556,7 @@ class TextMapPropagator {
         }
 
         const entry = `${header ? ',' : ''}${key}=${value}`
-        if (index >= requiredEntryCount && header.length + entry.length > maxLength) break
+        if (index >= firstOptionalTraceTagIndex && header.length + entry.length > maxLength) break
         header += entry
       }
     }
@@ -594,9 +625,10 @@ class TextMapPropagator {
    * @param {Record<string, string> | undefined} carrier
    * @param {boolean} injectTraceContext
    * @param {Array<string | undefined>} [traceTagReplacements]
+   * @param {TraceTagInjection} [traceTagInjection]
    * @returns {Record<string, string> | undefined}
    */
-  #injectTraceparent (spanContext, carrier, injectTraceContext, traceTagReplacements) {
+  #injectTraceparent (spanContext, carrier, injectTraceContext, traceTagReplacements, traceTagInjection) {
     if (!this.#hasPropagationStyle('inject', 'tracecontext')) return
 
     if (!injectTraceContext) {
@@ -651,8 +683,7 @@ class TextMapPropagator {
         const tagValueRaw = spanContext._trace.tags[key]
         if (!tagValueRaw || !key.startsWith('_dd.p.')) continue
 
-        const tagKey = 't.' + key.slice(6)
-          .replaceAll(tracestateTagKeyFilter, '_')
+        const tagKey = toTraceStateTagKey(key)
 
         const tagValue = tagValueRaw
           .toString()
@@ -667,8 +698,7 @@ class TextMapPropagator {
           const key = traceTagReplacements[index]
           if (!key.startsWith('_dd.p.')) continue
 
-          const tagKey = 't.' + key.slice(6)
-            .replaceAll(tracestateTagKeyFilter, '_')
+          const tagKey = toTraceStateTagKey(key)
           const tagValueRaw = traceTagReplacements[index + 1]
           if (!tagValueRaw) {
             state.delete(tagKey)
@@ -683,7 +713,7 @@ class TextMapPropagator {
           state.set(tagKey, tagValue)
         }
       }
-    })
+    }, isOptionalDatadogTraceStateField, traceTagInjection)
 
     writeTracestate(carrier, ts.toString())