[SVLS-4148] Inject trace context into AWS Stepfunction executions

.github/workflows/plugins.yml

@@ -55,7 +55,7 @@ jobs:
           majorVersion=$(echo "$version" | cut -d '.' -f 1)
           echo "Major Version: $majorVersion"
           echo "MAJOR_VERSION=$majorVersion" >> $GITHUB_ENV
-      - name: Install dependencies and run tests 
+      - name: Install dependencies and run tests
         if: env.MAJOR_VERSION == '3'
         run: |
           apt-get update && \
@@ -95,7 +95,7 @@ jobs:
           echo "Major Version: $majorVersion"
           echo "MAJOR_VERSION=$majorVersion" >> $GITHUB_ENV
       - uses: ./.github/actions/node/oldest
-      - name: Install dependencies and run tests 
+      - name: Install dependencies and run tests
         if: env.MAJOR_VERSION == '4'
         run: |
           yarn install --ignore-engines
@@ -139,7 +139,7 @@ jobs:
       - uses: actions/setup-node@v3
         with:
           node-version: ${{ matrix.node-version }}
-      - name: Install dependencies and run tests 
+      - name: Install dependencies and run tests
         if: env.MAJOR_VERSION >= '5'
         run: |
           yarn install --ignore-engines
@@ -227,7 +227,7 @@ jobs:
       localstack:
         image: localstack/localstack:3.0.2
         env:
-          LOCALSTACK_SERVICES: dynamodb,kinesis,s3,sqs,sns,redshift,route53,logs,serverless,lambda
+          LOCALSTACK_SERVICES: dynamodb,kinesis,s3,sqs,sns,redshift,route53,logs,serverless,lambda,stepfunctions,events
           EXTRA_CORS_ALLOWED_HEADERS: x-amz-request-id,x-amzn-requestid,x-amz-id-2
           EXTRA_CORS_EXPOSE_HEADERS: x-amz-request-id,x-amzn-requestid,x-amz-id-2
           AWS_DEFAULT_REGION: us-east-1

docker-compose.yml

@@ -92,7 +92,7 @@ services:
     ports:
       - "127.0.0.1:4566:4566" # Edge
     environment:
-      - LOCALSTACK_SERVICES=dynamodb,kinesis,s3,sqs,sns,redshift,route53,logs,serverless,lambda
+      - LOCALSTACK_SERVICES=dynamodb,kinesis,s3,sqs,sns,redshift,route53,logs,serverless,lambda,stepfunctions,events
       - EXTRA_CORS_ALLOWED_HEADERS=x-amz-request-id,x-amzn-requestid,x-amz-id-2
       - EXTRA_CORS_EXPOSE_HEADERS=x-amz-request-id,x-amzn-requestid,x-amz-id-2
       - AWS_DEFAULT_REGION=us-east-1

packages/datadog-instrumentations/src/aws-sdk.js

@@ -162,8 +162,11 @@ function getChannelSuffix (name) {
     'lambda',
     'redshift',
     's3',
+    'sfn',
     'sns',
-    'sqs'
+    'sqs',
+    'states',
+    'stepfunctions'
   ].includes(name)
     ? name
     : 'default'

packages/datadog-plugin-aws-sdk/src/services/index.js

@@ -7,6 +7,9 @@ exports.kinesis = require('./kinesis')
 exports.lambda = require('./lambda')
 exports.redshift = require('./redshift')
 exports.s3 = require('./s3')
+exports.sfn = require('./sfn')
 exports.sns = require('./sns')
 exports.sqs = require('./sqs')
+exports.states = require('./states')
+exports.stepfunctions = require('./stepfunctions')
 exports.default = require('./default')

packages/datadog-plugin-aws-sdk/src/services/sfn.js

@@ -0,0 +1,7 @@
+'use strict'
+const Stepfunctions = require('./stepfunctions')
+class Sfn extends Stepfunctions {
+  static get id () { return 'sfn' }
+}
+
+module.exports = Sfn

packages/datadog-plugin-aws-sdk/src/services/states.js

@@ -0,0 +1,7 @@
+'use strict'
+const Stepfunctions = require('./stepfunctions')
+class States extends Stepfunctions {
+  static get id () { return 'states' }
+}
+
+module.exports = States

packages/datadog-plugin-aws-sdk/src/services/stepfunctions.js

@@ -0,0 +1,64 @@
+'use strict'
+const log = require('../../../dd-trace/src/log')
+const BaseAwsSdkPlugin = require('../base')
+
+class Stepfunctions extends BaseAwsSdkPlugin {
+  static get id () { return 'stepfunctions' }
+
+  // This is the shape of StartExecutionInput, as defined in
+  // https://github.com/aws/aws-sdk-js/blob/master/apis/states-2016-11-23.normal.json
+  // "StartExecutionInput": {
+  //   "type": "structure",
+  //   "required": [
+  //     "stateMachineArn"
+  //   ],
+  //   "members": {
+  //     "stateMachineArn": {
+  //       "shape": "Arn",
+  //     },
+  //     "name": {
+  //       "shape": "Name",
+  //     },
+  //     "input": {
+  //       "shape": "SensitiveData",
+  //     },
+  //     "traceHeader": {
+  //       "shape": "TraceHeader",
+  //     }
+  //   }
+
+  generateTags (params, operation, response) {
+    if (!params) return {}
+    const tags = { 'resource.name': params.name ? `${operation} ${params.name}` : `${operation}` }
+    if (operation === 'startExecution' || operation === 'startSyncExecution') {
+      tags.statemachinearn = `${params.stateMachineArn}`
+    }
+    return tags
+  }
+
+  requestInject (span, request) {
+    const operation = request.operation
+    if (operation === 'startExecution' || operation === 'startSyncExecution') {
+      if (!request.params || !request.params.input) {
+        return
+      }
+
+      const input = request.params.input
+
+      try {
+        const inputObj = JSON.parse(input)
+        if (inputObj && typeof inputObj === 'object') {
+          // We've parsed the input JSON string
+          inputObj._datadog = {}
+          this.tracer.inject(span, 'text_map', inputObj._datadog)
+          const newInput = JSON.stringify(inputObj)
+          request.params.input = newInput
+        }
+      } catch (e) {
+        log.info('Unable to treat input as JSON')
+      }
+    }
+  }
+}
+
+module.exports = Stepfunctions

