add xray sdk language and version in subsegment in Lambda environment

aws · Jun 14, 2024 · b7706a8 · b7706a8
1 parent 63ace4a
commit b7706a8
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/...ray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java b/...ray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java
@@ -80,6 +80,7 @@ public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
                     ? new SubsegmentImpl(recorder, name, parentSegment)
                     : Subsegment.noOp(parentSegment, recorder);
             subsegment.setParent(parentSegment);
+            subsegment.putAllAws(recorder.getAwsRuntimeContext());
             // Enable FacadeSegment to keep track of its subsegments for subtree streaming
             parentSegment.addSubsegment(subsegment);
             setTraceEntity(subsegment);

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
@@ -341,7 +341,7 @@ public void testSubsegmentEmittedInLambdaContext() throws JSONException {
 
         JSONAssert.assertEquals(expectedLambdaSubsegment(
             header.getRootTraceId(), header.getParentId(), captured.getId(), captured.getStartTime(),
-            captured.getEndTime()).toString(), captured.streamSerialize(), JSONCompareMode.NON_EXTENSIBLE);
+            captured.getEndTime()).toString(), captured.streamSerialize(), JSONCompareMode.LENIENT);
     }
 
     @Test

diff --git a/...recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java b/...recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java
@@ -122,6 +122,18 @@ void testEndSubsegmentUsesContextMissing() {
         assertThatThrownBy(AWSXRay::endSubsegment).isInstanceOf(SubsegmentNotFoundException.class);
     }
 
+    @Test
+    @SetEnvironmentVariable(key = "_X_AMZN_TRACE_ID", value = TRACE_HEADER)
+    void testSampledSetsAwsXRaySdkVersionToSubsegment() {
+        LambdaSegmentContext lsc = new LambdaSegmentContext();
+        Subsegment subseg1 = lsc.beginSubsegment(AWSXRay.getGlobalRecorder(), "test");
+        Subsegment subseg2 = lsc.beginSubsegment(AWSXRay.getGlobalRecorder(), "test2");
+        assertThat(subseg1.getAws().get("xray")).isNotNull();
+        assertThat(subseg2.getAws().get("xray")).isNull();
+        lsc.endSubsegment(AWSXRay.getGlobalRecorder());
+        lsc.endSubsegment(AWSXRay.getGlobalRecorder());
+    }
+
     // We create segments twice with different environment variables for the same context, similar to how Lambda would invoke
     // a function.
     @Nested