Skip to content

Commit 00517eb

Browse files
retareta
authored
[BUG] The thread context is not properly cleared and messes up the traces (#10873)
* [BUG] The thread context is not properly cleared and messes up the traces Signed-off-by: Andriy Redko <andriy.redko@aiven.io> * Address code review comments Signed-off-by: Andriy Redko <andriy.redko@aiven.io> * Address code review comments Signed-off-by: Andriy Redko <andriy.redko@aiven.io> --------- Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
1 parent b0d6b3c commit 00517eb

File tree

22 files changed

+338
-141
lines changed

22 files changed

+338
-141
lines changed

‎CHANGELOG.md

+1
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
9090
- Fix 'org.apache.hc.core5.http.ParseException: Invalid protocol version' under JDK 16+ ([#4827](https://github.com/opensearch-project/OpenSearch/pull/4827))
9191
- Fix compression support for h2c protocol ([#4944](https://github.com/opensearch-project/OpenSearch/pull/4944))
9292
- Don't over-allocate in HeapBufferedAsyncEntityConsumer in order to consume the response ([#9993](https://github.com/opensearch-project/OpenSearch/pull/9993))
93+
- [BUG] Fix the thread context that is not properly cleared and messes up the traces ([#10873](https://github.com/opensearch-project/OpenSearch/pull/10873))
9394

9495
### Security
9596

‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultSpanScope.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -44,23 +44,23 @@ private DefaultSpanScope(Span span, SpanScope previousSpanScope, TracerContextSt
4444
public static SpanScope create(Span span, TracerContextStorage<String, Span> tracerContextStorage) {
4545
final SpanScope beforeSpanScope = spanScopeThreadLocal.get();
4646
SpanScope newSpanScope = new DefaultSpanScope(span, beforeSpanScope, tracerContextStorage);
47-
spanScopeThreadLocal.set(newSpanScope);
4847
return newSpanScope;
4948
}
5049

5150
@Override
5251
public void close() {
5352
detach();
54-
spanScopeThreadLocal.set(previousSpanScope);
5553
}
5654

5755
@Override
5856
public SpanScope attach() {
57+
spanScopeThreadLocal.set(this);
5958
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, this.span);
6059
return this;
6160
}
6261

6362
private void detach() {
63+
spanScopeThreadLocal.set(previousSpanScope);
6464
if (previousSpanScope != null) {
6565
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, previousSpanScope.getSpan());
6666
} else {

‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/DefaultTracer.java

+2-7
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
import java.io.Closeable;
1414
import java.io.IOException;
15-
import java.util.List;
15+
import java.util.Collection;
1616
import java.util.Map;
1717
import java.util.Optional;
1818

@@ -53,7 +53,6 @@ public Span startSpan(SpanCreationContext context) {
5353
parentSpan = getCurrentSpanInternal();
5454
}
5555
Span span = createSpan(context, parentSpan);
56-
setCurrentSpanInContext(span);
5756
addDefaultAttributes(span);
5857
return span;
5958
}
@@ -94,10 +93,6 @@ private Span createSpan(SpanCreationContext spanCreationContext, Span parentSpan
9493
return tracingTelemetry.createSpan(spanCreationContext, parentSpan);
9594
}
9695

97-
private void setCurrentSpanInContext(Span span) {
98-
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, span);
99-
}
100-
10196
/**
10297
* Adds default attributes in the span
10398
* @param span the current active span
@@ -107,7 +102,7 @@ protected void addDefaultAttributes(Span span) {
107102
}
108103

109104
@Override
110-
public Span startSpan(SpanCreationContext spanCreationContext, Map<String, List<String>> headers) {
105+
public Span startSpan(SpanCreationContext spanCreationContext, Map<String, Collection<String>> headers) {
111106
Optional<Span> propagatedSpan = tracingTelemetry.getContextPropagator().extractFromHeaders(headers);
112107
return startSpan(spanCreationContext.parent(propagatedSpan.map(SpanContext::new).orElse(null)));
113108
}

‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/Tracer.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
package org.opensearch.telemetry.tracing;
1010

1111
import org.opensearch.common.annotation.ExperimentalApi;
12-
import org.opensearch.telemetry.tracing.http.HttpTracer;
12+
import org.opensearch.telemetry.tracing.transport.TransportTracer;
1313

1414
import java.io.Closeable;
1515

@@ -22,7 +22,7 @@
2222
* @opensearch.experimental
2323
*/
2424
@ExperimentalApi
25-
public interface Tracer extends HttpTracer, Closeable {
25+
public interface Tracer extends TransportTracer, Closeable {
2626
/**
2727
* Starts the {@link Span} with given {@link SpanCreationContext}
2828
*

‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/TracingContextPropagator.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
import org.opensearch.common.annotation.ExperimentalApi;
1212

13-
import java.util.List;
13+
import java.util.Collection;
1414
import java.util.Map;
1515
import java.util.Optional;
1616
import java.util.function.BiConsumer;
@@ -36,7 +36,7 @@ public interface TracingContextPropagator {
3636
* @param headers request headers to extract the context from
3737
* @return current span
3838
*/
39-
Optional<Span> extractFromHeaders(Map<String, List<String>> headers);
39+
Optional<Span> extractFromHeaders(Map<String, Collection<String>> headers);
4040

4141
/**
4242
* Injects tracing context

‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/noop/NoopTracer.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
import org.opensearch.telemetry.tracing.SpanScope;
1717
import org.opensearch.telemetry.tracing.Tracer;
1818

19-
import java.util.List;
19+
import java.util.Collection;
2020
import java.util.Map;
2121

2222
/**
@@ -65,7 +65,7 @@ public void close() {
6565
}
6666

6767
@Override
68-
public Span startSpan(SpanCreationContext spanCreationContext, Map<String, List<String>> header) {
68+
public Span startSpan(SpanCreationContext spanCreationContext, Map<String, Collection<String>> header) {
6969
return NoopSpan.INSTANCE;
7070
}
7171
}

‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/http/HttpTracer.java ‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/transport/TransportTracer.java

+8-8
Original file line numberDiff line numberDiff line change
@@ -6,31 +6,31 @@
66
* compatible open source license.
77
*/
88

9-
package org.opensearch.telemetry.tracing.http;
9+
package org.opensearch.telemetry.tracing.transport;
1010

1111
import org.opensearch.common.annotation.ExperimentalApi;
1212
import org.opensearch.telemetry.tracing.Span;
1313
import org.opensearch.telemetry.tracing.SpanCreationContext;
1414

15-
import java.util.List;
15+
import java.util.Collection;
1616
import java.util.Map;
1717

1818
/**
19-
* HttpTracer helps in creating a {@link Span} which reads the incoming tracing information
20-
* from the HttpRequest header and propagate the span accordingly.
19+
* TransportTracer helps in creating a {@link Span} which reads the incoming tracing information
20+
* from the HTTP or TCP transport headers and propagate the span accordingly.
2121
* <p>
2222
* All methods on the Tracer object are multi-thread safe.
2323
*
2424
* @opensearch.experimental
2525
*/
2626
@ExperimentalApi
27-
public interface HttpTracer {
27+
public interface TransportTracer {
2828
/**
2929
* Start the span with propagating the tracing info from the HttpRequest header.
3030
*
3131
* @param spanCreationContext span name.
32-
* @param header http request header.
33-
* @return span.
32+
* @param headers transport headers
33+
* @return the span instance
3434
*/
35-
Span startSpan(SpanCreationContext spanCreationContext, Map<String, List<String>> header);
35+
Span startSpan(SpanCreationContext spanCreationContext, Map<String, Collection<String>> headers);
3636
}

‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/http/package-info.java ‎libs/telemetry/src/main/java/org/opensearch/telemetry/tracing/transport/package-info.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,6 @@
77
*/
88

99
/**
10-
* Contains No-op implementations
10+
* Contains HTTP or TCP transport related tracer capabilities
1111
*/
12-
package org.opensearch.telemetry.tracing.http;
12+
package org.opensearch.telemetry.tracing.transport;

‎libs/telemetry/src/test/java/org/opensearch/telemetry/tracing/DefaultTracerTests.java

+18-17
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
import java.util.concurrent.CompletableFuture;
2323
import java.util.concurrent.ExecutorService;
2424

25+
import static org.hamcrest.CoreMatchers.is;
26+
import static org.hamcrest.CoreMatchers.nullValue;
2527
import static org.mockito.ArgumentMatchers.eq;
2628
import static org.mockito.Mockito.mock;
2729
import static org.mockito.Mockito.never;
@@ -35,7 +37,6 @@ public class DefaultTracerTests extends OpenSearchTestCase {
3537
private Span mockSpan;
3638
private Span mockParentSpan;
3739

38-
private SpanScope mockSpanScope;
3940
private ThreadPool threadPool;
4041
private ExecutorService executorService;
4142
private SpanCreationContext spanCreationContext;
@@ -102,11 +103,11 @@ public void testCreateSpanWithAttributes() {
102103

103104
Span span = defaultTracer.startSpan(spanCreationContext);
104105

105-
assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
106-
assertEquals(1.0, ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key1"));
107-
assertEquals(2l, ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key2"));
108-
assertEquals(true, ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key3"));
109-
assertEquals("key4", ((MockSpan) defaultTracer.getCurrentSpan().getSpan()).getAttribute("key4"));
106+
assertThat(defaultTracer.getCurrentSpan(), is(nullValue()));
107+
assertEquals(1.0, ((MockSpan) span).getAttribute("key1"));
108+
assertEquals(2l, ((MockSpan) span).getAttribute("key2"));
109+
assertEquals(true, ((MockSpan) span).getAttribute("key3"));
110+
assertEquals("key4", ((MockSpan) span).getAttribute("key4"));
110111
span.endSpan();
111112
}
112113

@@ -121,16 +122,18 @@ public void testCreateSpanWithParent() {
121122

122123
Span span = defaultTracer.startSpan(spanCreationContext, null);
123124

124-
SpanContext parentSpan = defaultTracer.getCurrentSpan();
125-
126-
SpanCreationContext spanCreationContext1 = buildSpanCreationContext("span_name_1", Attributes.EMPTY, parentSpan.getSpan());
125+
try (final SpanScope scope = defaultTracer.withSpanInScope(span)) {
126+
SpanContext parentSpan = defaultTracer.getCurrentSpan();
127127

128-
Span span1 = defaultTracer.startSpan(spanCreationContext1);
128+
SpanCreationContext spanCreationContext1 = buildSpanCreationContext("span_name_1", Attributes.EMPTY, parentSpan.getSpan());
129129

130-
assertEquals("span_name_1", defaultTracer.getCurrentSpan().getSpan().getSpanName());
131-
assertEquals(parentSpan.getSpan(), defaultTracer.getCurrentSpan().getSpan().getParentSpan());
132-
span1.endSpan();
133-
span.endSpan();
130+
try (final ScopedSpan span1 = defaultTracer.startScopedSpan(spanCreationContext1)) {
131+
assertEquals("span_name_1", defaultTracer.getCurrentSpan().getSpan().getSpanName());
132+
assertEquals(parentSpan.getSpan(), defaultTracer.getCurrentSpan().getSpan().getParentSpan());
133+
}
134+
} finally {
135+
span.endSpan();
136+
}
134137
}
135138

136139
@SuppressWarnings("unchecked")
@@ -155,8 +158,7 @@ public void testCreateSpanWithNullParent() {
155158

156159
Span span = defaultTracer.startSpan(spanCreationContext);
157160

158-
assertEquals("span_name", defaultTracer.getCurrentSpan().getSpan().getSpanName());
159-
assertEquals(null, defaultTracer.getCurrentSpan().getSpan().getParentSpan());
161+
assertThat(defaultTracer.getCurrentSpan(), is(nullValue()));
160162
span.endSpan();
161163
}
162164

@@ -403,7 +405,6 @@ private void setupMocks() {
403405
mockTracingTelemetry = mock(TracingTelemetry.class);
404406
mockSpan = mock(Span.class);
405407
mockParentSpan = mock(Span.class);
406-
mockSpanScope = mock(SpanScope.class);
407408
mockTracerContextStorage = mock(TracerContextStorage.class);
408409
when(mockSpan.getSpanName()).thenReturn("span_name");
409410
when(mockSpan.getSpanId()).thenReturn("span_id");

‎plugins/telemetry-otel/src/internalClusterTest/java/org/opensearch/telemetry/tracing/TelemetryTracerEnabledSanityIT.java

+13-9
Original file line numberDiff line numberDiff line change
@@ -61,20 +61,22 @@ public void testSanityChecksWhenTracingEnabled() throws Exception {
6161

6262
// Create Index and ingest data
6363
String indexName = "test-index-11";
64-
Settings basicSettings = Settings.builder().put("number_of_shards", 3).put("number_of_replicas", 0).build();
64+
Settings basicSettings = Settings.builder()
65+
.put("number_of_shards", 2)
66+
.put("number_of_replicas", 0)
67+
.put("index.routing.allocation.total_shards_per_node", 1)
68+
.build();
6569
createIndex(indexName, basicSettings);
66-
indexRandom(true, client.prepareIndex(indexName).setId("1").setSource("field1", "the fox jumps in the well"));
67-
indexRandom(true, client.prepareIndex(indexName).setId("1").setSource("field2", "another fox did the same."));
70+
71+
indexRandom(false, client.prepareIndex(indexName).setId("1").setSource("field1", "the fox jumps in the well"));
72+
indexRandom(false, client.prepareIndex(indexName).setId("2").setSource("field2", "another fox did the same."));
6873

6974
ensureGreen();
7075
refresh();
7176

7277
// Make the search calls; adding the searchType and PreFilterShardSize to make the query path predictable across all the runs.
73-
client.prepareSearch().setSearchType("query_then_fetch").setPreFilterShardSize(3).setQuery(queryStringQuery("fox")).get();
74-
client.prepareSearch().setSearchType("query_then_fetch").setPreFilterShardSize(3).setQuery(queryStringQuery("jumps")).get();
75-
76-
ensureGreen();
77-
refresh();
78+
client.prepareSearch().setSearchType("dfs_query_then_fetch").setPreFilterShardSize(2).setQuery(queryStringQuery("fox")).get();
79+
client.prepareSearch().setSearchType("dfs_query_then_fetch").setPreFilterShardSize(2).setQuery(queryStringQuery("jumps")).get();
7880

7981
// Sleep for about 3s to wait for traces are published, delay is (the delay is 1s).
8082
Thread.sleep(3000);
@@ -88,8 +90,10 @@ public void testSanityChecksWhenTracingEnabled() throws Exception {
8890
)
8991
);
9092

93+
// See please https://github.com/opensearch-project/OpenSearch/issues/10291 till local transport is not instrumented,
94+
// capturing only the inter-nodes transport actions.
9195
InMemorySingletonSpanExporter exporter = InMemorySingletonSpanExporter.INSTANCE;
92-
validators.validate(exporter.getFinishedSpanItems(), 6);
96+
validators.validate(exporter.getFinishedSpanItems(), 4);
9397
}
9498

9599
private static void updateTelemetrySetting(Client client, boolean value) {

‎plugins/telemetry-otel/src/main/java/org/opensearch/telemetry/tracing/OTelTracingContextPropagator.java

+5-5
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,8 @@
1010

1111
import org.opensearch.core.common.Strings;
1212

13+
import java.util.Collection;
1314
import java.util.Collections;
14-
import java.util.List;
1515
import java.util.Map;
1616
import java.util.Optional;
1717
import java.util.function.BiConsumer;
@@ -51,7 +51,7 @@ private static OTelPropagatedSpan getPropagatedSpan(Context context) {
5151
}
5252

5353
@Override
54-
public Optional<Span> extractFromHeaders(Map<String, List<String>> headers) {
54+
public Optional<Span> extractFromHeaders(Map<String, Collection<String>> headers) {
5555
Context context = openTelemetry.getPropagators().getTextMapPropagator().extract(Context.current(), headers, HEADER_TEXT_MAP_GETTER);
5656
return Optional.ofNullable(getPropagatedSpan(context));
5757
}
@@ -87,9 +87,9 @@ public String get(Map<String, String> headers, String key) {
8787
}
8888
};
8989

90-
private static final TextMapGetter<Map<String, List<String>>> HEADER_TEXT_MAP_GETTER = new TextMapGetter<>() {
90+
private static final TextMapGetter<Map<String, Collection<String>>> HEADER_TEXT_MAP_GETTER = new TextMapGetter<>() {
9191
@Override
92-
public Iterable<String> keys(Map<String, List<String>> headers) {
92+
public Iterable<String> keys(Map<String, Collection<String>> headers) {
9393
if (headers != null) {
9494
return headers.keySet();
9595
} else {
@@ -98,7 +98,7 @@ public Iterable<String> keys(Map<String, List<String>> headers) {
9898
}
9999

100100
@Override
101-
public String get(Map<String, List<String>> headers, String key) {
101+
public String get(Map<String, Collection<String>> headers, String key) {
102102
if (headers != null && headers.containsKey(key)) {
103103
return Strings.collectionToCommaDelimitedString(headers.get(key));
104104
}

‎plugins/telemetry-otel/src/test/java/org/opensearch/telemetry/tracing/OTelTracingContextPropagatorTests.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@
1111
import org.opensearch.test.OpenSearchTestCase;
1212

1313
import java.util.Arrays;
14+
import java.util.Collection;
1415
import java.util.HashMap;
15-
import java.util.List;
1616
import java.util.Map;
1717

1818
import io.opentelemetry.api.OpenTelemetry;
@@ -57,7 +57,7 @@ public void testExtractTracerContextFromHeader() {
5757
}
5858

5959
public void testExtractTracerContextFromHttpHeader() {
60-
Map<String, List<String>> requestHeaders = new HashMap<>();
60+
Map<String, Collection<String>> requestHeaders = new HashMap<>();
6161
requestHeaders.put("traceparent", Arrays.asList("00-" + TRACE_ID + "-" + SPAN_ID + "-00"));
6262
OpenTelemetry mockOpenTelemetry = mock(OpenTelemetry.class);
6363
when(mockOpenTelemetry.getPropagators()).thenReturn(ContextPropagators.create(W3CTraceContextPropagator.getInstance()));

‎server/src/main/java/org/opensearch/http/AbstractHttpServerTransport.java

+8-1
Original file line numberDiff line numberDiff line change
@@ -69,9 +69,11 @@
6969
import java.nio.channels.CancelledKeyException;
7070
import java.util.ArrayList;
7171
import java.util.Arrays;
72+
import java.util.Collection;
7273
import java.util.Collections;
7374
import java.util.HashSet;
7475
import java.util.List;
76+
import java.util.Map;
7577
import java.util.Set;
7678
import java.util.concurrent.ConcurrentHashMap;
7779
import java.util.concurrent.atomic.AtomicLong;
@@ -362,7 +364,7 @@ protected void serverAcceptedChannel(HttpChannel httpChannel) {
362364
* @param httpChannel that received the http request
363365
*/
364366
public void incomingRequest(final HttpRequest httpRequest, final HttpChannel httpChannel) {
365-
final Span span = tracer.startSpan(SpanBuilder.from(httpRequest), httpRequest.getHeaders());
367+
final Span span = tracer.startSpan(SpanBuilder.from(httpRequest), extractHeaders(httpRequest.getHeaders()));
366368
try (final SpanScope httpRequestSpanScope = tracer.withSpanInScope(span)) {
367369
HttpChannel traceableHttpChannel = TraceableHttpChannel.create(httpChannel, span, tracer);
368370
handleIncomingRequest(httpRequest, traceableHttpChannel, httpRequest.getInboundException());
@@ -483,4 +485,9 @@ private static ActionListener<Void> earlyResponseListener(HttpRequest request, H
483485
return NO_OP;
484486
}
485487
}
488+
489+
@SuppressWarnings("unchecked")
490+
private static <Values extends Collection<String>> Map<String, Collection<String>> extractHeaders(Map<String, Values> headers) {
491+
return (Map<String, Collection<String>>) headers;
492+
}
486493
}

0 commit comments

Comments
 (0)