Skip to content

Commit 936aa5a

Browse files
ezhang6811ezhang6811
authored
Lambda Support Refactor + Enhance Unit Test Coverage for AWS Resources (#223)
1 parent 2b08d54 commit 936aa5a

File tree

3 files changed

+212
-19
lines changed

3 files changed

+212
-19
lines changed

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

Lines changed: 56 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,9 @@ internal class AwsMetricAttributeGenerator : IMetricAttributeGenerator
5454
private static readonly string NormalizedBedrockRuntimeServiceName = "AWS::BedrockRuntime";
5555
private static readonly string DbConnectionResourceType = "DB::Connection";
5656

57+
// Constants for Lambda operations
58+
private static readonly string LambdaInvokeOperation = "Invoke";
59+
5760
// Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
5861
private static readonly string GraphQL = "graphql";
5962

@@ -89,6 +92,7 @@ private ActivityTagsCollection GenerateDependencyMetricAttributes(Activity span,
8992
ActivityTagsCollection attributes = new ActivityTagsCollection();
9093
SetService(resource, span, attributes);
9194
SetEgressOperation(span, attributes);
95+
SetRemoteEnvironment(span, attributes);
9296
SetRemoteServiceAndOperation(span, attributes);
9397
SetRemoteResourceTypeAndIdentifier(span, attributes);
9498
SetSpanKindForDependency(span, attributes);
@@ -364,8 +368,6 @@ private static string NormalizeRemoteServiceName(Activity span, string serviceNa
364368
case "AmazonKinesis": // AWS SDK v1
365369
case "Kinesis": // AWS SDK v2
366370
return NormalizedKinesisServiceName;
367-
case "Lambda":
368-
return NormalizedLambdaServiceName;
369371
case "Amazon S3": // AWS SDK v1
370372
case "S3": // AWS SDK v2
371373
return NormalizedS3ServiceName;
@@ -384,6 +386,23 @@ private static string NormalizeRemoteServiceName(Activity span, string serviceNa
384386
return NormalizedBedrockServiceName;
385387
case "Bedrock Runtime":
386388
return NormalizedBedrockRuntimeServiceName;
389+
case "Lambda":
390+
if (IsLambdaInvokeOperation(span))
391+
{
392+
string? lambdaFunctionName = (string?)span.GetTagItem(AttributeAWSLambdaFunctionName);
393+
394+
// if Lambda function name is not present, use UnknownRemoteService
395+
// This is intentional - we want to clearly indicate when the Lambda function name
396+
// is missing rather than falling back to a generic service name
397+
return lambdaFunctionName != null
398+
? lambdaFunctionName
399+
: UnknownRemoteService;
400+
}
401+
else
402+
{
403+
return NormalizedLambdaServiceName;
404+
}
405+
387406
default:
388407
return "AWS::" + serviceName;
389408
}
@@ -414,23 +433,9 @@ private static void SetRemoteResourceTypeAndIdentifier(Activity span, ActivityTa
414433
}
415434
else if (IsKeyPresent(span, AttributeAWSLambdaFunctionName))
416435
{
417-
// Handling downstream Lambda as a service vs. an AWS resource:
418-
// - If the method call is "Invoke", we treat downstream Lambda as a service.
419-
// - Otherwise, we treat it as an AWS resource.
420-
//
421-
// This addresses a Lambda topology issue in Application Signals.
422-
// More context in PR: https://github.com/aws-observability/aws-otel-python-instrumentation/pull/319
423-
//
424-
// NOTE: The env var LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT was introduced as part of this fix.
425-
// It is optional and allow users to override the default value if needed.
426-
if (GetRemoteOperation(span, AttributeRpcMethod) == "Invoke")
427-
{
428-
attributes[AttributeAWSRemoteService] = EscapeDelimiters((string?)span.GetTagItem(AttributeAWSLambdaFunctionName));
429-
430-
string lambdaRemoteEnv = Environment.GetEnvironmentVariable("LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT") ?? "default";
431-
attributes.Add(AttributeAWSRemoteEnvironment, $"lambda:{lambdaRemoteEnv}");
432-
}
433-
else
436+
// For non-invoke Lambda operations, treat Lambda as a resource.
437+
// see NormalizeRemoteServiceName for more information.
438+
if (!IsLambdaInvokeOperation(span))
434439
{
435440
remoteResourceType = NormalizedLambdaServiceName + "::Function";
436441
remoteResourceIdentifier = EscapeDelimiters((string?)span.GetTagItem(AttributeAWSLambdaFunctionName));
@@ -541,6 +546,25 @@ private static void SetRemoteDbUser(Activity span, ActivityTagsCollection attrib
541546
}
542547
}
543548

549+
// Remote environment is used to identify the environment of downstream services.
550+
// Currently only set to "lambda:default" for Lambda Invoke operations when aws-api system is detected.
551+
private static void SetRemoteEnvironment(Activity span, ActivityTagsCollection attributes)
552+
{
553+
// We want to treat downstream Lambdas as a service rather than a resource because
554+
// Application Signals topology map gets disconnected due to conflicting Lambda Entity definitions
555+
// Additional context can be found in https://github.com/aws-observability/aws-otel-python-instrumentation/pull/319
556+
if (IsLambdaInvokeOperation(span))
557+
{
558+
var remoteEnvironment = Environment.GetEnvironmentVariable(Plugin.LambdaApplicationSignalsRemoteEnvironment);
559+
if (string.IsNullOrEmpty(remoteEnvironment))
560+
{
561+
remoteEnvironment = "default";
562+
}
563+
564+
attributes.Add(AttributeAWSRemoteEnvironment, "lambda:" + remoteEnvironment.Trim());
565+
}
566+
}
567+
544568
// Span kind is needed for differentiating metrics in the EMF exporter
545569
private static void SetSpanKindForService(Activity span, ActivityTagsCollection attributes)
546570
{
@@ -708,4 +732,17 @@ private static string BuildDbConnection(string? address, long? port)
708732

709733
return input.Replace("^", "^^").Replace("|", "^|");
710734
}
735+
736+
// Check if the span represents a Lambda Invoke operation.
737+
private static bool IsLambdaInvokeOperation(Activity span)
738+
{
739+
if (!IsAwsSDKSpan(span))
740+
{
741+
return false;
742+
}
743+
744+
string rpcService = GetRemoteService(span, AttributeRpcService);
745+
string rpcMethod = GetRemoteOperation(span, AttributeRpcMethod);
746+
return rpcService.Equals("Lambda") && rpcMethod.Equals(LambdaInvokeOperation);
747+
}
711748
}

‎src/AWS.Distro.OpenTelemetry.AutoInstrumentation/Plugin.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ public class Plugin
3939
/// OTEL_AWS_APPLICATION_SIGNALS_ENABLED
4040
/// </summary>
4141
public static readonly string ApplicationSignalsEnabledConfig = "OTEL_AWS_APPLICATION_SIGNALS_ENABLED";
42+
internal static readonly string LambdaApplicationSignalsRemoteEnvironment = "LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT";
4243
private static readonly string XRayOtlpEndpointPattern = "^https://xray\\.([a-z0-9-]+)\\.amazonaws\\.com/v1/traces$";
4344
private static readonly string SigV4EnabledConfig = "OTEL_AWS_SIG_V4_ENABLED";
4445
private static readonly string TracesExporterConfig = "OTEL_TRACES_EXPORTER";

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

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -879,6 +879,13 @@ public void TestSdkClientSpanWithRemoteResourceAttributes()
879879
attributesCombination[AttributeAWSSQSQueueUrl] = "invalidUrl";
880880
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::SQS::Queue", "aws_queue_name", "invalidUrl");
881881

882+
// Validate SQS behavior when QueueName isn't available
883+
attributesCombination = new Dictionary<string, object>
884+
{
885+
{ AttributeAWSSQSQueueUrl, "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue" },
886+
};
887+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::SQS::Queue", "MyQueue", "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue");
888+
882889
attributesCombination = new Dictionary<string, object>
883890
{
884891
{ AttributeAWSKinesisStreamName, "aws_stream_name" },
@@ -1043,6 +1050,128 @@ public void TestSdkClientSpanWithRemoteResourceAttributes()
10431050
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Bedrock::DataSource", "aws_data_source_^^id", "aws_knowledge_base_^^id|aws_data_source_^^id");
10441051
}
10451052

1053+
[Fact]
1054+
public void TestCloudformationPrimaryIdentifierFallbackToRemoteResourceIdentifier()
1055+
{
1056+
// Test case 1: S3 Bucket (no ARN available, should use bucket name for both)
1057+
Dictionary<string, object> attributesCombination = new Dictionary<string, object>
1058+
{
1059+
{ AttributeRpcService, "S3" },
1060+
{ AttributeAWSS3Bucket, "my-test-bucket" },
1061+
};
1062+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::S3::Bucket", "my-test-bucket", "my-test-bucket");
1063+
1064+
// Test S3 Bucket with speicial characters
1065+
attributesCombination[AttributeAWSS3Bucket] = "my-test|bucket^name";
1066+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::S3::Bucket", "my-test^|bucket^^name", "my-test^|bucket^^name");
1067+
1068+
// Test case 2: SQS Queue (no ARN, should use queue name for both)
1069+
attributesCombination = new Dictionary<string, object>
1070+
{
1071+
{ AttributeRpcService, "SQS" },
1072+
{ AttributeAWSSQSQueueName, "my-test-queue" },
1073+
};
1074+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::SQS::Queue", "my-test-queue", "my-test-queue");
1075+
1076+
// Test SQS Queue with special characters
1077+
attributesCombination[AttributeAWSSQSQueueName] = "my^queue|name";
1078+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::SQS::Queue", "my^^queue^|name", "my^^queue^|name");
1079+
1080+
// Test case 3: DynamoDB Table (no ARN, should use table name for both)
1081+
attributesCombination = new Dictionary<string, object>
1082+
{
1083+
{ AttributeRpcService, "DynamoDB" },
1084+
{ AttributeAWSDynamoTableName, "my-test-table" },
1085+
};
1086+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::DynamoDB::Table", "my-test-table", "my-test-table");
1087+
1088+
// Test DynamoDB Table with special characters
1089+
attributesCombination[AttributeAWSDynamoTableName] = "my|table^name";
1090+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::DynamoDB::Table", "my^|table^^name", "my^|table^^name");
1091+
1092+
// Test case 4: Kinesis Stream (no ARN, should use stream name for both)
1093+
attributesCombination = new Dictionary<string, object>
1094+
{
1095+
{ AttributeRpcService, "Kinesis" },
1096+
{ AttributeAWSKinesisStreamName, "my-test-stream" },
1097+
};
1098+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Kinesis::Stream", "my-test-stream", "my-test-stream");
1099+
1100+
// Test Kinesis Stream with special characters
1101+
attributesCombination[AttributeAWSKinesisStreamName] = "my|stream^name";
1102+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Kinesis::Stream", "my^|stream^^name", "my^|stream^^name");
1103+
1104+
// Test case 5: Lambda Function (non-invoke operation, no ARN)
1105+
attributesCombination = new Dictionary<string, object>
1106+
{
1107+
{ AttributeRpcService, "Lambda" },
1108+
{ AttributeRpcMethod, "GetFunction" },
1109+
{ AttributeAWSLambdaFunctionName, "my-test-function" },
1110+
};
1111+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Lambda::Function", "my-test-function", "my-test-function");
1112+
1113+
// Test Lambda Function with special characters
1114+
attributesCombination[AttributeAWSLambdaFunctionName] = "my|lambda^function";
1115+
this.ValidateRemoteResourceAttributes(attributesCombination, "AWS::Lambda::Function", "my^|lambda^^function", "my^|lambda^^function");
1116+
}
1117+
1118+
[Fact]
1119+
public void TestSetRemoteEnvironment()
1120+
{
1121+
// Test 1: Setting remote environment when all relevant attributes are present
1122+
Activity? spanDataMock = this.testSource.StartActivity("test", ActivityKind.Client);
1123+
spanDataMock.SetTag(AttributeRpcSystem, "aws-api");
1124+
spanDataMock.SetTag(AttributeRpcService, "Lambda");
1125+
spanDataMock.SetTag(AttributeRpcMethod, "Invoke");
1126+
spanDataMock.SetTag(AttributeAWSLambdaFunctionName, "testFunction");
1127+
1128+
this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
1129+
.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out ActivityTagsCollection? dependencyMetric);
1130+
dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out var remoteEnvironment);
1131+
Assert.Equal(remoteEnvironment, "lambda:default");
1132+
1133+
// Test 2: NOT setting remote environment when rpc.system is missing
1134+
spanDataMock.SetTag(AttributeRpcSystem, null);
1135+
this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
1136+
.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
1137+
dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
1138+
Assert.Null(remoteEnvironment);
1139+
spanDataMock.SetTag(AttributeRpcSystem, "aws-api");
1140+
1141+
// Test 3: NOT setting remote environment when rpc.method is missing
1142+
spanDataMock.SetTag(AttributeRpcMethod, null);
1143+
this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
1144+
.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
1145+
dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
1146+
Assert.Null(remoteEnvironment);
1147+
spanDataMock.SetTag(AttributeRpcMethod, "Invoke");
1148+
1149+
// Test 4: setting remote environment to lambda:default when FunctionName is missing
1150+
spanDataMock.SetTag(AttributeAWSLambdaFunctionName, null);
1151+
this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
1152+
.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
1153+
dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
1154+
Assert.Equal(remoteEnvironment, "lambda:default");
1155+
1156+
// Test 5: NOT setting remote environment for non-Lambda services
1157+
spanDataMock.SetTag(AttributeRpcService, "S3");
1158+
spanDataMock.SetTag(AttributeRpcMethod, "GetObject");
1159+
this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
1160+
.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
1161+
dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
1162+
Assert.Null(remoteEnvironment);
1163+
1164+
// Test 6: NOT setting remote environment for Lambda non-Invoke operations
1165+
spanDataMock.SetTag(AttributeRpcService, "Lambda");
1166+
spanDataMock.SetTag(AttributeRpcMethod, "GetFunction");
1167+
this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource)
1168+
.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
1169+
dependencyMetric.TryGetValue(AttributeAWSRemoteEnvironment, out remoteEnvironment);
1170+
Assert.Null(remoteEnvironment);
1171+
1172+
spanDataMock.Dispose();
1173+
}
1174+
10461175
[Fact]
10471176
public void TestNormalizeRemoteServiceName_NoNormalization()
10481177
{
@@ -1074,6 +1203,31 @@ public void TestNormalizeRemoteServiceName_AwsSdk()
10741203
this.TestAwsSdkServiceNormalization("Kinesis", "AWS::Kinesis");
10751204
this.TestAwsSdkServiceNormalization("S3", "AWS::S3");
10761205
this.TestAwsSdkServiceNormalization("Sqs", "AWS::SQS");
1206+
1207+
// Lambda: non-Invoke operations
1208+
this.TestAwsSdkServiceNormalization("Lambda", "AWS::Lambda");
1209+
1210+
// Lambda: Invoke with function name
1211+
Activity? spanDataMock = this.testSource.StartActivity("test", ActivityKind.Client);
1212+
spanDataMock.SetTag(AttributeRpcSystem, "aws-api");
1213+
spanDataMock.SetTag(AttributeRpcService, "Lambda");
1214+
1215+
spanDataMock.SetTag(AttributeRpcMethod, "Invoke");
1216+
spanDataMock.SetTag(AttributeAWSLambdaFunctionName, "testFunction");
1217+
1218+
var attributeMap = this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource);
1219+
attributeMap.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out ActivityTagsCollection? dependencyMetric);
1220+
dependencyMetric.TryGetValue(AttributeAWSRemoteService, out var actualServiceName);
1221+
Assert.Equal("testFunction", actualServiceName);
1222+
1223+
// Lambda: Invoke without function name - should fall back to UnknownRemoteService
1224+
spanDataMock.SetTag(AttributeAWSLambdaFunctionName, null);
1225+
attributeMap = this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource);
1226+
attributeMap.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out dependencyMetric);
1227+
dependencyMetric.TryGetValue(AttributeAWSRemoteService, out actualServiceName);
1228+
Assert.Equal(AutoInstrumentation.AwsSpanProcessingUtil.UnknownRemoteService, actualServiceName);
1229+
1230+
spanDataMock.Dispose();
10771231
}
10781232

10791233
[Fact]
@@ -1105,6 +1259,7 @@ private void TestAwsSdkServiceNormalization(string serviceName, string expectedR
11051259
Activity? spanDataMock = this.testSource.StartActivity("test", ActivityKind.Client);
11061260
spanDataMock.SetTag(AttributeRpcSystem, "aws-api");
11071261
spanDataMock.SetTag(AttributeRpcService, serviceName);
1262+
11081263
var attributeMap = this.generator.GenerateMetricAttributeMapFromSpan(spanDataMock, this.resource);
11091264
attributeMap.TryGetValue(MetricAttributeGeneratorConstants.DependencyMetric, out ActivityTagsCollection? dependencyMetric);
11101265
dependencyMetric.TryGetValue(AttributeAWSRemoteService, out var actualServiceName);

0 commit comments

Comments
 (0)