Skip to content

Commit 088c5e4

Browse files
mccullsmcculls
committed
WIP
1 parent c7533d0 commit 088c5e4

File tree

36 files changed

+206
-225
lines changed

36 files changed

+206
-225
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/AdviceUtils.java‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
package datadog.trace.bootstrap.instrumentation.java.concurrent;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
3+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeSpan;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
45
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled;
56

67
import datadog.trace.bootstrap.ContextStore;
78
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
9+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
810

911
/** Helper utils for Runnable/Callable instrumentation */
1012
public class AdviceUtils {
@@ -51,14 +53,14 @@ public static <T> void cancelTask(ContextStore<T, State> contextStore, final T t
5153
}
5254

5355
public static <T> void capture(ContextStore<T, State> contextStore, T task) {
54-
AgentScope activeScope = activeScope();
55-
if (null != activeScope && activeScope.isAsyncPropagating()) {
56+
AgentSpan span = activeSpan();
57+
if (span != null && isAsyncPropagationEnabled()) {
5658
State state = contextStore.get(task);
5759
if (null == state) {
5860
state = State.FACTORY.create();
5961
contextStore.put(task, state);
6062
}
61-
state.captureAndSetContinuation(activeScope);
63+
state.captureAndSetContinuation(span);
6264
}
6365
}
6466
}

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ConcurrentState.java‎

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package datadog.trace.bootstrap.instrumentation.java.concurrent;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureSpan;
34
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ContinuationClaim.CLAIMED;
45

56
import datadog.trace.bootstrap.ContextStore;
67
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
8+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
79
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
810
import org.slf4j.Logger;
911
import org.slf4j.LoggerFactory;
@@ -28,22 +30,13 @@ public final class ConcurrentState {
2830

2931
private ConcurrentState() {}
3032

31-
public static <K> ConcurrentState captureScope(
32-
ContextStore<K, ConcurrentState> contextStore, K key, AgentScope scope) {
33-
if (scope != null && scope.isAsyncPropagating()) {
34-
if (!scope.span().isValid()) {
35-
return null;
36-
}
37-
final ConcurrentState state = contextStore.putIfAbsent(key, FACTORY);
38-
if (!state.captureAndSetContinuation(scope) && log.isDebugEnabled()) {
39-
log.debug(
40-
"continuation was already set for {} in scope {}, no continuation captured.",
41-
key,
42-
scope);
43-
}
44-
return state;
33+
public static <K> void captureContinuation(
34+
ContextStore<K, ConcurrentState> contextStore, K key, AgentSpan span) {
35+
final ConcurrentState state = contextStore.putIfAbsent(key, FACTORY);
36+
if (!state.captureAndSetContinuation(span) && log.isDebugEnabled()) {
37+
log.debug(
38+
"continuation was already set for {} in span {}, no continuation captured.", key, span);
4539
}
46-
return null;
4740
}
4841

4942
public static <K> AgentScope activateAndContinueContinuation(
@@ -81,10 +74,10 @@ public static <K> void cancelAndClearContinuation(
8174
state.cancelAndClearContinuation();
8275
}
8376

84-
private boolean captureAndSetContinuation(final AgentScope scope) {
77+
private boolean captureAndSetContinuation(final AgentSpan span) {
8578
if (CONTINUATION.compareAndSet(this, null, CLAIMED)) {
8679
// lazy write is guaranteed to be seen by getAndSet
87-
CONTINUATION.lazySet(this, scope.capture().hold());
80+
CONTINUATION.lazySet(this, captureSpan(span).hold());
8881
return true;
8982
}
9083
return false;

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/ExecutorInstrumentationUtils.java‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import datadog.trace.bootstrap.ContextStore;
77
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
8+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
89
import java.util.concurrent.Executor;
910
import org.slf4j.Logger;
1011
import org.slf4j.LoggerFactory;
@@ -36,24 +37,22 @@ public static boolean shouldAttachStateToTask(final Object task, final Executor
3637
}
3738

3839
/**
39-
* Create task state given current scope.
40+
* Create task state given current span.
4041
*
4142
* @param contextStore context storage
4243
* @param task task instance
43-
* @param scope current scope
44+
* @param span current span
4445
* @param <T> task class type
4546
* @return new state
4647
*/
4748
public static <T> State setupState(
48-
final ContextStore<T, State> contextStore, final T task, final AgentScope scope) {
49+
final ContextStore<T, State> contextStore, final T task, final AgentSpan span) {
4950

5051
final State state = contextStore.putIfAbsent(task, State.FACTORY);
5152

52-
if (!state.captureAndSetContinuation(scope)) {
53+
if (!state.captureAndSetContinuation(span)) {
5354
log.debug(
54-
"continuation was already set for {} in scope {}, no continuation captured.",
55-
task,
56-
scope);
55+
"continuation was already set for {} in span {}, no continuation captured.", task, span);
5756
}
5857

5958
return state;

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/State.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.bootstrap.instrumentation.java.concurrent;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureSpan;
34
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ContinuationClaim.CLAIMED;
45

56
import datadog.trace.api.profiling.Timing;
@@ -25,15 +26,15 @@ public final class State {
2526

2627
private State() {}
2728

28-
public boolean captureAndSetContinuation(final AgentScope scope) {
29+
public boolean captureAndSetContinuation(final AgentSpan span) {
2930
if (CONTINUATION.compareAndSet(this, null, CLAIMED)) {
3031
// it's a real pain to do this twice, and this can actually
3132
// happen systematically - WITHOUT RACES - because of broken
3233
// instrumentation, e.g. SetExecuteRunnableStateAdvice
3334
// "double instruments" calls to ScheduledExecutorService.submit/schedule
3435
//
3536
// lazy write is guaranteed to be seen by getAndSet
36-
CONTINUATION.lazySet(this, scope.capture());
37+
CONTINUATION.lazySet(this, captureSpan(span));
3738
return true;
3839
}
3940
return false;

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/concurrent/Wrapper.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.trace.bootstrap.instrumentation.java.concurrent;
22

3-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
3+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation;
45
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
56
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;
67

@@ -17,12 +18,12 @@ public static <T extends Runnable> Runnable wrap(T task) {
1718
|| exclude(RUNNABLE, task)) {
1819
return task;
1920
}
20-
AgentScope scope = activeScope();
21-
if (null != scope) {
21+
AgentScope.Continuation continuation = captureActiveSpan();
22+
if (continuation != noopContinuation()) {
2223
if (task instanceof Comparable) {
23-
return new ComparableRunnable(task, scope.capture());
24+
return new ComparableRunnable(task, continuation);
2425
}
25-
return new Wrapper<>(task, scope.capture());
26+
return new Wrapper<>(task, continuation);
2627
}
2728
// don't wrap unless there is scope to propagate
2829
return task;

‎dd-java-agent/instrumentation/aerospike-4/src/main/java/datadog/trace/instrumentation/aerospike4/AerospikeClientInstrumentation.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
44
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
5+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureSpan;
56
import static datadog.trace.instrumentation.aerospike4.AerospikeClientDecorator.DECORATE;
67
import static net.bytebuddy.implementation.bytecode.assign.Assigner.Typing.DYNAMIC;
78
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
@@ -71,7 +72,7 @@ public static AgentScope beginRequest(
7172
AgentSpan clientSpan = DECORATE.startAerospikeSpan(methodName);
7273
AgentScope scope = activateSpan(clientSpan);
7374
// always want to wrap even when there's no listener so we get the true async time
74-
listener = new TracingListener(clientSpan, scope.capture(), listener);
75+
listener = new TracingListener(clientSpan, captureSpan(clientSpan), listener);
7576
return scope;
7677
}
7778

‎dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/ApacheHttpAsyncClientInstrumentation.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
5+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan;
66
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
77
import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
88
import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.HTTP_REQUEST;
@@ -90,13 +90,13 @@ public static AgentSpan methodEnter(
9090
@Advice.Argument(2) HttpContext context,
9191
@Advice.Argument(value = 3, readOnly = false) FutureCallback<?> futureCallback) {
9292

93-
final AgentScope parentScope = activeScope();
93+
final AgentScope.Continuation parentContinuation = captureActiveSpan();
9494
final AgentSpan clientSpan = startSpan(HTTP_REQUEST);
9595
DECORATE.afterStart(clientSpan);
9696

9797
requestProducer = new DelegatingRequestProducer(clientSpan, requestProducer);
9898
futureCallback =
99-
new TraceContinuedFutureCallback(parentScope, clientSpan, context, futureCallback);
99+
new TraceContinuedFutureCallback(parentContinuation, clientSpan, context, futureCallback);
100100

101101
return clientSpan;
102102
}

‎dd-java-agent/instrumentation/apache-httpasyncclient-4/src/main/java/datadog/trace/instrumentation/apachehttpasyncclient/TraceContinuedFutureCallback.java‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.instrumentation.apachehttpasyncclient;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled;
45
import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
56

@@ -15,15 +16,11 @@ public class TraceContinuedFutureCallback<T> implements FutureCallback<T> {
1516
private final FutureCallback<T> delegate;
1617

1718
public TraceContinuedFutureCallback(
18-
final AgentScope parentScope,
19+
final AgentScope.Continuation parentContinuation,
1920
final AgentSpan clientSpan,
2021
final HttpContext context,
2122
final FutureCallback<T> delegate) {
22-
if (parentScope != null) {
23-
parentContinuation = parentScope.capture();
24-
} else {
25-
parentContinuation = null;
26-
}
23+
this.parentContinuation = parentContinuation;
2724
this.clientSpan = clientSpan;
2825
this.context = context;
2926
// Note: this can be null in real life, so we have to handle this carefully
@@ -36,7 +33,7 @@ public void completed(final T result) {
3633
DECORATE.beforeFinish(clientSpan);
3734
clientSpan.finish(); // Finish span before calling delegate
3835

39-
if (parentContinuation == null) {
36+
if (parentContinuation == noopContinuation()) {
4037
completeDelegate(result);
4138
} else {
4239
try (final AgentScope scope = parentContinuation.activate()) {
@@ -53,7 +50,7 @@ public void failed(final Exception ex) {
5350
DECORATE.beforeFinish(clientSpan);
5451
clientSpan.finish(); // Finish span before calling delegate
5552

56-
if (parentContinuation == null) {
53+
if (parentContinuation == noopContinuation()) {
5754
failDelegate(ex);
5855
} else {
5956
try (final AgentScope scope = parentContinuation.activate()) {
@@ -69,7 +66,7 @@ public void cancelled() {
6966
DECORATE.beforeFinish(clientSpan);
7067
clientSpan.finish(); // Finish span before calling delegate
7168

72-
if (parentContinuation == null) {
69+
if (parentContinuation == noopContinuation()) {
7370
cancelDelegate();
7471
} else {
7572
try (final AgentScope scope = parentContinuation.activate()) {

‎dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/ApacheHttpAsyncClientInstrumentation.java‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface;
44
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
55
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
6-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.captureActiveSpan;
77
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
88
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
99
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.HTTP_REQUEST;
@@ -95,7 +95,7 @@ public static AgentScope methodEnter(
9595
@Advice.Argument(value = 3, readOnly = false) HttpContext context,
9696
@Advice.Argument(value = 4, readOnly = false) FutureCallback<?> futureCallback) {
9797

98-
final AgentScope parentScope = activeScope();
98+
final AgentScope.Continuation parentContinuation = captureActiveSpan();
9999
final AgentSpan clientSpan = startSpan(HTTP_REQUEST);
100100
final AgentScope clientScope = activateSpan(clientSpan);
101101
DECORATE.afterStart(clientSpan);
@@ -106,7 +106,8 @@ public static AgentScope methodEnter(
106106

107107
requestProducer = new DelegatingRequestProducer(clientSpan, requestProducer);
108108
futureCallback =
109-
new TraceContinuedFutureCallback<>(parentScope, clientSpan, context, futureCallback);
109+
new TraceContinuedFutureCallback<>(
110+
parentContinuation, clientSpan, context, futureCallback);
110111

111112
return clientScope;
112113
}

‎dd-java-agent/instrumentation/apache-httpclient-5/src/main/java/datadog/trace/instrumentation/apachehttpclient5/TraceContinuedFutureCallback.java‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package datadog.trace.instrumentation.apachehttpclient5;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.noopContinuation;
34
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled;
45
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
56

@@ -18,15 +19,11 @@ public class TraceContinuedFutureCallback<T> implements FutureCallback<T> {
1819
private final FutureCallback<T> delegate;
1920

2021
public TraceContinuedFutureCallback(
21-
final AgentScope parentScope,
22+
final AgentScope.Continuation parentContinuation,
2223
final AgentSpan clientSpan,
2324
final HttpContext context,
2425
final FutureCallback<T> delegate) {
25-
if (parentScope != null) {
26-
parentContinuation = parentScope.capture();
27-
} else {
28-
parentContinuation = null;
29-
}
26+
this.parentContinuation = parentContinuation;
3027
this.clientSpan = clientSpan;
3128
this.context = context;
3229
// Note: this can be null in real life, so we have to handle this carefully
@@ -39,7 +36,7 @@ public void completed(final T result) {
3936
DECORATE.beforeFinish(clientSpan);
4037
clientSpan.finish(); // Finish span before calling delegate
4138

42-
if (parentContinuation == null) {
39+
if (parentContinuation == noopContinuation()) {
4340
completeDelegate(result);
4441
} else {
4542
try (final AgentScope scope = parentContinuation.activate()) {
@@ -56,7 +53,7 @@ public void failed(final Exception ex) {
5653
DECORATE.beforeFinish(clientSpan);
5754
clientSpan.finish(); // Finish span before calling delegate
5855

59-
if (parentContinuation == null) {
56+
if (parentContinuation == noopContinuation()) {
6057
failDelegate(ex);
6158
} else {
6259
try (final AgentScope scope = parentContinuation.activate()) {
@@ -72,7 +69,7 @@ public void cancelled() {
7269
DECORATE.beforeFinish(clientSpan);
7370
clientSpan.finish(); // Finish span before calling delegate
7471

75-
if (parentContinuation == null) {
72+
if (parentContinuation == noopContinuation()) {
7673
cancelDelegate();
7774
} else {
7875
try (final AgentScope scope = parentContinuation.activate()) {

0 commit comments

Comments
 (0)