Skip to content

Commit d5592c6

Browse files
PerfectSlayerPerfectSlayer
authored
Move async propagation API from scope to tracer (#8231)
* feat(core): Move async propagation API from scope to tracer
1 parent 2b72cb5 commit d5592c6

File tree

47 files changed

+289
-244
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

47 files changed

+289
-244
lines changed

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

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

33
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled;
45

56
import datadog.trace.bootstrap.ContextStore;
67
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -26,7 +27,7 @@ public static AgentScope startTaskScope(State state) {
2627
final AgentScope.Continuation continuation = state.getAndResetContinuation();
2728
if (continuation != null) {
2829
final AgentScope scope = continuation.activate();
29-
scope.setAsyncPropagation(true);
30+
setAsyncPropagationEnabled(true);
3031
// important - stop timing after the scope has been activated so the time in the queue can
3132
// be attributed to the correct context without duplicating the propagated information
3233
state.stopTiming();

‎dd-java-agent/instrumentation/akka-concurrent/src/test/groovy/AkkaExecutorInstrumentationTest.groovy

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,8 @@ import java.util.concurrent.RejectedExecutionException
1616
import java.util.concurrent.ThreadPoolExecutor
1717
import java.util.concurrent.TimeUnit
1818

19-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
19+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled
20+
2021
/**
2122
* Test executor instrumentation for Akka specific classes.
2223
* This is to large extent a copy of ExecutorInstrumentationTest.
@@ -45,7 +46,7 @@ class AkkaExecutorInstrumentationTest extends AgentTestRunner {
4546
@Override
4647
@Trace(operationName = "parent")
4748
void run() {
48-
activeScope().setAsyncPropagation(true)
49+
setAsyncPropagationEnabled(true)
4950
// this child will have a span
5051
m(pool, new AkkaAsyncChild())
5152
// this child won't
@@ -101,7 +102,7 @@ class AkkaExecutorInstrumentationTest extends AgentTestRunner {
101102
@Override
102103
@Trace(operationName = "parent")
103104
void run() {
104-
activeScope().setAsyncPropagation(true)
105+
setAsyncPropagationEnabled(true)
105106
// this child will have a span
106107
dispatcher.execute(new AkkaAsyncChild())
107108
// this child won't
@@ -132,7 +133,7 @@ class AkkaExecutorInstrumentationTest extends AgentTestRunner {
132133
@Override
133134
@Trace(operationName = "parent")
134135
void run() {
135-
activeScope().setAsyncPropagation(true)
136+
setAsyncPropagationEnabled(true)
136137
try {
137138
for (int i = 0; i < 20; ++i) {
138139
// Our current instrumentation instrumentation does not behave very well

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

Lines changed: 4 additions & 3 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.setAsyncPropagationEnabled;
34
import static datadog.trace.instrumentation.apachehttpasyncclient.ApacheHttpAsyncClientDecorator.DECORATE;
45

56
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -39,7 +40,7 @@ public void completed(final T result) {
3940
completeDelegate(result);
4041
} else {
4142
try (final AgentScope scope = parentContinuation.activate()) {
42-
scope.setAsyncPropagation(true);
43+
setAsyncPropagationEnabled(true);
4344
completeDelegate(result);
4445
}
4546
}
@@ -56,7 +57,7 @@ public void failed(final Exception ex) {
5657
failDelegate(ex);
5758
} else {
5859
try (final AgentScope scope = parentContinuation.activate()) {
59-
scope.setAsyncPropagation(true);
60+
setAsyncPropagationEnabled(true);
6061
failDelegate(ex);
6162
}
6263
}
@@ -72,7 +73,7 @@ public void cancelled() {
7273
cancelDelegate();
7374
} else {
7475
try (final AgentScope scope = parentContinuation.activate()) {
75-
scope.setAsyncPropagation(true);
76+
setAsyncPropagationEnabled(true);
7677
cancelDelegate();
7778
}
7879
}

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

Lines changed: 4 additions & 3 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.setAsyncPropagationEnabled;
34
import static datadog.trace.instrumentation.apachehttpclient5.ApacheHttpClientDecorator.DECORATE;
45

56
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
@@ -42,7 +43,7 @@ public void completed(final T result) {
4243
completeDelegate(result);
4344
} else {
4445
try (final AgentScope scope = parentContinuation.activate()) {
45-
scope.setAsyncPropagation(true);
46+
setAsyncPropagationEnabled(true);
4647
completeDelegate(result);
4748
}
4849
}
@@ -59,7 +60,7 @@ public void failed(final Exception ex) {
5960
failDelegate(ex);
6061
} else {
6162
try (final AgentScope scope = parentContinuation.activate()) {
62-
scope.setAsyncPropagation(true);
63+
setAsyncPropagationEnabled(true);
6364
failDelegate(ex);
6465
}
6566
}
@@ -75,7 +76,7 @@ public void cancelled() {
7576
cancelDelegate();
7677
} else {
7778
try (final AgentScope scope = parentContinuation.activate()) {
78-
scope.setAsyncPropagation(true);
79+
setAsyncPropagationEnabled(true);
7980
cancelDelegate();
8081
}
8182
}

‎dd-java-agent/instrumentation/hystrix-1.4/src/main/java/datadog/trace/instrumentation/hystrix/HystrixThreadPoolInstrumentation.java

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

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
4+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
5+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled;
56
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
67
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
78

89
import com.google.auto.service.AutoService;
910
import datadog.trace.agent.tooling.Instrumenter;
1011
import datadog.trace.agent.tooling.InstrumenterModule;
11-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1212
import net.bytebuddy.asm.Advice;
1313

1414
@AutoService(InstrumenterModule.class)
@@ -35,23 +35,17 @@ public static class EnableAsyncAdvice {
3535

3636
@Advice.OnMethodEnter(suppress = Throwable.class)
3737
public static boolean enableAsyncTracking() {
38-
final AgentScope scope = activeScope();
39-
if (scope != null) {
40-
if (!scope.isAsyncPropagating()) {
41-
scope.setAsyncPropagation(true);
42-
return true;
43-
}
38+
if (!isAsyncPropagationEnabled()) {
39+
setAsyncPropagationEnabled(true);
40+
return true;
4441
}
4542
return false;
4643
}
4744

4845
@Advice.OnMethodExit(suppress = Throwable.class)
4946
public static void disableAsyncTracking(@Advice.Enter final boolean wasEnabled) {
5047
if (wasEnabled) {
51-
final AgentScope scope = activeScope();
52-
if (scope != null) {
53-
scope.setAsyncPropagation(false);
54-
}
48+
setAsyncPropagationEnabled(false);
5549
}
5650
}
5751
}

‎dd-java-agent/instrumentation/java-concurrent/java-completablefuture/src/test/groovy/CompletableFutureTest.groovy

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import java.util.function.Supplier
1111

1212
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
1313
import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
14-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
14+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled
1515

1616
/**
1717
* Note: ideally this should live with the rest of ExecutorInstrumentationTest,
@@ -45,7 +45,7 @@ class CompletableFutureTest extends AgentTestRunner {
4545
@Trace(operationName = "parent")
4646
CompletableFuture<String> get() {
4747
try {
48-
activeScope().setAsyncPropagation(true)
48+
setAsyncPropagationEnabled(true)
4949
return CompletableFuture.supplyAsync(supplier, pool)
5050
.thenCompose({ s -> CompletableFuture.supplyAsync(new AppendingSupplier(s), differentPool) })
5151
.thenApply(function)

‎dd-java-agent/instrumentation/java-concurrent/java-concurrent-21/src/test/groovy/VirtualThreadTest.groovy

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import java.util.concurrent.ExecutorCompletionService
88
import java.util.concurrent.Executors
99
import java.util.concurrent.TimeUnit
1010

11-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
11+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled
1212

1313
class VirtualThreadTest extends AgentTestRunner {
1414
@Shared
@@ -37,7 +37,7 @@ class VirtualThreadTest extends AgentTestRunner {
3737
@Override
3838
@Trace(operationName = "parent")
3939
void run() {
40-
activeScope().setAsyncPropagation(true)
40+
setAsyncPropagationEnabled(true)
4141
// this child will have a span
4242
m(pool, new JavaAsyncChild())
4343
// this child won't

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

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,15 @@
55
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.nameStartsWith;
66
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
77
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.namedOneOf;
8-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope;
8+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.isAsyncPropagationEnabled;
9+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled;
910
import static datadog.trace.instrumentation.java.concurrent.ConcurrentInstrumentationNames.EXECUTOR_INSTRUMENTATION_NAME;
1011
import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
1112
import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
1213

1314
import com.google.auto.service.AutoService;
1415
import datadog.trace.agent.tooling.Instrumenter;
1516
import datadog.trace.agent.tooling.InstrumenterModule;
16-
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
1717
import net.bytebuddy.asm.Advice;
1818
import net.bytebuddy.description.type.TypeDescription;
1919
import net.bytebuddy.matcher.ElementMatcher;
@@ -177,19 +177,18 @@ public void methodAdvice(MethodTransformer transformer) {
177177
public static class DisableAsyncAdvice {
178178

179179
@Advice.OnMethodEnter
180-
public static AgentScope before() {
181-
AgentScope scope = activeScope();
182-
if (null != scope && scope.isAsyncPropagating()) {
183-
scope.setAsyncPropagation(false);
184-
return scope;
180+
public static boolean before() {
181+
if (isAsyncPropagationEnabled()) {
182+
setAsyncPropagationEnabled(false);
183+
return true;
185184
}
186-
return null;
185+
return false;
187186
}
188187

189188
@Advice.OnMethodExit(onThrowable = Throwable.class)
190-
public static void after(@Advice.Enter AgentScope scope) {
191-
if (null != scope) {
192-
scope.setAsyncPropagation(true);
189+
public static void after(@Advice.Enter boolean wasDisabled) {
190+
if (wasDisabled) {
191+
setAsyncPropagationEnabled(true);
193192
}
194193
}
195194
}

‎dd-java-agent/instrumentation/java-concurrent/src/test/groovy/executor/ExecutorInstrumentationTest.groovy

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ import java.util.concurrent.ScheduledThreadPoolExecutor
2626
import java.util.concurrent.ThreadPoolExecutor
2727
import java.util.concurrent.TimeUnit
2828

29-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
29+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled
3030
import static org.junit.Assume.assumeTrue
3131

3232
abstract class ExecutorInstrumentationTest extends AgentTestRunner {
@@ -82,7 +82,7 @@ abstract class ExecutorInstrumentationTest extends AgentTestRunner {
8282
@Override
8383
@Trace(operationName = "parent")
8484
void run() {
85-
activeScope().setAsyncPropagation(true)
85+
setAsyncPropagationEnabled(true)
8686
// this child will have a span
8787
m(pool, new JavaAsyncChild())
8888
// this child won't
@@ -256,7 +256,7 @@ abstract class ExecutorInstrumentationTest extends AgentTestRunner {
256256
@Override
257257
@Trace(operationName = "parent")
258258
void run() {
259-
activeScope().setAsyncPropagation(true)
259+
setAsyncPropagationEnabled(true)
260260
def future = m(pool, task)
261261
sleep(500)
262262
future.cancel(true)
@@ -317,7 +317,7 @@ abstract class ExecutorInstrumentationTest extends AgentTestRunner {
317317
@Override
318318
@Trace(operationName = "parent")
319319
void run() {
320-
activeScope().setAsyncPropagation(true)
320+
setAsyncPropagationEnabled(true)
321321
// this child will have a span
322322
pool.execute(new JavaAsyncChild())
323323
// this child won't
@@ -368,7 +368,7 @@ abstract class ExecutorInstrumentationTest extends AgentTestRunner {
368368
@Override
369369
@Trace(operationName = "parent")
370370
void run() {
371-
activeScope().setAsyncPropagation(true)
371+
setAsyncPropagationEnabled(true)
372372
m(pool, w(child))
373373
}
374374
}.run()
@@ -406,7 +406,7 @@ abstract class ExecutorInstrumentationTest extends AgentTestRunner {
406406
@Override
407407
@Trace(operationName = "parent")
408408
void run() {
409-
activeScope().setAsyncPropagation(true)
409+
setAsyncPropagationEnabled(true)
410410
try {
411411
for (int i = 0; i < 20; ++i) {
412412
final JavaAsyncChild child = new JavaAsyncChild(false, true)

‎dd-java-agent/instrumentation/java-concurrent/src/test/groovy/executor/NettyExecutorInstrumentationTest.groovy

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ import java.util.concurrent.Future
1717
import java.util.concurrent.RejectedExecutionException
1818
import java.util.concurrent.TimeUnit
1919

20-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
20+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled
2121
import static org.junit.Assume.assumeTrue
2222

2323
class NettyExecutorInstrumentationTest extends AgentTestRunner {
@@ -65,7 +65,7 @@ class NettyExecutorInstrumentationTest extends AgentTestRunner {
6565
@Override
6666
@Trace(operationName = "parent")
6767
void run() {
68-
activeScope().setAsyncPropagation(true)
68+
setAsyncPropagationEnabled(true)
6969
// this child will have a span
7070
m(pool, new JavaAsyncChild())
7171
// this child won't
@@ -213,7 +213,7 @@ class NettyExecutorInstrumentationTest extends AgentTestRunner {
213213
@Override
214214
@Trace(operationName = "parent")
215215
void run() {
216-
activeScope().setAsyncPropagation(true)
216+
setAsyncPropagationEnabled(true)
217217
try {
218218
for (int i = 0; i < 20; ++i) {
219219
final JavaAsyncChild child = new JavaAsyncChild(false, true)

‎dd-java-agent/instrumentation/java-concurrent/src/test/groovy/executor/RejectedExecutionTest.groovy

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import java.util.concurrent.ThreadPoolExecutor
2121
import java.util.concurrent.TimeUnit
2222

2323
import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
24-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
24+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled
2525

2626
class RejectedExecutionTest extends AgentTestRunner {
2727

@@ -193,7 +193,7 @@ class RejectedExecutionTest extends AgentTestRunner {
193193

194194
when:
195195
runUnderTrace("parent") {
196-
activeScope().setAsyncPropagation(true)
196+
setAsyncPropagationEnabled(true)
197197
// must be rejected because the queue will be full until some
198198
// time after the first task is released
199199
executor.submit((Runnable) new JavaAsyncChild(true, false))
@@ -237,7 +237,7 @@ class RejectedExecutionTest extends AgentTestRunner {
237237

238238
return {
239239
runUnderTrace("parent") {
240-
activeScope().setAsyncPropagation(true)
240+
setAsyncPropagationEnabled(true)
241241
pool.submit({})
242242
}
243243
}
@@ -261,7 +261,7 @@ class RejectedExecutionTest extends AgentTestRunner {
261261

262262
return {
263263
runUnderTrace("parent") {
264-
activeScope().setAsyncPropagation(true)
264+
setAsyncPropagationEnabled(true)
265265
// must be rejected because the queue will be full until some
266266
// time after the first task is released
267267
def testTask = new JavaAsyncChild(true, false)

‎dd-java-agent/instrumentation/jetty-util/src/test/groovy/JettyExecutorInstrumentationTest.groovy

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import java.util.concurrent.Callable
1010
import java.util.concurrent.ExecutorService
1111
import java.util.concurrent.Executors
1212

13-
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activeScope
13+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.setAsyncPropagationEnabled
1414
import static org.junit.Assume.assumeTrue
1515

1616
class JettyExecutorInstrumentationTest extends AgentTestRunner {
@@ -46,7 +46,7 @@ class JettyExecutorInstrumentationTest extends AgentTestRunner {
4646
@Override
4747
@Trace(operationName = "parent")
4848
void run() {
49-
activeScope().setAsyncPropagation(true)
49+
setAsyncPropagationEnabled(true)
5050
// this child will have a span
5151
m(pool, new JavaAsyncChild())
5252
// this child won't

0 commit comments

Comments
 (0)