Fix: Lambda Topology Issue

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py

@@ -7,6 +7,7 @@
 AWS_REMOTE_DB_USER: str = "aws.remote.db.user"
 AWS_REMOTE_SERVICE: str = "aws.remote.service"
 AWS_REMOTE_OPERATION: str = "aws.remote.operation"
+AWS_REMOTE_ENVIRONMENT: str = "aws.remote.environment"
 AWS_REMOTE_RESOURCE_TYPE: str = "aws.remote.resource.type"
 AWS_REMOTE_RESOURCE_IDENTIFIER: str = "aws.remote.resource.identifier"
 AWS_SDK_DESCENDANT: str = "aws.sdk.descendant"

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py

@@ -1,5 +1,6 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
+import os
 import re
 from logging import DEBUG, Logger, getLogger
 from typing import Match, Optional
@@ -13,10 +14,13 @@
     AWS_BEDROCK_KNOWLEDGE_BASE_ID,
     AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
     AWS_KINESIS_STREAM_NAME,
+    AWS_LAMBDA_FUNCTION_ARN,
+    AWS_LAMBDA_FUNCTION_NAME,
     AWS_LAMBDA_RESOURCEMAPPING_ID,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
     AWS_REMOTE_DB_USER,
+    AWS_REMOTE_ENVIRONMENT,
     AWS_REMOTE_OPERATION,
     AWS_REMOTE_RESOURCE_IDENTIFIER,
     AWS_REMOTE_RESOURCE_TYPE,
@@ -445,6 +449,28 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
                 ":"
             )[-1]
             cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_ARN))
+        elif is_key_present(span, 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 span.attributes.get(_RPC_METHOD) == "Invoke":
+                attributes[AWS_REMOTE_SERVICE] = os.environ.get(
+                    "LAMBDA_APPLICATION_SIGNALS_REMOTE_SERVICE", span.attributes.get(AWS_LAMBDA_FUNCTION_NAME)
+                )
+                attributes[AWS_REMOTE_ENVIRONMENT] = (
+                    f'lambda:{os.environ.get("LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT", "default")}'
+                )
+            else:
+                remote_resource_type = _NORMALIZED_LAMBDA_SERVICE_NAME + "::Function"
+                remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_FUNCTION_NAME))
+                cloudformation_primary_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_FUNCTION_ARN))
         elif is_key_present(span, AWS_LAMBDA_RESOURCEMAPPING_ID):
             remote_resource_type = _NORMALIZED_LAMBDA_SERVICE_NAME + "::EventSourceMapping"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_LAMBDA_RESOURCEMAPPING_ID))

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py

@@ -3,6 +3,7 @@
 
 # pylint: disable=too-many-lines
 
+import os
 from typing import Dict, List, Optional
 from unittest import TestCase
 from unittest.mock import MagicMock
@@ -12,12 +13,16 @@
     AWS_BEDROCK_DATA_SOURCE_ID,
     AWS_BEDROCK_GUARDRAIL_ID,
     AWS_BEDROCK_KNOWLEDGE_BASE_ID,
+    AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
     AWS_CONSUMER_PARENT_SPAN_KIND,
     AWS_KINESIS_STREAM_NAME,
+    AWS_LAMBDA_FUNCTION_ARN,
+    AWS_LAMBDA_FUNCTION_NAME,
     AWS_LAMBDA_RESOURCEMAPPING_ID,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
     AWS_REMOTE_DB_USER,
+    AWS_REMOTE_ENVIRONMENT,
     AWS_REMOTE_OPERATION,
     AWS_REMOTE_RESOURCE_IDENTIFIER,
     AWS_REMOTE_RESOURCE_TYPE,
@@ -1152,6 +1157,78 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         self._validate_remote_resource_attributes("AWS::Lambda::EventSourceMapping", "aws_event_source_mapping_id")
         self._mock_attribute([AWS_LAMBDA_RESOURCEMAPPING_ID], [None])
 
