Skip to content

Commit 167a943

Browse files
moschemosche
authored
Improve newTraceContext to only modify the current thread context if a trace context exists. (#130989)
Improve newTraceContext to only modify the current thread context if a trace context exists. Otherwise, `newTraceContext` will behave exactly as #newStoredContextPreservingResponseHeaders as there's no reason to cleanup / move tracing headers. Relates to #ES-10969
1 parent 3716725 commit 167a943

File tree

3 files changed

+125
-27
lines changed

3 files changed

+125
-27
lines changed

‎server/src/main/java/org/elasticsearch/common/util/concurrent/ThreadContext.java‎

Lines changed: 49 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -189,37 +189,46 @@ public StoredContext newEmptySystemContext() {
189189
* moving tracing-related fields to different names so that a new child span can be started. This child span will pick up
190190
* the moved fields and use them to establish the parent-child relationship.
191191
*
192+
* Response headers will be propagated. If no parent span is in progress (meaning there's no trace context), this will behave exactly
193+
* the same way as {@link #newStoredContextPreservingResponseHeaders}.
194+
*
192195
* @return a stored context, which can be restored when this context is no longer needed.
193196
*/
194197
public StoredContext newTraceContext() {
195198
final ThreadContextStruct originalContext = threadLocal.get();
196-
final Map<String, String> newRequestHeaders = new HashMap<>(originalContext.requestHeaders);
197-
final Map<String, Object> newTransientHeaders = new HashMap<>(originalContext.transientHeaders);
198199

199-
final String previousTraceParent = newRequestHeaders.remove(Task.TRACE_PARENT_HTTP_HEADER);
200-
if (previousTraceParent != null) {
201-
newTransientHeaders.put(Task.PARENT_TRACE_PARENT_HEADER, previousTraceParent);
202-
}
200+
// this is the context when this method returns
201+
final ThreadContextStruct newContext;
202+
if (originalContext.hasTraceContext() == false) {
203+
newContext = originalContext;
204+
} else {
205+
final Map<String, String> newRequestHeaders = new HashMap<>(originalContext.requestHeaders);
206+
final Map<String, Object> newTransientHeaders = new HashMap<>(originalContext.transientHeaders);
203207

204-
final String previousTraceState = newRequestHeaders.remove(Task.TRACE_STATE);
205-
if (previousTraceState != null) {
206-
newTransientHeaders.put(Task.PARENT_TRACE_STATE, previousTraceState);
207-
}
208+
final String previousTraceParent = newRequestHeaders.remove(Task.TRACE_PARENT_HTTP_HEADER);
209+
if (previousTraceParent != null) {
210+
newTransientHeaders.put(Task.PARENT_TRACE_PARENT_HEADER, previousTraceParent);
211+
}
208212

209-
final Object previousTraceContext = newTransientHeaders.remove(Task.APM_TRACE_CONTEXT);
210-
if (previousTraceContext != null) {
211-
newTransientHeaders.put(Task.PARENT_APM_TRACE_CONTEXT, previousTraceContext);
212-
}
213+
final String previousTraceState = newRequestHeaders.remove(Task.TRACE_STATE);
214+
if (previousTraceState != null) {
215+
newTransientHeaders.put(Task.PARENT_TRACE_STATE, previousTraceState);
216+
}
213217

214-
// this is the context when this method returns
215-
final ThreadContextStruct newContext = new ThreadContextStruct(
216-
newRequestHeaders,
217-
originalContext.responseHeaders,
218-
newTransientHeaders,
219-
originalContext.isSystemContext,
220-
originalContext.warningHeadersSize
221-
);
222-
threadLocal.set(newContext);
218+
final Object previousTraceContext = newTransientHeaders.remove(Task.APM_TRACE_CONTEXT);
219+
if (previousTraceContext != null) {
220+
newTransientHeaders.put(Task.PARENT_APM_TRACE_CONTEXT, previousTraceContext);
221+
}
222+
223+
newContext = new ThreadContextStruct(
224+
newRequestHeaders,
225+
originalContext.responseHeaders,
226+
newTransientHeaders,
227+
originalContext.isSystemContext,
228+
originalContext.warningHeadersSize
229+
);
230+
threadLocal.set(newContext);
231+
}
223232
// Tracing shouldn't interrupt the propagation of response headers, so in the same as
224233
// #newStoredContextPreservingResponseHeaders(), pass on any potential changes to the response headers.
225234
return () -> {
@@ -233,10 +242,11 @@ public StoredContext newTraceContext() {
233242
}
234243

235244
public boolean hasTraceContext() {
236-
final ThreadContextStruct context = threadLocal.get();
237-
return context.requestHeaders.containsKey(Task.TRACE_PARENT_HTTP_HEADER)
238-
|| context.requestHeaders.containsKey(Task.TRACE_STATE)
239-
|| context.transientHeaders.containsKey(Task.APM_TRACE_CONTEXT);
245+
return threadLocal.get().hasTraceContext();
246+
}
247+
248+
public boolean hasParentTraceContext() {
249+
return threadLocal.get().hasParentTraceContext();
240250
}
241251

242252
/**
@@ -853,6 +863,18 @@ private ThreadContextStruct putResponseHeaders(Map<String, Set<String>> headers)
853863
return new ThreadContextStruct(requestHeaders, newResponseHeaders, transientHeaders, isSystemContext);
854864
}
855865

866+
private boolean hasTraceContext() {
867+
return requestHeaders.containsKey(Task.TRACE_PARENT_HTTP_HEADER)
868+
|| requestHeaders.containsKey(Task.TRACE_STATE)
869+
|| transientHeaders.containsKey(Task.APM_TRACE_CONTEXT);
870+
}
871+
872+
private boolean hasParentTraceContext() {
873+
return transientHeaders.containsKey(Task.PARENT_TRACE_PARENT_HEADER)
874+
|| transientHeaders.containsKey(Task.PARENT_TRACE_STATE)
875+
|| transientHeaders.containsKey(Task.PARENT_APM_TRACE_CONTEXT);
876+
}
877+
856878
private void logWarningHeaderThresholdExceeded(long threshold, Setting<?> thresholdSetting) {
857879
// If available, log some selected headers to help identifying the source of the request.
858880
// Note: Only Task.HEADERS_TO_COPY are guaranteed to be preserved at this point.

‎server/src/main/java/org/elasticsearch/tasks/Task.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ public class Task implements Traceable {
4444
* TRACE_PARENT once parsed in RestController.tryAllHandler is not preserved
4545
* has to be declared as a header copied over from http request.
4646
* May also be used internally when APM is enabled.
47+
* https://www.w3.org/TR/trace-context-1/#traceparent-header
4748
*/
4849
public static final String TRACE_PARENT_HTTP_HEADER = "traceparent";
4950

@@ -53,6 +54,10 @@ public class Task implements Traceable {
5354
*/
5455
public static final String TRACE_ID = "trace.id";
5556

57+
/**
58+
* Optional request header carrying vendor-specific trace information.
59+
* https://www.w3.org/TR/trace-context-1/#tracestate-header
60+
*/
5661
public static final String TRACE_STATE = "tracestate";
5762

5863
public static final String TRACE_START_TIME = "trace.starttime";

‎server/src/test/java/org/elasticsearch/common/util/concurrent/ThreadContextTests.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,12 +34,14 @@
3434

3535
import static com.carrotsearch.randomizedtesting.RandomizedTest.randomAsciiLettersOfLengthBetween;
3636
import static org.elasticsearch.tasks.Task.HEADERS_TO_COPY;
37+
import static org.hamcrest.Matchers.anEmptyMap;
3738
import static org.hamcrest.Matchers.contains;
3839
import static org.hamcrest.Matchers.containsInAnyOrder;
3940
import static org.hamcrest.Matchers.equalTo;
4041
import static org.hamcrest.Matchers.hasItem;
4142
import static org.hamcrest.Matchers.hasSize;
4243
import static org.hamcrest.Matchers.instanceOf;
44+
import static org.hamcrest.Matchers.is;
4345
import static org.hamcrest.Matchers.not;
4446
import static org.hamcrest.Matchers.nullValue;
4547
import static org.hamcrest.Matchers.sameInstance;
@@ -1161,6 +1163,75 @@ public void testNewEmptySystemContext() {
11611163
assertNotNull(threadContext.getHeader(header));
11621164
}
11631165

1166+
public void testNewTraceContext() {
1167+
final var threadContext = new ThreadContext(Settings.EMPTY);
1168+
1169+
var rootTraceContext = Map.of(Task.TRACE_PARENT_HTTP_HEADER, randomIdentifier(), Task.TRACE_STATE, randomIdentifier());
1170+
var apmTraceContext = new Object();
1171+
var responseKey = randomIdentifier();
1172+
var responseValue = randomAlphaOfLength(10);
1173+
1174+
threadContext.putHeader(rootTraceContext);
1175+
threadContext.putTransient(Task.APM_TRACE_CONTEXT, apmTraceContext);
1176+
1177+
assertThat(threadContext.hasTraceContext(), equalTo(true));
1178+
assertThat(threadContext.hasParentTraceContext(), equalTo(false));
1179+
1180+
try (var ignored = threadContext.newTraceContext()) {
1181+
assertThat(threadContext.hasTraceContext(), equalTo(false)); // no trace started yet
1182+
assertThat(threadContext.hasParentTraceContext(), equalTo(true));
1183+
1184+
assertThat(threadContext.getHeaders(), is(anEmptyMap()));
1185+
assertThat(
1186+
threadContext.getTransientHeaders(),
1187+
equalTo(
1188+
Map.of(
1189+
Task.PARENT_TRACE_PARENT_HEADER,
1190+
rootTraceContext.get(Task.TRACE_PARENT_HTTP_HEADER),
1191+
Task.PARENT_TRACE_STATE,
1192+
rootTraceContext.get(Task.TRACE_STATE),
1193+
Task.PARENT_APM_TRACE_CONTEXT,
1194+
apmTraceContext
1195+
)
1196+
)
1197+
);
1198+
// response headers shall be propagated
1199+
threadContext.addResponseHeader(responseKey, responseValue);
1200+
}
1201+
1202+
assertThat(threadContext.hasTraceContext(), equalTo(true));
1203+
assertThat(threadContext.hasParentTraceContext(), equalTo(false));
1204+
1205+
assertThat(threadContext.getHeaders(), equalTo(rootTraceContext));
1206+
assertThat(threadContext.getTransientHeaders(), equalTo(Map.of(Task.APM_TRACE_CONTEXT, apmTraceContext)));
1207+
assertThat(threadContext.getResponseHeaders(), equalTo(Map.of(responseKey, List.of(responseValue))));
1208+
}
1209+
1210+
public void testNewTraceContextWithoutParentTrace() {
1211+
final var threadContext = new ThreadContext(Settings.EMPTY);
1212+
1213+
var responseKey = randomIdentifier();
1214+
var responseValue = randomAlphaOfLength(10);
1215+
1216+
assertThat(threadContext.hasTraceContext(), equalTo(false));
1217+
assertThat(threadContext.hasParentTraceContext(), equalTo(false));
1218+
1219+
try (var ignored = threadContext.newTraceContext()) {
1220+
assertTrue(threadContext.isDefaultContext());
1221+
assertThat(threadContext.hasTraceContext(), equalTo(false));
1222+
assertThat(threadContext.hasParentTraceContext(), equalTo(false));
1223+
1224+
// discared, just making sure the context is isolated
1225+
threadContext.putTransient(randomIdentifier(), randomAlphaOfLength(10));
1226+
// response headers shall be propagated
1227+
threadContext.addResponseHeader(responseKey, responseValue);
1228+
}
1229+
1230+
assertThat(threadContext.getHeaders(), is(anEmptyMap()));
1231+
assertThat(threadContext.getTransientHeaders(), is(anEmptyMap()));
1232+
assertThat(threadContext.getResponseHeaders(), equalTo(Map.of(responseKey, List.of(responseValue))));
1233+
}
1234+
11641235
public void testRestoreExistingContext() {
11651236
final var threadContext = new ThreadContext(Settings.EMPTY);
11661237
final var header = randomIdentifier();

0 commit comments

Comments
 (0)