Fix: Lambda Topology Issue (#149)

**Issue #, if available:**
Lambda Topology Issue Context:
aws-observability/aws-otel-python-instrumentation#319

**Description of changes:**
- Apply fix for the Lambda Topology issue. The logic mimics the fix in
our ADOT Python SDK.
- Adding back logic to generate `aws.lambda.function.name` and
`aws.lambda.function.arn` span attributes for Lambda. Previously, we
removed these changes in the AWS Resources expansion due to the Lambda
Topology issue.
  - More context about adding support for Lambda as AWS Resource: 
-
aws-observability/aws-otel-python-instrumentation#265

**Test plan:**

The same cases in the test plan of this
[PR](aws-observability/aws-otel-python-instrumentation#319)
were validated by building a custom JavaScript lambda layer. Below is a
screenshot of all the test cases in a single topology.

We can observe the following:
- Service entity node for `Invoke` call to downstream lambda.
- AWS Resource node for `GetFunction` call to downstream lambda.
- AWS Resource node for `ListBuckets` call to downstream s3.
- Another AWS Resource node for `GetBucketLocation` call to downstream
s3.

<img width="1359" alt="Screenshot 2025-02-05 at 10 15 26 AM"
src="https://github.com/user-attachments/assets/14d38cad-32c7-4e2c-a987-c424c7fb7296"
/>

Here we see downstream lambda called with `Invoke` is correctly treated
as RemoteService entity when not instrumented:

![image](https://github.com/user-attachments/assets/7f1cfcec-8827-4dc7-beaf-2a1b9b98e4f4)

Additionally, I generated the spans locally to validate the new lambda
instrumentation patch behaves correctly.

**Generated span for `Invoke` call**

We see the new `aws.remote.environment` attribute is present in span so
the downstream lambda will be treated as a Service type in Application
Signals backend.

<img width="1124" alt="invoke"
src="https://github.com/user-attachments/assets/81cefb44-c951-4ad0-8aeb-9a2064d7d0ea"
/>


**Generated span for `GetFunction` call**

We see `aws.remote.resource.identifier` and
`aws.cloudformation.primary.identifier` are set so the downstream lambda
will be treated as an AWS resource type in Application Signals backend.

<img width="1124" alt="get-function"
src="https://github.com/user-attachments/assets/abfd700b-7966-4d93-bb47-868c31be9a34"
/>


By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Author: yiyuan-he

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-attribute-keys.ts

@@ -9,6 +9,7 @@ export const AWS_ATTRIBUTE_KEYS: { [key: string]: string } = {
   AWS_LOCAL_SERVICE: 'aws.local.service',
   AWS_LOCAL_OPERATION: 'aws.local.operation',
   AWS_REMOTE_SERVICE: 'aws.remote.service',
+  AWS_REMOTE_ENVIRONMENT: 'aws.remote.environment',
   AWS_REMOTE_OPERATION: 'aws.remote.operation',
   AWS_REMOTE_RESOURCE_TYPE: 'aws.remote.resource.type',
   AWS_REMOTE_RESOURCE_IDENTIFIER: 'aws.remote.resource.identifier',

aws-distro-opentelemetry-node-autoinstrumentation/src/aws-metric-attribute-generator.ts

@@ -409,6 +409,33 @@ export class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
           this.extractResourceNameFromArn(activityArn)
         );
         cloudFormationIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(activityArn);
+      } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_NAME)) {
+        // Handling downstream Lambda as a service vs. an AWS resource:
+        // - If the method call is "Invoke", we treat downstream Lambda as a service.
+        // - Otherwise, we treat it as an AWS resource.
+        //
+        // This addresses a Lambda topology issue in Application Signals.
+        // More context in PR: https://github.com/aws-observability/aws-otel-python-instrumentation/pull/319
+        //
+        // NOTE: The env vars LAMBDA_APPLICATION_SIGNALS_REMOTE_SERVICE and
+        // LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT were introduced as part of this fix.
+        // They are optional and allow users to override the default values if needed.
+        if (AwsMetricAttributeGenerator.getRemoteOperation(span, SEMATTRS_RPC_METHOD) === 'Invoke') {
+          attributes[AWS_ATTRIBUTE_KEYS.AWS_REMOTE_SERVICE] =
+            process.env.LAMBDA_APPLICATION_SIGNALS_REMOTE_SERVICE ||
+            span.attributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_NAME];
+          attributes[AWS_ATTRIBUTE_KEYS.AWS_REMOTE_ENVIRONMENT] = `lambda:${
+            process.env.LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT || 'default'
+          }`;
+        } else {
+          remoteResourceType = NORMALIZED_LAMBDA_SERVICE_NAME + '::Function';
+          remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+            span.attributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_NAME]
+          );
+          cloudFormationIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(
+            span.attributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_ARN]
+          );
+        }
       } else if (AwsSpanProcessingUtil.isKeyPresent(span, AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_RESOURCE_MAPPING_ID)) {
         remoteResourceType = NORMALIZED_LAMBDA_SERVICE_NAME + '::EventSourceMapping';
         remoteResourceIdentifier = AwsMetricAttributeGenerator.escapeDelimiters(

aws-distro-opentelemetry-node-autoinstrumentation/src/patches/instrumentation-patch.ts

@@ -10,9 +10,15 @@ import {
   ROOT_CONTEXT,
   TextMapGetter,
   trace,
+  Span,
+  Tracer,
 } from '@opentelemetry/api';
 import { Instrumentation } from '@opentelemetry/instrumentation';
-import { AwsSdkInstrumentationConfig, NormalizedRequest } from '@opentelemetry/instrumentation-aws-sdk';
+import {
+  AwsSdkInstrumentationConfig,
+  NormalizedRequest,
+  NormalizedResponse,
+} from '@opentelemetry/instrumentation-aws-sdk';
 import { AWSXRAY_TRACE_ID_HEADER, AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray';
 import { APIGatewayProxyEventHeaders, Context } from 'aws-lambda';
 import { AWS_ATTRIBUTE_KEYS } from '../aws-attribute-keys';
@@ -215,10 +221,41 @@ function patchLambdaServiceExtension(lambdaServiceExtension: any): void {
         if (resourceMappingId) {
           requestMetadata.spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_RESOURCE_MAPPING_ID] = resourceMappingId;
         }
+
+        const requestFunctionNameFormat = request.commandInput?.FunctionName;
+        let functionName = requestFunctionNameFormat;
+
+        if (requestFunctionNameFormat) {
+          if (requestFunctionNameFormat.startsWith('arn:aws:lambda')) {
+            const split = requestFunctionNameFormat.split(':');
+            functionName = split[split.length - 1];
+          }
+          requestMetadata.spanAttributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_NAME] = functionName;
+        }
       }
       return requestMetadata;
     };
 
     lambdaServiceExtension.requestPreSpanHook = patchedRequestPreSpanHook;
