-
Notifications
You must be signed in to change notification settings - Fork 303
Migrate OtelContext wrapper to new internal Context API #8645
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
175 changes: 57 additions & 118 deletions
175
...nt/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelContext.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,170 +1,109 @@ | ||
| package datadog.opentelemetry.shim.context; | ||
|
|
||
| import datadog.opentelemetry.shim.trace.OtelSpan; | ||
| import datadog.trace.bootstrap.instrumentation.api.AgentScope; | ||
| import datadog.trace.bootstrap.instrumentation.api.AgentSpan; | ||
| import datadog.trace.bootstrap.instrumentation.api.AgentTracer; | ||
| import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper; | ||
| import io.opentelemetry.api.trace.Span; | ||
| import io.opentelemetry.context.Context; | ||
| import io.opentelemetry.context.ContextKey; | ||
| import io.opentelemetry.context.Scope; | ||
| import java.util.Arrays; | ||
| import java.util.Map; | ||
| import java.util.concurrent.ConcurrentHashMap; | ||
| import javax.annotation.Nullable; | ||
| import javax.annotation.ParametersAreNonnullByDefault; | ||
|
|
||
| @SuppressWarnings({"rawtypes", "unchecked"}) | ||
| @ParametersAreNonnullByDefault | ||
| public class OtelContext implements Context { | ||
| private static final Object[] NO_ENTRIES = {}; | ||
|
|
||
| /** Overridden root context. */ | ||
| public static final Context ROOT = new OtelContext(OtelSpan.invalid(), OtelSpan.invalid()); | ||
| public static final Context ROOT = new OtelContext(datadog.context.Context.root()); | ||
|
|
||
| private static final String OTEL_CONTEXT_SPAN_KEY = "opentelemetry-trace-span-key"; | ||
| private static final String OTEL_CONTEXT_ROOT_SPAN_KEY = "opentelemetry-traces-local-root-span"; | ||
|
|
||
| /** Keep track of propagated context that has not been captured on the scope stack. */ | ||
| private static final ThreadLocal<OtelContext> lastPropagated = new ThreadLocal<>(); | ||
| /** Records the keys needed to access the delegate context, mapped by key name. */ | ||
| private static final Map<ContextKey<?>, datadog.context.ContextKey<?>> DELEGATE_KEYS = | ||
| new ConcurrentHashMap<>(); | ||
|
|
||
| private final Span currentSpan; | ||
| private final Span rootSpan; | ||
| private final datadog.context.Context delegate; | ||
|
|
||
| private final Object[] entries; | ||
| public OtelContext(datadog.context.Context delegate) { | ||
| this.delegate = delegate; | ||
| } | ||
|
|
||
| public OtelContext(Span currentSpan, Span rootSpan) { | ||
| this(currentSpan, rootSpan, NO_ENTRIES); | ||
| public static Context current() { | ||
| return new OtelContext(datadog.context.Context.current()); | ||
| } | ||
|
|
||
| public OtelContext(Span currentSpan, Span rootSpan, Object[] entries) { | ||
| this.currentSpan = currentSpan; | ||
| this.rootSpan = rootSpan; | ||
| this.entries = entries; | ||
| @Override | ||
| public Scope makeCurrent() { | ||
| return new OtelScope(Context.super.makeCurrent(), delegate.attach()); | ||
| } | ||
|
|
||
| @Nullable | ||
| @Override | ||
| @SuppressWarnings("unchecked") | ||
| public <V> V get(ContextKey<V> key) { | ||
| if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) { | ||
| return (V) this.currentSpan; | ||
| AgentSpan span = AgentSpan.fromContext(delegate); | ||
| if (span != null) { | ||
| return (V) toOtelSpan(span); | ||
| } | ||
| // fall-through and check for non-datadog span data | ||
| } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) { | ||
| return (V) this.rootSpan; | ||
| } | ||
| for (int i = 0; i < this.entries.length; i += 2) { | ||
| if (this.entries[i] == key) { | ||
| return (V) this.entries[i + 1]; | ||
| AgentSpan span = AgentSpan.fromContext(delegate); | ||
| if (span != null) { | ||
| return (V) toOtelSpan(span.getLocalRootSpan()); | ||
| } | ||
| // fall-through and check for non-datadog span data | ||
| } | ||
| return null; | ||
| return (V) delegate.get(delegateKey(key)); | ||
| } | ||
|
|
||
| @Override | ||
| public <V> Context with(ContextKey<V> key, V value) { | ||
| if (OTEL_CONTEXT_SPAN_KEY.equals(key.toString())) { | ||
| return new OtelContext((Span) value, this.rootSpan, this.entries); | ||
| } else if (OTEL_CONTEXT_ROOT_SPAN_KEY.equals(key.toString())) { | ||
| return new OtelContext(this.currentSpan, (Span) value, this.entries); | ||
| } | ||
| Object[] newEntries = null; | ||
| int oldEntriesLength = this.entries.length; | ||
| for (int i = 0; i < oldEntriesLength; i += 2) { | ||
| if (this.entries[i] == key) { | ||
| if (this.entries[i + 1] == value) { | ||
| return this; | ||
| } | ||
| newEntries = this.entries.clone(); | ||
| newEntries[i + 1] = value; | ||
| break; | ||
| if (value instanceof OtelSpan) { | ||
| AgentSpan span = ((OtelSpan) value).asAgentSpan(); | ||
| return new OtelContext(delegate.with(span)); | ||
| } | ||
| // fall-through and store as non-datadog span data | ||
| } | ||
| if (null == newEntries) { | ||
| newEntries = Arrays.copyOf(this.entries, oldEntriesLength + 2); | ||
| newEntries[oldEntriesLength] = key; | ||
| newEntries[oldEntriesLength + 1] = value; | ||
| } | ||
| return new OtelContext(this.currentSpan, this.rootSpan, newEntries); | ||
| return new OtelContext(delegate.with(delegateKey(key), value)); | ||
| } | ||
|
|
||
| @Override | ||
| public Scope makeCurrent() { | ||
| final Scope scope = Context.super.makeCurrent(); | ||
| if (this.currentSpan instanceof OtelSpan) { | ||
| // only keep propagated context until next span activation | ||
| lastPropagated.remove(); | ||
| AgentScope agentScope = ((OtelSpan) this.currentSpan).activate(); | ||
| return new OtelScope(scope, agentScope, this.entries); | ||
| } else { | ||
| // propagated context not on the scope stack, capture it here | ||
| lastPropagated.set(this); | ||
| return new Scope() { | ||
| @Override | ||
| public void close() { | ||
| lastPropagated.remove(); | ||
| scope.close(); | ||
| } | ||
| }; | ||
| } | ||
| } | ||
|
|
||
| public static Context current() { | ||
| // Check for propagated context not on the scope stack | ||
| Context context = lastPropagated.get(); | ||
| if (null != context) { | ||
| return context; | ||
| } | ||
| // Check empty context | ||
| AgentScope agentCurrentScope = AgentTracer.activeScope(); | ||
| if (null == agentCurrentScope) { | ||
| return OtelContext.ROOT; | ||
| } | ||
| // Get OTel current span | ||
| Span otelCurrentSpan = null; | ||
| AgentSpan agentCurrentSpan = agentCurrentScope.span(); | ||
| if (agentCurrentSpan instanceof AttachableWrapper) { | ||
| Object wrapper = ((AttachableWrapper) agentCurrentSpan).getWrapper(); | ||
| if (wrapper instanceof OtelSpan) { | ||
| otelCurrentSpan = (OtelSpan) wrapper; | ||
| } | ||
| public boolean equals(Object o) { | ||
| if (o == null || getClass() != o.getClass()) { | ||
| return false; | ||
| } | ||
| if (otelCurrentSpan == null) { | ||
| otelCurrentSpan = new OtelSpan(agentCurrentSpan); | ||
| } | ||
| // Get OTel root span | ||
| Span otelRootSpan = null; | ||
| AgentSpan agentRootSpan = agentCurrentSpan.getLocalRootSpan(); | ||
| if (agentRootSpan instanceof AttachableWrapper) { | ||
| Object wrapper = ((AttachableWrapper) agentRootSpan).getWrapper(); | ||
| if (wrapper instanceof OtelSpan) { | ||
| otelRootSpan = (OtelSpan) wrapper; | ||
| } | ||
| } | ||
| if (otelRootSpan == null) { | ||
| otelRootSpan = new OtelSpan(agentRootSpan); | ||
| } | ||
| // Get OTel custom context entries | ||
| Object[] contextEntries = NO_ENTRIES; | ||
| if (agentCurrentScope instanceof AttachableWrapper) { | ||
| Object wrapper = ((AttachableWrapper) agentCurrentScope).getWrapper(); | ||
| if (wrapper instanceof OtelScope) { | ||
| contextEntries = ((OtelScope) wrapper).contextEntries(); | ||
| } | ||
| } | ||
| return new OtelContext(otelCurrentSpan, otelRootSpan, contextEntries); | ||
| return delegate.equals(((OtelContext) o).delegate); | ||
| } | ||
|
|
||
| /** Last propagated context not on the scope stack; {@code null} if there's no such context. */ | ||
| @Nullable | ||
| public static Context lastPropagated() { | ||
| return lastPropagated.get(); | ||
| @Override | ||
| public int hashCode() { | ||
| return delegate.hashCode(); | ||
| } | ||
|
|
||
| @Override | ||
| public String toString() { | ||
| return "OtelContext{" | ||
| + "currentSpan=" | ||
| + this.currentSpan.getSpanContext() | ||
| + ", rootSpan=" | ||
| + this.rootSpan.getSpanContext() | ||
| + '}'; | ||
| return "OtelContext{" + "delegate=" + delegate + '}'; | ||
| } | ||
|
|
||
| private static datadog.context.ContextKey delegateKey(ContextKey key) { | ||
| return DELEGATE_KEYS.computeIfAbsent(key, OtelContext::mapByKeyName); | ||
| } | ||
|
|
||
| private static datadog.context.ContextKey mapByKeyName(ContextKey key) { | ||
| return datadog.context.ContextKey.named(key.toString()); | ||
| } | ||
|
|
||
| private static OtelSpan toOtelSpan(AgentSpan span) { | ||
| if (span instanceof AttachableWrapper) { | ||
| Object wrapper = ((AttachableWrapper) span).getWrapper(); | ||
| if (wrapper instanceof OtelSpan) { | ||
| return (OtelSpan) wrapper; | ||
| } | ||
| } | ||
| return new OtelSpan(span); | ||
| } | ||
| } | ||
21 changes: 5 additions & 16 deletions
21
...gent/agent-otel/otel-shim/src/main/java/datadog/opentelemetry/shim/context/OtelScope.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,31 +1,20 @@ | ||
| package datadog.opentelemetry.shim.context; | ||
|
|
||
| import datadog.trace.bootstrap.instrumentation.api.AgentScope; | ||
| import datadog.trace.bootstrap.instrumentation.api.AttachableWrapper; | ||
| import datadog.context.ContextScope; | ||
| import io.opentelemetry.context.Scope; | ||
|
|
||
| public class OtelScope implements Scope { | ||
| private final Scope scope; | ||
| private final AgentScope delegate; | ||
| private final Object[] contextEntries; | ||
| private final ContextScope delegate; | ||
|
|
||
| public OtelScope(Scope scope, AgentScope delegate, Object[] contextEntries) { | ||
| public OtelScope(Scope scope, ContextScope delegate) { | ||
| this.scope = scope; | ||
| this.delegate = delegate; | ||
| this.contextEntries = contextEntries; | ||
| if (delegate instanceof AttachableWrapper) { | ||
| ((AttachableWrapper) delegate).attachWrapper(this); | ||
| } | ||
| } | ||
|
|
||
| /** Context entries from {@link OtelContext}, captured when the context was made current. */ | ||
| Object[] contextEntries() { | ||
| return contextEntries; | ||
| } | ||
|
|
||
| @Override | ||
| public void close() { | ||
| this.delegate.close(); | ||
| this.scope.close(); | ||
| delegate.close(); | ||
| scope.close(); | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I imagine that this is not that big. However I did not see any cleanup or expiry but only entry insertion. Is that on purpose?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Correct, this is never going to grow that big because only a handful of context keys will ever be created and each one only takes up a small amount of space.
Context keys are also meant to be stored as constants, because each one is unique and must be available to allow querying - so even if we did have some kind of weak map here it wouldn't be evicting much because it would be rare for the mapped key to become unreachable during the application lifetime.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ok thank you for having given more "context for this context" :)