feat: Add auto-instrumentation support for Lambda

Author: yiyuan-he

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

@@ -28,3 +28,6 @@
 AWS_SNS_TOPIC_ARN: str = "aws.sns.topic.arn"
 AWS_STEPFUNCTIONS_STATEMACHINE_ARN: str = "aws.stepfunctions.state_machine.arn"
 AWS_STEPFUNCTIONS_ACTIVITY_ARN: str = "aws.stepfunctions.activity.arn"
+AWS_LAMBDA_FUNCTION_NAME: str = "aws.lambda.function.name"
+AWS_LAMBDA_RESOURCEMAPPING_ID: str = "aws.lambda.resource_mapping.id"
+AWS_LAMBDA_FUNCTION_ARN: str = "aws.lambda.function.arn"

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

@@ -13,6 +13,9 @@
     AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER,
     AWS_KINESIS_STREAM_CONSUMERNAME,
     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,
@@ -96,6 +99,7 @@
 _NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager"
 _NORMALIZED_SNS_SERVICE_NAME: str = "AWS::SNS"
 _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME: str = "AWS::StepFunctions"
+_NORMALIZED_LAMBDA_SERVICE_NAME: str = "AWS::Lambda"
 _DB_CONNECTION_STRING_TYPE: str = "DB::Connection"
 
 # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
@@ -442,6 +446,13 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
         elif is_key_present(span, AWS_STEPFUNCTIONS_ACTIVITY_ARN):
             remote_resource_type = _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME + "::Activity"
             remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_STEPFUNCTIONS_ACTIVITY_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))
+        elif is_key_present(span, AWS_LAMBDA_FUNCTION_NAME):
+            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_db_span(span):
         remote_resource_type = _DB_CONNECTION_STRING_TYPE
         remote_resource_identifier = _get_db_connection(span)

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py

@@ -6,6 +6,9 @@
 from amazon.opentelemetry.distro._aws_attribute_keys import (
     AWS_KINESIS_STREAM_CONSUMERNAME,
     AWS_KINESIS_STREAM_NAME,
+    AWS_LAMBDA_FUNCTION_ARN,
+    AWS_LAMBDA_FUNCTION_NAME,
+    AWS_LAMBDA_RESOURCEMAPPING_ID,
     AWS_SECRETSMANAGER_SECRET_ARN,
     AWS_SNS_TOPIC_ARN,
     AWS_SQS_QUEUE_NAME,
@@ -20,6 +23,7 @@
     _BedrockRuntimeExtension,
 )
 from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS
+from opentelemetry.instrumentation.botocore.extensions.lmbd import _LambdaExtension
 from opentelemetry.instrumentation.botocore.extensions.sns import _SnsExtension
 from opentelemetry.instrumentation.botocore.extensions.sqs import _SqsExtension
 from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension, _BotoResultT
@@ -39,6 +43,44 @@ def _apply_botocore_instrumentation_patches() -> None:
     _apply_botocore_secretsmanager_patch()
     _apply_botocore_sns_patch()
     _apply_botocore_stepfunctions_patch()
+    _apply_botocore_lambda_patch()
+
+
+def _apply_botocore_lambda_patch() -> None:
+    """Botocore instrumentation patch for Lambda
+
+    This patch adds an extension to the upstream's list of known extensions for Lambda.
+    Extensions allow for custom logic for adding service-specific information to spans,
+    such as attributes. Specifically, we are adding logic to add the 
+    `aws.lambda.function.name` and  `aws.lambda.resource_mapping.id` attributes
+
+    Sidenote: There exists SpanAttributes.FAAS_INVOKED_NAME for invoke operations
+    in upstream. However, we want to cover more operations to extract 'FunctionName',
+    so we define `aws.lambda.function.name` separately. Additionally, this helps
+    us maintain naming consistency with the other AWS resources.
+    """
+    old_extract_attributes = _LambdaExtension.extract_attributes
+
+    def patch_extract_attributes(self, attributes: _AttributeMapT):
+        old_extract_attributes(self, attributes)
+        function_name = self._call_context.params.get("FunctionName")
+        if function_name:
+            attributes[AWS_LAMBDA_FUNCTION_NAME] = function_name
+        resource_mapping_id = self._call_context.params.get("UUID")
+        if resource_mapping_id:
+            attributes[AWS_LAMBDA_RESOURCEMAPPING_ID] = resource_mapping_id
+
+    old_on_success  = _LambdaExtension.on_success 
+
+    def patch_on_success(self, span: Span, result: _BotoResultT):
+        old_on_success(self, span, result)
+        lambda_configuration = result.get("Configuration", {})
+        function_arn = lambda_configuration.get("FunctionArn")
+        if function_arn:
+            span.set_attribute(AWS_LAMBDA_FUNCTION_ARN, function_arn)
+
+    _LambdaExtension.extract_attributes = patch_extract_attributes
+    _LambdaExtension.on_success = patch_on_success
 
 
 def _apply_botocore_stepfunctions_patch() -> None:

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

@@ -15,6 +15,8 @@
     AWS_CONSUMER_PARENT_SPAN_KIND,
     AWS_KINESIS_STREAM_CONSUMERNAME,
     AWS_KINESIS_STREAM_NAME,
