dd-trace 6.9.0 → 6.10.0

package/packages/datadog-instrumentations/src/vitest-worker.js

@@ -3,6 +3,7 @@
 const path = require('node:path')
 const { performance } = require('node:perf_hooks')
 const { fileURLToPath } = require('node:url')
+const { isMainThread, parentPort } = require('node:worker_threads')
 
 const shimmer = require('../../datadog-shimmer')
 const log = require('../../dd-trace/src/log')
@@ -12,6 +13,8 @@ const {
   getTestSuitePath,
   recordAttemptToFixExecution,
   logAttemptToFixTestExecution,
+  VITEST_WORKER_EFD_SUITE_ADMISSION_REQUEST_CODE,
+  VITEST_WORKER_EFD_SUITE_ADMISSION_RESPONSE_CODE,
 } = require('../../dd-trace/src/plugins/util/test')
 const { getChannelPromise } = require('./helpers/channel')
 const { addHook } = require('./helpers/instrument')
@@ -36,12 +39,15 @@ const {
   getVitestTestProperties,
 } = require('./vitest-util')
 
+const EFD_SUITE_ADMISSION_TIMEOUT_MS = 5000
 const taskToCtx = new WeakMap()
 const taskToTestProperties = new WeakMap()
 const taskToStatuses = new WeakMap()
 const taskToReportedErrorCount = new WeakMap()
 const attemptToFixTaskToStatuses = new WeakMap()
 const fileToHasConcurrentTests = new WeakMap()
+const fileToEfdSuiteAdmission = new WeakMap()
+const pendingEfdSuiteAdmissionRequests = new Map()
 const originalHookFns = new WeakMap()
 const newTasks = new WeakSet()
 const dynamicNameTasks = new WeakSet()
@@ -62,6 +68,9 @@ let vitestGetFn = null
 let vitestSetFn = null
 let vitestGetHooks = null
 let preciseCoverageSession
+let didInitializeEfdSuiteAdmissionTransport = false
+let nextEfdSuiteAdmissionRequestId = 0
+let sendEfdSuiteAdmissionMessage
 let isPreciseCoverageUnavailable = false
 let vitestCoverageSnapshot
 const wrappedCoverageWorkerStates = new WeakSet()
@@ -348,6 +357,123 @@ function hasConcurrentTests (file) {
   return hasConcurrent
 }
 
