feat: Added entity linking attributes to aws-sdk v3 Lambda segments (#…

…2845)

Signed-off-by: mrickard <[email protected]>
Co-authored-by: Bob Evans <[email protected]>

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
+}

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],

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)

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()
     }
 

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
+      })
+    })
+  })
 })