+
+    if (typeof lambdaServiceExtension.responseHook === 'function') {
+      const originalResponseHook = lambdaServiceExtension.responseHook;
+
+      lambdaServiceExtension.responseHook = (
+        response: NormalizedResponse,
+        span: Span,
+        tracer: Tracer,
+        config: AwsSdkInstrumentationConfig
+      ): void => {
+        originalResponseHook.call(lambdaServiceExtension, response, span, tracer, config);
+
+        if (response.data && response.data.Configuration) {
+          const functionArn = response.data.Configuration.FunctionArn;
+          if (functionArn) {
+            span.setAttribute(AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_ARN, functionArn);
+          }
+        }
+      };
+    }
   }
 }

aws-distro-opentelemetry-node-autoinstrumentation/test/aws-metric-attribute-generator.test.ts

@@ -774,6 +774,16 @@ describe('AwsMetricAttributeGeneratorTest', () => {
     validateRemoteResourceAttributes('AWS::SecretsManager::Secret', 'testSecret');
     mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_SECRETSMANAGER_SECRET_ARN, undefined);
 
+    // Validate behaviour of AWS_LAMBDA_FUNCTION_NAME and AWS_LAMBDA_FUNCTION_ARN
+    mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_NAME, 'aws_lambda_function_name');
+    mockAttribute(
+      AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_ARN,
+      'arn:aws:lambda:us-east-1:123456789012:function:aws_lambda_function_name'
+    );
+    validateRemoteResourceAttributes('AWS::Lambda::Function', 'aws_lambda_function_name');
+    mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_NAME, undefined);
+    mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_ARN, undefined);
+
     // Validate behaviour of AWS_LAMBDA_RESOURCE_MAPPING_ID attribute then remove it.
     mockAttribute(AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_RESOURCE_MAPPING_ID, 'aws_lambda_resource_mapping_id');
     validateRemoteResourceAttributes('AWS::Lambda::EventSourceMapping', 'aws_lambda_resource_mapping_id');