+        # Test AWS Lambda Invoke scenario with default lambda remote environment
+        self.span_mock.kind = SpanKind.CLIENT
+        self._mock_attribute(
+            [AWS_LAMBDA_FUNCTION_NAME, SpanAttributes.RPC_METHOD],
+            ["test_downstream_lambda1", "Invoke"],
+            keys,
+            values,
+        )
+        dependency_attributes = _GENERATOR.generate_metric_attributes_dict_from_span(self.span_mock, self.resource).get(
+            DEPENDENCY_METRIC
+        )
+        self.assertEqual(dependency_attributes.get(AWS_REMOTE_SERVICE), "test_downstream_lambda1")
+        self.assertEqual(dependency_attributes.get(AWS_REMOTE_ENVIRONMENT), "lambda:default")
+        self.assertNotIn(AWS_REMOTE_RESOURCE_TYPE, dependency_attributes)
+        self.assertNotIn(AWS_REMOTE_RESOURCE_IDENTIFIER, dependency_attributes)
+        self.assertNotIn(AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER, dependency_attributes)
+        self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME, SpanAttributes.RPC_METHOD], [None, None])
+
+        # Test AWS Lambda Invoke scenario with user-configured lambda remote environment
+        os.environ["LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT"] = "test"
+        self.span_mock.kind = SpanKind.CLIENT
+        self._mock_attribute(
+            [AWS_LAMBDA_FUNCTION_NAME, SpanAttributes.RPC_METHOD],
+            ["testLambdaFunction", "Invoke"],
+            keys,
+            values,
+        )
+        dependency_attributes = _GENERATOR.generate_metric_attributes_dict_from_span(self.span_mock, self.resource).get(
+            DEPENDENCY_METRIC
+        )
+        self.assertEqual(dependency_attributes.get(AWS_REMOTE_SERVICE), "testLambdaFunction")
+        self.assertEqual(dependency_attributes.get(AWS_REMOTE_ENVIRONMENT), "lambda:test")
+        self.assertNotIn(AWS_REMOTE_RESOURCE_TYPE, dependency_attributes)
+        self.assertNotIn(AWS_REMOTE_RESOURCE_IDENTIFIER, dependency_attributes)
+        self.assertNotIn(AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER, dependency_attributes)
+        self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME, SpanAttributes.RPC_METHOD], [None, None])
+        os.environ.pop("LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT", None)
+
+        # Test AWS Lambda Invoke scenario with user-configured lambda remote service
+        os.environ["LAMBDA_APPLICATION_SIGNALS_REMOTE_SERVICE"] = "test_downstream_lambda2"
+        self.span_mock.kind = SpanKind.CLIENT
+        self._mock_attribute(
+            [AWS_LAMBDA_FUNCTION_NAME, SpanAttributes.RPC_METHOD],
+            ["testLambdaFunction", "Invoke"],
+            keys,
+            values,
+        )
+        dependency_attributes = _GENERATOR.generate_metric_attributes_dict_from_span(self.span_mock, self.resource).get(
+            DEPENDENCY_METRIC
+        )
+        self.assertEqual(dependency_attributes.get(AWS_REMOTE_SERVICE), "test_downstream_lambda2")
+        self.assertEqual(dependency_attributes.get(AWS_REMOTE_ENVIRONMENT), "lambda:default")
+        self.assertNotIn(AWS_REMOTE_RESOURCE_TYPE, dependency_attributes)
+        self.assertNotIn(AWS_REMOTE_RESOURCE_IDENTIFIER, dependency_attributes)
+        self.assertNotIn(AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER, dependency_attributes)
+        self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME, SpanAttributes.RPC_METHOD], [None, None])
+        os.environ.pop("LAMBDA_APPLICATION_SIGNALS_REMOTE_SERVICE", None)
+
+        # Test AWS Lambda non-Invoke scenario
+        self.span_mock.kind = SpanKind.CLIENT
+        lambda_arn = "arn:aws:lambda:us-east-1:123456789012:function:testLambda"
+        self._mock_attribute(
+            [AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_FUNCTION_ARN, SpanAttributes.RPC_METHOD],
+            ["testLambdaFunction", lambda_arn, "GetFunction"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes("AWS::Lambda::Function", "testLambdaFunction")
+        self._mock_attribute(
+            [AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_FUNCTION_ARN, SpanAttributes.RPC_METHOD], [None, None, None]
+        )
+
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
 
     def test_client_db_span_with_remote_resource_attributes(self):