update local operation of lambda span based on span attribute (#1106)

*Issue #, if available:*

*Description of changes:*

Currently the local operation of lambda span is hardcoded to
function_name//FunctionHandler. In some use cases, some server is
running inside lambda function and we should allow the setting of the
local operation dynamically at run time. For example, users can use span
api to set the attribute `aws.lambda.local.operation.override` and then
this value can be used for the local operation. We have an
Amazon-internal framework instrumentation that depends on this.

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

Author: pxaws

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

@@ -29,6 +29,19 @@ private AwsAttributeKeys() {}
   static final AttributeKey<String> AWS_LOCAL_OPERATION =
       AttributeKey.stringKey("aws.local.operation");
 
+  /*
+   * By default the local operation of a Lambda span is hard-coded to "<FunctionName>/FunctionHandler".
+   * To dynamically override this at runtime—such as when running a custom server inside your Lambda—
+   * you can set the span attribute "aws.lambda.local.operation.override" before ending the span. For example:
+   *
+   *   // Obtain the current Span and override its operation name
+   *   Span.current().setAttribute(
+   *       "aws.lambda.local.operation.override",
+   *       "MyService/handleRequest");
+   */
+  static final AttributeKey<String> AWS_LAMBDA_LOCAL_OPERATION_OVERRIDE =
+      AttributeKey.stringKey("aws.lambda.local.operation.override");
+
   static final AttributeKey<String> AWS_REMOTE_SERVICE =
       AttributeKey.stringKey("aws.remote.service");
 

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

@@ -27,6 +27,7 @@
 import static io.opentelemetry.semconv.SemanticAttributes.URL_PATH;
 import static software.amazon.opentelemetry.javaagent.providers.AwsApplicationSignalsCustomizerProvider.AWS_LAMBDA_FUNCTION_NAME_CONFIG;
 import static software.amazon.opentelemetry.javaagent.providers.AwsApplicationSignalsCustomizerProvider.isLambdaEnvironment;
+import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LAMBDA_LOCAL_OPERATION_OVERRIDE;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
 
 import com.fasterxml.jackson.core.type.TypeReference;
@@ -91,7 +92,23 @@ static List<String> getDialectKeywords() {
    */
   static String getIngressOperation(SpanData span) {
     if (isLambdaEnvironment()) {
-      return System.getenv(AWS_LAMBDA_FUNCTION_NAME_CONFIG) + "/FunctionHandler";
+      /*
+       * By default the local operation of a Lambda span is hard-coded to "<FunctionName>/FunctionHandler".
+       * To dynamically override this at runtime—such as when running a custom server inside your Lambda—
+       * you can set the span attribute "aws.lambda.local.operation.override" before ending the span. For example:
+       *
+       *   // Obtain the current Span and override its operation name
+       *   Span.current().setAttribute(
+       *       "aws.lambda.local.operation.override",
+       *       "MyServiceOperation");
+       *
+       * The code below will detect that override and use it instead of the default.
+       */
+      String operationOverride = span.getAttributes().get(AWS_LAMBDA_LOCAL_OPERATION_OVERRIDE);
+      if (operationOverride != null) {
+        return operationOverride;
+      }
+      return getFunctionNameFromEnv() + "/FunctionHandler";
     }
     String operation = span.getName();
     if (shouldUseInternalOperation(span)) {
@@ -102,6 +119,11 @@ static String getIngressOperation(SpanData span) {
     return operation;
   }
 
+  // define a function so that we can mock it in unit test
+  static String getFunctionNameFromEnv() {
+    return System.getenv(AWS_LAMBDA_FUNCTION_NAME_CONFIG);
+  }
+
   static String getEgressOperation(SpanData span) {
     if (shouldUseInternalOperation(span)) {
       return INTERNAL_OPERATION;

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

@@ -19,8 +19,12 @@
 import static io.opentelemetry.semconv.SemanticAttributes.MessagingOperationValues.PROCESS;
 import static io.opentelemetry.semconv.SemanticAttributes.MessagingOperationValues.RECEIVE;
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Answers.CALLS_REAL_METHODS;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
 import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
+import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LAMBDA_LOCAL_OPERATION_OVERRIDE;
 import static software.amazon.opentelemetry.javaagent.providers.AwsAttributeKeys.AWS_LOCAL_OPERATION;
 import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.MAX_KEYWORD_LENGTH;
 import static software.amazon.opentelemetry.javaagent.providers.AwsSpanProcessingUtil.getDialectKeywords;
@@ -33,6 +37,7 @@
 import java.util.List;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
 
 public class AwsSpanProcessingUtilTest {
   private static final String DEFAULT_PATH_VALUE = "/";
@@ -120,6 +125,49 @@ public void testGetIngressOperationInvalidNameAndValidTargetAndMethod() {
     assertThat(actualOperation).isEqualTo(validMethod + " " + validTarget);
   }
 
+  @Test
+  public void testGetIngressOperationLambdaOverride() {
+    try (MockedStatic<AwsApplicationSignalsCustomizerProvider> providerStatic =
+        mockStatic(
+            AwsApplicationSignalsCustomizerProvider.class,
+            withSettings().defaultAnswer(CALLS_REAL_METHODS))) {
+      // Force Lambda environment branch
+      providerStatic
+          .when(AwsApplicationSignalsCustomizerProvider::isLambdaEnvironment)
+          .thenReturn(true);
+      // Simulate an override attribute on the span
+      when(attributesMock.get(AWS_LAMBDA_LOCAL_OPERATION_OVERRIDE)).thenReturn("MyOverrideOp");
+
+      String actualOperation = AwsSpanProcessingUtil.getIngressOperation(spanDataMock);
+      assertThat(actualOperation).isEqualTo("MyOverrideOp");
+    }
+  }
+
+  @Test
+  public void testGetIngressOperationLambdaDefault() throws Exception {
+    try (
+    // Mock the AWS environment check
+    MockedStatic<AwsApplicationSignalsCustomizerProvider> providerStatic =
+            mockStatic(
+                AwsApplicationSignalsCustomizerProvider.class,
+                withSettings().defaultAnswer(CALLS_REAL_METHODS));
+        // Mock only getFunctionNameFromEnv, leave all other util logic untouched
+        MockedStatic<AwsSpanProcessingUtil> utilStatic =
+            mockStatic(
+                AwsSpanProcessingUtil.class, withSettings().defaultAnswer(CALLS_REAL_METHODS))) {
+      // force lambda branch and no override attribute
+      providerStatic
+          .when(AwsApplicationSignalsCustomizerProvider::isLambdaEnvironment)
+          .thenReturn(true);
+      when(attributesMock.get(AWS_LAMBDA_LOCAL_OPERATION_OVERRIDE)).thenReturn(null);
+      // Provide a deterministic function name
+      utilStatic.when(AwsSpanProcessingUtil::getFunctionNameFromEnv).thenReturn("MockFunction");
+
+      String actual = AwsSpanProcessingUtil.getIngressOperation(spanDataMock);
+      assertThat(actual).isEqualTo("MockFunction/FunctionHandler");
+    }
+  }
+
   @Test
   public void testGetEgressOperationUseInternalOperation() {
     String invalidName = null;