From 882026502dbc12b700b893485458176bcd03f68e Mon Sep 17 00:00:00 2001 From: Maurice Rickard Date: Wed, 18 Dec 2024 16:11:20 -0500 Subject: [PATCH] feat: Added entity linking attributes to aws-sdk v3 Lambda segments (#2845) Signed-off-by: mrickard Co-authored-by: Bob Evans --- lib/instrumentation/aws-sdk/v3/lambda.js | 57 ++++++++++ .../aws-sdk/v3/smithy-client.js | 2 + lib/serverless/aws-lambda.js | 4 +- test/unit/serverless/aws-lambda.test.js | 7 +- test/versioned/aws-sdk-v3/lambda.test.js | 107 +++++++++++++++++- 5 files changed, 167 insertions(+), 10 deletions(-) create mode 100644 lib/instrumentation/aws-sdk/v3/lambda.js diff --git a/lib/instrumentation/aws-sdk/v3/lambda.js b/lib/instrumentation/aws-sdk/v3/lambda.js new file mode 100644 index 0000000000..478c6be046 --- /dev/null +++ b/lib/instrumentation/aws-sdk/v3/lambda.js @@ -0,0 +1,57 @@ +/* + * Copyright 2021 New Relic Corporation. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +'use strict' +const InstrumentationDescriptor = require('../../../instrumentation-descriptor') + +/** + * Defines a deserialize middleware to add the + * cloud.resource_id segment attributes for the AWS command + * + * @param {Shim} shim New Relic agent shim + * @param {Object} config AWS command configuration + * @param {function} next next function in middleware chain + * @returns {function} wrapped version of middleware function + */ +function resourceIdMiddleware(shim, config, next) { + return async function wrappedResourceIdMiddleware(args) { + let result + try { + const region = await config.region() + result = await next(args) + const { response } = result + const segment = shim.getSegment(response.body.req) + // We can't derive account ID, so we have to consume it from config + const accountId = shim.agent.config.cloud.aws.account_id + const functionName = args?.input?.FunctionName // have to get function from params + if (accountId && functionName) { + segment.addAttribute( + 'cloud.resource_id', + `arn:aws:lambda:${region}:${accountId}:function:${functionName}` + ) + segment.addAttribute('cloud.platform', `aws_lambda`) + } + } catch (err) { + shim.logger.debug(err, 'Failed to add AWS cloud resource id to segment') + } finally { + return result + } + } +} + +const lambdaMiddlewareConfig = { + middleware: resourceIdMiddleware, + type: InstrumentationDescriptor.TYPE_GENERIC, + config: { + name: 'NewRelicGetResourceId', + step: 'deserialize', + priority: 'low', + override: true + } +} + +module.exports = { + lambdaMiddlewareConfig +} diff --git a/lib/instrumentation/aws-sdk/v3/smithy-client.js b/lib/instrumentation/aws-sdk/v3/smithy-client.js index a3ce0c3288..b6fe5404b3 100644 --- a/lib/instrumentation/aws-sdk/v3/smithy-client.js +++ b/lib/instrumentation/aws-sdk/v3/smithy-client.js @@ -9,12 +9,14 @@ const { middlewareConfig } = require('./common') const { snsMiddlewareConfig } = require('./sns') const { sqsMiddlewareConfig } = require('./sqs') const { dynamoMiddlewareConfig } = require('./dynamodb') +const { lambdaMiddlewareConfig } = require('./lambda') const { bedrockMiddlewareConfig } = require('./bedrock') const MIDDLEWARE = Symbol('nrMiddleware') const middlewareByClient = { Client: middlewareConfig, BedrockRuntime: [...middlewareConfig, bedrockMiddlewareConfig], + Lambda: [...middlewareConfig, lambdaMiddlewareConfig], SNS: [...middlewareConfig, snsMiddlewareConfig], SQS: [...middlewareConfig, sqsMiddlewareConfig], DynamoDB: [...middlewareConfig, ...dynamoMiddlewareConfig], diff --git a/lib/serverless/aws-lambda.js b/lib/serverless/aws-lambda.js index 8df11be080..3d3d2843e1 100644 --- a/lib/serverless/aws-lambda.js +++ b/lib/serverless/aws-lambda.js @@ -249,9 +249,7 @@ class AwsLambda { _getAwsAgentAttributes(event, context) { const attributes = { 'aws.lambda.arn': context.invokedFunctionArn, - 'aws.requestId': context.awsRequestId, - 'cloud.resource_id': context.invokedFunctionArn, - 'cloud.platform': 'aws_lambda' + 'aws.requestId': context.awsRequestId } const eventSourceInfo = this._detectEventType(event) diff --git a/test/unit/serverless/aws-lambda.test.js b/test/unit/serverless/aws-lambda.test.js index e96ece22b8..3f2d81b80d 100644 --- a/test/unit/serverless/aws-lambda.test.js +++ b/test/unit/serverless/aws-lambda.test.js @@ -25,8 +25,6 @@ const LAMBDA_ARN = 'aws.lambda.arn' const COLDSTART = 'aws.lambda.coldStart' const EVENTSOURCE_ARN = 'aws.lambda.eventSource.arn' const EVENTSOURCE_TYPE = 'aws.lambda.eventSource.eventType' -const PLATFORM = 'cloud.platform' -const RESOURCE_ID = 'cloud.resource_id' function getMetrics(agent) { return agent.metrics._metrics @@ -64,8 +62,7 @@ test('AwsLambda.patchLambdaHandler', async (t) => { functionVersion: 'TestVersion', invokedFunctionArn: 'arn:test:function', memoryLimitInMB: '128', - awsRequestId: 'testid', - platform: 'aws_lambda' + awsRequestId: 'testid' } ctx.nr.stubCallback = () => {} @@ -669,8 +666,6 @@ test('AwsLambda.patchLambdaHandler', async (t) => { assert.equal(txTrace[REQ_ID], stubContext.awsRequestId) assert.equal(txTrace[LAMBDA_ARN], stubContext.invokedFunctionArn) assert.equal(txTrace[COLDSTART], true) - assert.equal(txTrace[PLATFORM], stubContext.platform) - assert.equal(txTrace[RESOURCE_ID], stubContext.invokedFunctionArn) end() } diff --git a/test/versioned/aws-sdk-v3/lambda.test.js b/test/versioned/aws-sdk-v3/lambda.test.js index e52f61c65a..f163bc1e92 100644 --- a/test/versioned/aws-sdk-v3/lambda.test.js +++ b/test/versioned/aws-sdk-v3/lambda.test.js @@ -7,8 +7,76 @@ const test = require('node:test') const helper = require('../../lib/agent_helper') -const { afterEach, checkExternals } = require('./common') +const assert = require('node:assert') +const { + afterEach, + checkExternals, + checkAWSAttributes, + EXTERN_PATTERN, + SEGMENT_DESTINATION +} = require('./common') const { createEmptyResponseServer, FAKE_CREDENTIALS } = require('../../lib/aws-server-stubs') +const { match } = require('../../lib/custom-assertions') + +function checkEntityLinkingSegments({ operations, tx, end }) { + const root = tx.trace.root + + const segments = checkAWSAttributes(root, EXTERN_PATTERN) + const accountId = tx.agent.config.cloud.aws.account_id + const testFunctionName = 'funcName' + + assert(segments.length > 0, 'should have segments') + assert.ok(accountId, 'account id should be set on agent config') + + segments.forEach((segment) => { + const attrs = segment.attributes.get(SEGMENT_DESTINATION) + + match(attrs, { + 'aws.operation': operations[0], + 'aws.requestId': String, + 'aws.region': 'us-east-1', + 'aws.service': String, + 'cloud.resource_id': `arn:aws:lambda:${attrs['aws.region']}:${accountId}:function:${testFunctionName}`, + 'cloud.platform': `aws_lambda` + }) + }) + end() +} + +function checkNonLinkableSegments({ operations, tx, end }) { + // When no account ID or ARN is available, make sure not to set cloud resource id or platform + const root = tx.trace.root + + const segments = checkAWSAttributes(root, EXTERN_PATTERN) + const accountId = tx.agent.config?.cloud?.aws?.account_id + + assert(segments.length > 0, 'should have segments') + assert.equal(accountId, undefined, 'account id should not have been set for this test') + + segments.forEach((segment) => { + const attrs = segment.attributes.get(SEGMENT_DESTINATION) + + assert.equal( + attrs['cloud.resource_id'], + undefined, + 'if account Id has not been set, cloud.resource_id should not be set' + ) + assert.equal( + attrs['cloud.platform'], + undefined, + 'if account Id has not been set, cloud.platform should not be set' + ) + + // other attributes should be as expected + match(attrs, { + 'aws.operation': operations[0], + 'aws.requestId': String, + 'aws.region': 'us-east-1', + 'aws.service': String + }) + }) + end() +} test('LambdaClient', async (t) => { t.beforeEach(async (ctx) => { @@ -21,6 +89,7 @@ test('LambdaClient', async (t) => { ctx.nr.agent = helper.instrumentMockedAgent() const { LambdaClient, ...lib } = require('@aws-sdk/client-lambda') ctx.nr.AddLayerVersionPermissionCommand = lib.AddLayerVersionPermissionCommand + ctx.nr.InvokeCommand = lib.InvokeCommand const endpoint = `http://localhost:${server.address().port}` ctx.nr.service = new LambdaClient({ credentials: FAKE_CREDENTIALS, @@ -53,4 +122,40 @@ test('LambdaClient', async (t) => { }) }) }) + + await t.test('InvokeCommand', (t, end) => { + const { service, agent, InvokeCommand } = t.nr + agent.config.cloud.aws.account_id = 123456789123 + helper.runInTransaction(agent, async (tx) => { + const cmd = new InvokeCommand({ + FunctionName: 'funcName', + Payload: JSON.stringify({ prop1: 'test', prop2: 'test 2' }) + }) + await service.send(cmd) + tx.end() + setImmediate(checkEntityLinkingSegments, { + operations: ['InvokeCommand'], + tx, + end + }) + }) + }) + + await t.test('InvokeCommand without account ID defined', (t, end) => { + const { service, agent, InvokeCommand } = t.nr + agent.config.cloud.aws.account_id = null + helper.runInTransaction(agent, async (tx) => { + const cmd = new InvokeCommand({ + FunctionName: 'funcName', + Payload: JSON.stringify({ prop1: 'test', prop2: 'test 2' }) + }) + await service.send(cmd) + tx.end() + setImmediate(checkNonLinkableSegments, { + operations: ['InvokeCommand'], + tx, + end + }) + }) + }) })