+/**
+ * Returns whether a collected Vitest file contains a runnable new test eligible for EFD.
+ *
+ * @param {{ filepath: string, tasks?: object[] }} file
+ * @param {object} providedContext
+ * @returns {boolean}
+ */
+function hasRunnableNewTest (file, providedContext) {
+  for (const task of getTypeTasks(file.tasks)) {
+    if (task.mode === 'skip' || task.mode === 'todo') continue
+
+    const testProperties = getVitestTestProperties(providedContext, file.filepath, getTestName(task))
+    if (testProperties.isNew && !testProperties.isAttemptToFix && !testProperties.isDisabled) {
+      return true
+    }
+  }
+  return false
+}
+
+/**
+ * Resolves and removes one pending EFD suite admission request.
+ *
+ * @param {number} requestId
+ * @param {boolean} allowed
+ * @returns {void}
+ */
+function finishEfdSuiteAdmissionRequest (requestId, allowed) {
+  const request = pendingEfdSuiteAdmissionRequests.get(requestId)
+  if (!request) return
+
+  clearTimeout(request.timeout)
+  pendingEfdSuiteAdmissionRequests.delete(requestId)
+  request.resolve(allowed)
+}
+
+/**
+ * Handles an EFD suite admission response from the Vitest main process.
+ *
+ * @param {unknown} message
+ * @returns {void}
+ */
+function handleEfdSuiteAdmissionResponse (message) {
+  if (!Array.isArray(message) || message[0] !== VITEST_WORKER_EFD_SUITE_ADMISSION_RESPONSE_CODE) return
+
+  const { allowed, requestId } = message[1] || {}
+  if (Number.isSafeInteger(requestId)) {
+    finishEfdSuiteAdmissionRequest(requestId, allowed === true)
+  }
+}
+
+/**
+ * Returns the child-process or worker-thread sender used by the current Vitest worker.
+ *
+ * @returns {((message: unknown) => void)|undefined}
+ */
+function getEfdSuiteAdmissionSender () {
+  if (didInitializeEfdSuiteAdmissionTransport) return sendEfdSuiteAdmissionMessage
+  didInitializeEfdSuiteAdmissionTransport = true
+
+  if (typeof process.send === 'function') {
+    process.on('message', handleEfdSuiteAdmissionResponse)
+    sendEfdSuiteAdmissionMessage = process.send.bind(process)
+  } else if (!isMainThread && parentPort) {
+    parentPort.on('message', handleEfdSuiteAdmissionResponse)
+    sendEfdSuiteAdmissionMessage = parentPort.postMessage.bind(parentPort)
+  }
+  return sendEfdSuiteAdmissionMessage
+}
+
+/**
+ * Requests permission from the Vitest main process to schedule EFD retries for one suite.
+ *
+ * @param {string} testSuite
+ * @param {boolean} hasNewTest
+ * @returns {Promise<boolean>}
+ */
+function requestEfdSuiteAdmission (testSuite, hasNewTest) {
+  const sendMessage = getEfdSuiteAdmissionSender()
+  if (!sendMessage) return Promise.resolve(false)
+
+  const requestId = ++nextEfdSuiteAdmissionRequestId
+  return new Promise(resolve => {
+    const timeout = setTimeout(
+      () => finishEfdSuiteAdmissionRequest(requestId, false),
+      EFD_SUITE_ADMISSION_TIMEOUT_MS
+    )
+    timeout.unref?.()
+    pendingEfdSuiteAdmissionRequests.set(requestId, { resolve, timeout })
+
+    try {
+      sendMessage([VITEST_WORKER_EFD_SUITE_ADMISSION_REQUEST_CODE, { hasNewTest, requestId, testSuite }])
+    } catch {
+      finishEfdSuiteAdmissionRequest(requestId, false)
+    }
+  })
+}
+
+/**
+ * Requests a single EFD admission decision for all tests collected in a Vitest file.
+ *
+ * @param {{ file: { filepath: string, tasks?: object[] } }} task
+ * @param {object} providedContext
+ * @param {string|undefined} testSuite
+ * @returns {Promise<boolean>}
+ */
+function isEfdSuiteAdmissionAllowed (task, providedContext, testSuite) {
+  let admission = fileToEfdSuiteAdmission.get(task.file)
+  if (admission) return admission
+
+  admission = requestEfdSuiteAdmission(
+    testSuite || task.file.filepath,
+    hasRunnableNewTest(task.file, providedContext)
+  )
+  fileToEfdSuiteAdmission.set(task.file, admission)
+  return admission
+}
+
 /**
  * Gets the task associated with a Vitest hook invocation.
  *
@@ -383,7 +509,7 @@ function wrapSuiteHookFn (hookType, fn, fallbackTask) {
 function wrapVitestTestRunner (VitestTestRunner) {
   // `onBeforeRunTask` is run before any repetition or attempt is run
   // `onBeforeRunTask` is an async function
-  shimmer.wrap(VitestTestRunner.prototype, 'onBeforeRunTask', onBeforeRunTask => function (task) {
+  shimmer.wrap(VitestTestRunner.prototype, 'onBeforeRunTask', onBeforeRunTask => async function (task) {
     const testName = getTestName(task)
 
     const providedContext = getProvidedContext()
@@ -426,35 +552,38 @@ function wrapVitestTestRunner (VitestTestRunner) {
     }
 
     if (isImpactedTestsEnabled && testProperties.isModified) {
-      if (isEarlyFlakeDetectionEnabled) {
-        efdRetryTasks.add(task)
-        disableFrameworkRetries(task)
-        task.repeats = earlyFlakeDetectionRetryPolicy.schedulingRetryCount
-      }
       modifiedTasks.add(task)
-      taskToStatuses.set(task, [])
     }
 
     if (isKnownTestsEnabled && testProperties.isNew && !attemptToFixTasks.has(task)) {
-      if (isEarlyFlakeDetectionEnabled && !modifiedTasks.has(task)) {
-        efdRetryTasks.add(task)
-        disableFrameworkRetries(task)
-        task.repeats = earlyFlakeDetectionRetryPolicy.schedulingRetryCount
-      }
       newTasks.add(task)
-      taskToStatuses.set(task, [])
       if (DYNAMIC_NAME_RE.test(testName)) {
         dynamicNameTasks.add(task)
       }
     }
 
+    const isEfdCandidate = isEarlyFlakeDetectionEnabled &&
+      !attemptToFixTasks.has(task) &&
+      !disabledTasks.has(task) &&
+      (modifiedTasks.has(task) || newTasks.has(task))
+    if (isEfdCandidate) {
+      const isAdmissionAllowed = !providedContext.isEfdSuiteAdmissionEnabled ||
+        await isEfdSuiteAdmissionAllowed(task, providedContext, testProperties.testSuite)
+      if (isAdmissionAllowed) {
+        efdRetryTasks.add(task)
+        disableFrameworkRetries(task)
+        task.repeats = earlyFlakeDetectionRetryPolicy.schedulingRetryCount
+        taskToStatuses.set(task, [])
+      }
+    }
+
     return onBeforeRunTask.apply(this, arguments)
   })
 
   // `onAfterRunTask` is run after all repetitions or attempts are run
   // `onAfterRunTask` is an async function
   shimmer.wrap(VitestTestRunner.prototype, 'onAfterRunTask', onAfterRunTask => function (task) {
-    const { isEarlyFlakeDetectionEnabled, isTestManagementTestsEnabled } = getProvidedContext()
+    const { isTestManagementTestsEnabled } = getProvidedContext()
 
     if (isTestManagementTestsEnabled) {
       const isAttemptingToFix = attemptToFixTasks.has(task)
@@ -476,7 +605,7 @@ function wrapVitestTestRunner (VitestTestRunner) {
       }
     }
 
-    if (isEarlyFlakeDetectionEnabled && taskToStatuses.has(task) && !attemptToFixTasks.has(task)) {
+    if (efdRetryTasks.has(task)) {
       const statuses = taskToStatuses.get(task)
       // If the test has passed at least once, we consider it passed
       if (statuses.includes('pass')) {
@@ -501,7 +630,6 @@ function wrapVitestTestRunner (VitestTestRunner) {
     const providedContext = getProvidedContext()
     const {
       isKnownTestsEnabled,
-      isEarlyFlakeDetectionEnabled,
       isDiEnabled,
       earlyFlakeDetectionRetryPolicy,
     } = providedContext
@@ -512,7 +640,7 @@ function wrapVitestTestRunner (VitestTestRunner) {
 
     const { retry: numAttempt, repeats: numRepetition } = retryInfo
     const isFailedTestReplayAllowed = !hasConcurrentTests(task.file)
-    const isEfdManagedTask = isEarlyFlakeDetectionEnabled && taskToStatuses.has(task) && !attemptToFixTasks.has(task)
+    const isEfdManagedTask = efdRetryTasks.has(task)
 
     if (isEfdManagedTask && numRepetition > 0 && !efdDeterminedRetries.has(task)) {
       const previousExecutionStart = efdExecutionStartByTask.get(task)
@@ -574,8 +702,9 @@ function wrapVitestTestRunner (VitestTestRunner) {
 
     const lastExecutionStatus = task.result.state
     const isAtf = attemptToFixTasks.has(task)
-    const shouldTrackStatuses = isEarlyFlakeDetectionEnabled || isAtf
-    const shouldFlipStatus = isEarlyFlakeDetectionEnabled || isAtf
+    const isEfd = efdRetryTasks.has(task)
+    const shouldTrackStatuses = isEfd || isAtf
+    const shouldFlipStatus = isEfd || isAtf
     const statuses = isAtf ? attemptToFixTaskToStatuses.get(task) : taskToStatuses.get(task)
 
     // These clauses handle task.repeats, whether EFD is enabled or not
@@ -697,7 +826,6 @@ function wrapVitestTestRunner (VitestTestRunner) {
       const result = await onAfterTryTask.apply(this, arguments)
 
       const {
-        isEarlyFlakeDetectionEnabled,
         testManagementAttemptToFixRetries,
         earlyFlakeDetectionRetryPolicy,
       } = getProvidedContext()
@@ -719,10 +847,8 @@ function wrapVitestTestRunner (VitestTestRunner) {
       }
 
       if (
-        isEarlyFlakeDetectionEnabled &&
+        efdRetryTasks.has(task) &&
         (retryInfo.repeats ?? 0) === 0 &&
-        taskToStatuses.has(task) &&
-        !attemptToFixTasks.has(task) &&
         !efdDeterminedRetries.has(task)
       ) {
         const executionStart = efdExecutionStartByTask.get(task)
@@ -921,8 +1047,7 @@ addHook({
           }
 
           // Check if all EFD retries failed
-          const isEfdRetry =
-            providedContext.isEarlyFlakeDetectionEnabled && (newTasks.has(task) || modifiedTasks.has(task))
+          const isEfdRetry = efdRetryTasks.has(task)
           if (isEfdRetry) {
             const statuses = taskToStatuses.get(task)
             const efdRetryCount = efdDeterminedRetries.get(task) ??

package/packages/datadog-instrumentations/src/webdriverio.js

@@ -171,7 +171,7 @@ function createWorkerConfiguration () {
 }
 
 /**
- * Configures a Jasmine worker for reporting-only execution and notifies its coordinator.
+ * Configures a Jasmine worker and waits for its coordinator-owned optimization settings.
  *
  * @param {{_jrunner?: {env?: {addReporter?: (reporter: object) => void}}, _specs?: string[]}} adapter
  * @param {{
@@ -185,13 +185,7 @@ function initializeJasmineWorker (adapter, context) {
     return
   }
 
-  workerConfigurationCh.publish({
-    libraryConfig: createWorkerConfiguration(),
-    repositoryRoot: undefined,
-    specs: adapter?._specs || [],
-    testFramework: TEST_FRAMEWORK,
-    testFrameworkAdapter: JASMINE_FRAMEWORK_ADAPTER,
-  })
+  const specs = adapter?._specs || []
 
   if (jasmineDoneCh.hasSubscribers) {
     adapter?._jrunner?.env?.addReporter?.({
@@ -220,9 +214,10 @@ function initializeJasmineWorker (adapter, context) {
     /**
      * Releases Jasmine initialization exactly once.
      *
+     * @param {object} [configuration]
      * @returns {void}
      */
