Lambda Support Refactor + Enhance Unit Test Coverage for AWS Resources

src/AWS.Distro.OpenTelemetry.AutoInstrumentation/AwsMetricAttributeGenerator.cs

@@ -54,6 +54,10 @@ internal class AwsMetricAttributeGenerator : IMetricAttributeGenerator
     private static readonly string NormalizedBedrockRuntimeServiceName = "AWS::BedrockRuntime";
     private static readonly string DbConnectionResourceType = "DB::Connection";
 
+    // Constants for Lambda operations
+    private static readonly string LambdaApplicationSignalsRemoteEnvironment = "LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT";
+    private static readonly string LambdaInvokeOperation = "Invoke";
+
     // Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
     private static readonly string GraphQL = "graphql";
 
@@ -89,6 +93,7 @@ private ActivityTagsCollection GenerateDependencyMetricAttributes(Activity span,
         ActivityTagsCollection attributes = new ActivityTagsCollection();
         SetService(resource, span, attributes);
         SetEgressOperation(span, attributes);
+        SetRemoteEnvironment(span, attributes);
         SetRemoteServiceAndOperation(span, attributes);
         SetRemoteResourceTypeAndIdentifier(span, attributes);
         SetSpanKindForDependency(span, attributes);
@@ -364,8 +369,6 @@ private static string NormalizeRemoteServiceName(Activity span, string serviceNa
                 case "AmazonKinesis": // AWS SDK v1
                 case "Kinesis": // AWS SDK v2
                     return NormalizedKinesisServiceName;
-                case "Lambda":
-                    return NormalizedLambdaServiceName;
                 case "Amazon S3": // AWS SDK v1
                 case "S3": // AWS SDK v2
                     return NormalizedS3ServiceName;
@@ -384,6 +387,23 @@ private static string NormalizeRemoteServiceName(Activity span, string serviceNa
                     return NormalizedBedrockServiceName;
                 case "Bedrock Runtime":
                     return NormalizedBedrockRuntimeServiceName;
+                case "Lambda":
+                    if (IsLambdaInvokeOperation(span))
+                    {
+                        string? lambdaFunctionName = (string?)span.GetTagItem(AttributeAWSLambdaFunctionName);
+
+                        // if Lambda function name is not present, use UnknownRemoteService
+                        // This is intentional - we want to clearly indicate when the Lambda function name
+                        // is missing rather than falling back to a generic service name
+                        return lambdaFunctionName != null
+                            ? lambdaFunctionName
+                            : UnknownRemoteService;
+                    }
+                    else
+                    {
+                        return NormalizedLambdaServiceName;
+                    }
+
                 default:
                     return "AWS::" + serviceName;
             }
@@ -414,23 +434,9 @@ private static void SetRemoteResourceTypeAndIdentifier(Activity span, ActivityTa
             }
             else if (IsKeyPresent(span, AttributeAWSLambdaFunctionName))
             {
-                // 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 var LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT was introduced as part of this fix.
-                // It is optional and allow users to override the default value if needed.
-                if (GetRemoteOperation(span, AttributeRpcMethod) == "Invoke")
-                {
-                    attributes[AttributeAWSRemoteService] = EscapeDelimiters((string?)span.GetTagItem(AttributeAWSLambdaFunctionName));
-
-                    string lambdaRemoteEnv = Environment.GetEnvironmentVariable("LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT") ?? "default";
-                    attributes.Add(AttributeAWSRemoteEnvironment, $"lambda:{lambdaRemoteEnv}");
-                }
-                else
+                // For non-invoke Lambda operations, treat Lambda as a resource.
+                // see NormalizeRemoteServiceName for more information.
+                if (!IsLambdaInvokeOperation(span))
                 {
                     remoteResourceType = NormalizedLambdaServiceName + "::Function";
                     remoteResourceIdentifier = EscapeDelimiters((string?)span.GetTagItem(AttributeAWSLambdaFunctionName));
@@ -541,6 +547,25 @@ private static void SetRemoteDbUser(Activity span, ActivityTagsCollection attrib
         }
     }
 
+    // Remote environment is used to identify the environment of downstream services.
+    // Currently only set to "lambda:default" for Lambda Invoke operations when aws-api system is detected.
+    private static void SetRemoteEnvironment(Activity span, ActivityTagsCollection attributes)
+    {
+        // We want to treat downstream Lambdas as a service rather than a resource because
+        // Application Signals topology map gets disconnected due to conflicting Lambda Entity definitions
+        // Additional context can be found in https://github.com/aws-observability/aws-otel-python-instrumentation/pull/319
+        if (IsLambdaInvokeOperation(span))
+        {
+            var remoteEnvironment = Environment.GetEnvironmentVariable(LambdaApplicationSignalsRemoteEnvironment);
+            if (string.IsNullOrEmpty(remoteEnvironment))
+            {
+                remoteEnvironment = "default";
+            }
+
+            attributes.Add(AttributeAWSRemoteEnvironment, "lambda:" + remoteEnvironment.Trim());
+        }
+    }
+
     // Span kind is needed for differentiating metrics in the EMF exporter
     private static void SetSpanKindForService(Activity span, ActivityTagsCollection attributes)
     {
@@ -708,4 +733,17 @@ private static string BuildDbConnection(string? address, long? port)
 
         return input.Replace("^", "^^").Replace("|", "^|");
     }
+
+    // Check if the span represents a Lambda Invoke operation.
+    private static bool IsLambdaInvokeOperation(Activity span)
+    {
+        if (!IsAwsSDKSpan(span))
+        {
+            return false;
+        }
+
+        string rpcService = GetRemoteService(span, AttributeRpcService);
+        string rpcMethod = GetRemoteOperation(span, AttributeRpcMethod);
+        return rpcService.Equals("Lambda") && rpcMethod.Equals(LambdaInvokeOperation);
+    }
 }

test/AWS.Distro.OpenTelemetry.AutoInstrumentation.Tests/AwsMetricAttributesGeneratorTest.cs

@@ -879,6 +879,13 @@ public void TestSdkClientSpanWithRemoteResourceAttributes()
         attributesCombination[AttributeAWSSQSQueueUrl] = "invalidUrl";
         this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::SQS::Queue", "aws_queue_name", "invalidUrl");
 
+        // Validate SQS behavior when QueueName isn't available
+        attributesCombination = new Dictionary<string, object>
+        {
+            { AttributeAWSSQSQueueUrl, "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue" },
+        };
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::SQS::Queue", "MyQueue", "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue");
+
         attributesCombination = new Dictionary<string, object>
         {
             { AttributeAWSKinesisStreamName, "aws_stream_name" },
@@ -1043,6 +1050,128 @@ public void TestSdkClientSpanWithRemoteResourceAttributes()
         this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Bedrock::DataSource", "aws_data_source_^^id", "aws_knowledge_base_^^id|aws_data_source_^^id");
     }
 
+    [Fact]
+    public void TestCloudformationPrimaryIdentifierFallbackToRemoteResourceIdentifier()
+    {
+        // Test case 1: S3 Bucket (no ARN available, should use bucket name for both)
+        Dictionary<string, object> attributesCombination = new Dictionary<string, object>
+        {
+            { AttributeRpcService, "S3" },
+            { AttributeAWSS3Bucket, "my-test-bucket" },
+        };
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::S3::Bucket", "my-test-bucket", "my-test-bucket");
+
+        // Test S3 Bucket with speicial characters
+        attributesCombination[AttributeAWSS3Bucket] = "my-test|bucket^name";
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::S3::Bucket", "my-test^|bucket^^name", "my-test^|bucket^^name");
+
+        // Test case 2: SQS Queue (no ARN, should use queue name for both)
+        attributesCombination = new Dictionary<string, object>
+        {
+            { AttributeRpcService, "SQS" },
+            { AttributeAWSSQSQueueName, "my-test-queue" },
+        };
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::SQS::Queue", "my-test-queue", "my-test-queue");
+
+        // Test SQS Queue with special characters
+        attributesCombination[AttributeAWSSQSQueueName] = "my^queue|name";
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::SQS::Queue", "my^^queue^|name", "my^^queue^|name");
+
+        // Test case 3: DynamoDB Table (no ARN, should use table name for both)
+        attributesCombination = new Dictionary<string, object>
+        {
+            { AttributeRpcService, "DynamoDB" },
+            { AttributeAWSDynamoTableName, "my-test-table" },
+        };
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::DynamoDB::Table", "my-test-table", "my-test-table");
+
+        // Test DynamoDB Table with special characters
+        attributesCombination[AttributeAWSDynamoTableName] = "my|table^name";
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::DynamoDB::Table", "my^|table^^name", "my^|table^^name");
+
+        // Test case 4: Kinesis Stream (no ARN, should use stream name for both)
+        attributesCombination = new Dictionary<string, object>
+        {
+            { AttributeRpcService, "Kinesis" },
+            { AttributeAWSKinesisStreamName, "my-test-stream" },
+        };
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Kinesis::Stream", "my-test-stream", "my-test-stream");
+
+        // Test Kinesis Stream with special characters
+        attributesCombination[AttributeAWSKinesisStreamName] = "my|stream^name";
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Kinesis::Stream", "my^|stream^^name", "my^|stream^^name");
+
+        // Test case 5: Lambda Function (non-invoke operation, no ARN)
+        attributesCombination = new Dictionary<string, object>
+        {
+            { AttributeRpcService, "Lambda" },
+            { AttributeRpcMethod, "GetFunction" },
+            { AttributeAWSLambdaFunctionName, "my-test-function" },
+        };
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Lambda::Function", "my-test-function", "my-test-function");
+
+        // Test Lambda Function with special characters
+        attributesCombination[AttributeAWSLambdaFunctionName] = "my|lambda^function";
+        this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Lambda::Function", "my^|lambda^^function", "my^|lambda^^function");
+    }
+
+    [Fact]
+    public void TestSetRemoteEnvironmentForLambdaInvoke()
+    {
+        // Test 1: Setting remote environment when all relevant attributes are present
+        Activity? spanDataMock = this.testSource.StartActivity("test", ActivityKind.Client);
+        spanDataMock.SetTag(AttributeRpcSystem, "aws-api");
+        spanDataMock.SetTag(AttributeRpcService, "Lambda");
+        spanDataMock.SetTag(AttributeRpcMethod, "Invoke");
+        spanDataMock.SetTag(AttributeAWSLambdaFunctionName, "testFunction");
+
+        this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
+            .TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out ActivityTagsCollection? dependencyMetric);
+        dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out var remoteEnvironment);
+        Assert.Equal(remoteEnvironment, "lambda:default");
+
+        // Test 2: NOT setting remote environment when rpc.system is missing
+        spanDataMock.SetTag(AttributeRpcSystem, null);
+        this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
+            .TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
+        dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
+        Assert.Null(remoteEnvironment);
+        spanDataMock.SetTag(AttributeRpcSystem, "aws-api");
+
+        // Test 3: NOT setting remote environment when rpc.method is missing
+        spanDataMock.SetTag(AttributeRpcMethod, null);
+        this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
+            .TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
+        dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
+        Assert.Null(remoteEnvironment);
+        spanDataMock.SetTag(AttributeRpcMethod, "Invoke");
+
+        // Test 4: setting remote environment to lambda:default when FunctionName is missing
+        spanDataMock.SetTag(AttributeAWSLambdaFunctionName, null);
+        this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
+            .TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
+        dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
+        Assert.Equal(remoteEnvironment, "lambda:default");
+
+        // Test 5: NOT setting remote environment for non-Lambda services
+        spanDataMock.SetTag(AttributeRpcService, "S3");
+        spanDataMock.SetTag(AttributeRpcMethod, "GetObject");
+        this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
+            .TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
+        dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
+        Assert.Null(remoteEnvironment);
+
+        // Test 6: NOT setting remote environment for Lambda non-Invoke operations
+        spanDataMock.SetTag(AttributeRpcService, "Lambda");
+        spanDataMock.SetTag(AttributeRpcMethod, "GetFunction");
+        this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
+            .TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
+        dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
+        Assert.Null(remoteEnvironment);
+
+        spanDataMock.Dispose();
+    }
+
     [Fact]
     public void TestNormalizeRemoteServiceName_NoNormalization()
     {
@@ -1063,6 +1192,12 @@ public void TestNormalizeRemoteServiceName_AwsSdk()
         this.TestAwsSdkServiceNormalization("Bedrock Agent", "AWS::Bedrock");
         this.TestAwsSdkServiceNormalization("Bedrock Agent Runtime", "AWS::Bedrock");
 
+        // Test Lambda
+        // For Invoke operations, remote service should be Lambda function name or UnknownRemoteService
+        this.TestAwsSdkServiceNormalization("Lambda", "AWS::Lambda");
+        this.TestAwsSdkServiceNormalization("Lambda", "testFunction", "Invoke", "testFunction");
+        this.TestAwsSdkServiceNormalization("Lambda", "UnknownRemoteService", "Invoke", null);
+
         // AWS SDK V2
         this.TestAwsSdkServiceNormalization("AmazonDynamoDBv2", "AWS::DynamoDB");
         this.TestAwsSdkServiceNormalization("AmazonKinesis", "AWS::Kinesis");
@@ -1100,11 +1235,19 @@ public void TestBothMetricsWhenLocalRootConsumerProcess()
         Assert.Equal(2, attributeMap.Count);
     }
 
-    private void TestAwsSdkServiceNormalization(string serviceName, string expectedRemoteService)
+    private void TestAwsSdkServiceNormalization(string serviceName, string expectedRemoteService, string? rpcMethod = null, string? lambdaFunctionName = null)
     {
         Activity? spanDataMock = this.testSource.StartActivity("test", ActivityKind.Client);
         spanDataMock.SetTag(AttributeRpcSystem, "aws-api");
         spanDataMock.SetTag(AttributeRpcService, serviceName);
+
+        // For Lambda, if rpc.method is Invoke, depend on function name for remote service
+        if (serviceName == "Lambda")
+        {
+            spanDataMock.SetTag(AttributeRpcMethod, rpcMethod);
+            spanDataMock.SetTag(AttributeAWSLambdaFunctionName, lambdaFunctionName);
+        }
+
         var attributeMap = this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource);
         attributeMap.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out ActivityTagsCollection? dependencyMetric);
         dependencyMetric.TryGetValue(AttributeAWSRemoteService, out var actualServiceName);