Skip to content

Commit f40ef19

Browse files
GaganjunejaGagan Juneja
authored and
kaushalmahi12
committed
[Tracing Framework] Redefine telemetry context restoration and propagation (opensearch-project#9617)
* Add SpanBuilder support Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Refactor code Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Redefine telemetry context restoration Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Update changelog Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Stores the SpanScope in ThreadLocal Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Revert the context name changes Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Change the span::endSpan and SpanScope::close behaviour Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Supressed warnings Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Add more test cases Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Address review comment Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Address review comment Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Fix java doc Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Address review comment Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Fix failing test Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Empty-Commit Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> * Empty-Commit Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> --------- Signed-off-by: Gagan Juneja <gjjuneja@amazon.com> Signed-off-by: Gagan Juneja <gagandeepjuneja@gmail.com> Co-authored-by: Gagan Juneja <gjjuneja@amazon.com> Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>
1 parent 261eeaa commit f40ef19

File tree

27 files changed

+1236
-232
lines changed

27 files changed

+1236
-232
lines changed

‎CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
172172
- Separate request-based and settings-based concurrent segment search controls and introduce AggregatorFactory method to determine concurrent search support ([#9469](https://github.com/opensearch-project/OpenSearch/pull/9469))
173173
- [Remote Store] Rate limiter integration for remote store uploads and downloads([#9448](https://github.com/opensearch-project/OpenSearch/pull/9448/))
174174
- [Remote Store] Implicitly use replication type SEGMENT for remote store clusters ([#9264](https://github.com/opensearch-project/OpenSearch/pull/9264))
175+
- Redefine telemetry context restoration and propagation ([#9617](https://github.com/opensearch-project/OpenSearch/pull/9617))
175176
- Use non-concurrent path for sort request on timeseries index and field([#9562](https://github.com/opensearch-project/OpenSearch/pull/9562))
176177
- Added sampler based on `Blanket Probabilistic Sampling rate` and `Override for on demand` ([#9621](https://github.com/opensearch-project/OpenSearch/issues/9621))
177178

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

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.telemetry.tracing;
10+
11+
import java.util.Objects;
12+
13+
/**
14+
* Default implementation of Scope
15+
*
16+
* @opensearch.internal
17+
*/
18+
final class DefaultScopedSpan implements ScopedSpan {
19+
20+
private final Span span;
21+
22+
private final SpanScope spanScope;
23+
24+
/**
25+
* Creates Scope instance for the given span
26+
*
27+
* @param span underlying span
28+
* @param spanScope span scope.
29+
*/
30+
public DefaultScopedSpan(Span span, SpanScope spanScope) {
31+
this.span = Objects.requireNonNull(span);
32+
this.spanScope = Objects.requireNonNull(spanScope);
33+
}
34+
35+
@Override
36+
public void addAttribute(String key, String value) {
37+
span.addAttribute(key, value);
38+
}
39+
40+
@Override
41+
public void addAttribute(String key, long value) {
42+
span.addAttribute(key, value);
43+
}
44+
45+
@Override
46+
public void addAttribute(String key, double value) {
47+
span.addAttribute(key, value);
48+
}
49+
50+
@Override
51+
public void addAttribute(String key, boolean value) {
52+
span.addAttribute(key, value);
53+
}
54+
55+
@Override
56+
public void addEvent(String event) {
57+
span.addEvent(event);
58+
}
59+
60+
@Override
61+
public void setError(Exception exception) {
62+
span.setError(exception);
63+
}
64+
65+
/**
66+
* Executes the runnable to end the scope
67+
*/
68+
@Override
69+
public void close() {
70+
span.endSpan();
71+
spanScope.close();
72+
}
73+
74+
/**
75+
* Returns span.
76+
* @return
77+
*/
78+
Span getSpan() {
79+
return span;
80+
}
81+
82+
/**
83+
* Returns {@link SpanScope}
84+
* @return spanScope
85+
*/
86+
SpanScope getSpanScope() {
87+
return spanScope;
88+
}
89+
}

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

Lines changed: 40 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -8,65 +8,68 @@
88

99
package org.opensearch.telemetry.tracing;
1010

11-
import java.util.function.Consumer;
11+
import java.util.Objects;
1212

1313
/**
14-
* Default implementation of Scope
15-
*
16-
* @opensearch.internal
14+
* Default implementation for {@link SpanScope}
1715
*/
18-
final class DefaultSpanScope implements SpanScope {
19-
16+
public class DefaultSpanScope implements SpanScope {
2017
private final Span span;
21-
22-
private final Consumer<Span> onCloseConsumer;
18+
private final SpanScope previousSpanScope;
19+
private static final ThreadLocal<SpanScope> spanScopeThreadLocal = new ThreadLocal<>();
20+
private final TracerContextStorage<String, Span> tracerContextStorage;
2321

2422
/**
25-
* Creates Scope instance for the given span
26-
*
27-
* @param span underlying span
28-
* @param onCloseConsumer consumer to execute on scope close
23+
* Constructor
24+
* @param span span
25+
* @param previousSpanScope before attached span scope.
2926
*/
30-
public DefaultSpanScope(Span span, Consumer<Span> onCloseConsumer) {
31-
this.span = span;
32-
this.onCloseConsumer = onCloseConsumer;
27+
private DefaultSpanScope(Span span, SpanScope previousSpanScope, TracerContextStorage<String, Span> tracerContextStorage) {
28+
this.span = Objects.requireNonNull(span);
29+
this.previousSpanScope = previousSpanScope;
30+
this.tracerContextStorage = tracerContextStorage;
3331
}
3432

35-
@Override
36-
public void addSpanAttribute(String key, String value) {
37-
span.addAttribute(key, value);
33+
/**
34+
* Creates the SpanScope object.
35+
* @param span span.
36+
* @param tracerContextStorage tracer context storage.
37+
* @return SpanScope spanScope
38+
*/
39+
public static SpanScope create(Span span, TracerContextStorage<String, Span> tracerContextStorage) {
40+
final SpanScope beforeSpanScope = spanScopeThreadLocal.get();
41+
SpanScope newSpanScope = new DefaultSpanScope(span, beforeSpanScope, tracerContextStorage);
42+
spanScopeThreadLocal.set(newSpanScope);
43+
return newSpanScope;
3844
}
3945

4046
@Override
41-
public void addSpanAttribute(String key, long value) {
42-
span.addAttribute(key, value);
47+
public void close() {
48+
detach();
49+
spanScopeThreadLocal.set(previousSpanScope);
4350
}
4451

4552
@Override
46-
public void addSpanAttribute(String key, double value) {
47-
span.addAttribute(key, value);
53+
public SpanScope attach() {
54+
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, this.span);
55+
return this;
4856
}
4957

50-
@Override
51-
public void addSpanAttribute(String key, boolean value) {
52-
span.addAttribute(key, value);
58+
private void detach() {
59+
if (previousSpanScope != null) {
60+
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, previousSpanScope.getSpan());
61+
} else {
62+
tracerContextStorage.put(TracerContextStorage.CURRENT_SPAN, null);
63+
}
5364
}
5465

5566
@Override
56-
public void addSpanEvent(String event) {
57-
span.addEvent(event);
67+
public Span getSpan() {
68+
return span;
5869
}
5970

60-
@Override
61-
public void setError(Exception exception) {
62-
span.setError(exception);
71+
static SpanScope getCurrentSpanScope() {
72+
return spanScopeThreadLocal.get();
6373
}
6474

65-
/**
66-
* Executes the runnable to end the scope
67-
*/
68-
@Override
69-
public void close() {
70-
onCloseConsumer.accept(span);
71-
}
7275
}

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

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -41,17 +41,22 @@ public DefaultTracer(TracingTelemetry tracingTelemetry, TracerContextStorage<Str
4141
}
4242

4343
@Override
44-
public SpanScope startSpan(String spanName) {
44+
public Span startSpan(SpanCreationContext context) {
45+
return startSpan(context.getSpanName(), context.getAttributes());
46+
}
47+
48+
@Override
49+
public Span startSpan(String spanName) {
4550
return startSpan(spanName, Attributes.EMPTY);
4651
}
4752

4853
@Override
49-
public SpanScope startSpan(String spanName, Attributes attributes) {
54+
public Span startSpan(String spanName, Attributes attributes) {
5055
return startSpan(spanName, (SpanContext) null, attributes);
5156
}
5257

5358
@Override
54-
public SpanScope startSpan(String spanName, SpanContext parentSpan, Attributes attributes) {
59+
public Span startSpan(String spanName, SpanContext parentSpan, Attributes attributes) {
5560
Span span = null;
5661
if (parentSpan != null) {
5762
span = createSpan(spanName, parentSpan.getSpan(), attributes);
@@ -60,7 +65,7 @@ public SpanScope startSpan(String spanName, SpanContext parentSpan, Attributes a
6065
}
6166
setCurrentSpanInContext(span);
6267
addDefaultAttributes(span);
63-
return new DefaultSpanScope(span, (scopeSpan) -> endSpan(scopeSpan));
68+
return span;
6469
}
6570

6671
@Override
@@ -77,11 +82,21 @@ public SpanContext getCurrentSpan() {
7782
return (currentSpan == null) ? null : new SpanContext(currentSpan);
7883
}
7984

80-
private void endSpan(Span span) {
81-
if (span != null) {
82-
span.endSpan();
83-
setCurrentSpanInContext(span.getParentSpan());
84-
}
85+
@Override
86+
public ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext) {
87+
return startScopedSpan(spanCreationContext, null);
88+
}
89+
90+
@Override
91+
public ScopedSpan startScopedSpan(SpanCreationContext spanCreationContext, SpanContext parentSpan) {
92+
Span span = startSpan(spanCreationContext.getSpanName(), parentSpan, spanCreationContext.getAttributes());
93+
SpanScope spanScope = withSpanInScope(span);
94+
return new DefaultScopedSpan(span, spanScope);
95+
}
96+
97+
@Override
98+
public SpanScope withSpanInScope(Span span) {
99+
return DefaultSpanScope.create(span, tracerContextStorage).attach();
85100
}
86101

87102
private Span createSpan(String spanName, Span parentSpan, Attributes attributes) {
@@ -101,7 +116,7 @@ protected void addDefaultAttributes(Span span) {
101116
}
102117

103118
@Override
104-
public SpanScope startSpan(String spanName, Map<String, List<String>> headers, Attributes attributes) {
119+
public Span startSpan(String spanName, Map<String, List<String>> headers, Attributes attributes) {
105120
Optional<Span> propagatedSpan = tracingTelemetry.getContextPropagator().extractFromHeaders(headers);
106121
return startSpan(spanName, propagatedSpan.map(SpanContext::new).orElse(null), attributes);
107122
}

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

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.telemetry.tracing;
10+
11+
import org.opensearch.telemetry.tracing.noop.NoopScopedSpan;
12+
13+
/**
14+
* An auto-closeable that represents scoped span.
15+
* It provides interface for all the span operations.
16+
*/
17+
public interface ScopedSpan extends AutoCloseable {
18+
/**
19+
* No-op Scope implementation
20+
*/
21+
ScopedSpan NO_OP = new NoopScopedSpan();
22+
23+
/**
24+
* Adds string attribute to the {@link Span}.
25+
*
26+
* @param key attribute key
27+
* @param value attribute value
28+
*/
29+
void addAttribute(String key, String value);
30+
31+
/**
32+
* Adds long attribute to the {@link Span}.
33+
*
34+
* @param key attribute key
35+
* @param value attribute value
36+
*/
37+
void addAttribute(String key, long value);
38+
39+
/**
40+
* Adds double attribute to the {@link Span}.
41+
*
42+
* @param key attribute key
43+
* @param value attribute value
44+
*/
45+
void addAttribute(String key, double value);
46+
47+
/**
48+
* Adds boolean attribute to the {@link Span}.
49+
*
50+
* @param key attribute key
51+
* @param value attribute value
52+
*/
53+
void addAttribute(String key, boolean value);
54+
55+
/**
56+
* Adds an event to the {@link Span}.
57+
*
58+
* @param event event name
59+
*/
60+
void addEvent(String event);
61+
62+
/**
63+
* Records error in the span
64+
*
65+
* @param exception exception to be recorded
66+
*/
67+
void setError(Exception exception);
68+
69+
/**
70+
* closes the scope
71+
*/
72+
@Override
73+
void close();
74+
}

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

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.telemetry.tracing;
10+
11+
import org.opensearch.telemetry.tracing.attributes.Attributes;
12+
13+
/**
14+
* Context for span details.
15+
*/
16+
public final class SpanCreationContext {
17+
private final String spanName;
18+
private final Attributes attributes;
19+
20+
/**
21+
* Constructor.
22+
* @param spanName span name.
23+
* @param attributes attributes.
24+
*/
25+
public SpanCreationContext(String spanName, Attributes attributes) {
26+
this.spanName = spanName;
27+
this.attributes = attributes;
28+
}
29+
30+
/**
31+
* Returns the span name.
32+
* @return span name
33+
*/
34+
public String getSpanName() {
35+
return spanName;
36+
}
37+
38+
/**
39+
* Returns the span attributes.
40+
* @return attributes.
41+
*/
42+
public Attributes getAttributes() {
43+
return attributes;
44+
}
45+
}

0 commit comments

Comments
 (0)