Skip to content

Commit 93a5877

Browse files
atshaw43atshaw43
authored
Refactoring TraceHeader Calls (#361)
* Refactoring TraceHeader
1 parent fd93c8a commit 93a5877

File tree

12 files changed

+80
-115
lines changed

12 files changed

+80
-115
lines changed

‎aws-xray-recorder-sdk-apache-http/src/main/java/com/amazonaws/xray/proxies/apache/http/TracedHttpClient.java

+1-5
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import com.amazonaws.xray.entities.Segment;
2222
import com.amazonaws.xray.entities.Subsegment;
2323
import com.amazonaws.xray.entities.TraceHeader;
24-
import com.amazonaws.xray.entities.TraceHeader.SampleDecision;
2524
import java.io.IOException;
2625
import java.net.URI;
2726
import java.net.URISyntaxException;
@@ -105,10 +104,7 @@ public static void addRequestInformation(Subsegment subsegment, HttpRequest requ
105104
Segment parentSegment = subsegment.getParentSegment();
106105

107106
if (subsegment.shouldPropagate()) {
108-
TraceHeader header = new TraceHeader(parentSegment.getTraceId(),
109-
parentSegment.isSampled() ? subsegment.getId() : null,
110-
parentSegment.isSampled() ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
111-
request.addHeader(TraceHeader.HEADER_KEY, header.toString());
107+
request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(subsegment).toString());
112108
}
113109

114110
Map<String, Object> requestInformation = new HashMap<>();

‎aws-xray-recorder-sdk-aws-sdk-v2/src/main/java/com/amazonaws/xray/interceptors/TracingInterceptor.java

+3-8
Original file line numberDiff line numberDiff line change
@@ -249,14 +249,9 @@ public SdkHttpRequest modifyHttpRequest(Context.ModifyHttpRequest context, Execu
249249
return httpRequest;
250250
}
251251

252-
boolean isSampled = subsegment.getParentSegment().isSampled();
253-
TraceHeader header = new TraceHeader(
254-
subsegment.getParentSegment().getTraceId(),
255-
isSampled ? subsegment.getId() : null,
256-
isSampled ? TraceHeader.SampleDecision.SAMPLED : TraceHeader.SampleDecision.NOT_SAMPLED
257-
);
258-
259-
return httpRequest.toBuilder().appendHeader(TraceHeader.HEADER_KEY, header.toString()).build();
252+
return httpRequest.toBuilder().appendHeader(
253+
TraceHeader.HEADER_KEY,
254+
TraceHeader.fromEntity(subsegment).toString()).build();
260255
}
261256

262257
@Override

‎aws-xray-recorder-sdk-aws-sdk/src/main/java/com/amazonaws/xray/handlers/TracingHandler.java

+1-11
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,6 @@
3333
import com.amazonaws.xray.entities.Namespace;
3434
import com.amazonaws.xray.entities.Subsegment;
3535
import com.amazonaws.xray.entities.TraceHeader;
36-
import com.amazonaws.xray.entities.TraceHeader.SampleDecision;
3736
import com.amazonaws.xray.handlers.config.AWSOperationHandler;
3837
import com.amazonaws.xray.handlers.config.AWSOperationHandlerManifest;
3938
import com.amazonaws.xray.handlers.config.AWSServiceHandlerManifest;
@@ -184,7 +183,6 @@ public void beforeRequest(Request<?> request) {
184183
recorder.setTraceEntity(entityContext);
185184
}
186185

187-
Optional<Subsegment> previousSubsegment = recorder.getCurrentSubsegmentOptional();
188186
Subsegment currentSubsegment = recorder.beginSubsegment(serviceName);
189187
currentSubsegment.putAllAws(extractRequestParameters(request));
190188
currentSubsegment.putAws(EntityDataKeys.AWS.OPERATION_KEY, operationName);
@@ -194,15 +192,7 @@ public void beforeRequest(Request<?> request) {
194192
currentSubsegment.setNamespace(Namespace.AWS.toString());
195193

196194
if (recorder.getCurrentSegment() != null && recorder.getCurrentSubsegment().shouldPropagate()) {
197-
boolean isSampled = previousSubsegment.isPresent() ?
198-
previousSubsegment.get().isSampled() :
199-
recorder.getCurrentSegment().isSampled();
200-
201-
TraceHeader header =
202-
new TraceHeader(recorder.getCurrentSegment().getTraceId(),
203-
isSampled ? currentSubsegment.getId() : null,
204-
isSampled ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
205-
request.addHeader(TraceHeader.HEADER_KEY, header.toString());
195+
request.addHeader(TraceHeader.HEADER_KEY, TraceHeader.fromEntity(currentSubsegment).toString());
206196
}
207197
}
208198

‎aws-xray-recorder-sdk-aws-sdk/src/test/java/com/amazonaws/xray/handlers/TracingHandlerLambdaTest.java

+33-66
Original file line numberDiff line numberDiff line change
@@ -65,22 +65,7 @@ public void testSamplingOverrideFalseInLambda() throws Exception {
6565
.withEmitter(mockedEmitted)
6666
.build();
6767

68-
recorder.beginSubsegmentWithoutSampling("Test");
69-
70-
Subsegment subsegment = ((Subsegment) recorder.getTraceEntity());
71-
assertThat(subsegment.shouldPropagate()).isTrue();
72-
73-
DefaultRequest<Void> request = new DefaultRequest<>(new InvokeRequest(), "Test");
74-
75-
TracingHandler tracingHandler = new TracingHandler(recorder);
76-
tracingHandler.beforeRequest(request);
77-
tracingHandler.afterResponse(request, new Response(new InvokeResult(), new HttpResponse(request, null)));
78-
79-
assertThat(TraceHeader.fromString(request.getHeaders().get(TraceHeader.HEADER_KEY)).getSampled())
80-
.isEqualTo(TraceHeader.SampleDecision.NOT_SAMPLED);
81-
82-
recorder.endSubsegment();
83-
68+
lambdaTestHelper(recorder, "testFalse", false);
8469
Mockito.verify(mockedEmitted, Mockito.times(0)).sendSubsegment(any());
8570
}
8671

@@ -102,16 +87,7 @@ public void testSamplingOverrideTrueInLambda() {
10287

10388
Mockito.doAnswer(invocation -> { return true; }).when(mockedEmitted).sendSubsegment(any());
10489

105-
recorder.beginSubsegment("test1");
106-
Subsegment subsegment = ((Subsegment) recorder.getTraceEntity());
107-
assertThat(subsegment.shouldPropagate()).isTrue();
108-
DefaultRequest<Void> request = new DefaultRequest<>(new InvokeRequest(), "Test");
109-
TracingHandler tracingHandler = new TracingHandler(recorder);
110-
tracingHandler.beforeRequest(request);
111-
assertThat(TraceHeader.fromString(request.getHeaders().get(TraceHeader.HEADER_KEY)).getSampled())
112-
.isEqualTo(TraceHeader.SampleDecision.SAMPLED);
113-
tracingHandler.afterResponse(request, new Response(new InvokeResult(), new HttpResponse(request, null)));
114-
recorder.endSubsegment();
90+
lambdaTestHelper(recorder, "testTrue", true);
11591
Mockito.verify(mockedEmitted, Mockito.times(1)).sendSubsegment(any());
11692
}
11793

@@ -133,52 +109,43 @@ public void testSamplingOverrideMixedInLambda() {
133109

134110
Mockito.doAnswer(invocation -> { return true; }).when(mockedEmitted).sendSubsegment(any());
135111

136-
recorder.beginSubsegment("test1");
137-
Subsegment subsegment1 = ((Subsegment) recorder.getTraceEntity());
138-
assertThat(subsegment1.shouldPropagate()).isTrue();
139-
DefaultRequest<Void> request1 = new DefaultRequest<>(new InvokeRequest(), "Test");
140-
TracingHandler tracingHandler1 = new TracingHandler(recorder);
141-
tracingHandler1.beforeRequest(request1);
142-
assertThat(TraceHeader.fromString(request1.getHeaders().get(TraceHeader.HEADER_KEY)).getSampled())
143-
.isEqualTo(TraceHeader.SampleDecision.SAMPLED);
144-
tracingHandler1.afterResponse(request1, new Response(new InvokeResult(), new HttpResponse(request1, null)));
145-
recorder.endSubsegment();
112+
lambdaTestHelper(recorder, "test1", true);
146113
Mockito.verify(mockedEmitted, Mockito.times(1)).sendSubsegment(any());
147114

148-
recorder.beginSubsegmentWithoutSampling("test2");
149-
Subsegment subsegment2 = ((Subsegment) recorder.getTraceEntity());
150-
assertThat(subsegment2.shouldPropagate()).isTrue();
151-
DefaultRequest<Void> request2 = new DefaultRequest<>(new InvokeRequest(), "Test");
152-
TracingHandler tracingHandler2 = new TracingHandler(recorder);
153-
tracingHandler2.beforeRequest(request2);
154-
assertThat(TraceHeader.fromString(request2.getHeaders().get(TraceHeader.HEADER_KEY)).getSampled())
155-
.isEqualTo(TraceHeader.SampleDecision.NOT_SAMPLED);
156-
tracingHandler2.afterResponse(request2, new Response(new InvokeResult(), new HttpResponse(request2, null)));
157-
recorder.endSubsegment();
115+
lambdaTestHelper(recorder, "test2", false);
158116
Mockito.verify(mockedEmitted, Mockito.times(1)).sendSubsegment(any());
159117

160-
recorder.beginSubsegment("test3");
161-
Subsegment subsegment3 = ((Subsegment) recorder.getTraceEntity());
162-
assertThat(subsegment3.shouldPropagate()).isTrue();
163-
DefaultRequest<Void> request3 = new DefaultRequest<>(new InvokeRequest(), "Test");
164-
TracingHandler tracingHandler3 = new TracingHandler(recorder);
165-
tracingHandler3.beforeRequest(request3);
166-
assertThat(TraceHeader.fromString(request3.getHeaders().get(TraceHeader.HEADER_KEY)).getSampled())
167-
.isEqualTo(TraceHeader.SampleDecision.SAMPLED);
168-
tracingHandler3.afterResponse(request3, new Response(new InvokeResult(), new HttpResponse(request3, null)));
169-
recorder.endSubsegment();
118+
lambdaTestHelper(recorder, "test3", true);
170119
Mockito.verify(mockedEmitted, Mockito.times(2)).sendSubsegment(any());
171120

172-
recorder.beginSubsegmentWithoutSampling("test4");
173-
Subsegment subsegment4 = ((Subsegment) recorder.getTraceEntity());
174-
assertThat(subsegment4.shouldPropagate()).isTrue();
175-
DefaultRequest<Void> request4 = new DefaultRequest<>(new InvokeRequest(), "Test");
176-
TracingHandler tracingHandler4 = new TracingHandler(recorder);
177-
tracingHandler4.beforeRequest(request4);
178-
assertThat(TraceHeader.fromString(request4.getHeaders().get(TraceHeader.HEADER_KEY)).getSampled())
179-
.isEqualTo(TraceHeader.SampleDecision.NOT_SAMPLED);
180-
tracingHandler4.afterResponse(request4, new Response(new InvokeResult(), new HttpResponse(request4, null)));
181-
recorder.endSubsegment();
121+
lambdaTestHelper(recorder, "test4", false);
182122
Mockito.verify(mockedEmitted, Mockito.times(2)).sendSubsegment(any());
183123
}
124+
125+
public static void lambdaTestHelper(AWSXRayRecorder recorder, String name, boolean sampled) {
126+
if (sampled) {
127+
recorder.beginSubsegment(name);
128+
} else {
129+
recorder.beginSubsegmentWithoutSampling(name);
130+
}
131+
132+
Subsegment subsegment = ((Subsegment) recorder.getTraceEntity());
133+
assertThat(subsegment.shouldPropagate()).isTrue();
134+
DefaultRequest<Void> request = new DefaultRequest<>(new InvokeRequest(), "Test");
135+
TracingHandler tracingHandler = new TracingHandler(recorder);
136+
tracingHandler.beforeRequest(request);
137+
TraceHeader traceHeader = TraceHeader.fromString(request.getHeaders().get(TraceHeader.HEADER_KEY));
138+
String serviceEntityId = recorder.getCurrentSubsegment().getId();
139+
140+
assertThat(traceHeader.getSampled()).isEqualTo(
141+
subsegment.isSampled() ?
142+
TraceHeader.SampleDecision.SAMPLED :
143+
TraceHeader.SampleDecision.NOT_SAMPLED);
144+
assertThat(traceHeader.getRootTraceId()).isEqualTo(subsegment.getTraceId());
145+
assertThat(traceHeader.getParentId()).isEqualTo(subsegment.isSampled() ? serviceEntityId : null);
146+
147+
tracingHandler.afterResponse(request, new Response(new InvokeResult(), new HttpResponse(request, null)));
148+
149+
recorder.endSubsegment();
150+
}
184151
}

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/contexts/LambdaSegmentContext.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -91,11 +91,11 @@ public Subsegment beginSubsegment(AWSXRayRecorder recorder, String name) {
9191
}
9292
Segment parentSegment = parentSubsegment.getParentSegment();
9393

94-
boolean isRecording = parentSegment.isRecording();
94+
boolean isRecording = parentSubsegment.isRecording();
9595

9696
Subsegment subsegment = isRecording
9797
? new SubsegmentImpl(recorder, name, parentSegment)
98-
: Subsegment.noOp(parentSegment, recorder);
98+
: Subsegment.noOp(parentSegment, recorder, name);
9999
subsegment.setParent(parentSubsegment);
100100
parentSubsegment.addSubsegment(subsegment);
101101
setTraceEntity(subsegment);

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Entity.java

+2
Original file line numberDiff line numberDiff line change
@@ -600,6 +600,8 @@ default void run(Runnable runnable, AWSXRayRecorder recorder) {
600600

601601
boolean isEmitted();
602602

603+
boolean isSampled();
604+
603605
/**
604606
* Sets emitted on the entity.
605607
*/

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/NoOpSubSegment.java

+20-7
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@
2626
import org.checkerframework.checker.nullness.qual.Nullable;
2727

2828
class NoOpSubSegment implements Subsegment {
29-
3029
private final Segment parentSegment;
3130
private final AWSXRayRecorder creator;
3231
private final boolean shouldPropagate;
@@ -39,9 +38,21 @@ class NoOpSubSegment implements Subsegment {
3938
@JsonIgnore
4039
private boolean isRecording;
4140

41+
@JsonIgnore
42+
private String name;
43+
44+
@JsonIgnore
45+
@Nullable
46+
private String namespace;
47+
4248
@JsonIgnore
4349
private SamplingStrategyOverride samplingStrategyOverride;
4450

51+
NoOpSubSegment(Segment parentSegment, AWSXRayRecorder creator, String name) {
52+
this(parentSegment, creator);
53+
this.name = name;
54+
}
55+
4556
NoOpSubSegment(Segment parentSegment, AWSXRayRecorder creator) {
4657
this(parentSegment, creator, true);
4758
}
@@ -65,11 +76,11 @@ class NoOpSubSegment implements Subsegment {
6576
this.creator = creator;
6677
this.shouldPropagate = shouldPropagate;
6778
parent = parentSegment;
79+
name = "";
80+
namespace = "";
6881

69-
this.isSampled = samplingStrategyOverride == SamplingStrategyOverride.DISABLED ?
70-
parentSegment.isSampled() :
71-
false;
72-
this.isRecording = isSampled;
82+
this.isSampled = false;
83+
this.isRecording = false;
7384
this.samplingStrategyOverride = samplingStrategyOverride;
7485
}
7586

@@ -80,7 +91,7 @@ public boolean end() {
8091

8192
@Override
8293
public String getName() {
83-
return "";
94+
return name;
8495
}
8596

8697
@Override
@@ -129,12 +140,14 @@ public void setError(boolean error) {
129140
}
130141

131142
@Override
143+
@Nullable
132144
public String getNamespace() {
133-
return "";
145+
return namespace;
134146
}
135147

136148
@Override
137149
public void setNamespace(String namespace) {
150+
this.namespace = namespace;
138151
}
139152

140153
@Override

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Segment.java

-5
Original file line numberDiff line numberDiff line change
@@ -40,11 +40,6 @@ static Segment noOp(TraceID traceId, AWSXRayRecorder recorder) {
4040
@JsonIgnore
4141
boolean isRecording();
4242

43-
/**
44-
* @return the sampled
45-
*/
46-
boolean isSampled();
47-
4843
/**
4944
* @param sampled
5045
* the sampled to set

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/Subsegment.java

+4-3
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,10 @@ static Subsegment noOp(AWSXRayRecorder recorder, boolean shouldPropagate) {
4040
return new NoOpSubSegment(Segment.noOp(TraceID.invalid(), recorder), recorder, shouldPropagate);
4141
}
4242

43+
static Subsegment noOp(Segment parent, AWSXRayRecorder recorder, String name) {
44+
return new NoOpSubSegment(parent, recorder, name);
45+
}
46+
4347
static Subsegment noOp(Segment parent, AWSXRayRecorder recorder) {
4448
return new NoOpSubSegment(parent, recorder);
4549
}
@@ -124,9 +128,6 @@ static Subsegment noOp(Segment parent, AWSXRayRecorder recorder) {
124128
@Override
125129
void close();
126130

127-
@JsonIgnore
128-
boolean isSampled();
129-
130131
@JsonIgnore
131132
boolean isRecording();
132133

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/SubsegmentImpl.java

+6
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,12 @@ public boolean end() {
9393
return shouldEmit;
9494
}
9595

96+
@Override
97+
@JsonIgnore
98+
public TraceID getTraceId() {
99+
return parentSegment.getTraceId();
100+
}
101+
96102
@Override
97103
@Nullable
98104
public String getNamespace() {

‎aws-xray-recorder-sdk-core/src/main/java/com/amazonaws/xray/entities/TraceHeader.java

+7
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,13 @@ public TraceHeader(@Nullable TraceID rootTraceId, @Nullable String parentId, Sam
9191
}
9292
}
9393

94+
public static TraceHeader fromEntity(Entity entity) {
95+
return new TraceHeader(
96+
entity.getTraceId(),
97+
entity.isSampled() ? entity.getId() : null,
98+
entity.isSampled() ? SampleDecision.SAMPLED : SampleDecision.NOT_SAMPLED);
99+
}
100+
94101
/**
95102
* Creates a TraceHeader object from a String. Note that this will silently ignore any "Self=" trace ids injected from ALB.
96103
*

‎aws-xray-recorder-sdk-core/src/test/java/com/amazonaws/xray/AWSXRayRecorderTest.java

+1-8
Original file line numberDiff line numberDiff line change
@@ -507,9 +507,6 @@ public void testBeginSubsegmentWhenMissingContext() {
507507
Subsegment subsegment = recorder.beginSubsegment("hello");
508508
assertThat(subsegment).isNotNull();
509509
assertThat(subsegment.getNamespace()).isEmpty();
510-
// No-op
511-
subsegment.setNamespace("foo");
512-
assertThat(subsegment.getNamespace()).isEmpty();
513510
assertThat(subsegment.shouldPropagate()).isFalse();
514511
}
515512

@@ -801,8 +798,6 @@ public void testNoOpSubsegment() throws Exception {
801798
// Semantically relevant
802799
assertThat(subsegment.end()).isFalse();
803800
assertThat(subsegment.getName()).isEmpty();
804-
subsegment.setId("foo");
805-
assertThat(subsegment.getId()).isEmpty();
806801
subsegment.setStartTime(100);
807802
assertThat(subsegment.getStartTime()).isZero();
808803
subsegment.setEndTime(100);
@@ -811,7 +806,6 @@ public void testNoOpSubsegment() throws Exception {
811806
assertThat(subsegment.isFault()).isFalse();
812807
subsegment.setError(true);
813808
assertThat(subsegment.isError()).isFalse();
814-
subsegment.setNamespace("foo");
815809
assertThat(subsegment.getNamespace()).isEmpty();
816810
subsegment.setSubsegmentsLock(new ReentrantLock());
817811
Thread thread = new Thread(() -> subsegment.getSubsegmentsLock().lock());
@@ -990,9 +984,8 @@ public void testMalformedTraceId() {
990984
recorder.beginSubsegment("Test");
991985

992986
// Sanity checks that this is a NoOpSubsegment. (We cannot compare the instance of the private class directly)
993-
assertThat(recorder.getTraceEntity().getId()).isEqualTo("");
994987
assertThat(recorder.getTraceEntity().getName()).isEqualTo("");
995-
assertThat(recorder.getTraceEntity().getNamespace()).isEqualTo("");
988+
assertThat(recorder.getTraceEntity().getNamespace()).isEmpty();
996989
assertThat(recorder.getTraceEntity().getStartTime()).isEqualTo(0);
997990
assertThat(recorder.getTraceEntity().getEndTime()).isEqualTo(0);
998991

0 commit comments

Comments
 (0)