packages/datadog-plugin-aws-sdk/test/stepfunctions.spec.js

@@ -0,0 +1,128 @@
+/* eslint-disable max-len */
+'use strict'
+
+const semver = require('semver')
+const agent = require('../../dd-trace/test/plugins/agent')
+const { setup } = require('./spec_helpers')
+
+const helloWorldSMD = {
+  Comment: 'A Hello World example of the Amazon States Language using a Pass state',
+  StartAt: 'HelloWorld',
+  States: {
+    HelloWorld: {
+      Type: 'Pass',
+      Result: 'Hello World!',
+      End: true
+    }
+  }
+}
+
+describe('Sfn', () => {
+  let tracer
+
+  withVersions('aws-sdk', ['aws-sdk', '@aws-sdk/smithy-client'], (version, moduleName) => {
+    let stateMachineArn
+    let client
+
+    setup()
+
+    before(() => {
+      client = getClient()
+    })
+
+    function getClient () {
+      const params = { endpoint: 'http://127.0.0.1:4566', region: 'us-east-1' }
+      if (moduleName === '@aws-sdk/smithy-client') {
+        const lib = require(`../../../versions/@aws-sdk/client-sfn@${version}`).get()
+        const client = new lib.SFNClient(params)
+        return {
+          client,
+          createStateMachine: function () {
+            const req = new lib.CreateStateMachineCommand(...arguments)
+            return client.send(req)
+          },
+          deleteStateMachine: function () {
+            const req = new lib.DeleteStateMachineCommand(...arguments)
+            return client.send(req)
+          },
+          startExecution: function () {
+            const req = new lib.StartExecutionCommand(...arguments)
+            return client.send(req)
+          },
+          describeExecution: function () {
+            const req = new lib.DescribeExecutionCommand(...arguments)
+            return client.send(req)
+          }
+        }
+      } else {
+        const { StepFunctions } = require(`../../../versions/aws-sdk@${version}`).get()
+        const client = new StepFunctions(params)
+        return {
+          client,
+          createStateMachine: function () { return client.createStateMachine(...arguments).promise() },
+          deleteStateMachine: function () {
+            return client.deleteStateMachine(...arguments).promise()
+          },
+          startExecution: function () { return client.startExecution(...arguments).promise() },
+          describeExecution: function () { return client.describeExecution(...arguments).promise() }
+        }
+      }
+    }
+
+    async function createStateMachine (name, definition, xargs) {
+      return client.createStateMachine({
+        definition: JSON.stringify(definition),
+        name: name,
+        roleArn: 'arn:aws:iam::123456:role/test',
+        ...xargs
+      })
+    }
+
+    async function deleteStateMachine (arn) {
+      return client.deleteStateMachine({ stateMachineArn: arn })
+    }
+
+    describe('Traces', () => {
+      before(() => {
+        tracer = require('../../dd-trace')
+        tracer.use('aws-sdk')
+      })
+      // aws-sdk v2 doesn't support StepFunctions below 2.7.10
+      // https://github.com/aws/aws-sdk-js/blob/5dba638fd/CHANGELOG.md?plain=1#L18
+      if (moduleName !== 'aws-sdk' || semver.intersects(version, '>=2.7.10')) {
+        beforeEach(() => { return agent.load('aws-sdk') })
+        beforeEach(async () => {
+          const data = await createStateMachine('helloWorld', helloWorldSMD, {})
+          stateMachineArn = data.stateMachineArn
+        })
+
+        afterEach(() => { return agent.close({ ritmReset: false }) })
+
+        afterEach(async () => {
+          await deleteStateMachine(stateMachineArn)
+        })
+
+        it('is instrumented', async function () {
+          const startExecInput = {
+            stateMachineArn,
+            input: JSON.stringify({ moduleName })
+          }
+          const expectSpanPromise = agent.use(traces => {
+            const span = traces[0][0]
+            expect(span).to.have.property('resource', 'startExecution')
+            expect(span.meta).to.have.property('statemachinearn', stateMachineArn)
+          })
+
+          const resp = await client.startExecution(startExecInput)
+
+          const result = await client.describeExecution({ executionArn: resp.executionArn })
+          const sfInput = JSON.parse(result.input)
+          expect(sfInput).to.have.property('_datadog')
+          expect(sfInput._datadog).to.have.property('x-datadog-trace-id')
+          expect(sfInput._datadog).to.have.property('x-datadog-parent-id')
+          return expectSpanPromise.then(() => {})
+        })
+      }
+    })
+  })
+})

packages/dd-trace/test/plugins/externals.json

@@ -30,6 +30,10 @@
       "name": "@aws-sdk/client-s3",
       "versions": [">=3"]
     },
+    {
+      "name": "@aws-sdk/client-sfn",
+      "versions": [">=3"]
+    },
     {
       "name": "@aws-sdk/client-sns",
       "versions": [">=3"]