From b7706a84c823bb381862c934daf0ab8d9e59e335 Mon Sep 17 00:00:00 2001 From: Wang Date: Thu, 13 Jun 2024 17:10:40 -0700 Subject: [PATCH] add xray sdk language and version in subsegment in Lambda environment --- .../xray/contexts/LambdaSegmentContext.java | 1 + .../java/com/amazonaws/xray/AWSXRayRecorderTest.java | 2 +- .../xray/contexts/LambdaSegmentContextTest.java | 12 ++++++++++++ 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java b/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java index bb92a166..b03f1f43 100644 --- a/aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java +++ b/aws-xray-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 index 812c03a5..200558f4 100644 --- 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/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java b/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java index 2f7a728b..04f1debf 100644 --- a/aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/contexts/LambdaSegmentContextTest.java +++ b/aws-xray-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