diff --git a/package-lock.json b/package-lock.json index c6a1547a..abe81b2c 100644 --- a/package-lock.json +++ b/package-lock.json @@ -76,6 +76,14 @@ "npm": ">= 2.x" } }, + "node_modules/@aws/lambda-invoke-store": { + "version": "0.0.1", + "resolved": "https://registry.npmjs.org/@aws/lambda-invoke-store/-/lambda-invoke-store-0.0.1.tgz", + "integrity": "sha512-ORHRQ2tmvnBXc8t/X9Z8IcSbBA4xTLKuN873FopzklHMeqBst7YG0d+AX97inkvDX+NChYtSr+qGfcqGFaI8Zw==", + "engines": { + "node": ">=18.0.0" + } + }, "node_modules/@babel/code-frame": { "version": "7.12.11", "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.12.11.tgz", @@ -19341,6 +19349,7 @@ "license": "Apache-2.0", "dependencies": { "@aws-sdk/types": "^3.4.1", + "@aws/lambda-invoke-store": "^0.0.1", "@smithy/service-error-classification": "^2.0.4", "@types/cls-hooked": "^4.3.3", "atomic-batcher": "^1.0.2", @@ -19747,6 +19756,11 @@ } }, "dependencies": { + "@aws/lambda-invoke-store": { + "version": "0.0.1", + "resolved": "https://registry.npmjs.org/@aws/lambda-invoke-store/-/lambda-invoke-store-0.0.1.tgz", + "integrity": "sha512-ORHRQ2tmvnBXc8t/X9Z8IcSbBA4xTLKuN873FopzklHMeqBst7YG0d+AX97inkvDX+NChYtSr+qGfcqGFaI8Zw==" + }, "@babel/code-frame": { "version": "7.12.11", "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.12.11.tgz", @@ -23840,6 +23854,7 @@ "version": "file:packages/core", "requires": { "@aws-sdk/types": "^3.4.1", + "@aws/lambda-invoke-store": "^0.0.1", "@smithy/service-error-classification": "^2.0.4", "@types/cls-hooked": "^4.3.3", "atomic-batcher": "^1.0.2", @@ -24133,6 +24148,11 @@ "upath": "^1.2.0" }, "dependencies": { + "@aws/lambda-invoke-store": { + "version": "0.0.1", + "resolved": "https://registry.npmjs.org/@aws/lambda-invoke-store/-/lambda-invoke-store-0.0.1.tgz", + "integrity": "sha512-ORHRQ2tmvnBXc8t/X9Z8IcSbBA4xTLKuN873FopzklHMeqBst7YG0d+AX97inkvDX+NChYtSr+qGfcqGFaI8Zw==" + }, "@babel/code-frame": { "version": "7.12.11", "resolved": "https://registry.npmjs.org/@babel/code-frame/-/code-frame-7.12.11.tgz", @@ -28226,6 +28246,7 @@ "version": "file:packages/core", "requires": { "@aws-sdk/types": "^3.4.1", + "@aws/lambda-invoke-store": "^0.0.1", "@smithy/service-error-classification": "^2.0.4", "@types/cls-hooked": "^4.3.3", "atomic-batcher": "^1.0.2", diff --git a/packages/core/lib/env/aws_lambda.js b/packages/core/lib/env/aws_lambda.js index 18ed0f12..fcdd10a5 100644 --- a/packages/core/lib/env/aws_lambda.js +++ b/packages/core/lib/env/aws_lambda.js @@ -8,6 +8,8 @@ var SegmentUtils = require('../segments/segment_utils'); var logger = require('../logger'); const TraceID = require('../segments/attributes/trace_id'); +const { InvokeStore } = require('@aws/lambda-invoke-store'); + /** * @namespace * @ignore @@ -80,7 +82,9 @@ var facadeSegment = function facadeSegment() { }; segment.resolveLambdaTraceData = function resolveLambdaTraceData() { - var xAmznLambda = process.env._X_AMZN_TRACE_ID; + const traceIdFromInvokeStore = InvokeStore.getXRayTraceId(); + const traceIdFromEnv = process.env._X_AMZN_TRACE_ID; + var xAmznLambda = traceIdFromInvokeStore ?? traceIdFromEnv; if (xAmznLambda) { diff --git a/packages/core/package.json b/packages/core/package.json index dd9148fd..1b49dd60 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -24,6 +24,7 @@ "//": "@types/cls-hooked is exposed in API so must be in dependencies, not devDependencies", "dependencies": { "@aws-sdk/types": "^3.4.1", + "@aws/lambda-invoke-store": "^0.0.1", "@smithy/service-error-classification": "^2.0.4", "@types/cls-hooked": "^4.3.3", "atomic-batcher": "^1.0.2", @@ -55,4 +56,4 @@ ], "license": "Apache-2.0", "repository": "https://github.com/aws/aws-xray-sdk-node/tree/master/packages/core" -} \ No newline at end of file +} diff --git a/packages/core/test/unit/env/aws_lambda.test.js b/packages/core/test/unit/env/aws_lambda.test.js index eb7e6fb1..d2ba41d9 100644 --- a/packages/core/test/unit/env/aws_lambda.test.js +++ b/packages/core/test/unit/env/aws_lambda.test.js @@ -3,6 +3,7 @@ var chai = require('chai'); var sinon = require('sinon'); var sinonChai = require('sinon-chai'); +chai.should(); chai.use(sinonChai); var contextUtils = require('../../../lib/context_utils'); @@ -13,6 +14,7 @@ var Segment = require('../../../lib/segments/segment'); var SegmentUtils = require('../../../lib/segments/segment_utils'); var SegmentEmitter = require('../../../lib/segment_emitter'); const TraceID = require('../../../lib/segments/attributes/trace_id'); +const { InvokeStore } = require('@aws/lambda-invoke-store'); describe('AWSLambda', function() { var sandbox; @@ -147,6 +149,8 @@ describe('AWSLambda', function() { describe('#resolveLambdaTraceData', function() { var sandbox, traceId; + var INVOKE_STORE_TRACE_ID = 'Root=1-12345678-12345678901234567890abcd;Parent=abcdef0123456789;Sampled=1'; + var ENV_TRACE_ID = 'Root=1-87654321-09876543210987654321dcba;Parent=fedcba9876543210;Sampled=0'; beforeEach(function() { sandbox = sinon.createSandbox(); @@ -197,6 +201,47 @@ describe('AWSLambda', function() { facade.resolveLambdaTraceData(); populateStub.should.have.not.been.called; }); + + it('should prioritize InvokeStore trace ID over environment variable if both are defined', function() { + process.env._X_AMZN_TRACE_ID = ENV_TRACE_ID; + Lambda.init(); + + sandbox.stub(InvokeStore, 'getXRayTraceId').returns(INVOKE_STORE_TRACE_ID); + populateStub.reset(); + + var facade = setSegmentStub.args[0][0]; + facade.resolveLambdaTraceData(); + + populateStub.should.have.been.calledWith(facade, INVOKE_STORE_TRACE_ID); + }); + + it('should use InvokeStore trace ID when environment variable is undefined', function() { + delete process.env._X_AMZN_TRACE_ID; + + Lambda.init(); + + sandbox.stub(InvokeStore, 'getXRayTraceId').returns(INVOKE_STORE_TRACE_ID); + populateStub.reset(); + + var facade = setSegmentStub.args[0][0]; + facade.resolveLambdaTraceData(); + + populateStub.should.have.been.calledWith(facade, INVOKE_STORE_TRACE_ID); + }); + + it('should use environment variable when InvokeStore returns undefined', function() { + process.env._X_AMZN_TRACE_ID = INVOKE_STORE_TRACE_ID; + Lambda.init(); + + process.env._X_AMZN_TRACE_ID = ENV_TRACE_ID; + sandbox.stub(InvokeStore, 'getXRayTraceId').returns(undefined); + populateStub.reset(); + + var facade = setSegmentStub.args[0][0]; + facade.resolveLambdaTraceData(); + + populateStub.should.have.been.calledWith(facade, ENV_TRACE_ID); + }); }); });