dd-trace 5.2.0 → 5.4.0

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

@@ -24,7 +24,8 @@ const {
   TEST_SKIPPED_BY_ITR,
   TEST_ITR_UNSKIPPABLE,
   TEST_ITR_FORCED_RUN,
-  ITR_CORRELATION_ID
+  ITR_CORRELATION_ID,
+  TEST_SOURCE_FILE
 } = require('../../dd-trace/src/plugins/util/test')
 const { ORIGIN_KEY, COMPONENT } = require('../../dd-trace/src/constants')
 const log = require('../../dd-trace/src/log')
@@ -45,7 +46,8 @@ const {
   GIT_REPOSITORY_URL,
   GIT_COMMIT_SHA,
   GIT_BRANCH,
-  CI_PROVIDER_NAME
+  CI_PROVIDER_NAME,
+  CI_WORKSPACE_PATH
 } = require('../../dd-trace/src/plugins/util/tags')
 const {
   OS_VERSION,
@@ -113,20 +115,21 @@ function getSuiteStatus (suiteStats) {
   if (suiteStats.failures !== undefined && suiteStats.failures > 0) {
     return 'fail'
   }
-  if (suiteStats.tests !== undefined && suiteStats.tests === suiteStats.pending) {
+  if (suiteStats.tests !== undefined &&
+    (suiteStats.tests === suiteStats.pending || suiteStats.tests === suiteStats.skipped)) {
     return 'skip'
   }
   return 'pass'
 }
 
-function getItrConfig (tracer, testConfiguration) {
+function getLibraryConfiguration (tracer, testConfiguration) {
   return new Promise(resolve => {
-    if (!tracer._tracer._exporter || !tracer._tracer._exporter.getItrConfiguration) {
+    if (!tracer._tracer._exporter?.getLibraryConfiguration) {
       return resolve({ err: new Error('CI Visibility was not initialized correctly') })
     }
 
-    tracer._tracer._exporter.getItrConfiguration(testConfiguration, (err, itrConfig) => {
-      resolve({ err, itrConfig })
+    tracer._tracer._exporter.getLibraryConfiguration(testConfiguration, (err, libraryConfig) => {
+      resolve({ err, libraryConfig })
     })
   })
 }
@@ -136,7 +139,7 @@ function getSkippableTests (isSuitesSkippingEnabled, tracer, testConfiguration)
     return Promise.resolve({ skippableTests: [] })
   }
   return new Promise(resolve => {
-    if (!tracer._tracer._exporter || !tracer._tracer._exporter.getItrConfiguration) {
+    if (!tracer._tracer._exporter?.getLibraryConfiguration) {
       return resolve({ err: new Error('CI Visibility was not initialized correctly') })
     }
     tracer._tracer._exporter.getSkippableSuites(testConfiguration, (err, skippableTests, correlationId) => {
@@ -186,7 +189,8 @@ module.exports = (on, config) => {
     [RUNTIME_NAME]: runtimeName,
     [RUNTIME_VERSION]: runtimeVersion,
     [GIT_BRANCH]: branch,
-    [CI_PROVIDER_NAME]: ciProviderName
+    [CI_PROVIDER_NAME]: ciProviderName,
+    [CI_WORKSPACE_PATH]: repositoryRoot
   } = testEnvironmentMetadata
 
   const isUnsupportedCIProvider = !ciProviderName
@@ -205,7 +209,7 @@ module.exports = (on, config) => {
     testLevel: 'test'
   }
 
-  const codeOwnersEntries = getCodeOwnersFileEntries()
+  const codeOwnersEntries = getCodeOwnersFileEntries(repositoryRoot)
 
   let activeSpan = null
   let testSessionSpan = null
@@ -243,6 +247,10 @@ module.exports = (on, config) => {
     if (testSessionSpan && testModuleSpan) {
       testSuiteTags[TEST_SESSION_ID] = testSessionSpan.context().toTraceId()
       testSuiteTags[TEST_MODULE_ID] = testModuleSpan.context().toSpanId()
+      // If testSuiteSpan couldn't be created, we'll use the testModuleSpan as the parent
+      if (!testSuiteSpan) {
+        testSuiteTags[TEST_SUITE_ID] = testModuleSpan.context().toSpanId()
+      }
     }
 
     const {
@@ -283,13 +291,26 @@ module.exports = (on, config) => {
     })
   }
 
+  function getTestSuiteSpan (suite) {
+    const testSuiteSpanMetadata = getTestSuiteCommonTags(command, frameworkVersion, suite, TEST_FRAMEWORK_NAME)
+    ciVisEvent(TELEMETRY_EVENT_CREATED, 'suite')
+    return tracer.startSpan(`${TEST_FRAMEWORK_NAME}.test_suite`, {
+      childOf: testModuleSpan,
+      tags: {
+        [COMPONENT]: TEST_FRAMEWORK_NAME,
+        ...testEnvironmentMetadata,
+        ...testSuiteSpanMetadata
+      }
+    })
+  }
+
   on('before:run', (details) => {
-    return getItrConfig(tracer, testConfiguration).then(({ err, itrConfig }) => {
+    return getLibraryConfiguration(tracer, testConfiguration).then(({ err, libraryConfig }) => {
       if (err) {
         log.error(err)
       } else {
-        isSuitesSkippingEnabled = itrConfig.isSuitesSkippingEnabled
-        isCodeCoverageEnabled = itrConfig.isCodeCoverageEnabled
+        isSuitesSkippingEnabled = libraryConfig.isSuitesSkippingEnabled
+        isCodeCoverageEnabled = libraryConfig.isCodeCoverageEnabled
       }
 
       return getSkippableTests(isSuitesSkippingEnabled, tracer, testConfiguration)
@@ -346,6 +367,13 @@ module.exports = (on, config) => {
     const cypressTests = tests || []
     const finishedTests = finishedTestsByFile[spec.relative] || []
 
+    if (!testSuiteSpan) {
+      // dd:testSuiteStart hasn't been triggered for whatever reason
+      // We will create the test suite span on the spot if that's the case
+      log.warn('There was an error creating the test suite event.')
+      testSuiteSpan = getTestSuiteSpan(spec.relative)
+    }
+
     // Get tests that didn't go through `dd:afterEach`
     // and create a skipped test span for each of them
     cypressTests.filter(({ title }) => {
@@ -359,6 +387,11 @@ module.exports = (on, config) => {
         cypressTestName === test.name && spec.relative === test.suite
       )
       const skippedTestSpan = getTestSpan(cypressTestName, spec.relative)
+      if (spec.absolute && repositoryRoot) {
+        skippedTestSpan.setTag(TEST_SOURCE_FILE, getTestSuitePath(spec.absolute, repositoryRoot))
+      } else {
+        skippedTestSpan.setTag(TEST_SOURCE_FILE, spec.relative)
+      }
       skippedTestSpan.setTag(TEST_STATUS, 'skip')
       if (isSkippedByItr) {
         skippedTestSpan.setTag(TEST_SKIPPED_BY_ITR, 'true')
@@ -390,6 +423,11 @@ module.exports = (on, config) => {
       if (itrCorrelationId) {
         finishedTest.testSpan.setTag(ITR_CORRELATION_ID, itrCorrelationId)
       }
+      if (spec.absolute && repositoryRoot) {
+        finishedTest.testSpan.setTag(TEST_SOURCE_FILE, getTestSuitePath(spec.absolute, repositoryRoot))
+      } else {
+        finishedTest.testSpan.setTag(TEST_SOURCE_FILE, spec.relative)
+      }
       finishedTest.testSpan.finish(finishedTest.finishTime)
     })
 
@@ -457,16 +495,7 @@ module.exports = (on, config) => {
       if (testSuiteSpan) {
         return null
       }
-      const testSuiteSpanMetadata = getTestSuiteCommonTags(command, frameworkVersion, suite, TEST_FRAMEWORK_NAME)
-      testSuiteSpan = tracer.startSpan(`${TEST_FRAMEWORK_NAME}.test_suite`, {
-        childOf: testModuleSpan,
-        tags: {
-          [COMPONENT]: TEST_FRAMEWORK_NAME,
-          ...testEnvironmentMetadata,
-          ...testSuiteSpanMetadata
-        }
-      })
-      ciVisEvent(TELEMETRY_EVENT_CREATED, 'suite')
+      testSuiteSpan = getTestSuiteSpan(suite)
       return null
     },
     'dd:beforeEach': (test) => {

package/packages/datadog-plugin-grpc/src/client.js

@@ -41,7 +41,6 @@ class GrpcClientPlugin extends ClientPlugin {
         'grpc.status.code': 0
       }
     }, false)
-
     // needed as precursor for peer.service
     if (method.service && method.package) {
       span.setTag('rpc.service', method.package + '.' + method.service)
@@ -68,7 +67,7 @@ class GrpcClientPlugin extends ClientPlugin {
     this.addError(error, span)
   }
 
-  finish ({ span, result }) {
+  finish ({ span, result, peer }) {
     if (!span) return
 
     const { code, metadata } = result || {}
@@ -80,6 +79,21 @@ class GrpcClientPlugin extends ClientPlugin {
       addMetadataTags(span, metadata, metadataFilter, 'response')
     }
 
+    if (peer) {
+      // The only scheme we want to support here is ipv[46]:port, although
+      // more are supported by the library
+      // https://github.com/grpc/grpc/blob/v1.60.0/doc/naming.md
+      const parts = peer.split(':')
+      if (parts[parts.length - 1].match(/^\d+/)) {
+        const port = parts[parts.length - 1]
+        const ip = parts.slice(0, -1).join(':')
+        span.setTag('network.destination.ip', ip)
+        span.setTag('network.destination.port', port)
+      } else {
+        span.setTag('network.destination.ip', peer)
+      }
+    }
+
     this.tagPeerService(span)
     span.finish()
   }

package/packages/datadog-plugin-http/src/client.js

@@ -122,7 +122,7 @@ class HttpClientPlugin extends ClientPlugin {
       // conditions for no error:
       // 1. not using a custom agent instance with custom timeout specified
       // 2. no invocation of `req.setTimeout`
-      if (!args.options.agent?.options.timeout && !customRequestTimeout) return
+      if (!args.options.agent?.options?.timeout && !customRequestTimeout) return
 
       span.setTag('error', 1)
     }

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

@@ -14,7 +14,11 @@ const {
   TEST_ITR_UNSKIPPABLE,
   TEST_ITR_FORCED_RUN,
   TEST_CODE_OWNERS,
-  ITR_CORRELATION_ID
+  ITR_CORRELATION_ID,
+  TEST_SOURCE_FILE,
+  TEST_IS_NEW,
+  TEST_EARLY_FLAKE_IS_RETRY,
+  TEST_EARLY_FLAKE_IS_ENABLED
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 const id = require('../../dd-trace/src/id')
@@ -81,7 +85,9 @@ class JestPlugin extends CiPlugin {
       numSkippedSuites,
       hasUnskippableSuites,
       hasForcedToRunSuites,
-      error
+      error,
+      isEarlyFlakeDetectionEnabled,
+      onDone
     }) => {
       this.testSessionSpan.setTag(TEST_STATUS, status)
       this.testModuleSpan.setTag(TEST_STATUS, status)
@@ -106,23 +112,35 @@ class JestPlugin extends CiPlugin {
         }
       )
 
+      if (isEarlyFlakeDetectionEnabled) {
+        this.testSessionSpan.setTag(TEST_EARLY_FLAKE_IS_ENABLED, 'true')
+      }
+
       this.testModuleSpan.finish()
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'module')
       this.testSessionSpan.finish()
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'session')
       finishAllTraceSpans(this.testSessionSpan)
-      this.tracer._exporter.flush()
+
+      this.tracer._exporter.flush(() => {
+        if (onDone) {
+          onDone()
+        }
+      })
     })
 
     // Test suites can be run in a different process from jest's main one.
     // This subscriber changes the configuration objects from jest to inject the trace id
-    // of the test session to the processes that run the test suites.
+    // of the test session to the processes that run the test suites, and other data.
     this.addSub('ci:jest:session:configuration', configs => {
       configs.forEach(config => {
         config._ddTestSessionId = this.testSessionSpan.context().toTraceId()
         config._ddTestModuleId = this.testModuleSpan.context().toSpanId()
         config._ddTestCommand = this.testSessionSpan.context()._tags[TEST_COMMAND]
         config._ddItrCorrelationId = this.itrCorrelationId
+        config._ddIsEarlyFlakeDetectionEnabled = !!this.libraryConfig?.isEarlyFlakeDetectionEnabled
+        config._ddEarlyFlakeDetectionNumRetries = this.libraryConfig?.earlyFlakeDetectionNumRetries ?? 0
+        config._ddRepositoryRoot = this.repositoryRoot
       })
     })
 
@@ -223,7 +241,7 @@ class JestPlugin extends CiPlugin {
     })
 
     /**
-     * This can't use `this.itrConfig` like `ci:mocha:test-suite:code-coverage`
+     * This can't use `this.libraryConfig` like `ci:mocha:test-suite:code-coverage`
      * because this subscription happens in a different process from the one
      * fetching the ITR config.
      */
@@ -286,7 +304,17 @@ class JestPlugin extends CiPlugin {
   }
 
   startTestSpan (test) {
-    const { suite, name, runner, testParameters, frameworkVersion, testStartLine } = test
+    const {
+      suite,
+      name,
+      runner,
+      testParameters,
+      frameworkVersion,
+      testStartLine,
+      testSourceFile,
+      isNew,
+      isEfdRetry
+    } = test
 
     const extraTags = {
       [JEST_TEST_RUNNER]: runner,
@@ -296,6 +324,15 @@ class JestPlugin extends CiPlugin {
     if (testStartLine) {
       extraTags[TEST_SOURCE_START] = testStartLine
     }
+    // If for whatever we don't have the source file, we'll fall back to the suite name
+    extraTags[TEST_SOURCE_FILE] = testSourceFile || suite
+
+    if (isNew) {
+      extraTags[TEST_IS_NEW] = 'true'
+      if (isEfdRetry) {
+        extraTags[TEST_EARLY_FLAKE_IS_RETRY] = 'true'
+      }
+    }
 
     return super.startTestSpan(name, suite, this.testSuiteSpan, extraTags)
   }

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

@@ -1,8 +1,12 @@
 'use strict'
 
+const dc = require('dc-polyfill')
 const { getMessageSize, CONTEXT_PROPAGATION_KEY } = require('../../dd-trace/src/datastreams/processor')
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
 
+const afterStartCh = dc.channel('dd-trace:kafkajs:consumer:afterStart')
+const beforeFinishCh = dc.channel('dd-trace:kafkajs:consumer:beforeFinish')
+
 class KafkajsConsumerPlugin extends ConsumerPlugin {
   static get id () { return 'kafkajs' }
   static get operation () { return 'consume' }
@@ -79,6 +83,18 @@ class KafkajsConsumerPlugin extends ConsumerPlugin {
       this.tracer
         .setCheckpoint(['direction:in', `group:${groupId}`, `topic:${topic}`, 'type:kafka'], span, payloadSize)
     }
+
+    if (afterStartCh.hasSubscribers) {
+      afterStartCh.publish({ topic, partition, message, groupId })
+    }
+  }
+
+  finish () {
+    if (beforeFinishCh.hasSubscribers) {
+      beforeFinishCh.publish()
+    }
+
+    super.finish()
   }
 }
 

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

@@ -15,7 +15,12 @@ const {
   TEST_ITR_UNSKIPPABLE,
   TEST_ITR_FORCED_RUN,
   TEST_CODE_OWNERS,
-  ITR_CORRELATION_ID
+  ITR_CORRELATION_ID,
+  TEST_SOURCE_FILE,
+  removeEfdStringFromTestName,
+  TEST_IS_NEW,
+  TEST_EARLY_FLAKE_IS_RETRY,
+  TEST_EARLY_FLAKE_IS_ENABLED
 } = require('../../dd-trace/src/plugins/util/test')
 const { COMPONENT } = require('../../dd-trace/src/constants')
 const {
@@ -38,11 +43,11 @@ class MochaPlugin extends CiPlugin {
     super(...args)
 
     this._testSuites = new Map()
-    this._testNameToParams = {}
+    this._testTitleToParams = {}
     this.sourceRoot = process.cwd()
 
     this.addSub('ci:mocha:test-suite:code-coverage', ({ coverageFiles, suiteFile }) => {
-      if (!this.itrConfig || !this.itrConfig.isCodeCoverageEnabled) {
+      if (!this.libraryConfig?.isCodeCoverageEnabled) {
         return
       }
       const testSuiteSpan = this._testSuites.get(suiteFile)
@@ -98,7 +103,7 @@ class MochaPlugin extends CiPlugin {
         }
       })
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_CREATED, 'suite')
-      if (this.itrConfig?.isCodeCoverageEnabled) {
+      if (this.libraryConfig?.isCodeCoverageEnabled) {
         this.telemetry.ciVisEvent(TELEMETRY_CODE_COVERAGE_STARTED, 'suite', { library: 'istanbul' })
       }
       if (itrCorrelationId) {
@@ -130,9 +135,9 @@ class MochaPlugin extends CiPlugin {
       }
     })
 
-    this.addSub('ci:mocha:test:start', ({ test, testStartLine }) => {
+    this.addSub('ci:mocha:test:start', (testInfo) => {
       const store = storage.getStore()
-      const span = this.startTestSpan(test, testStartLine)
+      const span = this.startTestSpan(testInfo)
 
       this.enter(span, store)
     })
@@ -155,12 +160,12 @@ class MochaPlugin extends CiPlugin {
       }
     })
 
-    this.addSub('ci:mocha:test:skip', (test) => {
+    this.addSub('ci:mocha:test:skip', (testInfo) => {
       const store = storage.getStore()
       // skipped through it.skip, so the span is not created yet
       // for this test
       if (!store) {
-        const testSpan = this.startTestSpan(test)
+        const testSpan = this.startTestSpan(testInfo)
         this.enter(testSpan, store)
       }
     })
@@ -178,8 +183,8 @@ class MochaPlugin extends CiPlugin {
       }
     })
 
-    this.addSub('ci:mocha:test:parameterize', ({ name, params }) => {
-      this._testNameToParams[name] = params
+    this.addSub('ci:mocha:test:parameterize', ({ title, params }) => {
+      this._testTitleToParams[title] = params
     })
 
     this.addSub('ci:mocha:session:finish', ({
@@ -189,10 +194,11 @@ class MochaPlugin extends CiPlugin {
       numSkippedSuites,
       hasForcedToRunSuites,
       hasUnskippableSuites,
-      error
+      error,
+      isEarlyFlakeDetectionEnabled
     }) => {
       if (this.testSessionSpan) {
-        const { isSuitesSkippingEnabled, isCodeCoverageEnabled } = this.itrConfig || {}
+        const { isSuitesSkippingEnabled, isCodeCoverageEnabled } = this.libraryConfig || {}
         this.testSessionSpan.setTag(TEST_STATUS, status)
         this.testModuleSpan.setTag(TEST_STATUS, status)
 
@@ -216,23 +222,34 @@ class MochaPlugin extends CiPlugin {
           }
         )
 
+        if (isEarlyFlakeDetectionEnabled) {
+          this.testSessionSpan.setTag(TEST_EARLY_FLAKE_IS_ENABLED, 'true')
+        }
+
         this.testModuleSpan.finish()
         this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'module')
         this.testSessionSpan.finish()
         this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'session')
         finishAllTraceSpans(this.testSessionSpan)
       }
-      this.itrConfig = null
+      this.libraryConfig = null
       this.tracer._exporter.flush()
     })
   }
 
-  startTestSpan (test, testStartLine) {
-    const testName = test.fullTitle()
-    const { file: testSuiteAbsolutePath, title } = test
+  startTestSpan (testInfo) {
+    const {
+      testSuiteAbsolutePath,
+      title,
+      isNew,
+      isEfdRetry,
+      testStartLine
+    } = testInfo
+
+    const testName = removeEfdStringFromTestName(testInfo.testName)
 
     const extraTags = {}
-    const testParametersString = getTestParametersString(this._testNameToParams, title)
+    const testParametersString = getTestParametersString(this._testTitleToParams, title)
     if (testParametersString) {
       extraTags[TEST_PARAMETERS] = testParametersString
     }
@@ -244,6 +261,19 @@ class MochaPlugin extends CiPlugin {
     const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.sourceRoot)
     const testSuiteSpan = this._testSuites.get(testSuiteAbsolutePath)
 
+    if (this.repositoryRoot !== this.sourceRoot && !!this.repositoryRoot) {
+      extraTags[TEST_SOURCE_FILE] = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
+    } else {
+      extraTags[TEST_SOURCE_FILE] = testSuite
+    }
+
+    if (isNew) {
+      extraTags[TEST_IS_NEW] = 'true'
+      if (isEfdRetry) {
+        extraTags[TEST_EARLY_FLAKE_IS_RETRY] = 'true'
+      }
+    }
+
     return super.startTestSpan(testName, testSuite, testSuiteSpan, extraTags)
   }
 }

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

@@ -9,7 +9,9 @@ const {
   getTestSuitePath,
   getTestSuiteCommonTags,
   TEST_SOURCE_START,
-  TEST_CODE_OWNERS
+  TEST_CODE_OWNERS,
+  TEST_SOURCE_FILE,
+  TEST_CONFIGURATION_BROWSER_NAME
 } = require('../../dd-trace/src/plugins/util/test')
 const { RESOURCE_NAME } = require('../../../ext/tags')
 const { COMPONENT } = require('../../dd-trace/src/constants')
@@ -76,10 +78,11 @@ class PlaywrightPlugin extends CiPlugin {
       this.telemetry.ciVisEvent(TELEMETRY_EVENT_FINISHED, 'suite')
     })
 
-    this.addSub('ci:playwright:test:start', ({ testName, testSuiteAbsolutePath, testSourceLine }) => {
+    this.addSub('ci:playwright:test:start', ({ testName, testSuiteAbsolutePath, testSourceLine, browserName }) => {
       const store = storage.getStore()
       const testSuite = getTestSuitePath(testSuiteAbsolutePath, this.rootDir)
-      const span = this.startTestSpan(testName, testSuite, testSourceLine)
+      const testSourceFile = getTestSuitePath(testSuiteAbsolutePath, this.repositoryRoot)
+      const span = this.startTestSpan(testName, testSuite, testSourceFile, testSourceLine, browserName)
 
       this.enter(span, store)
     })
@@ -126,9 +129,20 @@ class PlaywrightPlugin extends CiPlugin {
     })
   }
 
-  startTestSpan (testName, testSuite, testSourceLine) {
+  startTestSpan (testName, testSuite, testSourceFile, testSourceLine, browserName) {
     const testSuiteSpan = this._testSuites.get(testSuite)
-    return super.startTestSpan(testName, testSuite, testSuiteSpan, { [TEST_SOURCE_START]: testSourceLine })
+
+    const extraTags = {
+      [TEST_SOURCE_START]: testSourceLine
+    }
+    if (testSourceFile) {
+      extraTags[TEST_SOURCE_FILE] = testSourceFile || testSuite
+    }
+    if (browserName) {
+      extraTags[TEST_CONFIGURATION_BROWSER_NAME] = browserName
+    }
+
+    return super.startTestSpan(testName, testSuite, testSuiteSpan, extraTags)
   }
 }
 

package/packages/datadog-plugin-rhea/src/consumer.js

@@ -2,6 +2,7 @@
 
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
 const { storage } = require('../../datadog-core')
+const { getAmqpMessageSize, CONTEXT_PROPAGATION_KEY } = require('../../dd-trace/src/datastreams/processor')
 
 class RheaConsumerPlugin extends ConsumerPlugin {
   static get id () { return 'rhea' }
@@ -19,7 +20,7 @@ class RheaConsumerPlugin extends ConsumerPlugin {
     const name = getResourceNameFromMessage(msgObj)
     const childOf = extractTextMap(msgObj, this.tracer)
 
-    this.startSpan({
+    const span = this.startSpan({
       childOf,
       resource: name,
       type: 'worker',
@@ -29,6 +30,15 @@ class RheaConsumerPlugin extends ConsumerPlugin {
         'amqp.link.role': 'receiver'
       }
     })
+
+    if (this.config.dsmEnabled && msgObj.message) {
+      const payloadSize = getAmqpMessageSize(
+        { headers: msgObj.message.delivery_annotations, content: msgObj.message.body }
+      )
+      this.tracer.decodeDataStreamsContext(msgObj.message.delivery_annotations[CONTEXT_PROPAGATION_KEY])
+      this.tracer
+        .setCheckpoint(['direction:in', `topic:${name}`, 'type:rabbitmq'], span, payloadSize)
+    }
   }
 }
 

package/packages/datadog-plugin-rhea/src/producer.js

@@ -2,6 +2,8 @@
 
 const { CLIENT_PORT_KEY } = require('../../dd-trace/src/constants')
 const ProducerPlugin = require('../../dd-trace/src/plugins/producer')
+const { encodePathwayContext } = require('../../dd-trace/src/datastreams/pathway')
+const { getAmqpMessageSize, CONTEXT_PROPAGATION_KEY } = require('../../dd-trace/src/datastreams/processor')
 
 class RheaProducerPlugin extends ProducerPlugin {
   static get id () { return 'rhea' }
@@ -36,6 +38,15 @@ function addDeliveryAnnotations (msg, tracer, span) {
     msg.delivery_annotations = msg.delivery_annotations || {}
 
     tracer.inject(span, 'text_map', msg.delivery_annotations)
+
+    if (tracer._config.dsmEnabled) {
+      const targetName = span.context()._tags['amqp.link.target.address']
+      const payloadSize = getAmqpMessageSize({ content: msg.body, headers: msg.delivery_annotations })
+      const dataStreamsContext = tracer
+        .setCheckpoint(['direction:out', `exchange:${targetName}`, 'type:rabbitmq'], span, payloadSize)
+      const pathwayCtx = encodePathwayContext(dataStreamsContext)
+      msg.delivery_annotations[CONTEXT_PROPAGATION_KEY] = pathwayCtx
+    }
   }
 }
 

package/packages/dd-trace/src/appsec/addresses.js

@@ -15,6 +15,8 @@ module.exports = {
   HTTP_INCOMING_GRAPHQL_RESOLVERS: 'graphql.server.all_resolvers',
   HTTP_INCOMING_GRAPHQL_RESOLVER: 'graphql.server.resolver',
 
+  HTTP_OUTGOING_BODY: 'server.response.body',
+
   HTTP_CLIENT_IP: 'http.client_ip',
 
   USER_ID: 'usr.id',

package/packages/dd-trace/src/appsec/api_security_sampler.js

@@ -5,6 +5,8 @@ const log = require('../log')
 let enabled
 let requestSampling
 
+const sampledRequests = new WeakSet()
+
 function configure ({ apiSecurity }) {
   enabled = apiSecurity.enabled
   setRequestSampling(apiSecurity.requestSampling)
@@ -32,17 +34,28 @@ function parseRequestSampling (requestSampling) {
   return parsed
 }
 
-function sampleRequest () {
+function sampleRequest (req) {
   if (!enabled || !requestSampling) {
     return false
   }
 
-  return Math.random() <= requestSampling
+  const shouldSample = Math.random() <= requestSampling
+
+  if (shouldSample) {
+    sampledRequests.add(req)
+  }
+
+  return shouldSample
+}
+
+function isSampled (req) {
+  return sampledRequests.has(req)
 }
 
 module.exports = {
   configure,
   disable,
   setRequestSampling,
-  sampleRequest
+  sampleRequest,
+  isSampled
 }

package/packages/dd-trace/src/appsec/channels.js

@@ -16,5 +16,6 @@ module.exports = {
   queryParser: dc.channel('datadog:query:read:finish'),
   setCookieChannel: dc.channel('datadog:iast:set-cookie'),
   nextBodyParsed: dc.channel('apm:next:body-parsed'),
-  nextQueryParsed: dc.channel('apm:next:query-parsed')
+  nextQueryParsed: dc.channel('apm:next:query-parsed'),
+  responseBody: dc.channel('datadog:express:response:json:start')
 }

package/packages/dd-trace/src/appsec/iast/analyzers/command-injection-analyzer.js

@@ -8,7 +8,7 @@ class CommandInjectionAnalyzer extends InjectionAnalyzer {
   }
 
   onConfigure () {
-    this.addSub('datadog:child_process:execution:start', ({ command }) => this.analyze(command))
+    this.addSub('tracing:datadog:child_process:execution:start', ({ command }) => this.analyze(command))
   }
 }