Skip to content

Commit b0aa443

Browse files
srprashsrprash
authored
Propagate additional trace data into AWS requests on Lambda (#549)
1 parent 8a4f584 commit b0aa443

File tree

8 files changed

+85
-28
lines changed

8 files changed

+85
-28
lines changed

‎.github/workflows/pr-build.yml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ jobs:
1111
name: Build Node ${{ matrix.node-version }} on ${{ matrix.os }}
1212
runs-on: ${{ matrix.os }}
1313
strategy:
14+
fail-fast: false
1415
matrix:
1516
os:
1617
- macos-latest

‎packages/core/lib/patchers/aws_p.js‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,10 +85,17 @@ function captureAWSRequest(req) {
8585
}
8686

8787
var traceId = parent.segment ? parent.segment.trace_id : parent.trace_id;
88+
const data = parent.segment ? parent.segment.additionalTraceData : parent.additionalTraceData;
8889

8990
var buildListener = function(req) {
90-
req.httpRequest.headers['X-Amzn-Trace-Id'] = 'Root=' + traceId + ';Parent=' + subsegment.id +
91+
let traceHeader = 'Root=' + traceId + ';Parent=' + subsegment.id +
9192
';Sampled=' + (subsegment.notTraced ? '0' : '1');
93+
if (data != null) {
94+
for (const [key, value] of Object.entries(data)) {
95+
traceHeader += ';' + key +'=' + value;
96+
}
97+
}
98+
req.httpRequest.headers['X-Amzn-Trace-Id'] = traceHeader;
9299
};
93100

94101
var completeListener = function(res) {

‎packages/core/lib/segments/segment.d.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@ declare class Segment {
1313
subsegments?: Array<Subsegment>;
1414
notTraced?: boolean;
1515

16+
additionalTraceData?: object
17+
1618
constructor(name: string, rootId?: string | null, parentId?: string | null);
1719

1820
addIncomingRequestData(data: IncomingRequestData): void;

‎packages/core/lib/utils.js‎

Lines changed: 23 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -162,33 +162,37 @@ var utils = {
162162
*/
163163
populateTraceData: function(segment, xAmznTraceId) {
164164
logger.getLogger().debug('Lambda trace data found: ' + xAmznTraceId);
165-
var data = utils.processTraceData(xAmznTraceId);
165+
let traceData = utils.processTraceData(xAmznTraceId);
166166
var valid = false;
167167

168-
if (!data) {
169-
data = {};
168+
if (!traceData) {
169+
traceData = {};
170170
logger.getLogger().error('_X_AMZN_TRACE_ID is empty or has an invalid format');
171-
} else if (!data.root || !data.parent || !data.sampled) {
171+
} else if (!traceData.root || !traceData.parent || !traceData.sampled) {
172172
logger.getLogger().error('_X_AMZN_TRACE_ID is missing required information');
173173
} else {
174174
valid = true;
175175
}
176176

177-
segment.trace_id = TraceID.FromString(data.root).toString(); // Will always assign valid trace_id
178-
segment.id = data.parent || crypto.randomBytes(8).toString('hex');
177+
segment.trace_id = TraceID.FromString(traceData.root).toString(); // Will always assign valid trace_id
178+
segment.id = traceData.parent || crypto.randomBytes(8).toString('hex');
179179

180-
if (data.root && segment.trace_id !== data.root) {
180+
if (traceData.root && segment.trace_id !== traceData.root) {
181181
logger.getLogger().error('_X_AMZN_TRACE_ID contains invalid trace ID');
182182
valid = false;
183183
}
184184

185-
if (!parseInt(data.sampled)) {
185+
if (!parseInt(traceData.sampled)) {
186186
segment.notTraced = true;
187187
} else {
188188
delete segment.notTraced;
189189
}
190190

191-
logger.getLogger().debug('Segment started: ' + JSON.stringify(data));
191+
if (traceData.data) {
192+
segment.userData = traceData.data;
193+
}
194+
195+
logger.getLogger().debug('Segment started: ' + JSON.stringify(traceData));
192196
return valid;
193197
}
194198
},
@@ -202,6 +206,7 @@ var utils = {
202206

203207
processTraceData: function processTraceData(traceData) {
204208
var amznTraceData = {};
209+
var data = {};
205210
var reservedKeywords = ['root', 'parent', 'sampled', 'self'];
206211
var remainingBytes = 256;
207212

@@ -217,19 +222,22 @@ var utils = {
217222
var pair = header.split('=');
218223

219224
if (pair[0] && pair[1]) {
220-
var key = pair[0].trim().toLowerCase();
221-
var value = pair[1].trim().toLowerCase();
222-
var reserved = reservedKeywords.indexOf(key) !== -1;
225+
let key = pair[0].trim();
226+
let value = pair[1].trim();
227+
let lowerCaseKey = key.toLowerCase();
228+
let reserved = reservedKeywords.indexOf(lowerCaseKey) !== -1;
223229

224230
if (reserved) {
225-
amznTraceData[key] = value;
226-
} else if (!reserved && remainingBytes - (key.length + value.length) >= 0) {
227-
amznTraceData[key] = value;
231+
amznTraceData[lowerCaseKey] = value;
232+
} else if (!reserved && remainingBytes - (lowerCaseKey.length + value.length) >= 0) {
233+
data[key] = value;
228234
remainingBytes -= (key.length + value.length);
229235
}
230236
}
231237
});
232238

239+
amznTraceData['data'] = data;
240+
233241
return amznTraceData;
234242
},
235243

‎packages/core/test/unit/env/aws_lambda.test.js‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,4 +199,31 @@ describe('AWSLambda', function() {
199199
});
200200
});
201201
});
202+
203+
describe('PopulateAdditionalTraceData', function() {
204+
var sandbox, setSegmentStub;
205+
206+
beforeEach(function() {
207+
sandbox = sinon.createSandbox();
208+
sandbox.stub(SegmentEmitter, 'disableReusableSocket');
209+
sandbox.stub(LambdaUtils, 'validTraceData').returns(true);
210+
211+
setSegmentStub = sandbox.stub(contextUtils, 'setSegment');
212+
});
213+
214+
afterEach(function() {
215+
delete process.env._X_AMZN_TRACE_ID;
216+
sandbox.restore();
217+
});
218+
219+
it('should populate additional trace data', function() {
220+
process.env._X_AMZN_TRACE_ID = 'Root=traceId;Lineage=1234abcd:4|3456abcd:6';
221+
Lambda.init();
222+
223+
var facade = setSegmentStub.args[0][0];
224+
facade.resolveLambdaTraceData();
225+
var userData = facade.userData;
226+
assert.equal(userData['Lineage'], '1234abcd:4|3456abcd:6');
227+
});
228+
});
202229
});

