Fix: Lambda Topology Issue (#1016) (#1085)

## What does this pull request do?
Cherry picks the Lambda Topology fix from our Lambda Layer v1 branch.

Original PR:
#1016

## Test strategy
Unit tests and followed e2e test strategy used in the [original
PR](#1016)
to sanity check correct behavior in service topology.

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 license.

Author: yiyuan-he

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

@@ -77,6 +77,8 @@
 public final class AwsApplicationSignalsCustomizerProvider
     implements AutoConfigurationCustomizerProvider {
   static final String AWS_LAMBDA_FUNCTION_NAME_CONFIG = "AWS_LAMBDA_FUNCTION_NAME";
+  static final String LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT =
+      "LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT";
 
   private static final Duration DEFAULT_METRIC_EXPORT_INTERVAL = Duration.ofMinutes(1);
   private static final Logger logger =

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

@@ -45,19 +45,23 @@
 import static io.opentelemetry.semconv.SemanticAttributes.SERVER_SOCKET_ADDRESS;
 import static io.opentelemetry.semconv.SemanticAttributes.SERVER_SOCKET_PORT;
 import static io.opentelemetry.semconv.SemanticAttributes.URL_FULL;
+import static software.amazon.opentelemetry.javaagent.providers.AwsApplicationSignalsCustomizerProvider.LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_AGENT_ID;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_BUCKET_NAME;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_CLOUDFORMATION_PRIMARY_IDENTIFIER;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_DATA_SOURCE_ID;
 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;
@@ -129,6 +133,9 @@ final class AwsMetricAttributeGenerator implements MetricAttributeGenerator {
   private static final String NORMALIZED_SECRETSMANAGER_SERVICE_NAME = "AWS::SecretsManager";
   private static final String NORMALIZED_LAMBDA_SERVICE_NAME = "AWS::Lambda";
 
+  // Constants for Lambda operations
+  private static final String LAMBDA_INVOKE_OPERATION = "Invoke";
+
   // Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
   private static final String GRAPHQL = "graphql";
 
@@ -167,6 +174,7 @@ private Attributes generateDependencyMetricAttributes(SpanData span, Resource re
     setService(resource, span, builder);
     setEgressOperation(span, builder);
     setRemoteServiceAndOperation(span, builder);
+    setRemoteEnvironment(span, builder);
     setRemoteResourceTypeAndIdentifier(span, builder);
     setSpanKindForDependency(span, builder);
     setHttpStatus(span, builder);
@@ -291,6 +299,27 @@ private static void setRemoteServiceAndOperation(SpanData span, AttributesBuilde
     builder.put(AWS_REMOTE_OPERATION, remoteOperation);
   }
 
+  /**
+   * 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(SpanData span, AttributesBuilder builder) {
+    // 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)) {
+      // TODO: This should be passed via ConfigProperties from
+      // AwsApplicationSignalsCustomizerProvider
+      String remoteEnvironment =
+          Optional.ofNullable(System.getenv(LAMBDA_APPLICATION_SIGNALS_REMOTE_ENVIRONMENT))
+              .filter(s -> !s.isEmpty())
+              .orElse("default");
+      builder.put(AWS_REMOTE_ENVIRONMENT, "lambda:" + remoteEnvironment);
+    }
+  }
+
   /**
    * When the remote call operation is undetermined for http use cases, will try to extract the
    * remote operation name from http url string
@@ -373,6 +402,15 @@ private static String generateRemoteService(SpanData span) {
     return remoteService;
   }
 
+  private static boolean isLambdaInvokeOperation(SpanData span) {
+    if (!isAwsSDKSpan(span)) {
+      return false;
+    }
+    String rpcService = getRemoteService(span, RPC_SERVICE);
+    return ("AWSLambda".equals(rpcService) || "Lambda".equals(rpcService))
+        && LAMBDA_INVOKE_OPERATION.equals(span.getAttributes().get(RPC_METHOD));
+  }
+
   /**
    * If the span is an AWS SDK span, normalize the name to align with <a
    * href="https://docs.aws.amazon.com/cloudcontrolapi/latest/userguide/supported-resources.html">AWS
@@ -420,7 +458,19 @@ private static String normalizeRemoteServiceName(SpanData span, String serviceNa
           return NORMALIZED_SECRETSMANAGER_SERVICE_NAME;
         case "AWSLambda": // AWS SDK v1
         case "Lambda": // AWS SDK v2
-          return NORMALIZED_LAMBDA_SERVICE_NAME;
+          if (isLambdaInvokeOperation(span)) {
+            // AWS_LAMBDA_NAME can contain either a function name or function ARN since Lambda AWS
+            // SDK calls accept both formats
+            Optional<String> lambdaFunctionName =
+                getLambdaFunctionNameFromArn(
+                    Optional.ofNullable(span.getAttributes().get(AWS_LAMBDA_NAME)));
+            // If Lambda 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.orElse(UNKNOWN_REMOTE_SERVICE);
+          } else {
+            return NORMALIZED_LAMBDA_SERVICE_NAME;
+          }
         default:
           return "AWS::" + serviceName;
       }
@@ -518,6 +568,19 @@ 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)) {
+        // For non-Invoke Lambda operations, treat Lambda as a resource,
+        // see normalizeRemoteServiceName for more information.
+        if (!isLambdaInvokeOperation(span)) {
+          remoteResourceType = Optional.of(NORMALIZED_LAMBDA_SERVICE_NAME + "::Function");
+          // AWS_LAMBDA_NAME can contain either a function name or function ARN since Lambda AWS SDK
+          // calls accept both formats
+          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 +602,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]);