Skip to content

Commit a374998

Browse files
mccullsmcculls
committed
Temporary bridge between contexts and spans/scopes
1 parent 769cec5 commit a374998

File tree

5 files changed

+93
-35
lines changed

5 files changed

+93
-35
lines changed

‎dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScope.java

+21-6
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.core.scopemanager;
22

3+
import datadog.context.Context;
34
import datadog.trace.api.Stateful;
45
import datadog.trace.api.scopemanager.ExtendedScopeListener;
56
import datadog.trace.api.scopemanager.ScopeListener;
@@ -15,10 +16,11 @@ class ContinuableScope implements AgentScope, AttachableWrapper {
1516
static final byte INSTRUMENTATION = 0;
1617
static final byte MANUAL = 1;
1718
static final byte ITERATION = 2;
19+
static final byte CONTEXT = 3;
1820

1921
private final ContinuableScopeManager scopeManager;
2022

21-
final AgentSpan span; // package-private so scopeManager can access it directly
23+
final Context context; // package-private so scopeManager can access it directly
2224

2325
/** Flag that this scope should be allowed to propagate across async boundaries. */
2426
private static final byte ASYNC_PROPAGATING = 1;
@@ -40,12 +42,12 @@ class ContinuableScope implements AgentScope, AttachableWrapper {
4042

4143
ContinuableScope(
4244
final ContinuableScopeManager scopeManager,
43-
final AgentSpan span,
45+
final Context context,
4446
final byte source,
4547
final boolean isAsyncPropagating,
4648
final Stateful scopeState) {
4749
this.scopeManager = scopeManager;
48-
this.span = span;
50+
this.context = context;
4951
this.source = source;
5052
this.flags = isAsyncPropagating ? ASYNC_PROPAGATING : 0;
5153
this.scopeState = scopeState;
@@ -65,7 +67,7 @@ public final void close() {
6567
byte source = source();
6668
scopeManager.healthMetrics.onScopeCloseError(source == MANUAL);
6769
if (source == MANUAL && scopeManager.strictMode) {
68-
throw new RuntimeException("Tried to close " + span + " scope when not on top");
70+
throw new RuntimeException("Tried to close " + context + " scope when not on top");
6971
}
7072
}
7173

@@ -131,7 +133,12 @@ public final boolean isAsyncPropagating() {
131133

132134
@Override
133135
public final AgentSpan span() {
134-
return span;
136+
return AgentSpan.fromContext(context);
137+
}
138+
139+
@Override
140+
public Context context() {
141+
return context;
135142
}
136143

137144
@Override
@@ -145,7 +152,7 @@ public final void setAsyncPropagation(final boolean value) {
145152

146153
@Override
147154
public final String toString() {
148-
return super.toString() + "->" + span;
155+
return super.toString() + "->" + context;
149156
}
150157

151158
public void checkpoint() {
@@ -162,6 +169,10 @@ public boolean rollback() {
162169
}
163170

164171
public final void beforeActivated() {
172+
AgentSpan span = span();
173+
if (span == null) {
174+
return;
175+
}
165176
try {
166177
scopeState.activate(span.context());
167178
} catch (Throwable e) {
@@ -171,6 +182,10 @@ public final void beforeActivated() {
171182
}
172183

173184
public final void afterActivated() {
185+
AgentSpan span = span();
186+
if (span == null) {
187+
return;
188+
}
174189
for (final ScopeListener listener : scopeManager.scopeListeners) {
175190
try {
176191
listener.afterScopeActivated();

‎dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuableScopeManager.java

+52-23
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,17 @@
33
import static datadog.trace.api.ConfigDefaults.DEFAULT_ASYNC_PROPAGATING;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopScope;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopSpan;
6+
import static datadog.trace.core.scopemanager.ContinuableScope.CONTEXT;
67
import static datadog.trace.core.scopemanager.ContinuableScope.INSTRUMENTATION;
78
import static datadog.trace.core.scopemanager.ContinuableScope.ITERATION;
89
import static datadog.trace.core.scopemanager.ContinuableScope.MANUAL;
910
import static java.util.concurrent.TimeUnit.MINUTES;
1011
import static java.util.concurrent.TimeUnit.NANOSECONDS;
1112
import static java.util.concurrent.TimeUnit.SECONDS;
1213

14+
import datadog.context.Context;
15+
import datadog.context.ContextManager;
16+
import datadog.context.ContextScope;
1317
import datadog.trace.api.Config;
1418
import datadog.trace.api.Stateful;
1519
import datadog.trace.api.scopemanager.ExtendedScopeListener;
@@ -40,7 +44,7 @@
4044
* from being reported even if all related spans are finished. It also delegates to other
4145
* ScopeInterceptors to provide additional functionality.
4246
*/
43-
public final class ContinuableScopeManager implements ScopeStateAware {
47+
public final class ContinuableScopeManager implements ScopeStateAware, ContextManager {
4448

4549
static final Logger log = LoggerFactory.getLogger(ContinuableScopeManager.class);
4650
static final RatelimitedLogger ratelimitedLog = new RatelimitedLogger(log, 1, MINUTES);
@@ -83,6 +87,8 @@ public ContinuableScopeManager(
8387
this.healthMetrics = healthMetrics;
8488
this.tlsScopeStack = new ScopeStackThreadLocal(profilingContextIntegration);
8589
this.profilingContextIntegration = profilingContextIntegration;
90+
91+
ContextManager.register(this);
8692
}
8793

8894
public AgentScope activateSpan(final AgentSpan span) {
@@ -96,10 +102,12 @@ public AgentScope activateManualSpan(final AgentSpan span) {
96102
public AgentScope.Continuation captureActiveSpan() {
97103
ContinuableScope activeScope = scopeStack().active();
98104
if (null != activeScope && activeScope.isAsyncPropagating()) {
99-
return captureSpan(activeScope.span(), activeScope.source());
100-
} else {
101-
return AgentTracer.noopContinuation();
105+
AgentSpan span = activeScope.span();
106+
if (span != null) {
107+
return captureSpan(span, activeScope.source());
108+
}
102109
}
110+
return AgentTracer.noopContinuation();
103111
}
104112

105113
public AgentScope.Continuation captureSpan(final AgentSpan span) {
@@ -111,14 +119,14 @@ private AgentScope.Continuation captureSpan(final AgentSpan span, byte source) {
111119
}
112120

113121
private AgentScope activate(
114-
final AgentSpan span,
122+
final Context context,
115123
final byte source,
116124
final boolean overrideAsyncPropagation,
117125
final boolean isAsyncPropagating) {
118126
ScopeStack scopeStack = scopeStack();
119127

120128
final ContinuableScope top = scopeStack.top;
121-
if (top != null && top.span.equals(span)) {
129+
if (top != null && top.context.equals(context)) {
122130
top.incrementReferences();
123131
return top;
124132
}
@@ -131,7 +139,7 @@ private AgentScope activate(
131139
return noopScope();
132140
}
133141

134-
assert span != null;
142+
assert context != null;
135143

136144
// Inherit the async propagation from the active scope unless the value is overridden
137145
boolean asyncPropagation =
@@ -140,7 +148,7 @@ private AgentScope activate(
140148
: top != null ? top.isAsyncPropagating() : DEFAULT_ASYNC_PROPAGATING;
141149

142150
final ContinuableScope scope =
143-
new ContinuableScope(this, span, source, asyncPropagation, createScopeState(span));
151+
new ContinuableScope(this, context, source, asyncPropagation, createScopeState(context));
144152
scopeStack.push(scope);
145153
healthMetrics.onActivateScope();
146154

@@ -153,26 +161,26 @@ private AgentScope activate(
153161
* @param continuation {@code null} if a continuation is re-used
154162
*/
155163
ContinuableScope continueSpan(
156-
final ScopeContinuation continuation, final AgentSpan span, final byte source) {
164+
final ScopeContinuation continuation, final Context context, final byte source) {
157165
ScopeStack scopeStack = scopeStack();
158166

159167
// optimization: if the top scope is already keeping the same span alive
160168
// then re-use that scope (avoids allocation) and cancel the continuation
161169
final ContinuableScope top = scopeStack.top;
162-
if (top != null && top.span.equals(span)) {
170+
if (top != null && top.context.equals(context)) {
163171
top.incrementReferences();
164172
if (continuation != null) {
165173
continuation.cancelFromContinuedScopeClose();
166174
}
167175
return top;
168176
}
169177

170-
Stateful scopeState = createScopeState(span);
178+
Stateful scopeState = createScopeState(context);
171179
final ContinuableScope scope;
172180
if (continuation != null) {
173-
scope = new ContinuingScope(this, span, source, true, continuation, scopeState);
181+
scope = new ContinuingScope(this, context, source, true, continuation, scopeState);
174182
} else {
175-
scope = new ContinuableScope(this, span, source, true, scopeState);
183+
scope = new ContinuableScope(this, context, source, true, scopeState);
176184
}
177185
scopeStack.push(scope);
178186

@@ -202,8 +210,9 @@ public void closePrevious(final boolean finishSpan) {
202210
}
203211
top.close();
204212
scopeStack.cleanup();
205-
if (finishSpan) {
206-
top.span.finishWithEndToEnd();
213+
AgentSpan span = top.span();
214+
if (finishSpan && span != null) {
215+
span.finishWithEndToEnd();
207216
}
208217
}
209218
}
@@ -261,7 +270,7 @@ public void rollbackActiveToCheckpoint() {
261270

262271
public AgentSpan activeSpan() {
263272
final ContinuableScope active = scopeStack().active();
264-
return active == null ? null : active.span;
273+
return active == null ? null : active.span();
265274
}
266275

267276
/** Attach a listener to scope activation events */
@@ -289,11 +298,12 @@ private void addExtendedScopeListener(final ExtendedScopeListener listener) {
289298
}
290299
}
291300

292-
private Stateful createScopeState(AgentSpan span) {
301+
private Stateful createScopeState(Context context) {
293302
// currently this just manages things the profiler has to do per scope, but could be expanded
294303
// to encapsulate other scope lifecycle activities
295304
// FIXME DDSpanContext is always a ProfilerContext anyway...
296-
if (span.context() instanceof ProfilerContext) {
305+
AgentSpan span = AgentSpan.fromContext(context);
306+
if (span != null && span.context() instanceof ProfilerContext) {
297307
return profilingContextIntegration.newScopeState((ProfilerContext) span.context());
298308
}
299309
return Stateful.DEFAULT;
@@ -308,6 +318,22 @@ public ScopeState newScopeState() {
308318
return new ContinuableScopeState();
309319
}
310320

321+
@Override
322+
public Context current() {
323+
final ContinuableScope active = scopeStack().active();
324+
return active == null ? Context.root() : active.context;
325+
}
326+
327+
@Override
328+
public ContextScope attach(Context context) {
329+
return activate(context, CONTEXT, false, true);
330+
}
331+
332+
@Override
333+
public Context swap(Context context) {
334+
throw new UnsupportedOperationException("Not yet implemented");
335+
}
336+
311337
private class ContinuableScopeState implements ScopeState {
312338

313339
private ScopeStack localScopeStack = tlsScopeStack.initialValue();
@@ -383,11 +409,14 @@ public void run(Map<ScopeStack, ContinuableScope> rootIterationScopes) {
383409

384410
if (!rootScope.alive()) { // no need to track this anymore
385411
itr.remove();
386-
} else if (NANOSECONDS.toMillis(rootScope.span.getStartTime()) < cutOff) {
387-
// mark scope as overdue to allow cleanup and avoid further spans being attached
388-
scopeStack.overdueRootScope = rootScope;
389-
rootScope.span.finishWithEndToEnd();
390-
itr.remove();
412+
} else {
413+
AgentSpan span = rootScope.span();
414+
if (span != null && NANOSECONDS.toMillis(span.getStartTime()) < cutOff) {
415+
// mark scope as overdue to allow cleanup and avoid further spans being attached
416+
scopeStack.overdueRootScope = rootScope;
417+
span.finishWithEndToEnd();
418+
itr.remove();
419+
}
391420
}
392421
}
393422
}

‎dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ContinuingScope.java

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
11
package datadog.trace.core.scopemanager;
22

3+
import datadog.context.Context;
34
import datadog.trace.api.Stateful;
4-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
55

66
final class ContinuingScope extends ContinuableScope {
77
/** Continuation that created this scope. */
88
private final ScopeContinuation continuation;
99

1010
ContinuingScope(
1111
final ContinuableScopeManager scopeManager,
12-
final AgentSpan span,
12+
final Context context,
1313
final byte source,
1414
final boolean isAsyncPropagating,
1515
final ScopeContinuation continuation,
1616
final Stateful scopeState) {
17-
super(scopeManager, span, source, isAsyncPropagating, scopeState);
17+
super(scopeManager, context, source, isAsyncPropagating, scopeState);
1818
this.continuation = continuation;
1919
}
2020

‎dd-trace-core/src/main/java/datadog/trace/core/scopemanager/ScopeStack.java

+9-2
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static datadog.trace.core.scopemanager.ContinuableScope.ITERATION;
44

5+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
56
import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
67
import java.util.ArrayDeque;
78

@@ -84,14 +85,20 @@ final boolean checkOverdueScopes(final ContinuableScope expectedScope) {
8485
// avoid calling close() as we're already in that method, instead just clear any
8586
// remaining references so the scope gets removed in the subsequent cleanup() call
8687
top.clearReferences();
87-
top.span.finishWithEndToEnd();
88+
AgentSpan span = top.span();
89+
if (span != null) {
90+
span.finishWithEndToEnd();
91+
}
8892
// now do the same for any previous iteration scopes ahead of the expected scope
8993
for (ContinuableScope scope : stack) {
9094
if (scope.source() != ITERATION) {
9195
return expectedScope.equals(scope);
9296
} else {
9397
scope.clearReferences();
94-
scope.span.finishWithEndToEnd();
98+
span = scope.span();
99+
if (span != null) {
100+
span.finishWithEndToEnd();
101+
}
95102
}
96103
}
97104
return false; // we didn't find the expected scope

‎internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/AgentScope.java

+8-1
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,18 @@
11
package datadog.trace.bootstrap.instrumentation.api;
22

3+
import datadog.context.Context;
4+
import datadog.context.ContextScope;
35
import datadog.trace.context.TraceScope;
46
import java.io.Closeable;
57

6-
public interface AgentScope extends TraceScope, Closeable {
8+
public interface AgentScope extends ContextScope, TraceScope, Closeable {
79
AgentSpan span();
810

11+
@Override
12+
default Context context() {
13+
return span();
14+
}
15+
916
@Override
1017
void close();
1118

0 commit comments

Comments
 (0)