dd-trace 5.17.0 → 5.18.0

package/packages/datadog-instrumentations/src/helpers/register.js

@@ -7,6 +7,7 @@ const Hook = require('./hook')
 const requirePackageJson = require('../../../dd-trace/src/require-package-json')
 const log = require('../../../dd-trace/src/log')
 const checkRequireCache = require('../check_require_cache')
+const telemetry = require('../../../dd-trace/src/telemetry/init-telemetry')
 
 const {
   DD_TRACE_DISABLED_INSTRUMENTATIONS = '',
@@ -35,22 +36,38 @@ if (DD_TRACE_DEBUG && DD_TRACE_DEBUG.toLowerCase() !== 'false') {
   setImmediate(checkRequireCache.checkForPotentialConflicts)
 }
 
+const seenCombo = new Set()
+
 // TODO: make this more efficient
 for (const packageName of names) {
   if (disabledInstrumentations.has(packageName)) continue
 
-  Hook([packageName], (moduleExports, moduleName, moduleBaseDir, moduleVersion) => {
+  const hookOptions = {}
+
+  let hook = hooks[packageName]
+
+  if (typeof hook === 'object') {
+    hookOptions.internals = hook.esmFirst
+    hook = hook.fn
+  }
+
+  Hook([packageName], hookOptions, (moduleExports, moduleName, moduleBaseDir, moduleVersion) => {
     moduleName = moduleName.replace(pathSepExpr, '/')
 
     // This executes the integration file thus adding its entries to `instrumentations`
-    hooks[packageName]()
+    hook()
 
     if (!instrumentations[packageName]) {
       return moduleExports
     }
 
-    for (const { name, file, versions, hook } of instrumentations[packageName]) {
+    const namesAndSuccesses = {}
+    for (const { name, file, versions, hook, filePattern } of instrumentations[packageName]) {
+      let fullFilePattern = filePattern
       const fullFilename = filename(name, file)
+      if (fullFilePattern) {
+        fullFilePattern = filename(name, fullFilePattern)
+      }
 
       // Create a WeakMap associated with the hook function so that patches on the same moduleExport only happens once
       // for example by instrumenting both dns and node:dns double the spans would be created
@@ -58,13 +75,29 @@ for (const packageName of names) {
       if (!hook[HOOK_SYMBOL]) {
         hook[HOOK_SYMBOL] = new WeakMap()
       }
+      let matchesFile = false
+
+      matchesFile = moduleName === fullFilename
 
-      if (moduleName === fullFilename) {
+      if (fullFilePattern) {
+        // Some libraries include a hash in their filenames when installed,
+        // so our instrumentation has to include a '.*' to match them for more than a single version.
+        matchesFile = matchesFile || new RegExp(fullFilePattern).test(moduleName)
+      }
+
+      if (matchesFile) {
         const version = moduleVersion || getVersion(moduleBaseDir)
+        if (!Object.hasOwnProperty(namesAndSuccesses, name)) {
+          namesAndSuccesses[name] = {
+            success: false,
+            version
+          }
+        }
 
         if (matchVersion(version, versions)) {
           // Check if the hook already has a set moduleExport
           if (hook[HOOK_SYMBOL].has(moduleExports)) {
+            namesAndSuccesses[name].success = true
             return moduleExports
           }
 
@@ -76,11 +109,29 @@ for (const packageName of names) {
             // Set the moduleExports in the hooks weakmap
             hook[HOOK_SYMBOL].set(moduleExports, name)
           } catch (e) {
-            log.error(e)
+            log.info('Error during ddtrace instrumentation of application, aborting.')
+            log.info(e)
+            telemetry('error', [
+              `error_type:${e.constructor.name}`,
+              `integration:${name}`,
+              `integration_version:${version}`
+            ])
           }
+          namesAndSuccesses[name].success = true
         }
       }
     }
+    for (const name of Object.keys(namesAndSuccesses)) {
+      const { success, version } = namesAndSuccesses[name]
+      if (!success && !seenCombo.has(`${name}@${version}`)) {
+        telemetry('abort.integration', [
+          `integration:${name}`,
+          `integration_version:${version}`
+        ])
+        log.info(`Found incompatible integration version: ${name}@${version}`)
+        seenCombo.add(`${name}@${version}`)
+      }
+    }
 
     return moduleExports
   })

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

@@ -21,6 +21,7 @@ const {
   runnableWrapper,
   getOnTestHandler,
   getOnTestEndHandler,
+  getOnTestRetryHandler,
   getOnHookEndHandler,
   getOnFailHandler,
   getOnPendingHandler,
@@ -37,10 +38,12 @@ let isSuitesSkipped = false
 let skippedSuites = []
 let isEarlyFlakeDetectionEnabled = false
 let isSuitesSkippingEnabled = false
+let isFlakyTestRetriesEnabled = false
 let earlyFlakeDetectionNumRetries = 0
 let knownTests = []
 let itrCorrelationId = ''
 let isForcedToRun = false
+const config = {}
 
 // We'll preserve the original coverage here
 const originalCoverageMap = createCoverageMap()
@@ -227,6 +230,12 @@ addHook({
       isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
       isSuitesSkippingEnabled = libraryConfig.isSuitesSkippingEnabled
       earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
+      isFlakyTestRetriesEnabled = libraryConfig.isFlakyTestRetriesEnabled
+
+      config.isEarlyFlakeDetectionEnabled = isEarlyFlakeDetectionEnabled
+      config.isSuitesSkippingEnabled = isSuitesSkippingEnabled
+      config.earlyFlakeDetectionNumRetries = earlyFlakeDetectionNumRetries
+      config.isFlakyTestRetriesEnabled = isFlakyTestRetriesEnabled
 
       if (isEarlyFlakeDetectionEnabled) {
         knownTestsCh.publish({
@@ -317,6 +326,8 @@ addHook({
 
     this.on('test end', getOnTestEndHandler())
 
+    this.on('retry', getOnTestRetryHandler())
+
     // If the hook passes, 'hook end' will be emitted. Otherwise, 'fail' will be emitted
     this.on('hook end', getOnHookEndHandler())
 
@@ -401,7 +412,7 @@ addHook({
   name: 'mocha',
   versions: ['>=5.2.0'],
   file: 'lib/runnable.js'
-}, runnableWrapper)
+}, (runnablePackage) => runnableWrapper(runnablePackage, config))
 
 // Only used in parallel mode (--parallel flag is passed)
 // Used to generate suite events and receive test payloads from workers

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

@@ -3,7 +3,8 @@
 const {
   getTestSuitePath,
   removeEfdStringFromTestName,
-  addEfdStringToTestName
+  addEfdStringToTestName,
+  NUM_FAILED_TEST_RETRIES
 } = require('../../../dd-trace/src/plugins/util/test')
 const { channel, AsyncResource } = require('../helpers/instrument')
 const shimmer = require('../../../datadog-shimmer')
@@ -11,6 +12,8 @@ const shimmer = require('../../../datadog-shimmer')
 // test channels
 const testStartCh = channel('ci:mocha:test:start')
 const testFinishCh = channel('ci:mocha:test:finish')
+// after a test has failed, we'll publish to this channel
+const testRetryCh = channel('ci:mocha:test:retry')
 const errorCh = channel('ci:mocha:test:error')
 const skipCh = channel('ci:mocha:test:skip')
 
@@ -70,6 +73,10 @@ function isMochaRetry (test) {
   return test._currentRetry !== undefined && test._currentRetry !== 0
 }
 
+function isLastRetry (test) {
+  return test._currentRetry === test._retries
+}
+
 function getTestFullName (test) {
   return `mocha.${getTestSuitePath(test.file, process.cwd())}.${removeEfdStringFromTestName(test.fullTitle())}`
 }
@@ -84,22 +91,34 @@ function getTestStatus (test) {
   return 'pass'
 }
 
-function getTestAsyncResource (test) {
+function getTestToArKey (test) {
   if (!test.fn) {
-    return testToAr.get(test)
+    return test
   }
   if (!wrappedFunctions.has(test.fn)) {
-    return testToAr.get(test.fn)
+    return test.fn
   }
   const originalFn = originalFns.get(test.fn)
-  return testToAr.get(originalFn)
+  return originalFn
+}
+
+function getTestAsyncResource (test) {
+  const key = getTestToArKey(test)
+  return testToAr.get(key)
 }
 
-function runnableWrapper (RunnablePackage) {
+function runnableWrapper (RunnablePackage, libraryConfig) {
   shimmer.wrap(RunnablePackage.prototype, 'run', run => function () {
     if (!testStartCh.hasSubscribers) {
       return run.apply(this, arguments)
     }
+    // Flaky test retries does not work in parallel mode
+    if (libraryConfig?.isFlakyTestRetriesEnabled) {
+      this.retries(NUM_FAILED_TEST_RETRIES)
+    }
+    // The reason why the wrapping logic is here is because we need to cover
+    // `afterEach` and `beforeEach` hooks as well.
+    // It can't be done in `getOnTestHandler` because it's only called for tests.
     const isBeforeEach = this.parent._beforeEach.includes(this)
     const isAfterEach = this.parent._afterEach.includes(this)
 
@@ -135,11 +154,16 @@ function runnableWrapper (RunnablePackage) {
 
 function getOnTestHandler (isMain, newTests) {
   return function (test) {
-    if (isMochaRetry(test)) {
-      return
-    }
     const testStartLine = testToStartLine.get(test)
     const asyncResource = new AsyncResource('bound-anonymous-fn')
+
+    // This may be a retry. If this is the case, `test.fn` is already wrapped,
+    // so we need to restore it.
+    if (wrappedFunctions.has(test.fn)) {
+      const originalFn = originalFns.get(test.fn)
+      test.fn = originalFn
+      wrappedFunctions.delete(test.fn)
+    }
     testToAr.set(test.fn, asyncResource)
 
     const {
@@ -186,7 +210,7 @@ function getOnTestEndHandler () {
     // if there are afterEach to be run, we don't finish the test yet
     if (asyncResource && !test.parent._afterEach.length) {
       asyncResource.runInAsyncScope(() => {
-        testFinishCh.publish(status)
+        testFinishCh.publish({ status, hasBeenRetried: isMochaRetry(test) })
       })
     }
   }
@@ -197,12 +221,17 @@ function getOnHookEndHandler () {
     const test = hook.ctx.currentTest
     if (test && hook.parent._afterEach.includes(hook)) { // only if it's an afterEach
       const isLastAfterEach = hook.parent._afterEach.indexOf(hook) === hook.parent._afterEach.length - 1
+      if (test._retries > 0 && !isLastRetry(test)) {
+        return
+      }
       if (isLastAfterEach) {
         const status = getTestStatus(test)
         const asyncResource = getTestAsyncResource(test)
-        asyncResource.runInAsyncScope(() => {
-          testFinishCh.publish(status)
-        })
+        if (asyncResource) {
+          asyncResource.runInAsyncScope(() => {
+            testFinishCh.publish({ status, hasBeenRetried: isMochaRetry(test) })
+          })
+        }
       }
     }
   }
@@ -226,7 +255,7 @@ function getOnFailHandler (isMain) {
           err.message = `${testOrHook.fullTitle()}: ${err.message}`
           errorCh.publish(err)
           // if it's a hook and it has failed, 'test end' will not be called
-          testFinishCh.publish('fail')
+          testFinishCh.publish({ status: 'fail', hasBeenRetried: isMochaRetry(test) })
         } else {
           errorCh.publish(err)
         }
@@ -250,6 +279,20 @@ function getOnFailHandler (isMain) {
   }
 }
 
+function getOnTestRetryHandler () {
+  return function (test) {
+    const asyncResource = getTestAsyncResource(test)
+    if (asyncResource) {
+      const isFirstAttempt = test._currentRetry === 0
+      asyncResource.runInAsyncScope(() => {
+        testRetryCh.publish(isFirstAttempt)
+      })
+    }
+    const key = getTestToArKey(test)
+    testToAr.delete(key)
+  }
+}
+
 function getOnPendingHandler () {
   return function (test) {
     const testStartLine = testToStartLine.get(test)
@@ -299,6 +342,7 @@ module.exports = {
   testToStartLine,
   getOnTestHandler,
   getOnTestEndHandler,
+  getOnTestRetryHandler,
   getOnHookEndHandler,
   getOnFailHandler,
   getOnPendingHandler,

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

@@ -49,3 +49,4 @@ addHook({
   versions: ['>=5.2.0'],
   file: 'lib/runnable.js'
 }, runnableWrapper)
+// TODO: parallel mode does not support flaky test retries, so no library config is passed.

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

@@ -249,7 +249,7 @@ function testEndHandler (test, annotations, testStatus, error, isTimeout) {
   testAsyncResource.runInAsyncScope(() => {
     testFinishCh.publish({
       testStatus,
-      steps: testResult.steps,
+      steps: testResult?.steps || [],
       error,
       extraTags: annotationTags,
       isNew: test._ddIsNew,

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

@@ -0,0 +1,303 @@
+const { addHook, channel, AsyncResource } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+// test hooks
+const testStartCh = channel('ci:vitest:test:start')
+const testFinishTimeCh = channel('ci:vitest:test:finish-time')
+const testPassCh = channel('ci:vitest:test:pass')
+const testErrorCh = channel('ci:vitest:test:error')
+const testSkipCh = channel('ci:vitest:test:skip')
+
+// test suite hooks
+const testSuiteStartCh = channel('ci:vitest:test-suite:start')
+const testSuiteFinishCh = channel('ci:vitest:test-suite:finish')
+const testSuiteErrorCh = channel('ci:vitest:test-suite:error')
+
+// test session hooks
+const testSessionStartCh = channel('ci:vitest:session:start')
+const testSessionFinishCh = channel('ci:vitest:session:finish')
+
+const taskToAsync = new WeakMap()
+
+const sessionAsyncResource = new AsyncResource('bound-anonymous-fn')
+
+function isReporterPackage (vitestPackage) {
+  return vitestPackage.B?.name === 'BaseSequencer'
+}
+
+// from 2.0.0
+function isReporterPackageNew (vitestPackage) {
+  return vitestPackage.e?.name === 'BaseSequencer'
+}
+
+function getSessionStatus (state) {
+  if (state.getCountOfFailedTests() > 0) {
+    return 'fail'
+  }
+  if (state.pathsSet.size === 0) {
+    return 'skip'
+  }
+  return 'pass'
+}
+
+// eslint-disable-next-line
+// From https://github.com/vitest-dev/vitest/blob/51c04e2f44d91322b334f8ccbcdb368facc3f8ec/packages/runner/src/run.ts#L243-L250
+function getVitestTestStatus (test, retryCount) {
+  if (test.result.state !== 'fail') {
+    if (!test.repeats) {
+      return 'pass'
+    } else if (test.repeats && (test.retry ?? 0) === retryCount) {
+      return 'pass'
+    }
+  }
+  return 'fail'
+}
+
+function getTypeTasks (fileTasks, type = 'test') {
+  const typeTasks = []
+
+  function getTasks (tasks) {
+    for (const task of tasks) {
+      if (task.type === type) {
+        typeTasks.push(task)
+      } else if (task.tasks) {
+        getTasks(task.tasks)
+      }
+    }
+  }
+
+  getTasks(fileTasks)
+
+  return typeTasks
+}
+
+function getTestName (task) {
+  let testName = task.name
+  let currentTask = task.suite
+
+  while (currentTask) {
+    if (currentTask.name) {
+      testName = `${currentTask.name} ${testName}`
+    }
+    currentTask = currentTask.suite
+  }
+
+  return testName
+}
+
+function getSortWrapper (sort) {
+  return async function () {
+    if (!testSessionFinishCh.hasSubscribers) {
+      return sort.apply(this, arguments)
+    }
+    shimmer.wrap(this.ctx, 'exit', exit => async function () {
+      let onFinish
+
+      const flushPromise = new Promise(resolve => {
+        onFinish = resolve
+      })
+      const failedSuites = this.state.getFailedFilepaths()
+      let error
+      if (failedSuites.length) {
+        error = new Error(`Test suites failed: ${failedSuites.length}.`)
+      }
+
+      sessionAsyncResource.runInAsyncScope(() => {
+        testSessionFinishCh.publish({
+          status: getSessionStatus(this.state),
+          onFinish,
+          error
+        })
+      })
+
+      await flushPromise
+
+      return exit.apply(this, arguments)
+    })
+
+    return sort.apply(this, arguments)
+  }
+}
+
+addHook({
+  name: 'vitest',
+  versions: ['>=1.6.0'],
+  file: 'dist/runners.js'
+}, (vitestPackage) => {
+  const { VitestTestRunner } = vitestPackage
+  // test start (only tests that are not marked as skip or todo)
+  shimmer.wrap(VitestTestRunner.prototype, 'onBeforeTryTask', onBeforeTryTask => async function (task) {
+    if (!testStartCh.hasSubscribers) {
+      return onBeforeTryTask.apply(this, arguments)
+    }
+    const asyncResource = new AsyncResource('bound-anonymous-fn')
+    taskToAsync.set(task, asyncResource)
+
+    asyncResource.runInAsyncScope(() => {
+      testStartCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
+    })
+    return onBeforeTryTask.apply(this, arguments)
+  })
+
+  // test finish (only passed tests)
+  shimmer.wrap(VitestTestRunner.prototype, 'onAfterTryTask', onAfterTryTask =>
+    async function (task, { retry: retryCount }) {
+      if (!testFinishTimeCh.hasSubscribers) {
+        return onAfterTryTask.apply(this, arguments)
+      }
+      const result = await onAfterTryTask.apply(this, arguments)
+
+      const status = getVitestTestStatus(task, retryCount)
+      const asyncResource = taskToAsync.get(task)
+
+      if (asyncResource) {
+        // We don't finish here because the test might fail in a later hook
+        asyncResource.runInAsyncScope(() => {
+          testFinishTimeCh.publish({ status, task })
+        })
+      }
+
+      return result
+    })
+
+  return vitestPackage
+})
+
+addHook({
+  name: 'vitest',
+  versions: ['>=2.0.0'],
+  filePattern: 'dist/vendor/index.*'
+}, (vitestPackage) => {
+  // there are multiple index* files so we have to check the exported values
+  if (isReporterPackageNew(vitestPackage)) {
+    shimmer.wrap(vitestPackage.e.prototype, 'sort', getSortWrapper)
+  }
+
+  return vitestPackage
+})
+
+addHook({
+  name: 'vitest',
+  versions: ['>=1.6.0'],
+  filePattern: 'dist/vendor/index.*'
+}, (vitestPackage) => {
+  // there are multiple index* files so we have to check the exported values
+  if (isReporterPackage(vitestPackage)) {
+    shimmer.wrap(vitestPackage.B.prototype, 'sort', getSortWrapper)
+  }
+
+  return vitestPackage
+})
+
+// Can't specify file because compiled vitest includes hashes in their files
+addHook({
+  name: 'vitest',
+  versions: ['>=1.6.0'],
+  filePattern: 'dist/vendor/cac.*'
+}, (vitestPackage, frameworkVersion) => {
+  shimmer.wrap(vitestPackage, 'c', oldCreateCli => function () {
+    if (!testSessionStartCh.hasSubscribers) {
+      return oldCreateCli.apply(this, arguments)
+    }
+    sessionAsyncResource.runInAsyncScope(() => {
+      const processArgv = process.argv.slice(2).join(' ')
+      testSessionStartCh.publish({ command: `vitest ${processArgv}`, frameworkVersion })
+    })
+    return oldCreateCli.apply(this, arguments)
+  })
+
+  return vitestPackage
+})
+
+// test suite start and finish
+// only relevant for workers
+addHook({
+  name: '@vitest/runner',
+  versions: ['>=1.6.0'],
+  file: 'dist/index.js'
+}, vitestPackage => {
+  shimmer.wrap(vitestPackage, 'startTests', startTests => async function (testPath) {
+    let testSuiteError = null
+    if (!testSuiteStartCh.hasSubscribers) {
+      return startTests.apply(this, arguments)
+    }
+
+    const testSuiteAsyncResource = new AsyncResource('bound-anonymous-fn')
+    testSuiteAsyncResource.runInAsyncScope(() => {
+      testSuiteStartCh.publish(testPath[0])
+    })
+    const startTestsResponse = await startTests.apply(this, arguments)
+
+    let onFinish = null
+    const onFinishPromise = new Promise(resolve => {
+      onFinish = resolve
+    })
+
+    const testTasks = getTypeTasks(startTestsResponse[0].tasks)
+
+    testTasks.forEach(task => {
+      const testAsyncResource = taskToAsync.get(task)
+      const { result } = task
+
+      if (result) {
+        const { state, duration, errors } = result
+        if (state === 'skip') { // programmatic skip
+          testSkipCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
+        } else if (state === 'pass') {
+          if (testAsyncResource) {
+            testAsyncResource.runInAsyncScope(() => {
+              testPassCh.publish({ task })
+            })
+          }
+        } else if (state === 'fail') {
+          // If it's failing, we have no accurate finish time, so we have to use `duration`
+          let testError
+
+          if (errors?.length) {
+            testError = errors[0]
+          }
+
+          if (testAsyncResource) {
+            testAsyncResource.runInAsyncScope(() => {
+              testErrorCh.publish({ duration, error: testError })
+            })
+          }
+          if (errors?.length) {
+            testSuiteError = testError // we store the error to bubble it up to the suite
+          }
+        }
+      } else { // test.skip or test.todo
+        testSkipCh.publish({ testName: getTestName(task), testSuiteAbsolutePath: task.suite.file.filepath })
+      }
+    })
+
+    const testSuiteResult = startTestsResponse[0].result
+
+    if (testSuiteResult.errors?.length) { // Errors from root level hooks
+      testSuiteError = testSuiteResult.errors[0]
+    } else if (testSuiteResult.state === 'fail') { // Errors from `describe` level hooks
+      const suiteTasks = getTypeTasks(startTestsResponse[0].tasks, 'suite')
+      const failedSuites = suiteTasks.filter(task => task.result?.state === 'fail')
+      if (failedSuites.length && failedSuites[0].result?.errors?.length) {
+        testSuiteError = failedSuites[0].result.errors[0]
+      }
+    }
+
+    if (testSuiteError) {
+      testSuiteAsyncResource.runInAsyncScope(() => {
+        testSuiteErrorCh.publish({ error: testSuiteError })
+      })
+    }
+
+    testSuiteAsyncResource.runInAsyncScope(() => {
+      testSuiteFinishCh.publish({ status: testSuiteResult.state, onFinish })
+    })
+
+    // TODO: fix too frequent flushes
+    await onFinishPromise
+
+    return startTestsResponse
+  })
+
+  return vitestPackage
+})

package/packages/datadog-plugin-aws-sdk/src/base.js

@@ -64,11 +64,17 @@ class BaseAwsSdkPlugin extends ClientPlugin {
       span.setTag('region', region)
     })
 
-    this.addSub(`apm:aws:request:complete:${this.serviceIdentifier}`, ({ response }) => {
+    this.addSub(`apm:aws:request:complete:${this.serviceIdentifier}`, ({ response, cbExists = false }) => {
       const store = storage.getStore()
       if (!store) return
       const { span } = store
       if (!span) return
+      // try to extract DSM context from response if no callback exists as extraction normally happens in CB
+      if (!cbExists && this.serviceIdentifier === 'sqs') {
+        const params = response.request.params
+        const operation = response.request.operation
+        this.responseExtractDSMContext(operation, params, response.data, span)
+      }
       this.addResponseTags(span, response)
       this.finish(span, response, response.error)
     })
@@ -159,6 +165,7 @@ function normalizeConfig (config, serviceIdentifier) {
 
   return Object.assign({}, config, specificConfig, {
     splitByAwsService: config.splitByAwsService !== false,
+    batchPropagationEnabled: config.batchPropagationEnabled !== false,
     hooks
   })
 }

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

@@ -52,7 +52,7 @@ class Kinesis extends BaseAwsSdkPlugin {
 
         // 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
+          request.operation, request.params, response, span || null, { streamName }
         )
       }
     })
@@ -100,7 +100,8 @@ class Kinesis extends BaseAwsSdkPlugin {
     }
   }
 
-  responseExtractDSMContext (operation, response, span, streamName) {
+  responseExtractDSMContext (operation, params, response, span, kwargs = {}) {
+    const { streamName } = kwargs
     if (!this.config.dsmEnabled) return
     if (operation !== 'getRecords') return
     if (!response || !response.Records || !response.Records[0]) return
@@ -151,7 +152,12 @@ class Kinesis extends BaseAwsSdkPlugin {
       case 'putRecords':
         stream = params.StreamArn ? params.StreamArn : (params.StreamName ? params.StreamName : '')
         for (let i = 0; i < params.Records.length; i++) {
-          this.injectToMessage(span, params.Records[i], stream, i === 0)
+          this.injectToMessage(
+            span,
+            params.Records[i],
+            stream,
+            i === 0 || (this.config.kinesis && this.config.kinesis.batchPropagationEnabled)
+          )
         }
     }
   }