Skip to content

Commit 1d0b64d

Browse files
amarzialidevflow.devflow-routing-intake
amarziali
and
devflow.devflow-routing-intake
authored
Move scala concurrent and promise advices to context tracking (#10537)
Use Context in ContextStore instead of Agentspan for scala promise Move to single instrumenter module and flag for context tracking Move scala concurrent advices to context tracking wip fixes fixes spotless Update dd-java-agent/instrumentation/scala/scala-promise/scala-promise-2.13/src/main/java/datadog/trace/instrumentation/scala213/concurrent/PromiseObjectInstrumentation.java Co-authored-by: Stuart McCulloch <stuart.mcculloch@datadoghq.com> suggestions spotless realign promise helper Co-authored-by: devflow.devflow-routing-intake <devflow.devflow-routing-intake@kubernetes.us1.ddbuild.io>
1 parent 307c79a commit 1d0b64d

File tree

13 files changed

+236
-272
lines changed

13 files changed

+236
-272
lines changed

‎dd-java-agent/instrumentation/scala/scala-concurrent-2.8/src/main/java/datadog/trace/instrumentation/scala/concurrent/ScalaConcurrentModule.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
package datadog.trace.instrumentation.scala.concurrent;
2+
3+
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE_FUTURE;
4+
import static java.util.Collections.singletonMap;
5+
6+
import com.google.auto.service.AutoService;
7+
import datadog.trace.agent.tooling.ExcludeFilterProvider;
8+
import datadog.trace.agent.tooling.Instrumenter;
9+
import datadog.trace.agent.tooling.InstrumenterModule;
10+
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
11+
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
12+
import java.util.Arrays;
13+
import java.util.Collection;
14+
import java.util.List;
15+
import java.util.Map;
16+
17+
@AutoService(InstrumenterModule.class)
18+
public final class ScalaConcurrentModule extends InstrumenterModule.ContextTracking
19+
implements ExcludeFilterProvider {
20+
21+
public ScalaConcurrentModule() {
22+
super("java_concurrent", "scala_concurrent");
23+
}
24+
25+
@Override
26+
public Map<String, String> contextStore() {
27+
return singletonMap("scala.concurrent.forkjoin.ForkJoinTask", State.class.getName());
28+
}
29+
30+
@Override
31+
public Map<ExcludeFilter.ExcludeType, ? extends Collection<String>> excludedClasses() {
32+
return singletonMap(
33+
RUNNABLE_FUTURE,
34+
Arrays.asList(
35+
"scala.concurrent.forkjoin.ForkJoinTask$AdaptedCallable",
36+
"scala.concurrent.forkjoin.ForkJoinTask$AdaptedRunnable",
37+
"scala.concurrent.forkjoin.ForkJoinTask$AdaptedRunnableAction"));
38+
}
39+
40+
@Override
41+
public List<Instrumenter> typeInstrumentations() {
42+
return Arrays.asList(
43+
new ScalaForkJoinTaskInstrumentation(), new ScalaForkJoinPoolInstrumentation());
44+
}
45+
}

‎dd-java-agent/instrumentation/scala/scala-concurrent-2.8/src/main/java/datadog/trace/instrumentation/scala/concurrent/ScalaForkJoinPoolInstrumentation.java‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -6,37 +6,23 @@
66
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
77
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.FORK_JOIN_TASK;
88
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;
9-
import static java.util.Collections.singletonMap;
109
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1110
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
1211

13-
import com.google.auto.service.AutoService;
1412
import datadog.trace.agent.tooling.Instrumenter;
15-
import datadog.trace.agent.tooling.InstrumenterModule;
1613
import datadog.trace.bootstrap.InstrumentationContext;
1714
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
18-
import java.util.Map;
1915
import net.bytebuddy.asm.Advice;
2016
import scala.concurrent.forkjoin.ForkJoinTask;
2117

22-
@AutoService(InstrumenterModule.class)
23-
public final class ScalaForkJoinPoolInstrumentation extends InstrumenterModule.Tracing
18+
public final class ScalaForkJoinPoolInstrumentation
2419
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
2520

26-
public ScalaForkJoinPoolInstrumentation() {
27-
super("java_concurrent", "scala_concurrent");
28-
}
29-
3021
@Override
3122
public String instrumentedType() {
3223
return "scala.concurrent.forkjoin.ForkJoinPool";
3324
}
3425

35-
@Override
36-
public Map<String, String> contextStore() {
37-
return singletonMap("scala.concurrent.forkjoin.ForkJoinTask", State.class.getName());
38-
}
39-
4026
@Override
4127
public void methodAdvice(MethodTransformer transformer) {
4228
transformer.applyAdvice(

‎dd-java-agent/instrumentation/scala/scala-concurrent-2.8/src/main/java/datadog/trace/instrumentation/scala/concurrent/ScalaForkJoinTaskInstrumentation.java‎

Lines changed: 2 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -10,22 +10,13 @@
1010
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope;
1111
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.startTaskScope;
1212
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.FORK_JOIN_TASK;
13-
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE_FUTURE;
1413
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.exclude;
15-
import static java.util.Collections.singletonMap;
1614
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
1715

18-
import com.google.auto.service.AutoService;
19-
import datadog.trace.agent.tooling.ExcludeFilterProvider;
2016
import datadog.trace.agent.tooling.Instrumenter;
21-
import datadog.trace.agent.tooling.InstrumenterModule;
2217
import datadog.trace.bootstrap.InstrumentationContext;
2318
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
24-
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
2519
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
26-
import java.util.Arrays;
27-
import java.util.Collection;
28-
import java.util.Map;
2920
import net.bytebuddy.asm.Advice;
3021
import net.bytebuddy.description.type.TypeDescription;
3122
import net.bytebuddy.matcher.ElementMatcher;
@@ -37,13 +28,8 @@
3728
* <p>Note: There are quite a few separate implementations of {@code ForkJoinTask}/{@code
3829
* ForkJoinPool}: JVM, Akka, Scala, Netty to name a few. This class handles Scala version.
3930
*/
40-
@AutoService(InstrumenterModule.class)
41-
public final class ScalaForkJoinTaskInstrumentation extends InstrumenterModule.Tracing
42-
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice, ExcludeFilterProvider {
43-
44-
public ScalaForkJoinTaskInstrumentation() {
45-
super("java_concurrent", "scala_concurrent");
46-
}
31+
public final class ScalaForkJoinTaskInstrumentation
32+
implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice {
4733

4834
@Override
4935
public String hierarchyMarkerType() {
@@ -59,11 +45,6 @@ public ElementMatcher<TypeDescription> hierarchyMatcher() {
5945
.and(extendsClass(named(hierarchyMarkerType())));
6046
}
6147

62-
@Override
63-
public Map<String, String> contextStore() {
64-
return singletonMap("scala.concurrent.forkjoin.ForkJoinTask", State.class.getName());
65-
}
66-
6748
@Override
6849
public void methodAdvice(MethodTransformer transformer) {
6950
transformer.applyAdvice(
@@ -72,16 +53,6 @@ public void methodAdvice(MethodTransformer transformer) {
7253
transformer.applyAdvice(isMethod().and(named("cancel")), getClass().getName() + "$Cancel");
7354
}
7455

75-
@Override
76-
public Map<ExcludeFilter.ExcludeType, ? extends Collection<String>> excludedClasses() {
77-
return singletonMap(
78-
RUNNABLE_FUTURE,
79-
Arrays.asList(
80-
"scala.concurrent.forkjoin.ForkJoinTask$AdaptedCallable",
81-
"scala.concurrent.forkjoin.ForkJoinTask$AdaptedRunnable",
82-
"scala.concurrent.forkjoin.ForkJoinTask$AdaptedRunnableAction"));
83-
}
84-
8556
public static final class Exec {
8657
@Advice.OnMethodEnter
8758
public static <T> AgentScope before(@Advice.This ForkJoinTask<T> task) {

‎dd-java-agent/instrumentation/scala/scala-promise/scala-promise-2.10/src/main/java/datadog/trace/instrumentation/scala210/concurrent/CallbackRunnableInstrumentation.java‎

Lines changed: 4 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -3,51 +3,28 @@
33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
44
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.capture;
55
import static datadog.trace.bootstrap.instrumentation.java.concurrent.AdviceUtils.endTaskScope;
6-
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.EXECUTOR;
7-
import static datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter.ExcludeType.RUNNABLE;
86
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
97
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
108

11-
import com.google.auto.service.AutoService;
12-
import datadog.trace.agent.tooling.ExcludeFilterProvider;
9+
import datadog.context.Context;
1310
import datadog.trace.agent.tooling.Instrumenter;
14-
import datadog.trace.agent.tooling.InstrumenterModule;
1511
import datadog.trace.bootstrap.ContextStore;
1612
import datadog.trace.bootstrap.InstrumentationContext;
1713
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
18-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
19-
import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
2014
import datadog.trace.bootstrap.instrumentation.java.concurrent.State;
2115
import datadog.trace.instrumentation.scala.PromiseHelper;
22-
import java.util.Collection;
23-
import java.util.Collections;
24-
import java.util.HashMap;
25-
import java.util.Map;
2616
import net.bytebuddy.asm.Advice;
2717
import scala.concurrent.impl.CallbackRunnable;
2818
import scala.util.Try;
2919

30-
@AutoService(InstrumenterModule.class)
31-
public class CallbackRunnableInstrumentation extends InstrumenterModule.Tracing
32-
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice, ExcludeFilterProvider {
33-
34-
public CallbackRunnableInstrumentation() {
35-
super("scala_concurrent");
36-
}
20+
public final class CallbackRunnableInstrumentation
21+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
3722

3823
@Override
3924
public String instrumentedType() {
4025
return "scala.concurrent.impl.CallbackRunnable";
4126
}
4227

43-
@Override
44-
public Map<String, String> contextStore() {
45-
Map<String, String> contextStore = new HashMap<>();
46-
contextStore.put("scala.concurrent.impl.CallbackRunnable", State.class.getName());
47-
contextStore.put("scala.util.Try", AgentSpan.class.getName());
48-
return contextStore;
49-
}
50-
5128
@Override
5229
public void methodAdvice(MethodTransformer transformer) {
5330
transformer.applyAdvice(isConstructor(), getClass().getName() + "$Construct");
@@ -56,21 +33,6 @@ public void methodAdvice(MethodTransformer transformer) {
5633
isMethod().and(named("executeWithValue")), getClass().getName() + "$ExecuteWithValue");
5734
}
5835

59-
@Override
60-
public Map<ExcludeFilter.ExcludeType, ? extends Collection<String>> excludedClasses() {
61-
// force other instrumentations (e.g. Runnable) not to deal with this type
62-
Map<ExcludeFilter.ExcludeType, Collection<String>> map = new HashMap<>();
63-
Collection<String> cbr = Collections.singleton("scala.concurrent.impl.CallbackRunnable");
64-
map.put(RUNNABLE, cbr);
65-
map.put(EXECUTOR, cbr);
66-
return map;
67-
}
68-
69-
@Override
70-
public String[] helperClassNames() {
71-
return new String[] {"datadog.trace.instrumentation.scala.PromiseHelper"};
72-
}
73-
7436
/** Capture the scope when the promise is created */
7537
public static final class Construct {
7638
@Advice.OnMethodExit(suppress = Throwable.class)
@@ -103,7 +65,7 @@ public static <T> void beforeExecute(
10365
if (PromiseHelper.completionPriority) {
10466
state =
10567
PromiseHelper.executeCaptureSpan(
106-
InstrumentationContext.get(Try.class, AgentSpan.class),
68+
InstrumentationContext.get(Try.class, Context.class),
10769
resolved,
10870
contextStore,
10971
task,

‎dd-java-agent/instrumentation/scala/scala-promise/scala-promise-2.10/src/main/java/datadog/trace/instrumentation/scala210/concurrent/FutureObjectInstrumentation.java‎

Lines changed: 3 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,13 @@
11
package datadog.trace.instrumentation.scala210.concurrent;
22

3-
import static java.util.Collections.singletonMap;
43
import static net.bytebuddy.matcher.ElementMatchers.isTypeInitializer;
54

6-
import com.google.auto.service.AutoService;
5+
import datadog.context.Context;
76
import datadog.trace.agent.tooling.Instrumenter;
8-
import datadog.trace.agent.tooling.InstrumenterModule;
97
import datadog.trace.bootstrap.InstrumentationContext;
10-
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
118
import java.lang.invoke.MethodHandle;
129
import java.lang.invoke.MethodHandles;
1310
import java.lang.invoke.MethodType;
14-
import java.util.Map;
1511
import net.bytebuddy.asm.Advice;
1612
import scala.concurrent.Future;
1713
import scala.concurrent.Future$;
@@ -24,25 +20,15 @@
2420
* every {@code Future} created via one of the static methods like {@code map} would always pick up
2521
* that context and propagate it forward, which is quite unexpected and not very relevant.
2622
*/
27-
@AutoService(InstrumenterModule.class)
28-
public class FutureObjectInstrumentation extends InstrumenterModule.Tracing
23+
public final class FutureObjectInstrumentation
2924
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
3025

31-
public FutureObjectInstrumentation() {
32-
super("scala_future_object", "scala_concurrent");
33-
}
34-
3526
@Override
3627
public String instrumentedType() {
3728
// The $ at the end is how Scala encodes a Scala object (as opposed to a class or trait)
3829
return "scala.concurrent.Future$";
3930
}
4031

41-
@Override
42-
public Map<String, String> contextStore() {
43-
return singletonMap("scala.util.Try", AgentSpan.class.getName());
44-
}
45-
4632
@Override
4733
public void methodAdvice(MethodTransformer transformer) {
4834
transformer.applyAdvice(isTypeInitializer(), getClass().getName() + "$ClassInit");
@@ -58,7 +44,7 @@ public static <T> void afterInit() {
5844
lookup.findVirtual(Future$.class, "unit", MethodType.methodType(Future.class));
5945
Future<?> unit = (Future<?>) mh.invoke(Future$.MODULE$);
6046
Try<?> result = unit.value().get();
61-
InstrumentationContext.get(Try.class, AgentSpan.class).put(result, null);
47+
InstrumentationContext.get(Try.class, Context.class).put(result, null);
6248
} catch (Throwable ignored) {
6349
}
6450
}

‎dd-java-agent/instrumentation/scala/scala-promise/scala-promise-2.10/src/main/java/datadog/trace/instrumentation/scala210/concurrent/PromiseObjectInstrumentation.java‎

Lines changed: 8 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,15 @@
11
package datadog.trace.instrumentation.scala210.concurrent;
22

33
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4-
import static java.util.Collections.singletonMap;
54
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
65

7-
import com.google.auto.service.AutoService;
6+
import datadog.context.Context;
87
import datadog.trace.agent.tooling.Instrumenter;
9-
import datadog.trace.agent.tooling.InstrumenterModule;
10-
import datadog.trace.api.InstrumenterConfig;
118
import datadog.trace.bootstrap.ContextStore;
129
import datadog.trace.bootstrap.InstrumentationContext;
1310
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
11+
import datadog.trace.bootstrap.instrumentation.api.Java8BytecodeBridge;
1412
import datadog.trace.instrumentation.scala.PromiseHelper;
15-
import java.util.Collections;
16-
import java.util.Map;
1713
import net.bytebuddy.asm.Advice;
1814
import scala.concurrent.impl.CallbackRunnable;
1915
import scala.util.Try;
@@ -23,55 +19,32 @@
2319
* to take priority over any spans captured while adding computations to a {@code Future} associated
2420
* with a {@code Promise}, then we capture the active span when the {@code Try} is resolved.
2521
*/
26-
@AutoService(InstrumenterModule.class)
27-
public class PromiseObjectInstrumentation extends InstrumenterModule.Tracing
22+
public final class PromiseObjectInstrumentation
2823
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
2924

30-
public PromiseObjectInstrumentation() {
31-
super("scala_promise_resolve", "scala_concurrent");
32-
}
33-
3425
@Override
3526
public String instrumentedType() {
3627
// The $ at the end is how Scala encodes a Scala object (as opposed to a class or trait)
3728
return "scala.concurrent.impl.Promise$";
3829
}
3930

40-
@Override
41-
public Map<String, String> contextStore() {
42-
return singletonMap("scala.util.Try", AgentSpan.class.getName());
43-
}
44-
4531
@Override
4632
public void methodAdvice(MethodTransformer transformer) {
4733
transformer.applyAdvice(
4834
isMethod().and(named("scala$concurrent$impl$Promise$$resolveTry")),
4935
getClass().getName() + "$ResolveTry");
5036
}
5137

52-
@Override
53-
public String[] helperClassNames() {
54-
return new String[] {"datadog.trace.instrumentation.scala.PromiseHelper"};
55-
}
56-
57-
@Override
58-
public boolean isEnabled() {
59-
// Only enable this if integrations have been enabled and the extra "integration"
60-
// scala_promise_completion_priority has been enabled specifically
61-
return super.isEnabled()
62-
&& InstrumenterConfig.get()
63-
.isIntegrationEnabled(
64-
Collections.singletonList("scala_promise_completion_priority"), false);
65-
}
66-
6738
public static final class ResolveTry {
6839
@Advice.OnMethodExit(suppress = Throwable.class)
6940
public static <T> void afterResolve(@Advice.Return(readOnly = false) Try<T> resolved) {
7041
AgentSpan span = PromiseHelper.getSpan();
7142
if (null != span) {
72-
ContextStore<Try, AgentSpan> contextStore =
73-
InstrumentationContext.get(Try.class, AgentSpan.class);
74-
Try<T> next = PromiseHelper.getTry(resolved, span, contextStore.get(resolved));
43+
ContextStore<Try, Context> contextStore =
44+
InstrumentationContext.get(Try.class, Context.class);
45+
final Context existing = contextStore.get(resolved);
46+
Try<T> next =
47+
PromiseHelper.getTry(resolved, span, Java8BytecodeBridge.spanFromContext(existing));
7548
if (next != resolved) {
7649
contextStore.put(next, span);
7750
resolved = next;

0 commit comments

Comments
 (0)