lambda topology issue fix

Author: yiyuan-he

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsAttributeKeys.java

@@ -32,6 +32,9 @@ private AwsAttributeKeys() {}
   static final AttributeKey<String> AWS_REMOTE_SERVICE =
       AttributeKey.stringKey("aws.remote.service");
 
+  static final AttributeKey<String> AWS_REMOTE_ENVIRONMENT =
+      AttributeKey.stringKey("aws.remote.environment");
+
   static final AttributeKey<String> AWS_REMOTE_OPERATION =
       AttributeKey.stringKey("aws.remote.operation");
 
@@ -64,6 +67,9 @@ private AwsAttributeKeys() {}
   static final AttributeKey<String> AWS_SECRET_ARN =
       AttributeKey.stringKey("aws.secretsmanager.secret.arn");
 
+  static final AttributeKey<String> AWS_LAMBDA_NAME =
+      AttributeKey.stringKey("aws.lambda.function.name");
+
   static final AttributeKey<String> AWS_LAMBDA_ARN =
       AttributeKey.stringKey("aws.lambda.function.arn");
 

awsagentprovider/src/main/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGenerator.java

@@ -52,12 +52,15 @@
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ARN;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGE_BASE_ID;
+import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LAMBDA_ARN;
+import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LAMBDA_NAME;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LAMBDA_RESOURCE_ID;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_SERVICE;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_QUEUE_NAME;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_QUEUE_URL;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_REMOTE_DB_USER;
+import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_REMOTE_ENVIRONMENT;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_REMOTE_OPERATION;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_REMOTE_RESOURCE_IDENTIFIER;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_REMOTE_RESOURCE_TYPE;
@@ -518,6 +521,29 @@ private static void setRemoteResourceTypeAndIdentifier(SpanData span, Attributes
                 Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_SECRET_ARN))));
         cloudformationPrimaryIdentifier =
             Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_SECRET_ARN)));
+      } else if (isKeyPresent(span, AWS_LAMBDA_NAME)) {
+        if ("Invoke".equals(getRemoteOperation(span, RPC_METHOD))) {
+          String remoteService =
+              Optional.ofNullable(System.getenv("LAMBDA_APPLICATION_SIGNALS_REMOTE_SERVICE"))
+                  .filter(s -> !s.isEmpty())
+                  .orElse(span.getAttributes().get(AWS_LAMBDA_NAME));
+          if (remoteService != null) {
+            builder.put(AWS_REMOTE_SERVICE, remoteService);
+          }
+
+          String remoteEnvironment =
+              Optional.ofNullable(System.getenv("LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT"))
+                  .filter(s -> !s.isEmpty())
+                  .orElse("default");
+          builder.put(AWS_REMOTE_ENVIRONMENT, "lambda:" + remoteEnvironment);
+        } else {
+          remoteResourceType = Optional.of(NORMALIZED_LAMBDA_SERVICE_NAME + "::Function");
+          remoteResourceIdentifier =
+              getLambdaFunctionNameFromArn(
+                  Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_LAMBDA_NAME))));
+          cloudformationPrimaryIdentifier =
+              Optional.ofNullable(escapeDelimiters(span.getAttributes().get(AWS_LAMBDA_ARN)));
+        }
       } else if (isKeyPresent(span, AWS_LAMBDA_RESOURCE_ID)) {
         remoteResourceType = Optional.of(NORMALIZED_LAMBDA_SERVICE_NAME + "::EventSourceMapping");
         remoteResourceIdentifier =
@@ -539,6 +565,14 @@ private static void setRemoteResourceTypeAndIdentifier(SpanData span, Attributes
     }
   }
 
+  private static Optional<String> getLambdaFunctionNameFromArn(Optional<String> stringArn) {
+    if (stringArn.isPresent() && stringArn.get().startsWith("arn:aws:lambda:")) {
+      Arn resourceArn = Arn.fromString(stringArn.get());
+      return Optional.of(resourceArn.getResource().toString().split(":")[1]);
+    }
+    return stringArn;
+  }
+
   private static Optional<String> getSecretsManagerResourceNameFromArn(Optional<String> stringArn) {
     Arn resourceArn = Arn.fromString(stringArn.get());
     return Optional.of(resourceArn.getResource().toString().split(":")[1]);

awsagentprovider/src/test/java/software/amazon/opentelemetry/javaagent/providers/AwsMetricAttributeGeneratorTest.java

@@ -26,6 +26,7 @@
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_DATA_SOURCE_ID;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_GUARDRAIL_ID;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_KNOWLEDGE_BASE_ID;
+import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LAMBDA_NAME;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LAMBDA_RESOURCE_ID;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_SERVICE;
@@ -874,6 +875,11 @@ public void testSdkClientSpanWithRemoteResourceAttributes() {
     validateRemoteResourceAttributes("AWS::SecretsManager::Secret", "secretName");
     mockAttribute(AWS_SECRET_ARN, null);
 
+    // Validate behaviour of AWS_LAMBDA_NAME, then remove it.
+    mockAttribute(AWS_LAMBDA_NAME, "testLambdaName");
+    validateRemoteResourceAttributes("AWS::Lambda::Function", "testLambdaName");
+    mockAttribute(AWS_LAMBDA_NAME, null);
+
     // Validate behaviour of AWS_LAMBDA_RESOURCE_ID
     mockAttribute(AWS_LAMBDA_RESOURCE_ID, "eventSourceId");
     validateRemoteResourceAttributes("AWS::Lambda::EventSourceMapping", "eventSourceId");