cache method for getXrayTraceId, and use MethodHandle instead

Author: jj22ee

lambda-layer/patches/opentelemetry-java-instrumentation.patch

@@ -147,10 +147,10 @@ index 9341bf6f79..f719c1ea93 100644
  }
 diff --git a/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ParentContextExtractor.java b/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ParentContextExtractor.java
 new file mode 100644
-index 0000000000..9bea9808a5
+index 0000000000..921a7b6b7f
 --- /dev/null
 +++ b/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/main/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ParentContextExtractor.java
-@@ -0,0 +1,76 @@
+@@ -0,0 +1,89 @@
 +/*
 + * Copyright The OpenTelemetry Authors
 + * SPDX-License-Identifier: Apache-2.0
@@ -162,6 +162,9 @@ index 0000000000..9bea9808a5
 +
 +import io.opentelemetry.context.Context;
 +import io.opentelemetry.context.propagation.TextMapGetter;
++import java.lang.invoke.MethodHandle;
++import java.lang.invoke.MethodHandles;
++import java.lang.invoke.MethodType;
 +import java.util.Locale;
 +import java.util.Map;
 +import java.util.logging.Logger;
@@ -177,6 +180,7 @@ index 0000000000..9bea9808a5
 +  private static final String AWS_TRACE_HEADER_PROP = "com.amazonaws.xray.traceHeader";
 +  // lower-case map getter used for extraction
 +  static final String AWS_TRACE_HEADER_PROPAGATOR_KEY = "x-amzn-trace-id";
++  static volatile MethodHandle getXrayTraceIdHandle;
 +
 +  static Context extract(
 +      Map<String, String> headers,
@@ -194,10 +198,19 @@ index 0000000000..9bea9808a5
 +      com.amazonaws.services.lambda.runtime.Context lambdaContext) {
 +    String traceHeader = null;
 +    try {
-+      // Use reflection to avoid muzzle issues with getXrayTraceId method
-+      traceHeader = (String) lambdaContext.getClass().getMethod("getXrayTraceId").invoke(lambdaContext);
-+    } catch (Exception e) {
-+      logger.fine("Failed to get X-Ray trace ID via reflection: " + e.getMessage());
++      MethodHandle handle = getXrayTraceIdHandle;
++      if (handle == null) {
++        handle =
++            MethodHandles.lookup()
++                .findVirtual(
++                    lambdaContext.getClass(),
++                    "getXrayTraceId",
++                    MethodType.methodType(String.class));
++        getXrayTraceIdHandle = handle;
++      }
++      traceHeader = (String) handle.invoke(lambdaContext);
++    } catch (Throwable t) {
++      logger.fine("Failed to get X-Ray trace ID via MethodHandle: " + t.getMessage());
 +    }
 +    if (traceHeader != null && !traceHeader.isEmpty()) {
 +      return traceHeader;
@@ -242,10 +255,10 @@ index cb19d1e568..12ed174bb2 100644
      assertThat(spanContext.getSpanId()).isEqualTo("00f067aa0ba902b7");
 diff --git a/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ParentContextExtractorTest.java b/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ParentContextExtractorTest.java
 new file mode 100644
-index 0000000000..edca3ce68c
+index 0000000000..e2a69ee6d8
 --- /dev/null
 +++ b/instrumentation/aws-lambda/aws-lambda-core-1.0/library/src/test/java/io/opentelemetry/instrumentation/awslambdacore/v1_0/internal/ParentContextExtractorTest.java
-@@ -0,0 +1,359 @@
+@@ -0,0 +1,418 @@
 +/*
 + * Copyright The OpenTelemetry Authors
 + * SPDX-License-Identifier: Apache-2.0
@@ -255,6 +268,8 @@ index 0000000000..edca3ce68c
 +
 +import static org.assertj.core.api.Assertions.assertThat;
 +import static org.mockito.Mockito.mock;
++import static org.mockito.Mockito.times;
++import static org.mockito.Mockito.verify;
 +import static org.mockito.Mockito.when;
 +
 +import com.amazonaws.services.lambda.runtime.Context;
@@ -266,6 +281,7 @@ index 0000000000..edca3ce68c
 +import io.opentelemetry.context.propagation.TextMapPropagator;
 +import io.opentelemetry.contrib.awsxray.propagator.AwsXrayPropagator;
 +import io.opentelemetry.extension.trace.propagation.B3Propagator;
++import java.lang.invoke.MethodHandle;
 +import java.util.Map;
 +import org.junit.jupiter.api.Test;
 +import org.junit.jupiter.api.extension.ExtendWith;
@@ -587,7 +603,8 @@ index 0000000000..edca3ce68c
 +  void shouldFallbackToSystemPropertyWhenReflectionFails() {
 +    // given
 +    Map<String, String> headers = ImmutableMap.of();
-+    // Create a mock context that doesn't have getXrayTraceId method (simulates older Lambda runtime)
++    // Create a mock context that doesn't have getXrayTraceId method (simulates older Lambda
++    // runtime)
 +    Context mockLambdaContextWithoutXrayMethod = mock(Context.class);
 +    systemProperties.set(
 +        "com.amazonaws.xray.traceHeader",
@@ -604,6 +621,61 @@ index 0000000000..edca3ce68c
 +    assertThat(spanContext.getSpanId()).isEqualTo("0000000000000789");
 +    assertThat(spanContext.getTraceId()).isEqualTo("8a3c60f7d188f8fa79d48a391a778fa7");
 +  }
++
++  @Test
++  void shouldFallbackToSystemPropertyWhenIllegalAccessException() {
++    // given
++    Map<String, String> headers = ImmutableMap.of();
++    Context mockLambdaContextWithAccessException = mock(Context.class);
++    when(mockLambdaContextWithAccessException.getXrayTraceId())
++        .thenThrow(new RuntimeException("Simulated IllegalAccessException"));
++    systemProperties.set(
++        "com.amazonaws.xray.traceHeader",
++        "Root=1-8a3c60f7-d188f8fa79d48a391a778fa7;Parent=0000000000000789;Sampled=1");
++
++    // when
++    io.opentelemetry.context.Context context =
++        ParentContextExtractor.extract(
++            headers, INSTRUMENTER_WITH_B3_XRAY_PROPAGATORS, mockLambdaContextWithAccessException);
++    // then
++    Span span = Span.fromContext(context);
++    SpanContext spanContext = span.getSpanContext();
++    assertThat(spanContext.isValid()).isTrue();
++    assertThat(spanContext.getSpanId()).isEqualTo("0000000000000789");
++    assertThat(spanContext.getTraceId()).isEqualTo("8a3c60f7d188f8fa79d48a391a778fa7");
++  }
++
++  @Test
++  void shouldCacheMethodHandleAcrossMultipleInvocations() {
++    // given
++    Context mockLambdaContext = mock(Context.class);
++    when(mockLambdaContext.getXrayTraceId())
++        .thenReturn("Root=1-4fd0b613-1f19f39af59518d127b0cafe;Parent=0000000000000123;Sampled=1");
++
++    // Clear any previously cached MethodHandle
++    ParentContextExtractor.getXrayTraceIdHandle = null;
++
++    // when - call extract multiple times (which internally calls getTraceHeader)
++    assertThat(ParentContextExtractor.getXrayTraceIdHandle).isNull(); // Initially null
++
++    ParentContextExtractor.extract(
++        ImmutableMap.of(), INSTRUMENTER_WITH_B3_XRAY_PROPAGATORS, mockLambdaContext);
++    assertThat(ParentContextExtractor.getXrayTraceIdHandle).isNotNull(); // Cached after first call
++
++    MethodHandle cachedHandle = ParentContextExtractor.getXrayTraceIdHandle;
++
++    ParentContextExtractor.extract(
++        ImmutableMap.of(), INSTRUMENTER_WITH_B3_XRAY_PROPAGATORS, mockLambdaContext);
++    ParentContextExtractor.extract(
++        ImmutableMap.of(), INSTRUMENTER_WITH_B3_XRAY_PROPAGATORS, mockLambdaContext);
++
++    // then - MethodHandle should remain the same (cached)
++    assertThat(ParentContextExtractor.getXrayTraceIdHandle).isSameAs(cachedHandle);
++
++    // verify getXrayTraceId was called for each invocation (MethodHandle caching doesn't affect
++    // this)
++    verify(mockLambdaContext, times(3)).getXrayTraceId();
++  }
 +}
 diff --git a/instrumentation/aws-lambda/aws-lambda-events-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awslambdaevents/v2_2/AwsLambdaRequestHandlerInstrumentation.java b/instrumentation/aws-lambda/aws-lambda-events-2.2/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/awslambdaevents/v2_2/AwsLambdaRequestHandlerInstrumentation.java
 index e059250807..1fa80c3735 100644