Skip to content

Commit f57faba

Browse files
tylerbensontylerbenson
committed
Allow closing of continuation explicitly
to allow a completed trace to be reported more timely. Continuation is closed when the returned scope is closed, but can also be closed directly.
1 parent 087b2e7 commit f57faba

3 files changed

Lines changed: 122 additions & 27 deletions

File tree

‎dd-trace-ot/src/main/java/datadog/opentracing/PendingTrace.java‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,11 +83,23 @@ public void addSpan(final DDSpan span) {
8383
* completed, so we need to wait till continuations are de-referenced before reporting.
8484
*/
8585
public void registerContinuation(final ContinuableScope.Continuation continuation) {
86-
weakReferences.add(
87-
new WeakReference<ContinuableScope.Continuation>(continuation, referenceQueue));
86+
continuation.ref =
87+
new WeakReference<ContinuableScope.Continuation>(continuation, referenceQueue);
88+
weakReferences.add(continuation.ref);
8889
pendingReferenceCount.incrementAndGet();
8990
}
9091

92+
public void cancelContinuation(final ContinuableScope.Continuation continuation) {
93+
synchronized (continuation) {
94+
if (continuation.ref != null) {
95+
weakReferences.remove(continuation.ref);
96+
continuation.ref.clear();
97+
continuation.ref = null;
98+
}
99+
}
100+
expireReference();
101+
}
102+
91103
private void expireReference() {
92104
if (pendingReferenceCount.decrementAndGet() == 0) {
93105
write();

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

Lines changed: 58 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,17 @@
11
package datadog.opentracing.scopemanager;
22

33
import datadog.opentracing.DDSpanContext;
4+
import datadog.opentracing.PendingTrace;
45
import io.opentracing.Scope;
56
import io.opentracing.Span;
7+
import io.opentracing.noop.NoopScopeManager;
8+
import java.io.Closeable;
9+
import java.lang.ref.WeakReference;
10+
import java.util.concurrent.atomic.AtomicBoolean;
611
import java.util.concurrent.atomic.AtomicInteger;
12+
import lombok.extern.slf4j.Slf4j;
713

14+
@Slf4j
815
public class ContinuableScope implements Scope {
916
final ContextualScopeManager scopeManager;
1017
final AtomicInteger refCount;
@@ -49,30 +56,75 @@ public Span span() {
4956
return wrapped;
5057
}
5158

59+
/**
60+
* The continuation returned should be closed after the associa
61+
*
62+
* @param finishOnClose
63+
* @return
64+
*/
5265
public Continuation capture(final boolean finishOnClose) {
5366
return new Continuation(this.finishOnClose && finishOnClose);
5467
}
5568

56-
public class Continuation {
69+
public class Continuation implements Closeable {
70+
public WeakReference<Continuation> ref;
5771

72+
private final AtomicBoolean used = new AtomicBoolean(false);
73+
private final PendingTrace trace;
5874
private final boolean finishSpanOnClose;
5975

6076
private Continuation(final boolean finishOnClose) {
6177
this.finishSpanOnClose = finishOnClose;
6278
refCount.incrementAndGet();
6379
if (wrapped.context() instanceof DDSpanContext) {
6480
final DDSpanContext context = (DDSpanContext) wrapped.context();
65-
context.getTrace().registerContinuation(this);
81+
trace = context.getTrace();
82+
trace.registerContinuation(this);
83+
} else {
84+
trace = null;
6685
}
6786
}
6887

6988
public Scope activate() {
70-
for (final ScopeContext context : scopeManager.scopeContexts) {
71-
if (context.inContext()) {
72-
return context.activate(wrapped, finishSpanOnClose);
89+
if (used.compareAndSet(false, true)) {
90+
for (final ScopeContext context : scopeManager.scopeContexts) {
91+
if (context.inContext()) {
92+
return new ClosingScope(context.activate(wrapped, finishSpanOnClose));
93+
}
7394
}
95+
return new ClosingScope(
96+
new ContinuableScope(scopeManager, refCount, wrapped, finishSpanOnClose));
97+
} else {
98+
log.debug("Reusing a continuation not allowed. Returning no-op scope.");
99+
return NoopScopeManager.NoopScope.INSTANCE;
100+
}
101+
}
102+
103+
@Override
104+
public void close() {
105+
used.getAndSet(true);
106+
if (trace != null) {
107+
trace.cancelContinuation(this);
108+
}
109+
}
110+
111+
private class ClosingScope implements Scope {
112+
private final Scope wrapped;
113+
114+
private ClosingScope(final Scope wrapped) {
115+
this.wrapped = wrapped;
116+
}
117+
118+
@Override
119+
public void close() {
120+
wrapped.close();
121+
Continuation.this.close();
122+
}
123+
124+
@Override
125+
public Span span() {
126+
return wrapped.span();
74127
}
75-
return new ContinuableScope(scopeManager, refCount, wrapped, finishSpanOnClose);
76128
}
77129
}
78130
}

‎dd-trace-ot/src/test/groovy/datadog/opentracing/scopemanager/ScopeManagerTest.groovy‎

Lines changed: 50 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -114,20 +114,34 @@ class ScopeManagerTest extends Specification {
114114

115115
when:
116116
continuation.activate()
117-
continuation = null // Continuation references also hold up traces.
118-
PendingTrace.awaitGC()
119-
((DDSpanContext) scopeManager.active().span().context()).trace.clean()
117+
if (forceGC) {
118+
continuation = null // Continuation references also hold up traces.
119+
PendingTrace.awaitGC()
120+
((DDSpanContext) scope.span().context()).trace.clean()
121+
}
122+
if (autoClose) {
123+
if (continuation != null) {
124+
continuation.close()
125+
}
126+
}
120127

121128
then:
122129
scopeManager.active() != null
123130

124131
when:
125132
scopeManager.active().close()
133+
writer.waitForTraces(1)
126134

127135
then:
128136
scopeManager.active() == null
129137
spanFinished(scope.span())
130138
writer == [[scope.span()]]
139+
140+
where:
141+
autoClose | forceGC
142+
true | true
143+
true | false
144+
false | true
131145
}
132146

133147
def "hard reference on continuation prevents trace from reporting"() {
@@ -145,13 +159,26 @@ class ScopeManagerTest extends Specification {
145159
writer == []
146160

147161
when:
148-
// remove hard reference and force GC
149-
continuation = null
150-
PendingTrace.awaitGC()
151-
span.context().trace.clean()
162+
if (forceGC) {
163+
continuation = null // Continuation references also hold up traces.
164+
PendingTrace.awaitGC()
165+
((DDSpanContext) span.context()).trace.clean()
166+
writer.waitForTraces(1)
167+
}
168+
if (autoClose) {
169+
if (continuation != null) {
170+
continuation.close()
171+
}
172+
}
152173

153174
then:
154175
writer == [[span]]
176+
177+
where:
178+
autoClose | forceGC
179+
true | true
180+
true | false
181+
false | true
155182
}
156183

157184
def "continuation restores trace"() {
@@ -161,7 +188,7 @@ class ScopeManagerTest extends Specification {
161188
ContinuableScope childScope = (ContinuableScope) tracer.buildSpan("parent").startActive(true)
162189
def childSpan = childScope.span()
163190

164-
def cont = childScope.capture(true)
191+
def continuation = childScope.capture(true)
165192
childScope.close()
166193

167194
expect:
@@ -181,13 +208,10 @@ class ScopeManagerTest extends Specification {
181208
writer == []
182209

183210
when:
184-
def newScope = cont.activate()
185-
cont = null // Continuation references also hold up traces.
186-
PendingTrace.awaitGC()
187-
((DDSpanContext) scopeManager.active().span().context()).trace.clean()
211+
def newScope = continuation.activate()
188212

189213
then:
190-
scopeManager.active() == newScope
214+
scopeManager.active() == newScope.wrapped
191215
newScope != childScope && newScope != parentScope
192216
newScope.span() == childSpan
193217
!spanFinished(childSpan)
@@ -217,7 +241,7 @@ class ScopeManagerTest extends Specification {
217241

218242
expect:
219243
newScope != scope
220-
scopeManager.active() == newScope
244+
scopeManager.active() == newScope.wrapped
221245
spanFinished(span)
222246
writer == []
223247

@@ -234,14 +258,21 @@ class ScopeManagerTest extends Specification {
234258
writer == []
235259

236260
when:
237-
// remove hard reference and force GC
238-
continuation = null
239-
PendingTrace.awaitGC()
240-
span.context().trace.clean()
241-
writer.waitForTraces(1)
261+
if (closeScope) {
262+
newScope.close()
263+
}
264+
if (closeContinuation) {
265+
continuation.close()
266+
}
242267

243268
then:
244269
writer == [[childSpan, span]]
270+
271+
where:
272+
closeScope | closeContinuation
273+
true | false
274+
false | true
275+
true | true
245276
}
246277

247278
@Unroll

0 commit comments

Comments
 (0)