aws-distro-opentelemetry-node-autoinstrumentation/test/patches/instrumentation-patch.test.ts

@@ -38,6 +38,8 @@ const _SECRETS_ARN: string = 'arn:aws:secretsmanager:us-east-1:123456789123:secr
 const _UUID: string = 'random-uuid';
 const _TOPIC_ARN: string = 'arn:aws:sns:us-east-1:123456789012:mystack-mytopic-NZJ5JSMVGFIE';
 const _QUEUE_URL: string = 'https://sqs.us-east-1.amazonaws.com/123412341234/queueName';
+const _FUNCTION_NAME: string = 'testFunction';
+const _FUNCTION_ARN: string = `arn:aws:lambda:us-east-1:123456789012:function:${_FUNCTION_NAME}`;
 const _BEDROCK_AGENT_ID: string = 'agentId';
 const _BEDROCK_DATASOURCE_ID: string = 'DataSourceId';
 const _BEDROCK_GUARDRAIL_ID: string = 'GuardrailId';
@@ -165,6 +167,8 @@ describe('InstrumentationPatchTest', () => {
 
     const lambdaAttributes: Attributes = doExtractLambdaAttributes(services);
     expect(lambdaAttributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_RESOURCE_MAPPING_ID]).toBeUndefined();
+    expect(lambdaAttributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_NAME]).toBeUndefined();
+    expect(lambdaAttributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_ARN]).toBeUndefined();
   });
 
   it('SFN without patching', () => {
@@ -228,6 +232,9 @@ describe('InstrumentationPatchTest', () => {
     const services: Map<string, any> = extractServicesFromAwsSdkInstrumentation(patchedAwsSdkInstrumentation);
     const requestLambdaAttributes: Attributes = doExtractLambdaAttributes(services);
     expect(requestLambdaAttributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_RESOURCE_MAPPING_ID]).toEqual(_UUID);
+    expect(requestLambdaAttributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_NAME]).toEqual(_FUNCTION_NAME);
+    const responseLambdaAttributes: Attributes = doResponseHookLambda(services);
+    expect(responseLambdaAttributes[AWS_ATTRIBUTE_KEYS.AWS_LAMBDA_FUNCTION_ARN]).toEqual(_FUNCTION_ARN);
   });
 
   it('SFN with patching', () => {
@@ -429,6 +436,7 @@ describe('InstrumentationPatchTest', () => {
       commandName: 'mockCommandName',
       commandInput: {
         UUID: _UUID,
+        FunctionName: _FUNCTION_NAME,
       },
     };
     return doExtractAttributes(services, serviceName, params);
@@ -507,6 +515,23 @@ describe('InstrumentationPatchTest', () => {
     return doResponseHook(services, 'SecretsManager', results as NormalizedResponse);
   }
 
+  function doResponseHookLambda(services: Map<string, ServiceExtension>): Attributes {
+    const results: Partial<NormalizedResponse> = {
+      data: {
+        Configuration: {
+          FunctionArn: _FUNCTION_ARN,
+        },
+      },
+      request: {
+        commandInput: {},
+        commandName: 'dummy_operation',
+        serviceName: 'Lambda',
+      },
+    };
+
+    return doResponseHook(services, 'Lambda', results as NormalizedResponse);
+  }
+
   function doResponseHookBedrock(
     services: Map<string, ServiceExtension>,
     serviceName: string,