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 <[email protected]>

Author: rohitnair

aws-xray-recorder-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);

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<>();

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)