-    function finish () {
+    function finish (configuration = createWorkerConfiguration()) {
       if (finished) {
         return
       }
@@ -230,6 +225,13 @@ function initializeJasmineWorker (adapter, context) {
       clearTimeout(timeout)
       process.off('message', onMessage)
       process.off('disconnect', finish)
+      workerConfigurationCh.publish({
+        libraryConfig: configuration,
+        repositoryRoot: configuration.repositoryRoot,
+        specs,
+        testFramework: TEST_FRAMEWORK,
+        testFrameworkAdapter: JASMINE_FRAMEWORK_ADAPTER,
+      })
       onDone()
     }
 
@@ -241,7 +243,7 @@ function initializeJasmineWorker (adapter, context) {
      */
     function onMessage (message) {
       if (message?.name === WORKER_READY_RESPONSE && message.content?.requestId === requestId) {
-        finish()
+        finish(message.content.configuration)
       }
     }
 
@@ -446,6 +448,18 @@ function getMochaFrameworkData (frameworkData) {
   return frameworkData?.[TEST_FRAMEWORK] ?? frameworkData?.mocha
 }
 
+/**
+ * Applies the worker configuration for an EFD session with faulty known-tests data.
+ *
+ * @param {object} configuration
+ * @returns {void}
+ */
+function setEarlyFlakeDetectionFaulty (configuration) {
+  configuration.isEarlyFlakeDetectionEnabled = false
+  configuration.isEarlyFlakeDetectionFaulty = true
+  configuration.isKnownTestsEnabled = false
+}
+
 /**
  * Applies settings and requests the enabled non-TIA datasets once for the whole run.
  *
@@ -498,9 +512,11 @@ function configureCoordinator (state, response) {
     pendingRequests++
     requestCoordinatorData(state, knownTestsCh, ({ err, knownTests } = {}) => {
       const mochaKnownTests = getMochaFrameworkData(knownTests)
-      if (err || !mochaKnownTests) {
+      if (err) {
         configuration.isEarlyFlakeDetectionEnabled = false
         configuration.isKnownTestsEnabled = false
+      } else if (mochaKnownTests === undefined) {
+        setEarlyFlakeDetectionFaulty(configuration)
       } else {
         configuration.knownTests = { mocha: mochaKnownTests }
       }
@@ -564,7 +580,6 @@ function initializeCoordinator (state, onDone) {
 
   try {
     libraryConfigurationCh.runStores({
-      basicReportingOnly: state.testFrameworkAdapter === JASMINE_FRAMEWORK_ADAPTER,
       disableTestImpactAnalysis: true,
       frameworkVersion: state.frameworkVersion,
       isParallel: state.maxActiveWorkers > 1,
@@ -727,9 +742,7 @@ function updateEarlyFlakeDetectionFaultyState (state, files) {
     configuration.earlyFlakeDetectionFaultyThreshold
   )
   if (isFaulty) {
-    configuration.isEarlyFlakeDetectionEnabled = false
-    configuration.isEarlyFlakeDetectionFaulty = true
-    configuration.isKnownTestsEnabled = false
+    setEarlyFlakeDetectionFaulty(configuration)
   }
 }
 
@@ -809,15 +822,16 @@ function handleWorkerMessage (state, workerRecord, message) {
   }
 
   if (message.name === WORKER_READY) {
-    initializeCoordinator(state, () => {
+    initializeCoordinator(state, configuration => {
       if (state.testFrameworkAdapter === JASMINE_FRAMEWORK_ADAPTER) {
+        updateEarlyFlakeDetectionFaultyState(state, workerRecord.specs)
         startWorkerSuites(workerRecord, workerRecord.specs)
       }
       if (message.content?.requestId) {
         sendWorkerMessage(workerRecord, {
           origin: 'datadog',
           name: WORKER_READY_RESPONSE,
-          content: { requestId: message.content.requestId },
+          content: { configuration, requestId: message.content.requestId },
         })
       }
     })

package/packages/datadog-plugin-aws-durable-execution-sdk-js/src/trace-checkpoint.js

@@ -1,6 +1,12 @@
 'use strict'
 
 const crypto = require('crypto')
+const {
+  deleteDatadogParentId,
+  readDatadogParentId,
+  readDatadogTraceId,
+  writeDatadogParentId,
+} = require('../../dd-trace/src/carrier')
 const log = require('../../dd-trace/src/log')
 const TextMapPropagator = require('../../dd-trace/src/opentracing/propagation/text_map')
 
@@ -37,7 +43,7 @@ function getDatadogOnlyPropagator (tracer) {
  * @returns {Record<string, string>}
  */
 function injectHeaders (tracer, span) {
-  const headers = {}
+  const headers = /** @type {Record<string, string>} */ ({})
   const propagator = getDatadogOnlyPropagator(tracer)
   const ctx = typeof span?.context === 'function' ? span.context() : span
   propagator.inject(ctx, headers)
@@ -51,8 +57,8 @@ function injectHeaders (tracer, span) {
  */
 function overrideParentId (headers, parentId) {
   if (!parentId) return
-  if (headers['x-datadog-trace-id']) {
-    headers['x-datadog-parent-id'] = String(parentId)
+  if (readDatadogTraceId(headers)) {
+    writeDatadogParentId(headers, String(parentId))
   }
 }
 
@@ -192,7 +198,7 @@ async function saveTraceContextCheckpointIfUpdated (
   if (typeof checkpointManager?.checkpoint !== 'function') return
 
   const currentHeaders = injectHeaders(tracer, span)
-  if (currentHeaders['x-datadog-trace-id'] === undefined) return
+  if (readDatadogTraceId(currentHeaders) === undefined) return
 
   const latest = findLastCheckpoint(event)
 
@@ -204,8 +210,8 @@ async function saveTraceContextCheckpointIfUpdated (
     // x-datadog-parent-id reflects the active span at save time and always differs, so exclude it
     // from the comparison. Capture the previous anchor first to carry it forward on a real update.
     // needsCheckpointUpdate only reads currentHeaders' keys, so deleting it from there is enough.
-    const anchoredSpanId = previousHeaders['x-datadog-parent-id']
-    delete currentHeaders['x-datadog-parent-id']
+    const anchoredSpanId = readDatadogParentId(previousHeaders)
+    deleteDatadogParentId(currentHeaders)
     if (!needsCheckpointUpdate(currentHeaders, previousHeaders)) return
 
     newNumber = latest.checkpointNumber + 1

package/packages/datadog-plugin-aws-sdk/src/services/kinesis.js

@@ -196,14 +196,11 @@ class Kinesis extends BaseAwsSdkPlugin {
       try {
         parsedAttributes = JSON.parse(recordDataAsString(record.Data))
       } catch {
-        // Non-JSON record. Skip DSM context for this entry; the
-        // checkpoint payload size below is still reported.
+        // Non-JSON record: it carries no context, so the checkpoint below starts a new pathway.
       }
 
       const payloadSize = getSizeOrZero(record.Data)
-      if (parsedAttributes?._datadog) {
-        this.tracer.decodeDataStreamsContext(parsedAttributes._datadog)
-      }
+      this.tracer.decodeDataStreamsContext(parsedAttributes?._datadog)
       this.tracer.setCheckpoint(tags, span, payloadSize)
     }
   }

package/packages/datadog-plugin-aws-sdk/src/services/sqs.js

@@ -275,11 +275,7 @@ class Sqs extends BaseAwsSdkPlugin {
     for (let i = 0; i < messages.length; i++) {
       const message = messages[i]
       const carrier = carriers === undefined ? this.parseMessageCarrier(message) : carriers[i]
-      if (carrier) {
-        // Inert for EventBridge until its producer emits a pathway (separate
-        // change) — no `dd-pathway-ctx-base64` to decode yet; SQS/SNS decode now.
-        this.tracer.decodeDataStreamsContext(carrier)
-      }
+      this.tracer.decodeDataStreamsContext(carrier)
       const payloadSize = getHeadersSize({
         Body: message.Body,
         MessageAttributes: message.MessageAttributes,

package/packages/datadog-plugin-cypress/src/cypress-plugin.js

@@ -78,6 +78,7 @@ const {
   TEST_HAS_DYNAMIC_NAME,
   getIsFaultyEarlyFlakeDetection,
   DYNAMIC_NAME_RE,
+  isMarkedAsUnskippable,
   recordAttemptToFixExecution,
   recordTestManagementExecution,
   logAttemptToFixTestExecution,
@@ -86,7 +87,6 @@ const {
   TEST_FINAL_STATUS,
   getTestOptimizationRequestResults,
 } = require('../../dd-trace/src/plugins/util/test')
-const { isMarkedAsUnskippable } = require('../../datadog-plugin-jest/src/util')
 const { ORIGIN_KEY, COMPONENT } = require('../../dd-trace/src/constants')
 const { RESOURCE_NAME } = require('../../../ext/tags')
 const getConfig = require('../../dd-trace/src/config')

package/packages/datadog-plugin-google-cloud-pubsub/src/producer.js

@@ -1,8 +1,17 @@
 'use strict'
 
-const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
+const {
+  readDatadogParentId,
+  readDatadogSamplingPriority,
+  readDatadogTraceId,
+  readTraceparent,
+  writeDatadogParentId,
+  writeDatadogSamplingPriority,
+  writeDatadogTraceId,
+} = require('../../dd-trace/src/carrier')
 const { DsmPathwayCodec, getHeadersSize } = require('../../dd-trace/src/datastreams')
 const id = require('../../dd-trace/src/id')
+const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
 
 class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
   static id = 'google-cloud-pubsub'
@@ -21,7 +30,7 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
 
   handleMessagePublish ({ attributes, pubsub, topicName }) {
     // Skip if message already has trace context from upstream
-    if (attributes['x-datadog-trace-id'] || attributes.traceparent) return
+    if (readDatadogTraceId(attributes) || readTraceparent(attributes)) return
 
     const activeSpan = this.tracer.scope().active()
     if (!activeSpan) return
@@ -62,7 +71,7 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
     const messages = request.messages || []
     const topic = request.topic
     const messageCount = messages.length
-    const hasTraceContext = messages[0]?.attributes?.['x-datadog-trace-id']
+    const hasTraceContext = messages[0]?.attributes && readDatadogTraceId(messages[0].attributes)
 
     /**
      * Batch Publishing Strategy:
@@ -82,12 +91,14 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
     }
 
     const firstAttrs = messages[0]?.attributes
-    const parentData = firstAttrs?.['x-datadog-trace-id'] && firstAttrs['x-datadog-parent-id']
+    const traceId = firstAttrs && readDatadogTraceId(firstAttrs)
+    const spanId = firstAttrs && readDatadogParentId(firstAttrs)
+    const parentData = traceId && spanId
       ? {
-          traceId: firstAttrs['x-datadog-trace-id'],
-          spanId: firstAttrs['x-datadog-parent-id'],
+          traceId,
+          spanId,
           traceIdUpper: firstAttrs['_dd.p.tid'],
-          samplingPriority: firstAttrs['x-datadog-sampling-priority'],
+          samplingPriority: readDatadogSamplingPriority(firstAttrs),
         }
       : null
 
@@ -174,33 +185,36 @@ class GoogleCloudPubsubProducerPlugin extends ProducerPlugin {
   }
 
   #extractSpanLink (attrs) {
-    if (!attrs?.['x-datadog-trace-id'] || !attrs['x-datadog-parent-id']) return null
+    if (!attrs) return null
+    const traceId = readDatadogTraceId(attrs)
+    const spanId = readDatadogParentId(attrs)
+    if (!traceId || !spanId) return null
 
     // Convert to hex strings
-    const lowerHex = id(attrs['x-datadog-trace-id']).toString(16)
-    const spanIdHex = id(attrs['x-datadog-parent-id']).toString(16)
+    const lowerHex = id(traceId).toString(16)
+    const spanIdHex = id(spanId).toString(16)
 
     // Build full 128-bit trace ID
     const traceIdHex = attrs['_dd.p.tid']
       ? attrs['_dd.p.tid'] + lowerHex
       : lowerHex.padStart(32, '0')
 
+    const samplingPriority = readDatadogSamplingPriority(attrs)
     return {
       traceId: traceIdHex,
       spanId: spanIdHex,
-      samplingPriority: attrs['x-datadog-sampling-priority']
-        ? Number.parseInt(attrs['x-datadog-sampling-priority'], 10)
+      samplingPriority: samplingPriority
+        ? Number.parseInt(samplingPriority, 10)
         : undefined,
     }
   }
 
   #extractParentContext (data) {
-    const carrier = {
-      'x-datadog-trace-id': data.traceId,
-      'x-datadog-parent-id': data.spanId,
-    }
+    const carrier = /** @type {Record<string, unknown>} */ ({})
+    writeDatadogTraceId(carrier, data.traceId)
+    writeDatadogParentId(carrier, data.spanId)
     if (data.traceIdUpper) carrier['_dd.p.tid'] = data.traceIdUpper
-    if (data.samplingPriority) carrier['x-datadog-sampling-priority'] = String(data.samplingPriority)
+    if (data.samplingPriority) writeDatadogSamplingPriority(carrier, String(data.samplingPriority))
 
     return this.tracer.extract('text_map', carrier)
   }

package/packages/datadog-plugin-jest/src/index.js

@@ -5,6 +5,7 @@ const realSetTimeout = setTimeout
 
 const CiPlugin = require('../../dd-trace/src/plugins/ci_plugin')
 const { storage } = require('../../datadog-core')
+const { writeDatadogParentId, writeDatadogTraceId } = require('../../dd-trace/src/carrier')
 const { getEnvironmentVariable } = require('../../dd-trace/src/config/helper')
 const { appClosing: appClosingTelemetry } = require('../../dd-trace/src/telemetry')
 
@@ -230,10 +231,10 @@ class JestPlugin extends CiPlugin {
         _ddItrSkippingEnabledTags: itrSkippingEnabledTags,
       } = testEnvironmentOptions
 
-      const testSessionSpanContext = this.tracer.extract('text_map', {
-        'x-datadog-trace-id': testSessionId,
-        'x-datadog-parent-id': testModuleId,
-      })
+      const carrier = /** @type {Record<string, unknown>} */ ({})
+      writeDatadogTraceId(carrier, testSessionId)
+      writeDatadogParentId(carrier, testModuleId)
+      const testSessionSpanContext = this.tracer.extract('text_map', carrier)
 
       const testSuiteMetadata = {
         ...getTestSuiteCommonTags(testCommand, frameworkVersion, testSuite, 'jest'),

package/packages/datadog-plugin-kafkajs/src/batch-consumer.js

@@ -1,8 +1,7 @@
 'use strict'
 
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
-const { getMessageSize } = require('../../dd-trace/src/datastreams')
-const { convertToTextMap } = require('./utils')
+const { convertToTextMap, getKafkaMessageSize } = require('./utils')
 
 class KafkajsBatchConsumerPlugin extends ConsumerPlugin {
   static id = 'kafkajs'
@@ -39,7 +38,7 @@ class KafkajsBatchConsumerPlugin extends ConsumerPlugin {
       }
 
       if (!this.config.dsmEnabled) continue
-      const payloadSize = getMessageSize(message)
+      const payloadSize = getKafkaMessageSize(message)
       this.tracer.decodeDataStreamsContext(headers)
       const edgeTags = ['direction:in', `group:${groupId}`, `topic:${topic}`, 'type:kafka']
       if (clusterId) {