dd-trace 5.3.0 → 5.5.0

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

@@ -9,7 +9,9 @@ const {
   JEST_WORKER_COVERAGE_PAYLOAD_CODE,
   getTestLineStart,
   getTestSuitePath,
-  getTestParametersString
+  getTestParametersString,
+  EFD_STRING,
+  removeEfdStringFromTestName
 } = require('../../dd-trace/src/plugins/util/test')
 const {
   getFormattedJestTestParameters,
@@ -56,9 +58,7 @@ 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')
+let hasFilteredSkippableSuites = false
 
 const sessionAsyncResource = new AsyncResource('bound-anonymous-fn')
 
@@ -105,10 +105,6 @@ 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) {
@@ -116,12 +112,17 @@ 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)
 
+      const repositoryRoot = this.testEnvironmentOptions._ddRepositoryRoot
+
+      if (repositoryRoot) {
+        this.testSourceFile = getTestSuitePath(context.testPath, repositoryRoot)
+      }
+
       this.isEarlyFlakeDetectionEnabled = this.testEnvironmentOptions._ddIsEarlyFlakeDetectionEnabled
 
       if (this.isEarlyFlakeDetectionEnabled) {
@@ -152,7 +153,7 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
     // we use its describe block to get the full name
     getTestNameFromAddTestEvent (event, state) {
       const describeSuffix = getJestTestName(state.currentDescribeBlock)
-      return removeEfdTestName(`${describeSuffix} ${event.testName}`).trim()
+      return removeEfdStringFromTestName(`${describeSuffix} ${event.testName}`).trim()
     }
 
     async handleTestEvent (event, state) {
@@ -186,7 +187,7 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
         const testName = getJestTestName(event.test)
 
         if (this.isEarlyFlakeDetectionEnabled) {
-          const originalTestName = removeEfdTestName(testName)
+          const originalTestName = removeEfdStringFromTestName(testName)
           isNewTest = retriedTestsToNumAttempts.has(originalTestName)
           if (isNewTest) {
             numEfdRetry = retriedTestsToNumAttempts.get(originalTestName)
@@ -196,9 +197,9 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
 
         asyncResource.runInAsyncScope(() => {
           testStartCh.publish({
-            name: removeEfdTestName(testName),
+            name: removeEfdStringFromTestName(testName),
             suite: this.testSuite,
-            testFileAbsolutePath: this.testFileAbsolutePath,
+            testSourceFile: this.testSourceFile,
             runner: 'jest-circus',
             testParameters,
             frameworkVersion: jestVersion,
@@ -249,7 +250,7 @@ function getWrappedEnvironment (BaseEnvironment, jestVersion) {
           testSkippedCh.publish({
             name: getJestTestName(event.test),
             suite: this.testSuite,
-            testFileAbsolutePath: this.testFileAbsolutePath,
+            testSourceFile: this.testSourceFile,
             runner: 'jest-circus',
             frameworkVersion: jestVersion,
             testStartLine: getTestLineStart(event.test.asyncError, this.testSuite)
@@ -270,6 +271,23 @@ function getTestEnvironment (pkg, jestVersion) {
   return getWrappedEnvironment(pkg, jestVersion)
 }
 
+function applySuiteSkipping (originalTests, rootDir, frameworkVersion) {
+  const jestSuitesToRun = getJestSuitesToRun(skippableSuites, originalTests, rootDir || process.cwd())
+  hasFilteredSkippableSuites = true
+  log.debug(
+    () => `${jestSuitesToRun.suitesToRun.length} out of ${originalTests.length} suites are going to run.`
+  )
+  hasUnskippableSuites = jestSuitesToRun.hasUnskippableSuites
+  hasForcedToRunSuites = jestSuitesToRun.hasForcedToRunSuites
+
+  isSuitesSkipped = jestSuitesToRun.suitesToRun.length !== originalTests.length
+  numSkippedSuites = jestSuitesToRun.skippedSuites.length
+
+  itrSkippedSuitesCh.publish({ skippedSuites: jestSuitesToRun.skippedSuites, frameworkVersion })
+  skippableSuites = []
+  return jestSuitesToRun.suitesToRun
+}
+
 addHook({
   name: 'jest-environment-node',
   versions: ['>=24.8.0']
@@ -280,6 +298,51 @@ addHook({
   versions: ['>=24.8.0']
 }, getTestEnvironment)
 
+function getWrappedScheduleTests (scheduleTests, frameworkVersion) {
+  return async function (tests) {
+    if (!isSuitesSkippingEnabled || hasFilteredSkippableSuites) {
+      return scheduleTests.apply(this, arguments)
+    }
+    const [test] = tests
+    const rootDir = test?.context?.config?.rootDir
+
+    arguments[0] = applySuiteSkipping(tests, rootDir, frameworkVersion)
+
+    return scheduleTests.apply(this, arguments)
+  }
+}
+
+addHook({
+  name: '@jest/core',
+  file: 'build/TestScheduler.js',
+  versions: ['>=27.0.0']
+}, (testSchedulerPackage, frameworkVersion) => {
+  const oldCreateTestScheduler = testSchedulerPackage.createTestScheduler
+  const newCreateTestScheduler = async function () {
+    if (!isSuitesSkippingEnabled || hasFilteredSkippableSuites) {
+      return oldCreateTestScheduler.apply(this, arguments)
+    }
+    // If suite skipping is enabled and has not filtered skippable suites yet, we'll attempt to do it
+    const scheduler = await oldCreateTestScheduler.apply(this, arguments)
+    shimmer.wrap(scheduler, 'scheduleTests', scheduleTests => getWrappedScheduleTests(scheduleTests, frameworkVersion))
+    return scheduler
+  }
+  testSchedulerPackage.createTestScheduler = newCreateTestScheduler
+  return testSchedulerPackage
+})
+
+addHook({
+  name: '@jest/core',
+  file: 'build/TestScheduler.js',
+  versions: ['>=24.8.0 <27.0.0']
+}, (testSchedulerPackage, frameworkVersion) => {
+  shimmer.wrap(
+    testSchedulerPackage.default.prototype,
+    'scheduleTests', scheduleTests => getWrappedScheduleTests(scheduleTests, frameworkVersion)
+  )
+  return testSchedulerPackage
+})
+
 addHook({
   name: '@jest/test-sequencer',
   versions: ['>=24.8.0']
@@ -287,29 +350,13 @@ addHook({
   shimmer.wrap(sequencerPackage.default.prototype, 'shard', shard => function () {
     const shardedTests = shard.apply(this, arguments)
 
-    if (!shardedTests.length) {
+    if (!shardedTests.length || !isSuitesSkippingEnabled || !skippableSuites.length) {
       return shardedTests
     }
-    // TODO: could we get the rootDir from each test?
     const [test] = shardedTests
     const rootDir = test?.context?.config?.rootDir
 
-    const jestSuitesToRun = getJestSuitesToRun(skippableSuites, shardedTests, rootDir || process.cwd())
-
-    log.debug(
-      () => `${jestSuitesToRun.suitesToRun.length} out of ${shardedTests.length} suites are going to run.`
-    )
-
-    hasUnskippableSuites = jestSuitesToRun.hasUnskippableSuites
-    hasForcedToRunSuites = jestSuitesToRun.hasForcedToRunSuites
-
-    isSuitesSkipped = jestSuitesToRun.suitesToRun.length !== shardedTests.length
-    numSkippedSuites = jestSuitesToRun.skippedSuites.length
-
-    itrSkippedSuitesCh.publish({ skippedSuites: jestSuitesToRun.skippedSuites, frameworkVersion })
-
-    skippableSuites = []
-    return jestSuitesToRun.suitesToRun
+    return applySuiteSkipping(shardedTests, rootDir, frameworkVersion)
   })
   return sequencerPackage
 })
@@ -635,6 +682,7 @@ addHook({
       _ddKnownTests,
       _ddIsEarlyFlakeDetectionEnabled,
       _ddEarlyFlakeDetectionNumRetries,
+      _ddRepositoryRoot,
       ...restOfTestEnvironmentOptions
     } = testEnvironmentOptions
 
@@ -659,13 +707,13 @@ addHook({
   const SearchSource = searchSourcePackage.default ? searchSourcePackage.default : searchSourcePackage
 
   shimmer.wrap(SearchSource.prototype, 'getTestPaths', getTestPaths => async function () {
-    if (!skippableSuites.length) {
+    if (!isSuitesSkippingEnabled || !skippableSuites.length) {
       return getTestPaths.apply(this, arguments)
     }
 
     const [{ rootDir, shard }] = arguments
 
-    if (shard && shard.shardIndex) {
+    if (shard?.shardCount > 1) {
       // If the user is using jest sharding, we want to apply the filtering of tests in the shard process.
       // The reason for this is the following:
       // The tests for different shards are likely being run in different CI jobs so
@@ -679,21 +727,8 @@ addHook({
     const testPaths = await getTestPaths.apply(this, arguments)
     const { tests } = testPaths
 
-    const jestSuitesToRun = getJestSuitesToRun(skippableSuites, tests, rootDir)
-
-    log.debug(() => `${jestSuitesToRun.suitesToRun.length} out of ${tests.length} suites are going to run.`)
-
-    hasUnskippableSuites = jestSuitesToRun.hasUnskippableSuites
-    hasForcedToRunSuites = jestSuitesToRun.hasForcedToRunSuites
-
-    isSuitesSkipped = jestSuitesToRun.suitesToRun.length !== tests.length
-    numSkippedSuites = jestSuitesToRun.skippedSuites.length
-
-    itrSkippedSuitesCh.publish({ skippedSuites: jestSuitesToRun.skippedSuites, frameworkVersion })
-
-    skippableSuites = []
-
-    return { ...testPaths, tests: jestSuitesToRun.suitesToRun }
+    const suitesToRun = applySuiteSkipping(tests, rootDir, frameworkVersion)
+    return { ...testPaths, tests: suitesToRun }
   })
 
   return searchSourcePackage

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

@@ -11,7 +11,9 @@ const {
   mergeCoverage,
   getTestSuitePath,
   fromCoverageMapToCoverage,
-  getCallSites
+  getCallSites,
+  addEfdStringToTestName,
+  removeEfdStringFromTestName
 } = require('../../dd-trace/src/plugins/util/test')
 
 const testStartCh = channel('ci:mocha:test:start')
@@ -21,6 +23,7 @@ const testFinishCh = channel('ci:mocha:test:finish')
 const parameterizedTestCh = channel('ci:mocha:test:parameterize')
 
 const libraryConfigurationCh = channel('ci:mocha:library-configuration')
+const knownTestsCh = channel('ci:mocha:known-tests')
 const skippableSuitesCh = channel('ci:mocha:test-suite:skippable')
 
 const testSessionStartCh = channel('ci:mocha:session:start')
@@ -40,6 +43,7 @@ const testToAr = new WeakMap()
 const originalFns = new WeakMap()
 const testFileToSuiteAr = new Map()
 const testToStartLine = new WeakMap()
+const newTests = {}
 
 // `isWorker` is true if it's a Mocha worker
 let isWorker = false
@@ -54,6 +58,10 @@ let skippedSuites = []
 const unskippableSuites = []
 let isForcedToRun = false
 let itrCorrelationId = ''
+let isEarlyFlakeDetectionEnabled = false
+let earlyFlakeDetectionNumRetries = 0
+let isSuitesSkippingEnabled = false
+let knownTests = []
 
 function getSuitesByTestFile (root) {
   const suitesByTestFile = {}
@@ -93,6 +101,26 @@ function isRetry (test) {
   return test._currentRetry !== undefined && test._currentRetry !== 0
 }
 
+function getTestFullName (test) {
+  return `mocha.${getTestSuitePath(test.file, process.cwd())}.${removeEfdStringFromTestName(test.fullTitle())}`
+}
+
+function isNewTest (test) {
+  return !knownTests.includes(getTestFullName(test))
+}
+
+function retryTest (test) {
+  const originalTestName = test.title
+  const suite = test.parent
+  for (let retryIndex = 0; retryIndex < earlyFlakeDetectionNumRetries; retryIndex++) {
+    const clonedTest = test.clone()
+    clonedTest.title = addEfdStringToTestName(originalTestName, retryIndex + 1)
+    suite.addTest(clonedTest)
+    clonedTest._ddIsNew = true
+    clonedTest._ddIsEfdRetry = true
+  }
+}
+
 function getTestAsyncResource (test) {
   if (!test.fn) {
     return testToAr.get(test)
@@ -123,6 +151,19 @@ function mochaHook (Runner) {
 
   patched.add(Runner)
 
+  shimmer.wrap(Runner.prototype, 'runTests', runTests => function (suite, fn) {
+    if (isEarlyFlakeDetectionEnabled) {
+      // by the time we reach `this.on('test')`, it is too late. We need to add retries here
+      suite.tests.forEach(test => {
+        if (!test.isPending() && isNewTest(test)) {
+          test._ddIsNew = true
+          retryTest(test)
+        }
+      })
+    }
+    return runTests.apply(this, arguments)
+  })
+
   shimmer.wrap(Runner.prototype, 'run', run => function () {
     if (!testStartCh.hasSubscribers || isWorker) {
       return run.apply(this, arguments)
@@ -144,6 +185,24 @@ function mochaHook (Runner) {
         status = 'fail'
       }
 
+      if (isEarlyFlakeDetectionEnabled) {
+        /**
+         * If Early Flake Detection (EFD) is enabled the logic is as follows:
+         * - If all attempts for a test are failing, the test has failed and we will let the test process fail.
+         * - If just a single attempt passes, we will prevent the test process from failing.
+         * The rationale behind is the following: you may still be able to block your CI pipeline by gating
+         * on flakiness (the test will be considered flaky), but you may choose to unblock the pipeline too.
+         */
+        for (const tests of Object.values(newTests)) {
+          const failingNewTests = tests.filter(test => test.isFailed())
+          const areAllNewTestsFailing = failingNewTests.length === tests.length
+          if (failingNewTests.length && !areAllNewTestsFailing) {
+            this.stats.failures -= failingNewTests.length
+            this.failures -= failingNewTests.length
+          }
+        }
+      }
+
       if (status === 'fail') {
         error = new Error(`Failed tests: ${this.failures}.`)
       }
@@ -168,7 +227,8 @@ function mochaHook (Runner) {
         numSkippedSuites: skippedSuites.length,
         hasForcedToRunSuites: isForcedToRun,
         hasUnskippableSuites: !!unskippableSuites.length,
-        error
+        error,
+        isEarlyFlakeDetectionEnabled
       })
     }))
 
@@ -253,8 +313,35 @@ function mochaHook (Runner) {
       const testStartLine = testToStartLine.get(test)
       const asyncResource = new AsyncResource('bound-anonymous-fn')
       testToAr.set(test.fn, asyncResource)
+
+      const {
+        file: testSuiteAbsolutePath,
+        title,
+        _ddIsNew: isNew,
+        _ddIsEfdRetry: isEfdRetry
+      } = test
+
+      const testInfo = {
+        testName: test.fullTitle(),
+        testSuiteAbsolutePath,
+        title,
+        isNew,
+        isEfdRetry,
+        testStartLine
+      }
+
+      // We want to store the result of the new tests
+      if (isNew) {
+        const testFullName = getTestFullName(test)
+        if (newTests[testFullName]) {
+          newTests[testFullName].push(test)
+        } else {
+          newTests[testFullName] = [test]
+        }
+      }
+
       asyncResource.runInAsyncScope(() => {
-        testStartCh.publish({ test, testStartLine })
+        testStartCh.publish(testInfo)
       })
     })
 
@@ -323,10 +410,23 @@ function mochaHook (Runner) {
     })
 
     this.on('pending', (test) => {
+      const testStartLine = testToStartLine.get(test)
+      const {
+        file: testSuiteAbsolutePath,
+        title
+      } = test
+
+      const testInfo = {
+        testName: test.fullTitle(),
+        testSuiteAbsolutePath,
+        title,
+        testStartLine
+      }
+
       const asyncResource = getTestAsyncResource(test)
       if (asyncResource) {
         asyncResource.runInAsyncScope(() => {
-          skipCh.publish(test)
+          skipCh.publish(testInfo)
         })
       } else {
         // if there is no async resource, the test has been skipped through `test.skip`
@@ -338,7 +438,7 @@ function mochaHook (Runner) {
           testToAr.set(test, skippedTestAsyncResource)
         }
         skippedTestAsyncResource.runInAsyncScope(() => {
-          skipCh.publish(test)
+          skipCh.publish(testInfo)
         })
       }
     })
@@ -358,8 +458,8 @@ function mochaEachHook (mochaEach) {
     const [params] = arguments
     const { it, ...rest } = mochaEach.apply(this, arguments)
     return {
-      it: function (name) {
-        parameterizedTestCh.publish({ name, params })
+      it: function (title) {
+        parameterizedTestCh.publish({ title, params })
         it.apply(this, arguments)
       },
       ...rest
@@ -425,17 +525,43 @@ addHook({
       global.run()
     }
 
-    const onReceivedConfiguration = ({ err }) => {
+    const onReceivedKnownTests = ({ err, knownTests: receivedKnownTests }) => {
       if (err) {
-        return global.run()
+        knownTests = []
+        isEarlyFlakeDetectionEnabled = false
+      } else {
+        knownTests = receivedKnownTests
+      }
+
+      if (isSuitesSkippingEnabled) {
+        skippableSuitesCh.publish({
+          onDone: mochaRunAsyncResource.bind(onReceivedSkippableSuites)
+        })
+      } else {
+        global.run()
       }
-      if (!skippableSuitesCh.hasSubscribers) {
+    }
+
+    const onReceivedConfiguration = ({ err, libraryConfig }) => {
+      if (err || !skippableSuitesCh.hasSubscribers || !knownTestsCh.hasSubscribers) {
         return global.run()
       }
 
-      skippableSuitesCh.publish({
-        onDone: mochaRunAsyncResource.bind(onReceivedSkippableSuites)
-      })
+      isEarlyFlakeDetectionEnabled = libraryConfig.isEarlyFlakeDetectionEnabled
+      isSuitesSkippingEnabled = libraryConfig.isSuitesSkippingEnabled
+      earlyFlakeDetectionNumRetries = libraryConfig.earlyFlakeDetectionNumRetries
+
+      if (isEarlyFlakeDetectionEnabled) {
+        knownTestsCh.publish({
+          onDone: mochaRunAsyncResource.bind(onReceivedKnownTests)
+        })
+      } else if (isSuitesSkippingEnabled) {
+        skippableSuitesCh.publish({
+          onDone: mochaRunAsyncResource.bind(onReceivedSkippableSuites)
+        })
+      } else {
+        global.run()
+      }
     }
 
     mochaRunAsyncResource.runInAsyncScope(() => {

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

@@ -28,9 +28,12 @@ class AmqplibConsumerPlugin extends ConsumerPlugin {
       }
     })
 
-    if (this.config.dsmEnabled && message) {
+    if (
+      this.config.dsmEnabled &&
+      message?.properties?.headers?.[CONTEXT_PROPAGATION_KEY]
+    ) {
       const payloadSize = getAmqpMessageSize({ headers: message.properties.headers, content: message.content })
-      const queue = fields.queue ?? fields.routingKey
+      const queue = fields.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-aws-sdk/src/services/kinesis.js

@@ -53,7 +53,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, response, span || null, streamName
         )
       }
     })
@@ -143,62 +143,72 @@ class Kinesis extends BaseAwsSdkPlugin {
   }
 
   requestInject (span, request) {
-    const operation = request.operation
-    if (operation === 'putRecord' || operation === 'putRecords') {
-      if (!request.params) {
-        return
-      }
-      const traceData = {}
-
-      // inject data with DD context
-      this.tracer.inject(span, 'text_map', traceData)
-      let injectPath
-      if (request.params.Records && request.params.Records.length > 0) {
-        injectPath = request.params.Records[0]
-      } else if (request.params.Data) {
-        injectPath = request.params
-      } else {
-        log.error('No valid payload passed, unable to pass trace context')
+    const { operation, params } = request
+    if (!params) return
+
+    let stream
+    switch (operation) {
+      case 'putRecord':
+        stream = params.StreamArn ? params.StreamArn : (params.StreamName ? params.StreamName : '')
+        this.injectToMessage(span, params, stream, true)
+        break
+      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)
+        }
+    }
+  }
+
+  injectToMessage (span, params, stream, injectTraceContext) {
+    if (!params) {
+      return
+    }
+
+    let parsedData
+    if (injectTraceContext || this.config.dsmEnabled) {
+      parsedData = this._tryParse(params.Data)
+      if (!parsedData) {
+        log.error('Unable to parse payload, unable to pass trace context or set DSM checkpoint (if enabled)')
         return
       }
+    }
 
-      const parsedData = this._tryParse(injectPath.Data)
-      if (parsedData) {
-        parsedData._datadog = traceData
-
-        // set DSM hash if enabled
-        if (this.config.dsmEnabled) {
-          // get payload size of request data
-          const payloadSize = Buffer.from(JSON.stringify(parsedData)).byteLength
-          let stream
-          // users can optionally use either stream name or stream arn
-          if (request.params && request.params.StreamArn) {
-            stream = request.params.StreamArn
-          } else if (request.params && request.params.StreamName) {
-            stream = request.params.StreamName
-          }
-          const dataStreamsContext = this.tracer
-            .setCheckpoint(['direction:out', `topic:${stream}`, 'type:kinesis'], span, payloadSize)
-          if (dataStreamsContext) {
-            const pathwayCtx = encodePathwayContext(dataStreamsContext)
-            parsedData._datadog[CONTEXT_PROPAGATION_KEY] = pathwayCtx.toJSON()
-          }
-        }
+    const ddInfo = {}
+    // for now, we only want to inject to the first message, this may change for batches in the future
+    if (injectTraceContext) { this.tracer.inject(span, 'text_map', ddInfo) }
+
+    // set DSM hash if enabled
+    if (this.config.dsmEnabled) {
+      parsedData._datadog = ddInfo
+      const dataStreamsContext = this.setDSMCheckpoint(span, parsedData, stream)
+      if (dataStreamsContext) {
+        const pathwayCtx = encodePathwayContext(dataStreamsContext)
+        ddInfo[CONTEXT_PROPAGATION_KEY] = pathwayCtx.toJSON()
+      }
+    }
 
-        const finalData = Buffer.from(JSON.stringify(parsedData))
-        const byteSize = finalData.length
-        // Kinesis max payload size is 1MB
-        // So we must ensure adding DD context won't go over that (512b is an estimate)
-        if (byteSize >= 1048576) {
-          log.info('Payload size too large to pass context')
-          return
-        }
-        injectPath.Data = finalData
-      } else {
-        log.error('Unable to parse payload, unable to pass trace context')
+    if (Object.keys(ddInfo).length !== 0) {
+      parsedData._datadog = ddInfo
+      const finalData = Buffer.from(JSON.stringify(parsedData))
+      const byteSize = finalData.length
+      // Kinesis max payload size is 1MB
+      // So we must ensure adding DD context won't go over that (512b is an estimate)
+      if (byteSize >= 1048576) {
+        log.info('Payload size too large to pass context')
+        return
       }
+      params.Data = finalData
     }
   }
+
+  setDSMCheckpoint (span, parsedData, stream) {
+    // get payload size of request data
+    const payloadSize = Buffer.from(JSON.stringify(parsedData)).byteLength
+    const dataStreamsContext = this.tracer
+      .setCheckpoint(['direction:out', `topic:${stream}`, 'type:kinesis'], span, payloadSize)
+    return dataStreamsContext
+  }
 }
 
 module.exports = Kinesis