Skip to content

Commit 431a7b2

Browse files
DaveCTurnerDaveCTurner
authored
Destringify APM tracer interface (#94864)
Today the APM `Tracer` interface identifies each span by a raw string, but in practice there is structure to these strings: task-related spans have IDs like `task-NNNN` and spans that relate to REST requests have IDs like `rest-NNNN`. This convention is distributed across the codebase a little too widely, so with this commit we centralise it into a `SpanId` class, and introduce specific overrides for `Task` and `RestRequest` to avoid callers needing to construct IDs themselves.
1 parent ce737d3 commit 431a7b2

File tree

11 files changed

+208
-85
lines changed

11 files changed

+208
-85
lines changed

‎modules/apm/src/main/java/org/elasticsearch/tracing/apm/APMTracer.java‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import org.elasticsearch.core.Nullable;
3535
import org.elasticsearch.core.Releasable;
3636
import org.elasticsearch.tasks.Task;
37+
import org.elasticsearch.tracing.SpanId;
3738

3839
import java.security.AccessController;
3940
import java.security.PrivilegedAction;
@@ -59,7 +60,7 @@ public class APMTracer extends AbstractLifecycleComponent implements org.elastic
5960
private static final Logger logger = LogManager.getLogger(APMTracer.class);
6061

6162
/** Holds in-flight span information. */
62-
private final Map<String, Context> spans = ConcurrentCollections.newConcurrentMap();
63+
private final Map<SpanId, Context> spans = ConcurrentCollections.newConcurrentMap();
6364

6465
private volatile boolean enabled;
6566
private volatile APMServices services;
@@ -158,7 +159,7 @@ private void destroyApmServices() {
158159
}
159160

160161
@Override
161-
public void startTrace(ThreadContext threadContext, String spanId, String spanName, @Nullable Map<String, Object> attributes) {
162+
public void startTrace(ThreadContext threadContext, SpanId spanId, String spanName, @Nullable Map<String, Object> attributes) {
162163
assert threadContext != null;
163164
assert spanId != null;
164165
assert spanName != null;
@@ -273,7 +274,7 @@ private Context getParentContext(ThreadContext threadContext) {
273274
* @return a method to close the scope when you are finished with it.
274275
*/
275276
@Override
276-
public Releasable withScope(String spanId) {
277+
public Releasable withScope(SpanId spanId) {
277278
final Context context = spans.get(spanId);
278279
if (context != null) {
279280
var scope = context.makeCurrent();
@@ -330,47 +331,47 @@ private void setSpanAttributes(ThreadContext threadContext, @Nullable Map<String
330331
}
331332

332333
@Override
333-
public void addError(String spanId, Throwable throwable) {
334+
public void addError(SpanId spanId, Throwable throwable) {
334335
final var span = Span.fromContextOrNull(spans.get(spanId));
335336
if (span != null) {
336337
span.recordException(throwable);
337338
}
338339
}
339340

340341
@Override
341-
public void setAttribute(String spanId, String key, boolean value) {
342+
public void setAttribute(SpanId spanId, String key, boolean value) {
342343
final var span = Span.fromContextOrNull(spans.get(spanId));
343344
if (span != null) {
344345
span.setAttribute(key, value);
345346
}
346347
}
347348

348349
@Override
349-
public void setAttribute(String spanId, String key, double value) {
350+
public void setAttribute(SpanId spanId, String key, double value) {
350351
final var span = Span.fromContextOrNull(spans.get(spanId));
351352
if (span != null) {
352353
span.setAttribute(key, value);
353354
}
354355
}
355356

356357
@Override
357-
public void setAttribute(String spanId, String key, long value) {
358+
public void setAttribute(SpanId spanId, String key, long value) {
358359
final var span = Span.fromContextOrNull(spans.get(spanId));
359360
if (span != null) {
360361
span.setAttribute(key, value);
361362
}
362363
}
363364

364365
@Override
365-
public void setAttribute(String spanId, String key, String value) {
366+
public void setAttribute(SpanId spanId, String key, String value) {
366367
final var span = Span.fromContextOrNull(spans.get(spanId));
367368
if (span != null) {
368369
span.setAttribute(key, value);
369370
}
370371
}
371372

372373
@Override
373-
public void stopTrace(String spanId) {
374+
public void stopTrace(SpanId spanId) {
374375
final var span = Span.fromContextOrNull(spans.remove(spanId));
375376
if (span != null) {
376377
logger.trace("Finishing trace [{}]", spanId);
@@ -387,7 +388,7 @@ public void stopTrace() {
387388
}
388389

389390
@Override
390-
public void addEvent(String spanId, String eventName) {
391+
public void addEvent(SpanId spanId, String eventName) {
391392
final var span = Span.fromContextOrNull(spans.get(spanId));
392393
if (span != null) {
393394
span.addEvent(eventName);
@@ -412,7 +413,7 @@ private static boolean isSupportedContextKey(String key) {
412413
}
413414

414415
// VisibleForTesting
415-
Map<String, Context> getSpans() {
416+
Map<SpanId, Context> getSpans() {
416417
return spans;
417418
}
418419

‎modules/apm/src/test/java/org/elasticsearch/tracing/apm/APMTracerTests.java‎

Lines changed: 25 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.elasticsearch.common.util.concurrent.ThreadContext;
1414
import org.elasticsearch.tasks.Task;
1515
import org.elasticsearch.test.ESTestCase;
16+
import org.elasticsearch.tracing.SpanId;
1617

1718
import java.util.List;
1819
import java.util.stream.Stream;
@@ -28,14 +29,18 @@
2829

2930
public class APMTracerTests extends ESTestCase {
3031

32+
private static final SpanId SPAN_ID1 = SpanId.forBareString("id1");
33+
private static final SpanId SPAN_ID2 = SpanId.forBareString("id2");
34+
private static final SpanId SPAN_ID3 = SpanId.forBareString("id3");
35+
3136
/**
3237
* Check that the tracer doesn't create spans when tracing is disabled.
3338
*/
3439
public void test_onTraceStarted_withTracingDisabled_doesNotStartTrace() {
3540
Settings settings = Settings.builder().put(APM_ENABLED_SETTING.getKey(), false).build();
3641
APMTracer apmTracer = buildTracer(settings);
3742

38-
apmTracer.startTrace(new ThreadContext(settings), "id1", "name1", null);
43+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID1, "name1", null);
3944

4045
assertThat(apmTracer.getSpans(), anEmptyMap());
4146
}
@@ -50,7 +55,7 @@ public void test_onTraceStarted_withSpanNameOmitted_doesNotStartTrace() {
5055
.build();
5156
APMTracer apmTracer = buildTracer(settings);
5257

53-
apmTracer.startTrace(new ThreadContext(settings), "id1", "name1", null);
58+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID1, "name1", null);
5459

5560
assertThat(apmTracer.getSpans(), anEmptyMap());
5661
}
@@ -62,10 +67,10 @@ public void test_onTraceStarted_startsTrace() {
6267
Settings settings = Settings.builder().put(APM_ENABLED_SETTING.getKey(), true).build();
6368
APMTracer apmTracer = buildTracer(settings);
6469

65-
apmTracer.startTrace(new ThreadContext(settings), "id1", "name1", null);
70+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID1, "name1", null);
6671

6772
assertThat(apmTracer.getSpans(), aMapWithSize(1));
68-
assertThat(apmTracer.getSpans(), hasKey("id1"));
73+
assertThat(apmTracer.getSpans(), hasKey(SPAN_ID1));
6974
}
7075

7176
/**
@@ -75,8 +80,8 @@ public void test_onTraceStopped_stopsTrace() {
7580
Settings settings = Settings.builder().put(APM_ENABLED_SETTING.getKey(), true).build();
7681
APMTracer apmTracer = buildTracer(settings);
7782

78-
apmTracer.startTrace(new ThreadContext(settings), "id1", "name1", null);
79-
apmTracer.stopTrace("id1");
83+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID1, "name1", null);
84+
apmTracer.stopTrace(SPAN_ID1);
8085

8186
assertThat(apmTracer.getSpans(), anEmptyMap());
8287
}
@@ -93,7 +98,7 @@ public void test_whenTraceStarted_threadContextIsPopulated() {
9398
APMTracer apmTracer = buildTracer(settings);
9499

95100
ThreadContext threadContext = new ThreadContext(settings);
96-
apmTracer.startTrace(threadContext, "id1", "name1", null);
101+
apmTracer.startTrace(threadContext, SPAN_ID1, "name1", null);
97102
assertThat(threadContext.getTransient(Task.APM_TRACE_CONTEXT), notNullValue());
98103
}
99104

@@ -114,13 +119,13 @@ public void test_whenTraceStarted_andSpanNameIncluded_thenSpanIsStarted() {
114119
.build();
115120
APMTracer apmTracer = buildTracer(settings);
116121

117-
apmTracer.startTrace(new ThreadContext(settings), "id1", "name-aaa", null);
118-
apmTracer.startTrace(new ThreadContext(settings), "id2", "name-bbb", null);
119-
apmTracer.startTrace(new ThreadContext(settings), "id3", "name-ccc", null);
122+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID1, "name-aaa", null);
123+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID2, "name-bbb", null);
124+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID3, "name-ccc", null);
120125

121-
assertThat(apmTracer.getSpans(), hasKey("id1"));
122-
assertThat(apmTracer.getSpans(), hasKey("id2"));
123-
assertThat(apmTracer.getSpans(), not(hasKey("id3")));
126+
assertThat(apmTracer.getSpans(), hasKey(SPAN_ID1));
127+
assertThat(apmTracer.getSpans(), hasKey(SPAN_ID2));
128+
assertThat(apmTracer.getSpans(), not(hasKey(SPAN_ID3)));
124129
}
125130

126131
/**
@@ -137,7 +142,7 @@ public void test_whenTraceStarted_andSpanNameIncludedAndExcluded_thenSpanIsNotSt
137142
.build();
138143
APMTracer apmTracer = buildTracer(settings);
139144

140-
apmTracer.startTrace(new ThreadContext(settings), "id1", "name-aaa", null);
145+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID1, "name-aaa", null);
141146

142147
assertThat(apmTracer.getSpans(), not(hasKey("id1")));
143148
}
@@ -159,13 +164,13 @@ public void test_whenTraceStarted_andSpanNameExcluded_thenSpanIsNotStarted() {
159164
.build();
160165
APMTracer apmTracer = buildTracer(settings);
161166

162-
apmTracer.startTrace(new ThreadContext(settings), "id1", "name-aaa", null);
163-
apmTracer.startTrace(new ThreadContext(settings), "id2", "name-bbb", null);
164-
apmTracer.startTrace(new ThreadContext(settings), "id3", "name-ccc", null);
167+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID1, "name-aaa", null);
168+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID2, "name-bbb", null);
169+
apmTracer.startTrace(new ThreadContext(settings), SPAN_ID3, "name-ccc", null);
165170

166-
assertThat(apmTracer.getSpans(), not(hasKey("id1")));
167-
assertThat(apmTracer.getSpans(), not(hasKey("id2")));
168-
assertThat(apmTracer.getSpans(), hasKey("id3"));
171+
assertThat(apmTracer.getSpans(), not(hasKey(SPAN_ID1)));
172+
assertThat(apmTracer.getSpans(), not(hasKey(SPAN_ID2)));
173+
assertThat(apmTracer.getSpans(), hasKey(SPAN_ID3));
169174
}
170175

171176
/**

‎server/src/main/java/org/elasticsearch/http/DefaultRestChannel.java‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import org.elasticsearch.rest.RestRequest;
2828
import org.elasticsearch.rest.RestResponse;
2929
import org.elasticsearch.rest.RestStatus;
30+
import org.elasticsearch.tracing.SpanId;
3031
import org.elasticsearch.tracing.Tracer;
3132

3233
import java.util.ArrayList;
@@ -91,13 +92,13 @@ public void sendResponse(RestResponse restResponse) {
9192
// We're sending a response so we know we won't be needing the request content again and release it
9293
httpRequest.release();
9394

94-
final String traceId = "rest-" + this.request.getRequestId();
95+
final SpanId spanId = SpanId.forRestRequest(request);
9596

9697
final ArrayList<Releasable> toClose = new ArrayList<>(4);
9798
if (HttpUtils.shouldCloseConnection(httpRequest)) {
9899
toClose.add(() -> CloseableChannel.closeChannel(httpChannel));
99100
}
100-
toClose.add(() -> tracer.stopTrace(traceId));
101+
toClose.add(() -> tracer.stopTrace(request));
101102

102103
boolean success = false;
103104
String opaque = null;
@@ -171,9 +172,9 @@ public void sendResponse(RestResponse restResponse) {
171172

172173
addCookies(httpResponse);
173174

174-
tracer.setAttribute(traceId, "http.status_code", restResponse.status().getStatus());
175+
tracer.setAttribute(spanId, "http.status_code", restResponse.status().getStatus());
175176
restResponse.getHeaders()
176-
.forEach((key, values) -> tracer.setAttribute(traceId, "http.response.headers." + key, String.join("; ", values)));
177+
.forEach((key, values) -> tracer.setAttribute(spanId, "http.response.headers." + key, String.join("; ", values)));
177178

178179
ActionListener<Void> listener = ActionListener.releasing(Releasables.wrap(toClose));
179180
if (httpLogger != null) {

‎server/src/main/java/org/elasticsearch/rest/RestController.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -490,11 +490,11 @@ private void startTrace(ThreadContext threadContext, RestChannel channel, String
490490
case HTTP_1_1 -> attributes.put("http.flavour", "1.1");
491491
}
492492

493-
tracer.startTrace(threadContext, "rest-" + channel.request().getRequestId(), name, attributes);
493+
tracer.startTrace(threadContext, channel.request(), name, attributes);
494494
}
495495

496496
private void traceException(RestChannel channel, Throwable e) {
497-
this.tracer.addError("rest-" + channel.request().getRequestId(), e);
497+
this.tracer.addError(channel.request(), e);
498498
}
499499

500500
private static void sendContentTypeErrorMessage(@Nullable List<String> contentTypeHeader, RestChannel channel) throws IOException {

‎server/src/main/java/org/elasticsearch/search/SearchService.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -445,7 +445,7 @@ public void executeDfsPhase(ShardSearchRequest request, SearchShardTask task, Ac
445445
private DfsSearchResult executeDfsPhase(ShardSearchRequest request, SearchShardTask task) throws IOException {
446446
ReaderContext readerContext = createOrGetReaderContext(request);
447447
try (
448-
Releasable scope = tracer.withScope("task-" + task.getId());
448+
Releasable scope = tracer.withScope(task);
449449
Releasable ignored = readerContext.markAsUsed(getKeepAlive(request));
450450
SearchContext context = createContext(readerContext, request, task, true)
451451
) {
@@ -626,7 +626,7 @@ private static <T> void runAsync(Executor executor, CheckedSupplier<T, Exception
626626
private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, SearchShardTask task) throws Exception {
627627
final ReaderContext readerContext = createOrGetReaderContext(request);
628628
try (
629-
Releasable scope = tracer.withScope("task-" + task.getId());
629+
Releasable scope = tracer.withScope(task);
630630
Releasable ignored = readerContext.markAsUsed(getKeepAlive(request));
631631
SearchContext context = createContext(readerContext, request, task, true)
632632
) {
@@ -668,7 +668,7 @@ private SearchPhaseResult executeQueryPhase(ShardSearchRequest request, SearchSh
668668

669669
private QueryFetchSearchResult executeFetchPhase(ReaderContext reader, SearchContext context, long afterQueryTime) {
670670
try (
671-
Releasable scope = tracer.withScope("task-" + context.getTask().getId());
671+
Releasable scope = tracer.withScope(context.getTask());
672672
SearchOperationListenerExecutor executor = new SearchOperationListenerExecutor(context, true, afterQueryTime)
673673
) {
674674
shortcutDocIdsToLoad(context);

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ void startTrace(ThreadContext threadContext, Task task) {
171171
Tracer.AttributeKeys.PARENT_TASK_ID,
172172
parentTask.toString()
173173
);
174-
tracer.startTrace(threadContext, "task-" + task.getId(), task.getAction(), attributes);
174+
tracer.startTrace(threadContext, task, task.getAction(), attributes);
175175
}
176176

177177
public <Request extends ActionRequest, Response extends ActionResponse> Task registerAndExecute(
@@ -288,7 +288,7 @@ public Task unregister(Task task) {
288288
return removedTask;
289289
}
290290
} finally {
291-
tracer.stopTrace("task-" + task.getId());
291+
tracer.stopTrace(task);
292292
for (RemovedTaskListener listener : removedTaskListeners) {
293293
listener.onRemoved(task);
294294
}

‎server/src/main/java/org/elasticsearch/tracing/SpanId.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0 and the Server Side Public License, v 1; you may not use this file except
5+
* in compliance with, at your election, the Elastic License 2.0 or the Server
6+
* Side Public License, v 1.
7+
*/
8+
9+
package org.elasticsearch.tracing;
10+
11+
import org.elasticsearch.rest.RestRequest;
12+
import org.elasticsearch.tasks.Task;
13+
14+
import java.util.Objects;
15+
16+
public class SpanId {
17+
private final String rawId;
18+
19+
private SpanId(String rawId) {
20+
this.rawId = Objects.requireNonNull(rawId);
21+
}
22+
23+
public String getRawId() {
24+
return rawId;
25+
}
26+
27+
@Override
28+
public String toString() {
29+
return "SpanId[" + rawId + "]";
30+
}
31+
32+
@Override
33+
public boolean equals(Object o) {
34+
if (this == o) return true;
35+
if (o == null || getClass() != o.getClass()) return false;
36+
SpanId spanId = (SpanId) o;
37+
return rawId.equals(spanId.rawId);
38+
}
39+
40+
@Override
41+
public int hashCode() {
42+
return Objects.hash(rawId);
43+
}
44+
45+
public static SpanId forTask(Task task) {
46+
return new SpanId("task-" + task.getId());
47+
}
48+
49+
public static SpanId forRestRequest(RestRequest restRequest) {
50+
return new SpanId("rest-" + restRequest.getRequestId());
51+
}
52+
53+
public static SpanId forBareString(String rawId) {
54+
return new SpanId(rawId);
55+
}
56+
}

0 commit comments

Comments
 (0)