Optimize default resolver chain construction (#387)

* Optimize default resolver chain construction We can figure out the right SegmentContextResolver at init time since it should remain static throughout the lifetime of an application. Trying to resolve through the resolver chain each time at runtime seems unnecessary. --------- Co-authored-by: atshaw43 <108552302+atshaw43@users.noreply.github.com>
aws · Jul 24, 2023 · 398ec05 · 398ec05
1 parent babdde3
commit 398ec05
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 17 deletions.
diff --git a/...order-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java b/...order-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java
@@ -73,18 +73,18 @@ public void testSamplingOverrideFalseInLambda() throws Exception {
     public void testSamplingOverrideTrueInLambda() {
         Emitter mockedEmitted = Mockito.mock(DefaultEmitter.class);
 
-        AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard()
-                .withSamplingStrategy(new NoSamplingStrategy())
-                .withEmitter(mockedEmitted)
-                .build();
-
         TraceHeader header = TraceHeader.fromString(TRACE_HEADER);
 
         PowerMockito.stub(PowerMockito.method(
                 LambdaSegmentContext.class, "getTraceHeaderFromEnvironment")).toReturn(header);
         PowerMockito.stub(PowerMockito.method(
                 LambdaSegmentContextResolver.class, "getLambdaTaskRoot")).toReturn("/var/task");
 
+        AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard()
+                .withSamplingStrategy(new NoSamplingStrategy())
+                .withEmitter(mockedEmitted)
+                .build();
+
         Mockito.doAnswer(invocation -> { return true; }).when(mockedEmitted).sendSubsegment(any());
 
         lambdaTestHelper(recorder, "testTrue", true);
@@ -95,18 +95,18 @@ public void testSamplingOverrideTrueInLambda() {
     public void testSamplingOverrideMixedInLambda() {
         Emitter mockedEmitted = Mockito.mock(DefaultEmitter.class);
 
-        AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard()
-                .withSamplingStrategy(new NoSamplingStrategy())
-                .withEmitter(mockedEmitted)
-                .build();
-
         TraceHeader header = TraceHeader.fromString(TRACE_HEADER);
 
         PowerMockito.stub(PowerMockito.method(
                 LambdaSegmentContext.class, "getTraceHeaderFromEnvironment")).toReturn(header);
         PowerMockito.stub(PowerMockito.method(
                 LambdaSegmentContextResolver.class, "getLambdaTaskRoot")).toReturn("/var/task");
 
+        AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard()
+                .withSamplingStrategy(new NoSamplingStrategy())
+                .withEmitter(mockedEmitted)
+                .build();
+
         Mockito.doAnswer(invocation -> { return true; }).when(mockedEmitted).sendSubsegment(any());
 
         lambdaTestHelper(recorder, "test1", true);

diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/AWSXRayRecorder.java
@@ -164,9 +164,13 @@ public AWSXRayRecorder() {
         }
 
         segmentContextResolverChain = new SegmentContextResolverChain();
-        segmentContextResolverChain.addResolver(new LambdaSegmentContextResolver());
-        segmentContextResolverChain.addResolver(new ThreadLocalSegmentContextResolver());
-
+        LambdaSegmentContextResolver lambdaSegmentContextResolver = new LambdaSegmentContextResolver();
+        if (lambdaSegmentContextResolver.resolve() != null) {
+            segmentContextResolverChain.addResolver(lambdaSegmentContextResolver);
+        } else {
+            segmentContextResolverChain.addResolver(new ThreadLocalSegmentContextResolver());
+        }
+
         segmentListeners = new ArrayList<>();
 
         awsRuntimeContext = new ConcurrentHashMap<>();

diff --git a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java
@@ -967,10 +967,6 @@ public void testDefaultContextMissingBehaviorEnd() {
 
     @Test
     public void testMalformedTraceId() {
-        AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard()
-            .withSamplingStrategy(new NoSamplingStrategy())
-            .build();
-
         TraceHeader malformedHeader = TraceHeader.fromString("malformedTraceID");
 
         PowerMockito.stub(PowerMockito.method(
@@ -981,6 +977,10 @@ public void testMalformedTraceId() {
                 LambdaSegmentContextResolver.class, "getLambdaTaskRoot"))
                 .toReturn("/var/task");
 
+        AWSXRayRecorder recorder = AWSXRayRecorderBuilder.standard()
+            .withSamplingStrategy(new NoSamplingStrategy())
+            .build();
+
         recorder.beginSubsegment("Test");
 
         // Sanity checks that this is a NoOpSubsegment. (We cannot compare the instance of the private class directly)