dd-trace 5.2.0 → 5.3.0

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

@@ -37,11 +37,16 @@ const testRunFinishCh = channel('ci:jest:test:finish')
 const testErrCh = channel('ci:jest:test:err')
 
 const skippableSuitesCh = channel('ci:jest:test-suite:skippable')
-const jestItrConfigurationCh = channel('ci:jest:itr-configuration')
+const libraryConfigurationCh = channel('ci:jest:library-configuration')
+const knownTestsCh = channel('ci:jest:known-tests')
 
 const itrSkippedSuitesCh = channel('ci:jest:itr:skipped-suites')
 
+// Maximum time we'll wait for the tracer to flush
+const FLUSH_TIMEOUT = 10000
+
 let skippableSuites = []
+let knownTests = []
 let isCodeCoverageEnabled = false
 let isSuitesSkippingEnabled = false
 let isUserCodeCoverageEnabled = false
@@ -49,6 +54,11 @@ let isSuitesSkipped = false
 let numSkippedSuites = 0
 let hasUnskippableSuites = false
 let hasForcedToRunSuites = false
+let isEarlyFlakeDetectionEnabled = false
+let earlyFlakeDetectionNumRetries = 0
+
+const EFD_STRING = "Retried by Datadog's Early Flake Detection"
+const EFD_TEST_NAME_REGEX = new RegExp(EFD_STRING + ' \\(#\\d+\\): ', 'g')
 
 const sessionAsyncResource = new AsyncResource('bound-anonymous-fn')
 
