Skip to content

Commit bd52fc0

Browse files
agocsJordi Bertran de Balandatlhunter
authored andcommitted
[SVLS-4148] Inject trace context into AWS Stepfunction executions (#4069)
Adds a plugin for AWS Stepfunctions that injects the datadog trace context into the input parameter of a startExecution or startSyncExecution request, if the input parameter is already a JSON object. SVLS-4148 Serverless Integrations has added the ability to trace Stepfunction executions, and we want to be able to link Stepfunction executions with upstream traces. We do so by inspecting input._datadog in the Logs to Traces Reducer. --------- Co-authored-by: Jordi Bertran de Balanda <[email protected]> Co-authored-by: Thomas Hunter II <[email protected]>
1 parent b3a10bf commit bd52fc0

File tree

9 files changed

+222
-6
lines changed

9 files changed

+222
-6
lines changed

.github/workflows/plugins.yml

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ jobs:
5555
majorVersion=$(echo "$version" | cut -d '.' -f 1)
5656
echo "Major Version: $majorVersion"
5757
echo "MAJOR_VERSION=$majorVersion" >> $GITHUB_ENV
58-
- name: Install dependencies and run tests
58+
- name: Install dependencies and run tests
5959
if: env.MAJOR_VERSION == '3'
6060
run: |
6161
apt-get update && \
@@ -95,7 +95,7 @@ jobs:
9595
echo "Major Version: $majorVersion"
9696
echo "MAJOR_VERSION=$majorVersion" >> $GITHUB_ENV
9797
- uses: ./.github/actions/node/oldest
98-
- name: Install dependencies and run tests
98+
- name: Install dependencies and run tests
9999
if: env.MAJOR_VERSION == '4'
100100
run: |
101101
yarn install --ignore-engines
@@ -139,7 +139,7 @@ jobs:
139139
- uses: actions/setup-node@v3
140140
with:
141141
node-version: ${{ matrix.node-version }}
142-
- name: Install dependencies and run tests
142+
- name: Install dependencies and run tests
143143
if: env.MAJOR_VERSION >= '5'
144144
run: |
145145
yarn install --ignore-engines
@@ -227,7 +227,7 @@ jobs:
227227
localstack:
228228
image: localstack/localstack:3.0.2
229229
env:
230-
LOCALSTACK_SERVICES: dynamodb,kinesis,s3,sqs,sns,redshift,route53,logs,serverless,lambda
230+
LOCALSTACK_SERVICES: dynamodb,kinesis,s3,sqs,sns,redshift,route53,logs,serverless,lambda,stepfunctions,events
231231
EXTRA_CORS_ALLOWED_HEADERS: x-amz-request-id,x-amzn-requestid,x-amz-id-2
232232
EXTRA_CORS_EXPOSE_HEADERS: x-amz-request-id,x-amzn-requestid,x-amz-id-2
233233
AWS_DEFAULT_REGION: us-east-1

docker-compose.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,7 @@ services:
9292
ports:
9393
- "127.0.0.1:4566:4566" # Edge
9494
environment:
95-
- LOCALSTACK_SERVICES=dynamodb,kinesis,s3,sqs,sns,redshift,route53,logs,serverless,lambda
95+
- LOCALSTACK_SERVICES=dynamodb,kinesis,s3,sqs,sns,redshift,route53,logs,serverless,lambda,stepfunctions,events
9696
- EXTRA_CORS_ALLOWED_HEADERS=x-amz-request-id,x-amzn-requestid,x-amz-id-2
9797
- EXTRA_CORS_EXPOSE_HEADERS=x-amz-request-id,x-amzn-requestid,x-amz-id-2
9898
- AWS_DEFAULT_REGION=us-east-1

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,8 +162,11 @@ function getChannelSuffix (name) {
162162
'lambda',
163163
'redshift',
164164
's3',
165+
'sfn',
165166
'sns',
166-
'sqs'
167+
'sqs',
168+
'states',
169+
'stepfunctions'
167170
].includes(name)
168171
? name
169172
: 'default'

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@ exports.kinesis = require('./kinesis')
77
exports.lambda = require('./lambda')
88
exports.redshift = require('./redshift')
99
exports.s3 = require('./s3')
10+
exports.sfn = require('./sfn')
1011
exports.sns = require('./sns')
1112
exports.sqs = require('./sqs')
13+
exports.states = require('./states')
14+
exports.stepfunctions = require('./stepfunctions')
1215
exports.default = require('./default')
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
'use strict'
2+
const Stepfunctions = require('./stepfunctions')
3+
class Sfn extends Stepfunctions {
4+
static get id () { return 'sfn' }
5+
}
6+
7+
module.exports = Sfn
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
'use strict'
2+
const Stepfunctions = require('./stepfunctions')
3+
class States extends Stepfunctions {
4+
static get id () { return 'states' }
5+
}
6+
7+
module.exports = States
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
'use strict'
2+
const log = require('../../../dd-trace/src/log')
3+
const BaseAwsSdkPlugin = require('../base')
4+
5+
class Stepfunctions extends BaseAwsSdkPlugin {
6+
static get id () { return 'stepfunctions' }
7+
8+
// This is the shape of StartExecutionInput, as defined in
9+
// https://github.com/aws/aws-sdk-js/blob/master/apis/states-2016-11-23.normal.json
10+
// "StartExecutionInput": {
11+
// "type": "structure",
12+
// "required": [
13+
// "stateMachineArn"
14+
// ],
15+
// "members": {
16+
// "stateMachineArn": {
17+
// "shape": "Arn",
18+
// },
19+
// "name": {
20+
// "shape": "Name",
21+
// },
22+
// "input": {
23+
// "shape": "SensitiveData",
24+
// },
25+
// "traceHeader": {
26+
// "shape": "TraceHeader",
27+
// }
28+
// }
29+
30+
generateTags (params, operation, response) {
31+
if (!params) return {}
32+
const tags = { 'resource.name': params.name ? `${operation} ${params.name}` : `${operation}` }
33+
if (operation === 'startExecution' || operation === 'startSyncExecution') {
34+
tags.statemachinearn = `${params.stateMachineArn}`
35+
}
36+
return tags
37+
}
38+
39+
requestInject (span, request) {
40+
const operation = request.operation
41+
if (operation === 'startExecution' || operation === 'startSyncExecution') {
42+
if (!request.params || !request.params.input) {
43+
return
44+
}
45+
46+
const input = request.params.input
47+
48+
try {
49+
const inputObj = JSON.parse(input)
50+
if (inputObj && typeof inputObj === 'object') {
51+
// We've parsed the input JSON string
52+
inputObj._datadog = {}
53+
this.tracer.inject(span, 'text_map', inputObj._datadog)
54+
const newInput = JSON.stringify(inputObj)
55+
request.params.input = newInput
56+
}
57+
} catch (e) {
58+
log.info('Unable to treat input as JSON')
59+
}
60+
}
61+
}
62+
}
63+
64+
module.exports = Stepfunctions
Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
/* eslint-disable max-len */
2+
'use strict'
3+
4+
const semver = require('semver')
5+
const agent = require('../../dd-trace/test/plugins/agent')
6+
const { setup } = require('./spec_helpers')
7+
8+
const helloWorldSMD = {
9+
Comment: 'A Hello World example of the Amazon States Language using a Pass state',
10+
StartAt: 'HelloWorld',
11+
States: {
12+
HelloWorld: {
13+
Type: 'Pass',
14+
Result: 'Hello World!',
15+
End: true
16+
}
17+
}
18+
}
19+
20+
describe('Sfn', () => {
21+
let tracer
22+
23+
withVersions('aws-sdk', ['aws-sdk', '@aws-sdk/smithy-client'], (version, moduleName) => {
24+
let stateMachineArn
25+
let client
26+
27+
setup()
28+
29+
before(() => {
30+
client = getClient()
31+
})
32+
33+
function getClient () {
34+
const params = { endpoint: 'http://127.0.0.1:4566', region: 'us-east-1' }
35+
if (moduleName === '@aws-sdk/smithy-client') {
36+
const lib = require(`../../../versions/@aws-sdk/client-sfn@${version}`).get()
37+
const client = new lib.SFNClient(params)
38+
return {
39+
client,
40+
createStateMachine: function () {
41+
const req = new lib.CreateStateMachineCommand(...arguments)
42+
return client.send(req)
43+
},
44+
deleteStateMachine: function () {
45+
const req = new lib.DeleteStateMachineCommand(...arguments)
46+
return client.send(req)
47+
},
48+
startExecution: function () {
49+
const req = new lib.StartExecutionCommand(...arguments)
50+
return client.send(req)
51+
},
52+
describeExecution: function () {
53+
const req = new lib.DescribeExecutionCommand(...arguments)
54+
return client.send(req)
55+
}
56+
}
57+
} else {
58+
const { StepFunctions } = require(`../../../versions/aws-sdk@${version}`).get()
59+
const client = new StepFunctions(params)
60+
return {
61+
client,
62+
createStateMachine: function () { return client.createStateMachine(...arguments).promise() },
63+
deleteStateMachine: function () {
64+
return client.deleteStateMachine(...arguments).promise()
65+
},
66+
startExecution: function () { return client.startExecution(...arguments).promise() },
67+
describeExecution: function () { return client.describeExecution(...arguments).promise() }
68+
}
69+
}
70+
}
71+
72+
async function createStateMachine (name, definition, xargs) {
73+
return client.createStateMachine({
74+
definition: JSON.stringify(definition),
75+
name: name,
76+
roleArn: 'arn:aws:iam::123456:role/test',
77+
...xargs
78+
})
79+
}
80+
81+
async function deleteStateMachine (arn) {
82+
return client.deleteStateMachine({ stateMachineArn: arn })
83+
}
84+
85+
describe('Traces', () => {
86+
before(() => {
87+
tracer = require('../../dd-trace')
88+
tracer.use('aws-sdk')
89+
})
90+
// aws-sdk v2 doesn't support StepFunctions below 2.7.10
91+
// https://github.com/aws/aws-sdk-js/blob/5dba638fd/CHANGELOG.md?plain=1#L18
92+
if (moduleName !== 'aws-sdk' || semver.intersects(version, '>=2.7.10')) {
93+
beforeEach(() => { return agent.load('aws-sdk') })
94+
beforeEach(async () => {
95+
const data = await createStateMachine('helloWorld', helloWorldSMD, {})
96+
stateMachineArn = data.stateMachineArn
97+
})
98+
99+
afterEach(() => { return agent.close({ ritmReset: false }) })
100+
101+
afterEach(async () => {
102+
await deleteStateMachine(stateMachineArn)
103+
})
104+
105+
it('is instrumented', async function () {
106+
const startExecInput = {
107+
stateMachineArn,
108+
input: JSON.stringify({ moduleName })
109+
}
110+
const expectSpanPromise = agent.use(traces => {
111+
const span = traces[0][0]
112+
expect(span).to.have.property('resource', 'startExecution')
113+
expect(span.meta).to.have.property('statemachinearn', stateMachineArn)
114+
})
115+
116+
const resp = await client.startExecution(startExecInput)
117+
118+
const result = await client.describeExecution({ executionArn: resp.executionArn })
119+
const sfInput = JSON.parse(result.input)
120+
expect(sfInput).to.have.property('_datadog')
121+
expect(sfInput._datadog).to.have.property('x-datadog-trace-id')
122+
expect(sfInput._datadog).to.have.property('x-datadog-parent-id')
123+
return expectSpanPromise.then(() => {})
124+
})
125+
}
126+
})
127+
})
128+
})

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,10 @@
3030
"name": "@aws-sdk/client-s3",
3131
"versions": [">=3"]
3232
},
33+
{
34+
"name": "@aws-sdk/client-sfn",
35+
"versions": [">=3"]
36+
},
3337
{
3438
"name": "@aws-sdk/client-sns",
3539
"versions": [">=3"]

0 commit comments

Comments
 (0)