Combine Threadlocal and RefCounting into single Scope
Also added additional tests.
MODIFY
dd-java-agent/instrumentation/aws-sdk/src/test/groovy/AWSClientTest.groovy
dd-java-agent/instrumentation/aws-sdk/src/test/groovy/AWSClientTest.groovy
@@ -17,7 +17,7 @@ import java.util.concurrent.atomic.AtomicReference
import static ratpack.groovy.test.embed.GroovyEmbeddedApp.ratpack
-@Timeout(5)
+@Timeout(20)
class AWSClientTest extends AgentTestRunner {
def "request handler is hooked up with builder"() {
@@ -61,7 +61,6 @@ class AWSClientTest extends AgentTestRunner {
false | 1
}
- @Timeout(10)
def "send request with mocked back end"() {
setup:
def receivedHeaders = new AtomicReference<Headers>()

MODIFY
dd-trace-ot/src/main/java/datadog/opentracing/DDSpanContext.java
dd-trace-ot/src/main/java/datadog/opentracing/DDSpanContext.java
@@ -58,8 +58,6 @@ public class DDSpanContext implements io.opentracing.SpanContext {
/** Tags are associated to the current span, they will not propagate to the children span */
private Map<String, Object> tags;
- @JsonIgnore public final boolean useRefCounting;
-
public DDSpanContext(
final long traceId,
final long spanId,
@@ -73,8 +71,7 @@ public class DDSpanContext implements io.opentracing.SpanContext {
final String spanType,
final Map<String, Object> tags,
final Queue<DDSpan> trace,
- final DDTracer tracer,
- final boolean useRefCounting) {
+ final DDTracer tracer) {
this.traceId = traceId;
this.spanId = spanId;
@@ -95,8 +92,6 @@ public class DDSpanContext implements io.opentracing.SpanContext {
this.tags = tags;
- this.useRefCounting = useRefCounting;
-
if (trace == null) {
// TODO: figure out better concurrency model.
this.trace = new ConcurrentLinkedQueue<>();

MODIFY
dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java
dd-trace-ot/src/main/java/datadog/opentracing/DDTracer.java
@@ -279,7 +279,6 @@ public class DDTracer implements io.opentracing.Tracer {
private boolean errorFlag;
private String spanType;
private boolean ignoreScope = false;
- private boolean useRefCounting = false;
public DDSpanBuilder(final String operationName, final ScopeManager scopeManager) {
this.operationName = operationName;
@@ -341,12 +340,8 @@ public class DDTracer implements io.opentracing.Tracer {
@Override
public DDSpanBuilder withTag(final String tag, final boolean bool) {
- if (bool && tag.equals("dd.use.ref.counting")) {
- return withReferenceCounting();
- } else {
return withTag(tag, (Object) bool);
}
- }
@Override
public DDSpanBuilder withStartTimestamp(final long timestampMillis) {
@@ -369,11 +364,6 @@ public class DDTracer implements io.opentracing.Tracer {
return this;
}
- public DDSpanBuilder withReferenceCounting() {
- this.useRefCounting = true;
- return this;
- }
-
public DDSpanBuilder withSpanType(final String spanType) {
this.spanType = spanType;
return this;
@@ -481,8 +471,7 @@ public class DDTracer implements io.opentracing.Tracer {
spanType,
this.tags,
parentTrace,
- DDTracer.this,
- useRefCounting);
+ DDTracer.this);
return context;
}

MODIFY
dd-trace-ot/src/main/java/datadog/opentracing/propagation/HTTPCodec.java
dd-trace-ot/src/main/java/datadog/opentracing/propagation/HTTPCodec.java
@@ -72,8 +72,7 @@ public class HTTPCodec implements Codec<TextMap> {
null,
null,
null,
- null,
- false);
+ null);
context.lockSamplingPriority();
log.debug("{} - Parent context extracted", context);

MODIFY
dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContextualScopeManager.java
dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/ContextualScopeManager.java
package datadog.opentracing.scopemanager;
-import datadog.opentracing.DDSpan;
import io.opentracing.Scope;
import io.opentracing.ScopeManager;
import io.opentracing.Span;
-import java.util.Set;
-import java.util.concurrent.CopyOnWriteArraySet;
+import java.util.Deque;
+import java.util.concurrent.ConcurrentLinkedDeque;
import java.util.concurrent.atomic.AtomicInteger;
public class ContextualScopeManager implements ScopeManager {
final ThreadLocal<Scope> tlsScope = new ThreadLocal<>();
- final Set<ScopeContext> scopeContexts = new CopyOnWriteArraySet<>();
+ final Deque<ScopeContext> scopeContexts = new ConcurrentLinkedDeque<>();
@Override
public Scope activate(final Span span, final boolean finishOnClose) {
@@ -19,11 +18,7 @@ public class ContextualScopeManager implements ScopeManager {
return context.activate(span, finishOnClose);
}
}
- if (span instanceof DDSpan && ((DDSpan) span).context().useRefCounting) {
- return new RefCountingScope(this, new AtomicInteger(1), span);
- } else {
- return new ThreadLocalScope(this, span, finishOnClose);
- }
+ return new RefCountingScope(this, new AtomicInteger(1), span, finishOnClose);
}
@Override
@@ -37,91 +32,6 @@ public class ContextualScopeManager implements ScopeManager {
}
public void addScopeContext(final ScopeContext context) {
- scopeContexts.add(context);
- }
-
- class ThreadLocalScope implements Scope {
- private final ContextualScopeManager scopeManager;
- private final Span wrapped;
- private final boolean finishOnClose;
- private final Scope toRestore;
-
- ThreadLocalScope(
- final ContextualScopeManager scopeManager,
- final Span wrapped,
- final boolean finishOnClose) {
- this.scopeManager = scopeManager;
- this.wrapped = wrapped;
- this.finishOnClose = finishOnClose;
- this.toRestore = scopeManager.tlsScope.get();
- scopeManager.tlsScope.set(this);
- }
-
- @Override
- public void close() {
- if (scopeManager.tlsScope.get() != this) {
- // This shouldn't happen if users call methods in the expected order. Bail out.
- return;
- }
-
- if (finishOnClose) {
- wrapped.finish();
- }
-
- scopeManager.tlsScope.set(toRestore);
- }
-
- @Override
- public Span span() {
- return wrapped;
- }
- }
-
- public class RefCountingScope implements Scope {
- final ContextualScopeManager manager;
- final AtomicInteger refCount;
- private final Span wrapped;
- private final Scope toRestore;
-
- RefCountingScope(
- final ContextualScopeManager manager, final AtomicInteger refCount, final Span wrapped) {
- this.manager = manager;
- this.refCount = refCount;
- this.wrapped = wrapped;
- this.toRestore = manager.tlsScope.get();
- manager.tlsScope.set(this);
- }
-
- public class Continuation {
- public Continuation() {
- refCount.incrementAndGet();
- }
-
- public RefCountingScope activate() {
- return new RefCountingScope(manager, refCount, wrapped);
- }
- }
-
- public RefCountingScope.Continuation capture() {
- return new RefCountingScope.Continuation();
- }
-
- @Override
- public void close() {
- if (manager.tlsScope.get() != this) {
- return;
- }
-
- if (refCount.decrementAndGet() == 0) {
- wrapped.finish();
- }
-
- manager.tlsScope.set(toRestore);
- }
-
- @Override
- public Span span() {
- return wrapped;
- }
+ scopeContexts.addFirst(context);
}
}

ADD
dd-trace-ot/src/main/java/datadog/opentracing/scopemanager/RefCountingScope.java
+package datadog.opentracing.scopemanager;
+
+import io.opentracing.Scope;
+import io.opentracing.Span;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class RefCountingScope implements Scope {
+ final ContextualScopeManager scopeManager;
+ final AtomicInteger refCount;
+ private final Span wrapped;
+ private final boolean finishOnClose;
+ private final Scope toRestore;
+
+ RefCountingScope(
+ final ContextualScopeManager scopeManager,
+ final AtomicInteger refCount,
+ final Span wrapped,
+ final boolean finishOnClose) {
+ this.scopeManager = scopeManager;
+ this.refCount = refCount;
+ this.wrapped = wrapped;
+ this.finishOnClose = finishOnClose;
+ this.toRestore = scopeManager.tlsScope.get();
+ scopeManager.tlsScope.set(this);
+ }
+
+ @Override
+ public void close() {
+ if (scopeManager.tlsScope.get() != this) {
+ return;
+ }
+
+ if (refCount.decrementAndGet() == 0 && finishOnClose) {
+ wrapped.finish();
+ }
+
+ scopeManager.tlsScope.set(toRestore);
+ }
+
+ @Override
+ public Span span() {
+ return wrapped;
+ }
+
+ public Continuation capture() {
+ return new Continuation();
+ }
+
+ public class Continuation {
+ public Continuation() {
+ refCount.incrementAndGet();
+ }
+
+ public Scope activate() {
+ for (final ScopeContext context : scopeManager.scopeContexts) {
+ if (context.inContext()) {
+ return context.activate(wrapped, finishOnClose);
+ }
+ }
+ return new RefCountingScope(scopeManager, refCount, wrapped, finishOnClose);
+ }
+ }
+}

MODIFY
dd-trace-ot/src/test/groovy/datadog/opentracing/DDSpanSerializationTest.groovy
dd-trace-ot/src/test/groovy/datadog/opentracing/DDSpanSerializationTest.groovy
@@ -49,8 +49,7 @@ class DDSpanSerializationTest extends Specification {
"type",
tags,
null,
- null,
- false)
+ null)
baggage.put(DDTags.THREAD_NAME, Thread.currentThread().getName())
baggage.put(DDTags.THREAD_ID, String.valueOf(Thread.currentThread().getId()))

MODIFY
dd-trace-ot/src/test/groovy/datadog/opentracing/decorators/URLAsResourceNameTest.groovy
dd-trace-ot/src/test/groovy/datadog/opentracing/decorators/URLAsResourceNameTest.groovy
@@ -96,8 +96,7 @@ class URLAsResourceNameTest extends Specification {
"fakeType",
tags,
null,
- null,
- false)
+ null)
then:
decorator.afterSetTag(context, Tags.HTTP_URL.getKey(), value)

MODIFY
dd-trace-ot/src/test/groovy/datadog/opentracing/propagation/HTTPCodecTest.groovy
dd-trace-ot/src/test/groovy/datadog/opentracing/propagation/HTTPCodecTest.groovy
@@ -42,8 +42,7 @@ class HTTPCodecTest extends Specification {
"fakeType",
null,
null,
- null,
- false)
+ null)
final Map<String, String> carrier = new HashMap<>()

MODIFY
dd-trace-ot/src/test/groovy/datadog/opentracing/scopemanager/ScopeManagerTest.groovy
dd-trace-ot/src/test/groovy/datadog/opentracing/scopemanager/ScopeManagerTest.groovy
package datadog.opentracing.scopemanager
-import datadog.opentracing.DDSpan
import datadog.opentracing.DDTracer
import datadog.trace.common.writer.ListWriter
import io.opentracing.Scope
import io.opentracing.Span
import spock.lang.Specification
import spock.lang.Subject
+import spock.lang.Timeout
+import spock.lang.Unroll
+import java.util.concurrent.atomic.AtomicBoolean
import java.util.concurrent.atomic.AtomicReference
+@Timeout(1)
class ScopeManagerTest extends Specification {
def writer = new ListWriter()
def tracer = new DDTracer(writer)
@@ -34,78 +37,70 @@ class ScopeManagerTest extends Specification {
def "threadlocal is active"() {
when:
def builder = tracer.buildSpan("test")
- builder.withTag("dd.use.ref.counting", false)
- def scope = builder.startActive(true)
+ def scope = builder.startActive(finishSpan)
then:
!spanReported(scope.span())
scopeManager.active() == scope
- scope instanceof ContextualScopeManager.ThreadLocalScope
+ scope instanceof RefCountingScope
writer.empty
when:
scope.close()
then:
- spanReported(scope.span())
- writer == [[scope.span()]]
+ spanReported(scope.span()) == finishSpan
+ writer == [[scope.span()]] || !finishSpan
scopeManager.active() == null
+
+ where:
+ finishSpan << [true, false]
}
- def "threadlocal is active with ref counting scope"() {
+ def "sets parent as current upon close"() {
setup:
- def builder = tracer.buildSpan("test")
- builder.withReferenceCounting()
- def scope = builder.startActive(true)
+ def parentScope = tracer.buildSpan("parent").startActive(finishSpan)
+ def childScope = tracer.buildSpan("parent").startActive(finishSpan)
expect:
- !spanReported(scope.span())
- scopeManager.active() == scope
- scope instanceof ContextualScopeManager.RefCountingScope
+ scopeManager.active() == childScope
+ childScope.span().context().parentId == parentScope.span().context().spanId
+ childScope.span().context().trace == parentScope.span().context().trace
when:
- scope.close()
+ childScope.close()
then:
- spanReported(scope.span())
- writer == [[scope.span()]]
- scopeManager.active() == null
- }
-
- def "threadlocal is active with ref counting scope using tag"() {
- setup:
- def builder = tracer.buildSpan("test")
- builder.withTag("dd.use.ref.counting", true)
- def scope = builder.startActive(true)
-
- expect:
- !spanReported(scope.span())
- scopeManager.active() == scope
- scope instanceof ContextualScopeManager.RefCountingScope
+ scopeManager.active() == parentScope
+ spanReported(childScope.span()) == finishSpan
+ !spanReported(parentScope.span())
+ writer == []
when:
- scope.close()
+ parentScope.close()
then:
- spanReported(scope.span())
- writer == [[scope.span()]]
+ spanReported(childScope.span()) == finishSpan
+ spanReported(parentScope.span()) == finishSpan
+ writer == [[parentScope.span(), childScope.span()]] || !finishSpan
scopeManager.active() == null
+
+ where:
+ finishSpan << [true, false]
}
def "ref counting scope doesn't close if non-zero"() {
setup:
def builder = tracer.buildSpan("test")
- builder.withReferenceCounting()
- def scope = (ContextualScopeManager.RefCountingScope) builder.startActive(true)
+ def scope = (RefCountingScope) builder.startActive(true)
def continuation = scope.capture()
expect:
!spanReported(scope.span())
scopeManager.active() == scope
- scope instanceof ContextualScopeManager.RefCountingScope
+ scope instanceof RefCountingScope
writer.empty
-
when:
scope.close()
@@ -129,7 +124,64 @@ class ScopeManagerTest extends Specification {
writer == [[scope.span()]]
}
- def "context takes control"() {
+ def "continuation restores trace"() {
+ setup:
+ def parentScope = tracer.buildSpan("parent").startActive(false) //false or trace is reported early
+ def parentSpan = parentScope.span()
+ RefCountingScope childScope = (RefCountingScope) tracer.buildSpan("parent").startActive(true)
+ def childSpan = childScope.span()
+
+ def cont = childScope.capture()
+ childScope.close()
+
+ expect:
+ parentSpan.context().trace == childSpan.context().trace
+ scopeManager.active() == parentScope
+ !spanReported(childSpan)
+ !spanReported(parentSpan)
+
+ when:
+ parentScope.close()
+ // If finishSpanOnClose was true for the parent, the trace would get reported even though the child was not done.
+
+ then:
+ scopeManager.active() == null
+ !spanReported(childSpan)
+ !spanReported(parentSpan)
+ writer == []
+
+ when:
+ def newScope = cont.activate()
+
+ then:
+ scopeManager.active() == newScope
+ newScope != childScope && newScope != parentScope
+ newScope.span() == childSpan
+ !spanReported(childSpan)
+ !spanReported(parentSpan)
+ writer == []
+
+ when:
+ newScope.close()
+
+ then:
+ scopeManager.active() == null
+ spanReported(childSpan)
+ !spanReported(parentSpan)
+ writer == []
+
+ when:
+ // Since finishSpanOnClose was false, we must manually finish the span.
+ parentSpan.finish()
+
+ then:
+ spanReported(childSpan)
+ spanReported(parentSpan)
+ writer == [[parentSpan, childSpan]]
+ }
+
+ @Unroll
+ def "context takes control (#active)"() {
setup:
contexts.each {
scopeManager.addScopeContext(it)
@@ -140,18 +192,18 @@ class ScopeManagerTest extends Specification {
expect:
scopeManager.tlsScope.get() == null
scopeManager.active() == scope
- contexts[activeIndex].get() == scope.get()
+ contexts[active].get() == scope.get()
writer.empty
where:
- activeIndex | contexts
+ active | contexts
0 | [new AtomicReferenceScope(true)]
- 0 | [new AtomicReferenceScope(true), new AtomicReferenceScope(true)]
- 1 | [new AtomicReferenceScope(false), new AtomicReferenceScope(true), new AtomicReferenceScope(true), new AtomicReferenceScope(false)]
- 2 | [new AtomicReferenceScope(false), new AtomicReferenceScope(false), new AtomicReferenceScope(true)]
+ 1 | [new AtomicReferenceScope(true), new AtomicReferenceScope(true)]
+ 3 | [new AtomicReferenceScope(false), new AtomicReferenceScope(true), new AtomicReferenceScope(false), new AtomicReferenceScope(true)]
}
- def "disabled context is ignored"() {
+ @Unroll
+ def "disabled context is ignored (#contexts.size)"() {
setup:
contexts.each {
scopeManager.addScopeContext(it)
@@ -160,12 +212,13 @@ class ScopeManagerTest extends Specification {
def scope = builder.startActive(true)
expect:
+ contexts.findAll {
+ it.get() != null
+ } == []
+
scopeManager.tlsScope.get() == scope
scopeManager.active() == scope
writer.empty
- contexts.each {
- assert it.get() == null
- } == contexts
where:
contexts | _
@@ -175,20 +228,86 @@ class ScopeManagerTest extends Specification {
[new AtomicReferenceScope(false), new AtomicReferenceScope(false), new AtomicReferenceScope(false)] | _
}
- boolean spanReported(DDSpan span) {
+ @Unroll
+ def "threadlocal to context with capture (#active)"() {
+ setup:
+ contexts.each {
+ scopeManager.addScopeContext(it)
+ }
+ RefCountingScope scope = (RefCountingScope) tracer.buildSpan("parent").startActive(true)
+
+ expect:
+ scopeManager.tlsScope.get() == scope
+
+ when:
+ def cont = scope.capture()
+ scope.close()
+
+ then:
+ scopeManager.tlsScope.get() == null
+
+ when:
+ active.each {
+ ((AtomicBoolean) contexts[it].enabled).set(true)
+ }
+ cont.activate()
+
+ then:
+ scopeManager.tlsScope.get() == null
+
+ where:
+ active | contexts
+ [0] | [new AtomicReferenceScope(false)]
+ [0] | [new AtomicReferenceScope(false), new AtomicReferenceScope(false)]
+ [1] | [new AtomicReferenceScope(false), new AtomicReferenceScope(false), new AtomicReferenceScope(false)]
+ [2] | [new AtomicReferenceScope(false), new AtomicReferenceScope(false), new AtomicReferenceScope(false)]
+ [0, 2] | [new AtomicReferenceScope(false), new AtomicReferenceScope(false), new AtomicReferenceScope(false)]
+ }
+
+ @Unroll
+ def "context to threadlocal (#contexts.size)"() {
+ setup:
+ contexts.each {
+ scopeManager.addScopeContext(it)
+ }
+ def scope = tracer.buildSpan("parent").startActive(false)
+ def span = scope.span()
+
+ expect:
+ scope instanceof AtomicReferenceScope
+ scopeManager.tlsScope.get() == null
+
+ when:
+ scope.close()
+ contexts.each {
+ ((AtomicBoolean) it.enabled).set(false)
+ }
+ scope = scopeManager.activate(span, true)
+
+ then:
+ scope instanceof RefCountingScope
+ scopeManager.tlsScope.get() == scope
+
+ where:
+ contexts | _
+ [new AtomicReferenceScope(true)] | _
+ [new AtomicReferenceScope(true), new AtomicReferenceScope(true)] | _
+ }
+
+ boolean spanReported(Span span) {
return span.durationNano != 0
}
class AtomicReferenceScope extends AtomicReference<Span> implements ScopeContext, Scope {
- final boolean enabled
+ final AtomicBoolean enabled
AtomicReferenceScope(boolean enabled) {
- this.enabled = enabled
+ this.enabled = new AtomicBoolean(enabled)
}
@Override
boolean inContext() {
- return enabled
+ return enabled.get()
}
@Override
@@ -211,5 +330,9 @@ class ScopeManagerTest extends Specification {
Scope active() {
return get() == null ? null : this
}
+
+ String toString() {
+ return "Ref: " + super.toString()
+ }
}
}

MODIFY
dd-trace-ot/src/test/groovy/datadog/trace/SpanFactory.groovy
dd-trace-ot/src/test/groovy/datadog/trace/SpanFactory.groovy
@@ -20,8 +20,7 @@ class SpanFactory {
"fakeType",
Collections.emptyMap(),
null,
- new DDTracer(),
- false)
+ new DDTracer())
return new DDSpan(timestampMicro, context)
}
@@ -39,8 +38,7 @@ class SpanFactory {
"fakeType",
Collections.emptyMap(),
null,
- tracer,
- false)
+ tracer)
return new DDSpan(1, context)
}
}

MODIFY
dd-trace-ot/src/test/groovy/datadog/trace/api/sampling/RateByServiceSamplerTest.groovy
dd-trace-ot/src/test/groovy/datadog/trace/api/sampling/RateByServiceSamplerTest.groovy
@@ -66,8 +66,7 @@ class RateByServiceSamplerTest extends Specification {
"fakeType",
Collections.emptyMap(),
null,
- new DDTracer(),
- false)
+ new DDTracer())
context.setTag("env", envName)
return new DDSpan(0l, context)
}

Combine Threadlocal and RefCounting into single Scope

Also added additional tests.