‎packages/core/test/unit/middleware/mw_utils.test.js‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -72,23 +72,23 @@ describe('Middleware utils', function() {
7272
req.headers[XRAY_HEADER] = 'Root=' + traceId;
7373
var headers = MWUtils.processHeaders(req);
7474

75-
assert.deepEqual(headers, {root: traceId});
75+
assert.deepEqual(headers, {root: traceId, data: {}});
7676
});
7777

7878
it('should return a split array on an request with an "x-amzn-trace-id" header with a root ID and parent ID', function() {
7979
var req = { headers: {}};
8080
req.headers[XRAY_HEADER] = 'Root=' + traceId + '; Parent=' + parentId;
8181
var headers = MWUtils.processHeaders(req);
8282

83-
assert.deepEqual(headers, {root: traceId, parent: parentId});
83+
assert.deepEqual(headers, {root: traceId, parent: parentId, data: {}});
8484
});
8585

8686
it('should return a split array on an request with an "x-amzn-trace-id" header with a root ID, parent ID and sampling', function() {
8787
var req = { headers: {}};
8888
req.headers[XRAY_HEADER] = 'Root=' + traceId + '; Parent=' + parentId + '; Sampled=0';
8989
var headers = MWUtils.processHeaders(req);
9090

91-
assert.deepEqual(headers, {root: traceId, parent: parentId, sampled: '0'});
91+
assert.deepEqual(headers, {root: traceId, parent: parentId, sampled: '0', data: {}});
9292
});
9393
});
9494