@@ -62,6 +72,7 @@ const specStatusToTestStatus = {
 
 const asyncResources = new WeakMap()
 const originalTestFns = new WeakMap()
+const retriedTestsToNumAttempts = new Map()
 
 // based on https://github.com/facebook/jest/blob/main/packages/jest-circus/src/formatNodeAssertErrors.ts#L41
 function formatJestError (errors) {
@@ -90,6 +101,14 @@ function getTestEnvironmentOptions (config) {
   return {}
 }
 
+function getEfdTestName (testName, numAttempt) {
+  return `${EFD_STRING} (#${numAttempt}): ${testName}`
+}
+
+function removeEfdTestName (testName) {
+  return testName.replace(EFD_TEST_NAME_REGEX, '')
+}
+
 function getWrappedEnvironment (BaseEnvironment, jestVersion) {
   return class DatadogEnvironment extends BaseEnvironment {
     constructor (config, context) {
@@ -97,10 +116,43 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
       const rootDir = config.globalConfig ? config.globalConfig.rootDir : config.rootDir
       this.rootDir = rootDir
       this.testSuite = getTestSuitePath(context.testPath, rootDir)
+      this.testFileAbsolutePath = context.testPath
       this.nameToParams = {}
       this.global._ddtrace = global._ddtrace
 
       this.testEnvironmentOptions = getTestEnvironmentOptions(config)
+
+      this.isEarlyFlakeDetectionEnabled = this.testEnvironmentOptions._ddIsEarlyFlakeDetectionEnabled
+
+      if (this.isEarlyFlakeDetectionEnabled) {
+        earlyFlakeDetectionNumRetries = this.testEnvironmentOptions._ddEarlyFlakeDetectionNumRetries
+        try {
+          this.knownTestsForThisSuite = this.getKnownTestsForSuite(this.testEnvironmentOptions._ddKnownTests)
+        } catch (e) {
+          // If there has been an error parsing the tests, we'll disable Early Flake Deteciton
+          this.isEarlyFlakeDetectionEnabled = false
+        }
+      }
+    }
+
+    // Function that receives a list of known tests for a test service and
+    // returns the ones that belong to the current suite
+    getKnownTestsForSuite (knownTests) {
+      let knownTestsForSuite = knownTests
+      // If jest runs in band, the known tests are not serialized, so they're an array.
+      if (!Array.isArray(knownTests)) {
+        knownTestsForSuite = JSON.parse(knownTestsForSuite)
+      }
+      return knownTestsForSuite
+        .filter(test => test.includes(this.testSuite))
+        .map(test => test.replace(`jest.${this.testSuite}.`, '').trim())
+    }
+
+    // Add the `add_test` event we don't have the test object yet, so
+    // we use its describe block to get the full name
+    getTestNameFromAddTestEvent (event, state) {
+      const describeSuffix = getJestTestName(state.currentDescribeBlock)
+      return removeEfdTestName(`${describeSuffix} ${event.testName}`).trim()
     }
 
     async handleTestEvent (event, state) {
@@ -124,23 +176,56 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
         }
       }
       if (event.name === 'test_start') {
+        let isNewTest = false
+        let numEfdRetry = null
         const testParameters = getTestParametersString(this.nameToParams, event.test.name)
         // Async resource for this test is created here
         // It is used later on by the test_done handler
         const asyncResource = new AsyncResource('bound-anonymous-fn')
         asyncResources.set(event.test, asyncResource)
+        const testName = getJestTestName(event.test)
+
+        if (this.isEarlyFlakeDetectionEnabled) {
+          const originalTestName = removeEfdTestName(testName)
+          isNewTest = retriedTestsToNumAttempts.has(originalTestName)
+          if (isNewTest) {
+            numEfdRetry = retriedTestsToNumAttempts.get(originalTestName)
+            retriedTestsToNumAttempts.set(originalTestName, numEfdRetry + 1)
+          }
+        }
+
         asyncResource.runInAsyncScope(() => {
           testStartCh.publish({
-            name: getJestTestName(event.test),
+            name: removeEfdTestName(testName),
             suite: this.testSuite,
+            testFileAbsolutePath: this.testFileAbsolutePath,
             runner: 'jest-circus',
             testParameters,
-            frameworkVersion: jestVersion
+            frameworkVersion: jestVersion,
+            isNew: isNewTest,
+            isEfdRetry: numEfdRetry > 0
           })
           originalTestFns.set(event.test, event.test.fn)
           event.test.fn = asyncResource.bind(event.test.fn)
         })
       }
+      if (event.name === 'add_test') {
+        if (this.isEarlyFlakeDetectionEnabled) {
+          const testName = this.getTestNameFromAddTestEvent(event, state)
+          const isNew = !this.knownTestsForThisSuite?.includes(testName)
+          const isSkipped = event.mode === 'todo' || event.mode === 'skip'
+          if (isNew && !isSkipped && !retriedTestsToNumAttempts.has(testName)) {
+            retriedTestsToNumAttempts.set(testName, 0)
+            for (let retryIndex = 0; retryIndex < earlyFlakeDetectionNumRetries; retryIndex++) {
+              if (this.global.test) {
+                this.global.test(getEfdTestName(event.testName, retryIndex), event.fn, event.timeout)
+              } else {
+                log.error('Early flake detection could not retry test because global.test is undefined')
+              }
+            }
+          }
+        }
+      }
       if (event.name === 'test_done') {
         const asyncResource = asyncResources.get(event.test)
         asyncResource.runInAsyncScope(() => {
@@ -164,6 +249,7 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
           testSkippedCh.publish({
             name: getJestTestName(event.test),
             suite: this.testSuite,
+            testFileAbsolutePath: this.testFileAbsolutePath,
             runner: 'jest-circus',
             frameworkVersion: jestVersion,
             testStartLine: getTestLineStart(event.test.asyncError, this.testSuite)
@@ -206,7 +292,7 @@ addHook({
     }
     // TODO: could we get the rootDir from each test?
     const [test] = shardedTests
-    const rootDir = test && test.context && test.context.config && test.context.config.rootDir
+    const rootDir = test?.context?.config?.rootDir
 
     const jestSuitesToRun = getJestSuitesToRun(skippableSuites, shardedTests, rootDir || process.cwd())
 
@@ -234,24 +320,45 @@ function cliWrapper (cli, jestVersion) {
     const configurationPromise = new Promise((resolve) => {
       onDone = resolve
     })
-    if (!jestItrConfigurationCh.hasSubscribers) {
+    if (!libraryConfigurationCh.hasSubscribers) {
       return runCLI.apply(this, arguments)
     }
 
     sessionAsyncResource.runInAsyncScope(() => {
-      jestItrConfigurationCh.publish({ onDone })
+      libraryConfigurationCh.publish({ onDone })
     })
 
     try {
-      const { err, itrConfig } = await configurationPromise
+      const { err, libraryConfig } = await configurationPromise
       if (!err) {
-        isCodeCoverageEnabled = itrConfig.isCodeCoverageEnabled
-        isSuitesSkippingEnabled = itrConfig.isSuitesSkippingEnabled
+        isCodeCoverageEnabled = libraryConfig.isCodeCoverageEnabled
+        isSuitesSkippingEnabled = libraryConfig.isSuitesSkippingEnabled
+        isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
+        earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
       }
     } catch (err) {
       log.error(err)
     }
 
+    if (isEarlyFlakeDetectionEnabled) {
+      const knownTestsPromise = new Promise((resolve) => {
+        onDone = resolve
+      })
+
+      sessionAsyncResource.runInAsyncScope(() => {
+        knownTestsCh.publish({ onDone })
+      })
+
+      try {
+        const { err, knownTests: receivedKnownTests } = await knownTestsPromise
+        if (!err) {
+          knownTests = receivedKnownTests
+        }
+      } catch (err) {
+        log.error(err)
+      }
+    }
+
     if (isSuitesSkippingEnabled) {
       const skippableSuitesPromise = new Promise((resolve) => {
         onDone = resolve
@@ -311,6 +418,21 @@ function cliWrapper (cli, jestVersion) {
       status = 'fail'
       error = new Error(`Failed test suites: ${numFailedTestSuites}. Failed tests: ${numFailedTests}`)
     }
+    let timeoutId
+
+    // Pass the resolve callback to defer it to DC listener
+    const flushPromise = new Promise((resolve) => {
+      onDone = () => {
+        clearTimeout(timeoutId)
+        resolve()
+      }
+    })
+
+    const timeoutPromise = new Promise((resolve) => {
+      timeoutId = setTimeout(() => {
+        resolve('timeout')
+      }, FLUSH_TIMEOUT).unref()
+    })
 
     sessionAsyncResource.runInAsyncScope(() => {
       testSessionFinishCh.publish({
@@ -322,9 +444,16 @@ function cliWrapper (cli, jestVersion) {
         numSkippedSuites,
         hasUnskippableSuites,
         hasForcedToRunSuites,
-        error
+        error,
+        isEarlyFlakeDetectionEnabled,
+        onDone
       })
     })
+    const waitingResult = await Promise.race([flushPromise, timeoutPromise])
+
+    if (waitingResult === 'timeout') {
+      log.error('Timeout waiting for the tracer to flush')
+    }
 
     numSkippedSuites = 0
 
@@ -438,10 +567,15 @@ function configureTestEnvironment (readConfigsResult) {
   // because `jestAdapterWrapper` runs in a different process. We have to go through `testEnvironmentOptions`
   configs.forEach(config => {
     config.testEnvironmentOptions._ddTestCodeCoverageEnabled = isCodeCoverageEnabled
+    config.testEnvironmentOptions._ddKnownTests = knownTests
   })
 
   isUserCodeCoverageEnabled = !!readConfigsResult.globalConfig.collectCoverage
 
+  if (readConfigsResult.globalConfig.forceExit) {
+    log.warn("Jest's '--forceExit' flag has been passed. This may cause loss of data.")
+  }
+
   if (isCodeCoverageEnabled) {
     const globalConfig = {
       ...readConfigsResult.globalConfig,
@@ -498,6 +632,9 @@ addHook({
       _ddForcedToRun,
       _ddUnskippable,
       _ddItrCorrelationId,
+      _ddKnownTests,
+      _ddIsEarlyFlakeDetectionEnabled,
+      _ddEarlyFlakeDetectionNumRetries,
       ...restOfTestEnvironmentOptions
     } = testEnvironmentOptions
 

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

@@ -20,7 +20,7 @@ const skipCh = channel('ci:mocha:test:skip')
 const testFinishCh = channel('ci:mocha:test:finish')
 const parameterizedTestCh = channel('ci:mocha:test:parameterize')
 
-const itrConfigurationCh = channel('ci:mocha:itr-configuration')
+const libraryConfigurationCh = channel('ci:mocha:library-configuration')
 const skippableSuitesCh = channel('ci:mocha:test-suite:skippable')
 
 const testSessionStartCh = channel('ci:mocha:session:start')
@@ -384,7 +384,7 @@ addHook({
       return run.apply(this, arguments)
     }
 
-    if (!itrConfigurationCh.hasSubscribers || this.isWorker) {
+    if (!libraryConfigurationCh.hasSubscribers || this.isWorker) {
       if (this.isWorker) {
         isWorker = true
       }
@@ -439,7 +439,7 @@ addHook({
     }
 
     mochaRunAsyncResource.runInAsyncScope(() => {
-      itrConfigurationCh.publish({
+      libraryConfigurationCh.publish({
         onDone: mochaRunAsyncResource.bind(onReceivedConfiguration)
       })
     })

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

@@ -79,21 +79,26 @@ addHook({
         })
 
         let callbackWrapped = false
-        const lastArgumentIndex = arguments.length - 1
 
-        if (typeof arguments[lastArgumentIndex] === 'function') {
-          // is a callback, wrap it to execute finish()
-          shimmer.wrap(arguments, lastArgumentIndex, originalCb => {
-            return function () {
-              finish()
+        const wrapCallbackIfExist = (args) => {
+          const lastArgumentIndex = args.length - 1
 
-              return originalCb.apply(this, arguments)
-            }
-          })
+          if (typeof args[lastArgumentIndex] === 'function') {
+            // is a callback, wrap it to execute finish()
+            shimmer.wrap(args, lastArgumentIndex, originalCb => {
+              return function () {
+                finish()
+
+                return originalCb.apply(this, arguments)
+              }
+            })
 
-          callbackWrapped = true
+            callbackWrapped = true
+          }
         }
 
+        wrapCallbackIfExist(arguments)
+
         return asyncResource.runInAsyncScope(() => {
           startCh.publish({
             filters,
@@ -106,8 +111,16 @@ addHook({
           if (!callbackWrapped) {
             shimmer.wrap(res, 'exec', originalExec => {
               return function wrappedExec () {
+                if (!callbackWrapped) {
+                  wrapCallbackIfExist(arguments)
+                }
+
                 const execResult = originalExec.apply(this, arguments)
 
+                if (callbackWrapped || typeof execResult?.then !== 'function') {
+                  return execResult
+                }
+
                 // wrap them method, wrap resolve and reject methods
                 shimmer.wrap(execResult, 'then', originalThen => {
                   return function wrappedThen () {

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

@@ -290,9 +290,23 @@ addHook({
   shimmer.massWrap(request.NextRequest.prototype, ['text', 'json'], function (originalMethod) {
     return async function wrappedJson () {
       const body = await originalMethod.apply(this, arguments)
-      bodyParsedChannel.publish({
-        body
-      })
+
+      bodyParsedChannel.publish({ body })
+
+      return body
+    }
+  })
+
+  shimmer.wrap(request.NextRequest.prototype, 'formData', function (originalFormData) {
+    return async function wrappedFormData () {
+      const body = await originalFormData.apply(this, arguments)
+
+      let normalizedBody = body
+      if (typeof body.entries === 'function') {
+        normalizedBody = Object.fromEntries(body.entries())
+      }
+      bodyParsedChannel.publish({ body: normalizedBody })
+
       return body
     }
   })

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

@@ -73,14 +73,41 @@ function getRootDir (playwrightRunner) {
   if (playwrightRunner._configDir) {
     return playwrightRunner._configDir
   }
-  if (playwrightRunner._config && playwrightRunner._config.config) {
-    return playwrightRunner._config.config.rootDir
+  if (playwrightRunner._config) {
+    return playwrightRunner._config.config?.rootDir || process.cwd()
   }
   return process.cwd()
 }
 
-function testBeginHandler (test) {
-  const { _requireFile: testSuiteAbsolutePath, title: testName, _type, location: { line: testSourceLine } } = test
+function getProjectsFromRunner (runner) {
+  const config = getPlaywrightConfig(runner)
+  return config.projects?.map(({ project }) => project)
+}
+
+function getProjectsFromDispatcher (dispatcher) {
+  const newConfig = dispatcher._config?.config?.projects
+  if (newConfig) {
+    return newConfig
+  }
+  // old
+  return dispatcher._loader?.fullConfig()?.projects
+}
+
+function getBrowserNameFromProjects (projects, projectId) {
+  if (!projects) {
+    return null
+  }
+  return projects.find(project =>
+    project.__projectId === projectId || project._id === projectId
+  )?.name
+}
+
+function testBeginHandler (test, browserName) {
+  const {
+    _requireFile: testSuiteAbsolutePath,
+    title: testName, _type,
+    location: { line: testSourceLine }
+  } = test
 
   if (_type === 'beforeAll' || _type === 'afterAll') {
     return
@@ -100,7 +127,7 @@ function testBeginHandler (test) {
   const testAsyncResource = new AsyncResource('bound-anonymous-fn')
   testToAr.set(test, testAsyncResource)
   testAsyncResource.runInAsyncScope(() => {
-    testStartCh.publish({ testName, testSuiteAbsolutePath, testSourceLine })
+    testStartCh.publish({ testName, testSuiteAbsolutePath, testSourceLine, browserName })
   })
 }
 
@@ -166,11 +193,12 @@ function dispatcherHook (dispatcherExport) {
   shimmer.wrap(dispatcherExport.Dispatcher.prototype, '_createWorker', createWorker => function () {
     const dispatcher = this
     const worker = createWorker.apply(this, arguments)
-
     worker.process.on('message', ({ method, params }) => {
       if (method === 'testBegin') {
         const { test } = dispatcher._testById.get(params.testId)
-        testBeginHandler(test)
+        const projects = getProjectsFromDispatcher(dispatcher)
+        const browser = getBrowserNameFromProjects(projects, test._projectId)
+        testBeginHandler(test, browser)
       } else if (method === 'testEnd') {
         const { test } = dispatcher._testById.get(params.testId)
 
@@ -203,7 +231,9 @@ function dispatcherHookNew (dispatcherExport, runWrapper) {
 
     worker.on('testBegin', ({ testId }) => {
       const test = getTestByTestId(dispatcher, testId)
-      testBeginHandler(test)
+      const projects = getProjectsFromDispatcher(dispatcher)
+      const browser = getBrowserNameFromProjects(projects, test._projectId)
+      testBeginHandler(test, browser)
     })
     worker.on('testEnd', ({ testId, status, errors, annotations }) => {
       const test = getTestByTestId(dispatcher, testId)
@@ -226,6 +256,7 @@ function runnerHook (runnerExport, playwrightVersion) {
     testSessionAsyncResource.runInAsyncScope(() => {
       testSessionStartCh.publish({ command, frameworkVersion: playwrightVersion, rootDir })
     })
+    const projects = getProjectsFromRunner(this)
 
     const runAllTestsReturn = await runAllTests.apply(this, arguments)
 
@@ -234,7 +265,8 @@ function runnerHook (runnerExport, playwrightVersion) {
       // there were tests that did not go through `testBegin` or `testEnd`,
       // because they were skipped
       tests.forEach(test => {
-        testBeginHandler(test)
+        const browser = getBrowserNameFromProjects(projects, test._projectId)
+        testBeginHandler(test, browser)
         testEndHandler(test, [], 'skip')
       })
     })

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

@@ -2,6 +2,7 @@
 
 const { TEXT_MAP } = require('../../../ext/formats')
 const ConsumerPlugin = require('../../dd-trace/src/plugins/consumer')
+const { getAmqpMessageSize, CONTEXT_PROPAGATION_KEY } = require('../../dd-trace/src/datastreams/processor')
 const { getResourceName } = require('./util')
 
 class AmqplibConsumerPlugin extends ConsumerPlugin {
@@ -13,7 +14,7 @@ class AmqplibConsumerPlugin extends ConsumerPlugin {
 
     const childOf = extract(this.tracer, message)
 
-    this.startSpan({
+    const span = this.startSpan({
       childOf,
       resource: getResourceName(method, fields),
       type: 'worker',
@@ -26,6 +27,14 @@ class AmqplibConsumerPlugin extends ConsumerPlugin {
         'amqp.destination': fields.destination
       }
     })
+
+    if (this.config.dsmEnabled && message) {
+      const payloadSize = getAmqpMessageSize({ headers: message.properties.headers, content: message.content })
+      const queue = fields.queue ?? fields.routingKey
+      this.tracer.decodeDataStreamsContext(message.properties.headers[CONTEXT_PROPAGATION_KEY])
+      this.tracer
+        .setCheckpoint(['direction:in', `topic:${queue}`, 'type:rabbitmq'], span, payloadSize)
+    }
   }
 }
 

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

@@ -3,13 +3,15 @@
 const { TEXT_MAP } = require('../../../ext/formats')
 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')
 const { getResourceName } = require('./util')
 
 class AmqplibProducerPlugin extends ProducerPlugin {
   static get id () { return 'amqplib' }
   static get operation () { return 'command' }
 
-  start ({ channel = {}, method, fields }) {
+  start ({ channel = {}, method, fields, message }) {
     if (method !== 'basic.publish') return
 
     const stream = (channel.connection && channel.connection.stream) || {}
@@ -30,6 +32,17 @@ class AmqplibProducerPlugin extends ProducerPlugin {
     fields.headers = fields.headers || {}
 
     this.tracer.inject(span, TEXT_MAP, fields.headers)
+
+    if (this.config.dsmEnabled) {
+      const hasRoutingKey = fields.routingKey != null
+      const payloadSize = getAmqpMessageSize({ content: message, headers: fields.headers })
+      const dataStreamsContext = this.tracer
+        .setCheckpoint(
+          ['direction:out', `exchange:${fields.exchange}`, `has_routing_key:${hasRoutingKey}`, 'type:rabbitmq']
+          , span, payloadSize)
+      const pathwayCtx = encodePathwayContext(dataStreamsContext)
+      fields.headers[CONTEXT_PROPAGATION_KEY] = pathwayCtx
+    }
   }
 }
 

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

@@ -1,15 +1,69 @@
 'use strict'
 const {
-  CONTEXT_PROPAGATION_KEY
+  CONTEXT_PROPAGATION_KEY,
+  getSizeOrZero
 } = require('../../../dd-trace/src/datastreams/processor')
 const { encodePathwayContext } = require('../../../dd-trace/src/datastreams/pathway')
 const log = require('../../../dd-trace/src/log')
 const BaseAwsSdkPlugin = require('../base')
+const { storage } = require('../../../datadog-core')
 
 class Kinesis extends BaseAwsSdkPlugin {
   static get id () { return 'kinesis' }
   static get peerServicePrecursors () { return ['streamname'] }
 
+  constructor (...args) {
+    super(...args)
+
+    // TODO(bengl) Find a way to create the response span tags without this WeakMap being populated
+    // in the base class
+    this.requestTags = new WeakMap()
+
+    this.addSub('apm:aws:response:start:kinesis', obj => {
+      const { request, response } = obj
+      const store = storage.getStore()
+      const plugin = this
+
+      // if we have either of these operations, we want to store the streamName param
+      // since it is not typically available during get/put records requests
+      if (request.operation === 'getShardIterator' || request.operation === 'listShards') {
+        this.storeStreamName(request.params, request.operation, store)
+        return
+      }
+
+      if (request.operation === 'getRecords') {
+        let span
+        const responseExtraction = this.responseExtract(request.params, request.operation, response)
+        if (responseExtraction && responseExtraction.maybeChildOf) {
+          obj.needsFinish = true
+          const options = {
+            childOf: responseExtraction.maybeChildOf,
+            tags: Object.assign(
+              {},
+              this.requestTags.get(request) || {},
+              { 'span.kind': 'server' }
+            )
+          }
+          span = plugin.tracer.startSpan('aws.response', options)
+          this.enter(span, store)
+        }
+
+        // get the stream name that should have been stored previously
+        const { streamName } = storage.getStore()
+
+        // extract DSM context after as we might not have a parent-child but may have a DSM context
+        this.responseExtractDSMContext(
+          request.operation, response, span ?? null, streamName
+        )
+      }
+    })
+
+    this.addSub('apm:aws:response:finish:kinesis', err => {
+      const { span } = storage.getStore()
+      this.finish(span, null, err)
+    })
+  }
+
   generateTags (params, operation, response) {
     if (!params || !params.StreamName) return {}
 
@@ -20,6 +74,58 @@ class Kinesis extends BaseAwsSdkPlugin {
     }
   }
 
+  storeStreamName (params, operation, store) {
+    if (!operation || (operation !== 'getShardIterator' && operation !== 'listShards')) return
+    if (!params || !params.StreamName) return
+
+    const streamName = params.StreamName
+    storage.enterWith({ ...store, streamName })
+  }
+
+  responseExtract (params, operation, response) {
+    if (operation !== 'getRecords') return
+    if (params.Limit && params.Limit !== 1) return
+    if (!response || !response.Records || !response.Records[0]) return
+
+    const record = response.Records[0]
+
+    try {
+      const decodedData = JSON.parse(Buffer.from(record.Data).toString())
+
+      return {
+        maybeChildOf: this.tracer.extract('text_map', decodedData._datadog),
+        parsedAttributes: decodedData._datadog
+      }
+    } catch (e) {
+      log.error(e)
+    }
+  }
+
+  responseExtractDSMContext (operation, response, span, streamName) {
+    if (!this.config.dsmEnabled) return
+    if (operation !== 'getRecords') return
+    if (!response || !response.Records || !response.Records[0]) return
+
+    // we only want to set the payloadSize on the span if we have one message, not repeatedly
+    span = response.Records.length > 1 ? null : span
+
+    response.Records.forEach(record => {
+      const parsedAttributes = JSON.parse(Buffer.from(record.Data).toString())
+
+      if (
+        parsedAttributes &&
+        parsedAttributes._datadog &&
+        parsedAttributes._datadog[CONTEXT_PROPAGATION_KEY] &&
+        streamName
+      ) {
+        const payloadSize = getSizeOrZero(record.Data)
+        this.tracer.decodeDataStreamsContext(Buffer.from(parsedAttributes._datadog[CONTEXT_PROPAGATION_KEY]))
+        this.tracer
+          .setCheckpoint(['direction:in', `topic:${streamName}`, 'type:kinesis'], span, payloadSize)
+      }
+    })
+  }
+
   // AWS-SDK will b64 kinesis payloads
   // or will accept an already b64 encoded payload
   // This method handles both