+    AWS_LAMBDA_FUNCTION_NAME,
+    AWS_LAMBDA_RESOURCEMAPPING_ID,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
     AWS_REMOTE_DB_USER,
@@ -1167,6 +1169,27 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         self._validate_remote_resource_attributes("AWS::Kinesis::Stream", "aws_stream_name")
         self._mock_attribute([AWS_KINESIS_STREAM_NAME, AWS_KINESIS_STREAM_CONSUMERNAME], [None, None])
 
+        # Validate behaviour of AWS_LAMBDA_FUNCTION_NAME attribute, then remove it.
+        self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME], ["aws_lambda_function_name"], keys, values)
+        self._validate_remote_resource_attributes("AWS::Lambda::Function", "aws_lambda_function_name")
+        self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME], [None])
+
+        # Validate behaviour of AWS_LAMBDA_RESOURCEMAPPING_ID attribute, then remove it.
+        self._mock_attribute([AWS_LAMBDA_RESOURCEMAPPING_ID], ["aws_event_source_mapping_id"], keys, values)
+        self._validate_remote_resource_attributes("AWS::Lambda::EventSourceMapping", "aws_event_source_mapping_id")
+        self._mock_attribute([AWS_LAMBDA_RESOURCEMAPPING_ID], [None])
+
+        # Validate behaviour of both AWS_LAMBDA_FUNCTION_NAME and AWS_LAMBDA_RESOURCE_MAPPING_ID,
+        # then remove it.
+        self._mock_attribute(
+            [AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_RESOURCEMAPPING_ID],
+            ["aws_lambda_function_name", "aws_event_source_mapping_id"],
+            keys,
+            values
+        )
+        self._validate_remote_resource_attributes("AWS::Lambda::EventSourceMapping", "aws_event_source_mapping_id")
+        self._mock_attribute([AWS_LAMBDA_FUNCTION_NAME, AWS_LAMBDA_RESOURCEMAPPING_ID], [None, None])
+
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
 
     def test_client_db_span_with_remote_resource_attributes(self):

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

@@ -31,6 +31,8 @@
 _TOPIC_ARN: str = "topicArn"
 _STATE_MACHINE_ARN: str = "arn:aws:states:us-west-2:000000000000:stateMachine:testStateMachine"
 _ACTIVITY_ARN: str = "arn:aws:states:us-east-1:007003123456789012:activity:testActivity"
+_LAMBDA_FUNCTION_NAME: str = "lambdaFunctionName"
+_LAMBDA_SOURCE_MAPPING_ID: str = "lambdaEventSourceMappingID"
 
 # Patch names
 GET_DISTRIBUTION_PATCH: str = (
@@ -155,6 +157,9 @@ def _test_unpatched_botocore_instrumentation(self):
         # StepFunctions
         self.assertFalse("stepfunctions" in _KNOWN_EXTENSIONS, "Upstream has added a StepFunctions extension")
 
+        # Lambda
+        self.assertTrue("lambda" in _KNOWN_EXTENSIONS, "Upstream has removed the Lambda extension")
+
     def _test_unpatched_gevent_instrumentation(self):
         self.assertFalse(gevent.monkey.is_module_patched("os"), "gevent os module has been patched")
         self.assertFalse(gevent.monkey.is_module_patched("thread"), "gevent thread module has been patched")
@@ -239,6 +244,14 @@ def _test_patched_botocore_instrumentation(self):
         self.assertTrue("aws.stepfunctions.activity.arn" in stepfunctions_attributes)
         self.assertEqual(stepfunctions_attributes["aws.stepfunctions.activity.arn"], _ACTIVITY_ARN)
 
+        # Lambda
+        self.assertTrue("lambda" in _KNOWN_EXTENSIONS)
+        lambda_attributes: Dict[str, str] = _do_extract_lambda_attributes()
+        self.assertTrue("aws.lambda.function.name" in lambda_attributes)
+        self.assertEqual(lambda_attributes["aws.lambda.function.name"], _LAMBDA_FUNCTION_NAME)
+        self.assertTrue("aws.lambda.resource_mapping.id" in lambda_attributes)
+        self.assertEqual(lambda_attributes["aws.lambda.resource_mapping.id"], _LAMBDA_SOURCE_MAPPING_ID)
+
     def _test_patched_gevent_os_ssl_instrumentation(self):
         # Only ssl and os module should have been patched since the environment variable was set to 'os, ssl'
         self.assertTrue(gevent.monkey.is_module_patched("ssl"), "gevent ssl module has not been patched")
@@ -421,6 +434,12 @@ def _do_extract_stepfunctions_attributes() -> Dict[str, str]:
     return _do_extract_attributes(service_name, params)
 
 
+def _do_extract_lambda_attributes() -> Dict[str, str]:
+    service_name: str = "lambda"
+    params: Dict[str, str] = {"FunctionName": _LAMBDA_FUNCTION_NAME, "UUID": _LAMBDA_SOURCE_MAPPING_ID}
+    return _do_extract_attributes(service_name, params)
+
+
 def _do_extract_attributes(service_name: str, params: Dict[str, Any], operation: str = None) -> Dict[str, str]:
     mock_call_context: MagicMock = MagicMock()
     mock_call_context.params = params