‎packages/core/test/unit/patchers/aws_p.test.js‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,7 @@ describe('AWS patcher', function() {
124124
awsRequest.emitter = new MyEmitter();
125125

126126
segment = new Segment('testSegment', traceId);
127+
segment.additionalTraceData = {'Foo': 'bar'};
127128
sub = segment.addNewSubsegment('subseg');
128129

129130
stubResolveManual = sandbox.stub(contextUtils, 'resolveManualSegmentParams');
@@ -161,7 +162,7 @@ describe('AWS patcher', function() {
161162
awsRequest.emitter.emit('build');
162163

163164
setTimeout(function() {
164-
var expected = new RegExp('^Root=' + traceId + ';Parent=' + sub.id + ';Sampled=1$');
165+
var expected = new RegExp('^Root=' + traceId + ';Parent=' + sub.id + ';Sampled=1' + ';Foo=bar$');
165166
assert.match(awsRequest.httpRequest.headers['X-Amzn-Trace-Id'], expected);
166167
done();
167168
}, 50);
@@ -307,6 +308,7 @@ describe('AWS patcher', function() {
307308
awsRequest.emitter = new MyEmitter();
308309

309310
segment = new Segment('testSegment', traceId);
311+
segment.additionalTraceData = {'Foo': 'bar'};
310312
sub = segment.addNewSubsegmentWithoutSampling('subseg');
311313
service = sub.addNewSubsegmentWithoutSampling('service');
312314

@@ -339,7 +341,7 @@ describe('AWS patcher', function() {
339341
awsRequest.emitter.emit('build');
340342

341343
setTimeout(function() {
342-
var expected = new RegExp('^Root=' + traceId + ';Parent=' + service.id + ';Sampled=0$');
344+
var expected = new RegExp('^Root=' + traceId + ';Parent=' + service.id + ';Sampled=0' + ';Foo=bar$');
343345
assert.match(awsRequest.httpRequest.headers['X-Amzn-Trace-Id'], expected);
344346
done();
345347
}, 50);

‎packages/core/test/unit/utils.test.js‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -75,48 +75,58 @@ describe('Utils', function() {
7575

7676
it('should handle trace header values with excess semicolons correctly', function() {
7777
assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;'), {
78-
root: '1-58ed6027-14afb2e09172c337713486c0'
78+
root: '1-58ed6027-14afb2e09172c337713486c0',
79+
data: {}
7980
});
8081
});
8182

8283
it('should handle malformed key=value pairs correctly (missing value)', function() {
8384
assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;Parent'), {
84-
root: '1-58ed6027-14afb2e09172c337713486c0'
85+
root: '1-58ed6027-14afb2e09172c337713486c0',
86+
data: {}
8587
});
8688
});
8789

8890
it('should handle malformed key=value pairs correctly (empty key)', function() {
8991
assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;=48af77592b6dd73f'), {
90-
root: '1-58ed6027-14afb2e09172c337713486c0'
92+
root: '1-58ed6027-14afb2e09172c337713486c0',
93+
data: {}
9194
});
9295
});
9396

9497
it('should handle malformed key=value pairs correctly (empty value)', function() {
9598
assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;Parent='), {
96-
root: '1-58ed6027-14afb2e09172c337713486c0'
99+
root: '1-58ed6027-14afb2e09172c337713486c0',
100+
data: {}
97101
});
98102
});
99103

100104
it('should accept arbitrary key=value pairs', function() {
101105
assert.deepEqual(Utils.processTraceData('Root=1-58ed6027-14afb2e09172c337713486c0;Foo=bar'), {
102106
root: '1-58ed6027-14afb2e09172c337713486c0',
103-
foo: 'bar'
107+
data: {
108+
Foo: 'bar'
109+
}
104110
});
105111
});
106112

107113
it('should not accept arbitrary key=value pairs that exceed the 256 byte limit', function() {
108114
var longVal = 'a'.repeat(251);
109115
assert.deepEqual(Utils.processTraceData(`Root=1-58ed6027-14afb2e09172c337713486c0;Foo=bar;Baz=${longVal}`), {
110116
root: '1-58ed6027-14afb2e09172c337713486c0',
111-
foo: 'bar'
117+
data: {
118+
Foo: 'bar'
119+
}
112120
});
113121
});
114122

115123
it('should always accept reserved keywords even if unreserved capacity exceeded', function() {
116124
var longVal = 'a'.repeat(251);
117125
assert.deepEqual(Utils.processTraceData(`Baz=${longVal};Root=1-58ed6027-14afb2e09172c337713486c0;Foo=bar`), {
118126
root: '1-58ed6027-14afb2e09172c337713486c0',
119-
baz: longVal
127+
data: {
128+
Baz: longVal
129+
}
120130
});
121131
});
122132
});

0 commit comments

Comments
 (0)