Skip to content

Commit 207f770

Browse files
mccullsmcculls
authored
Fix native-image generation of reactive applications (#8012)
* Refactor registration of AsyncResult extensions so it doesn't pull in AsyncResultDecorator * Rebuild AsyncResultExtensions at native-image runtime
1 parent f7a6771 commit 207f770

File tree

17 files changed

+152
-136
lines changed

17 files changed

+152
-136
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/AsyncResultDecorator.java‎

Lines changed: 9 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -1,47 +1,27 @@
11
package datadog.trace.bootstrap.instrumentation.decorator;
22

3-
import static java.util.Collections.singletonList;
4-
53
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
6-
import java.util.concurrent.CompletableFuture;
7-
import java.util.concurrent.CompletionException;
8-
import java.util.concurrent.CompletionStage;
9-
import java.util.concurrent.CopyOnWriteArrayList;
10-
import java.util.concurrent.ExecutionException;
11-
import java.util.function.BiConsumer;
4+
import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtension;
5+
import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions;
126

137
/**
148
* This decorator handles asynchronous result types, finishing spans only when the async calls are
15-
* complete. The different async types are supported using {@link AsyncResultSupportExtension} that
16-
* should be registered using {@link #registerExtension(AsyncResultSupportExtension)} first.
9+
* complete. The different async types are supported using {@link AsyncResultExtension} that should
10+
* be registered using {@link AsyncResultExtensions#register(AsyncResultExtension)} first.
1711
*/
1812
public abstract class AsyncResultDecorator extends BaseDecorator {
19-
private static final CopyOnWriteArrayList<AsyncResultSupportExtension> EXTENSIONS =
20-
new CopyOnWriteArrayList<>(
21-
singletonList(new JavaUtilConcurrentAsyncResultSupportExtension()));
2213

23-
private static final ClassValue<AsyncResultSupportExtension> EXTENSION_CLASS_VALUE =
24-
new ClassValue<AsyncResultSupportExtension>() {
14+
private static final ClassValue<AsyncResultExtension> EXTENSION_CLASS_VALUE =
15+
new ClassValue<AsyncResultExtension>() {
2516
@Override
26-
protected AsyncResultSupportExtension computeValue(Class<?> type) {
27-
return EXTENSIONS.stream()
17+
protected AsyncResultExtension computeValue(Class<?> type) {
18+
return AsyncResultExtensions.registered().stream()
2819
.filter(extension -> extension.supports(type))
2920
.findFirst()
3021
.orElse(null);
3122
}
3223
};
3324

34-
/**
35-
* Registers an extension to add supported async types.
36-
*
37-
* @param extension The extension to register.
38-
*/
39-
public static void registerExtension(AsyncResultSupportExtension extension) {
40-
if (extension != null) {
41-
EXTENSIONS.add(extension);
42-
}
43-
}
44-
4525
/**
4626
* Look for asynchronous result and decorate it with span finisher. If the result is not
4727
* asynchronous, it will be return unmodified and span will be finished.
@@ -53,7 +33,7 @@ public static void registerExtension(AsyncResultSupportExtension extension) {
5333
*/
5434
public Object wrapAsyncResultOrFinishSpan(
5535
final Object result, final Class<?> methodReturnType, final AgentSpan span) {
56-
AsyncResultSupportExtension extension;
36+
AsyncResultExtension extension;
5737
if (result != null && (extension = EXTENSION_CLASS_VALUE.get(methodReturnType)) != null) {
5838
Object applied = extension.apply(result, span);
5939
if (applied != null) {
@@ -64,63 +44,4 @@ public Object wrapAsyncResultOrFinishSpan(
6444
span.finish();
6545
return result;
6646
}
67-
68-
/**
69-
* This interface defines asynchronous result type support extension. It allows deferring the
70-
* support implementations where types are available on classpath.
71-
*/
72-
public interface AsyncResultSupportExtension {
73-
/**
74-
* Checks whether this extensions support a result type.
75-
*
76-
* @param result The result type to check.
77-
* @return {@code true} if the type is supported by this extension, {@code false} otherwise.
78-
*/
79-
boolean supports(Class<?> result);
80-
81-
/**
82-
* Applies the extension to the async result.
83-
*
84-
* @param result The async result.
85-
* @param span The related span.
86-
* @return The result object to return (can be the original result if not modified), or {@code
87-
* null} if the extension could not be applied.
88-
*/
89-
Object apply(Object result, AgentSpan span);
90-
}
91-
92-
private static class JavaUtilConcurrentAsyncResultSupportExtension
93-
implements AsyncResultSupportExtension {
94-
@Override
95-
public boolean supports(Class<?> result) {
96-
return CompletableFuture.class.isAssignableFrom(result)
97-
|| CompletionStage.class.isAssignableFrom(result);
98-
}
99-
100-
@Override
101-
public Object apply(Object result, AgentSpan span) {
102-
if (result instanceof CompletableFuture<?>) {
103-
CompletableFuture<?> completableFuture = (CompletableFuture<?>) result;
104-
if (!completableFuture.isDone() && !completableFuture.isCancelled()) {
105-
return completableFuture.whenComplete(finishSpan(span));
106-
}
107-
} else if (result instanceof CompletionStage<?>) {
108-
CompletionStage<?> completionStage = (CompletionStage<?>) result;
109-
return completionStage.whenComplete(finishSpan(span));
110-
}
111-
return null;
112-
}
113-
114-
private <T> BiConsumer<T, Throwable> finishSpan(AgentSpan span) {
115-
return (o, throwable) -> {
116-
if (throwable != null) {
117-
span.addThrowable(
118-
throwable instanceof ExecutionException || throwable instanceof CompletionException
119-
? throwable.getCause()
120-
: throwable);
121-
}
122-
span.finish();
123-
};
124-
}
125-
}
12647
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package datadog.trace.bootstrap.instrumentation.java.concurrent;
2+
3+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
4+
5+
/**
6+
* This interface defines asynchronous result type support extension. It allows deferring the
7+
* support implementations where types are available on classpath.
8+
*/
9+
public interface AsyncResultExtension {
10+
/**
11+
* Checks whether this extensions support a result type.
12+
*
13+
* @param result The result type to check.
14+
* @return {@code true} if the type is supported by this extension, {@code false} otherwise.
15+
*/
16+
boolean supports(Class<?> result);
17+
18+
/**
19+
* Applies the extension to the async result.
20+
*
21+
* @param result The async result.
22+
* @param span The related span.
23+
* @return The result object to return (can be the original result if not modified), or {@code
24+
* null} if the extension could not be applied.
25+
*/
26+
Object apply(Object result, AgentSpan span);
27+
}

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

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
package datadog.trace.bootstrap.instrumentation.java.concurrent;
2+
3+
import static java.util.Collections.singletonList;
4+
5+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
6+
import java.util.List;
7+
import java.util.concurrent.CompletableFuture;
8+
import java.util.concurrent.CompletionException;
9+
import java.util.concurrent.CompletionStage;
10+
import java.util.concurrent.CopyOnWriteArrayList;
11+
import java.util.concurrent.ExecutionException;
12+
import java.util.function.BiConsumer;
13+
14+
public final class AsyncResultExtensions {
15+
private static final List<AsyncResultExtension> EXTENSIONS =
16+
new CopyOnWriteArrayList<>(singletonList(new CompletableAsyncResultExtension()));
17+
18+
/**
19+
* Registers an extension to add supported async types.
20+
*
21+
* @param extension The extension to register.
22+
*/
23+
public static void register(AsyncResultExtension extension) {
24+
if (extension != null) {
25+
EXTENSIONS.add(extension);
26+
}
27+
}
28+
29+
/** Returns the list of currently registered extensions. */
30+
public static List<AsyncResultExtension> registered() {
31+
return EXTENSIONS;
32+
}
33+
34+
static final class CompletableAsyncResultExtension implements AsyncResultExtension {
35+
@Override
36+
public boolean supports(Class<?> result) {
37+
return CompletableFuture.class.isAssignableFrom(result)
38+
|| CompletionStage.class.isAssignableFrom(result);
39+
}
40+
41+
@Override
42+
public Object apply(Object result, AgentSpan span) {
43+
if (result instanceof CompletableFuture<?>) {
44+
CompletableFuture<?> completableFuture = (CompletableFuture<?>) result;
45+
if (!completableFuture.isDone() && !completableFuture.isCancelled()) {
46+
return completableFuture.whenComplete(finishSpan(span));
47+
}
48+
} else if (result instanceof CompletionStage<?>) {
49+
CompletionStage<?> completionStage = (CompletionStage<?>) result;
50+
return completionStage.whenComplete(finishSpan(span));
51+
}
52+
return null;
53+
}
54+
55+
private <T> BiConsumer<T, Throwable> finishSpan(AgentSpan span) {
56+
return (o, throwable) -> {
57+
if (throwable != null) {
58+
span.addThrowable(
59+
throwable instanceof ExecutionException || throwable instanceof CompletionException
60+
? throwable.getCause()
61+
: throwable);
62+
}
63+
span.finish();
64+
};
65+
}
66+
}
67+
}

‎dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/NativeImageGeneratorRunnerInstrumentation.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ public static void onEnter(@Advice.Argument(value = 0, readOnly = false) String[
9696
+ "datadog.trace.bootstrap.benchmark.StaticEventLogger:build_time,"
9797
+ "datadog.trace.bootstrap.blocking.BlockingExceptionHandler:build_time,"
9898
+ "datadog.trace.bootstrap.InstrumentationErrors:build_time,"
99+
+ "datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions:rerun,"
99100
+ "datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState:build_time,"
100101
+ "datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter:build_time,"
101102
+ "datadog.trace.bootstrap.instrumentation.java.concurrent.QueueTimeHelper:build_time,"

‎dd-java-agent/instrumentation/guava-10/src/main/java/datadog/trace/instrumentation/guava10/GuavaAsyncResultSupportExtension.java‎ renamed to ‎dd-java-agent/instrumentation/guava-10/src/main/java/datadog/trace/instrumentation/guava10/GuavaAsyncResultExtension.java‎

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

33
import com.google.common.util.concurrent.ListenableFuture;
44
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
5-
import datadog.trace.bootstrap.instrumentation.decorator.AsyncResultDecorator;
5+
import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtension;
6+
import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions;
67
import java.util.concurrent.CancellationException;
78
import java.util.concurrent.ExecutionException;
89

9-
public class GuavaAsyncResultSupportExtension
10-
implements AsyncResultDecorator.AsyncResultSupportExtension {
10+
public class GuavaAsyncResultExtension implements AsyncResultExtension {
1111
static {
12-
AsyncResultDecorator.registerExtension(new GuavaAsyncResultSupportExtension());
12+
AsyncResultExtensions.register(new GuavaAsyncResultExtension());
1313
}
1414

1515
/**
16-
* Register the extension as an {@link AsyncResultDecorator.AsyncResultSupportExtension} using
17-
* static class initialization.<br>
16+
* Register the extension as an {@link AsyncResultExtension} using static class initialization.
17+
* <br>
1818
* It uses an empty static method call to ensure the class loading and the one-time-only static
19-
* class initialization. This will ensure this extension will only be registered once to the
20-
* {@link AsyncResultDecorator}.
19+
* class initialization. This will ensure this extension will only be registered once under {@link
20+
* AsyncResultExtensions}.
2121
*/
2222
public static void initialize() {}
2323

‎dd-java-agent/instrumentation/guava-10/src/main/java/datadog/trace/instrumentation/guava10/ListenableFutureInstrumentation.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ public String instrumentedType() {
3636
@Override
3737
public String[] helperClassNames() {
3838
return new String[] {
39-
this.packageName + ".GuavaAsyncResultSupportExtension",
39+
this.packageName + ".GuavaAsyncResultExtension",
4040
};
4141
}
4242

@@ -57,7 +57,7 @@ public void methodAdvice(MethodTransformer transformer) {
5757
public static class AbstractFutureAdvice {
5858
@Advice.OnMethodExit(suppress = Throwable.class)
5959
public static void init() {
60-
GuavaAsyncResultSupportExtension.initialize();
60+
GuavaAsyncResultExtension.initialize();
6161
}
6262
}
6363

‎dd-java-agent/instrumentation/guava-10/src/test/groovy/GuavaAsyncResultSupportExtensionTest.groovy‎ renamed to ‎dd-java-agent/instrumentation/guava-10/src/test/groovy/GuavaAsyncResultExtensionTest.groovy‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import java.util.concurrent.ExecutionException
88
import java.util.concurrent.ExecutorService
99
import java.util.concurrent.Executors
1010

11-
class GuavaAsyncResultSupportExtensionTest extends AgentTestRunner {
11+
class GuavaAsyncResultExtensionTest extends AgentTestRunner {
1212
@Override
1313
void configurePreAgent() {
1414
super.configurePreAgent()

‎dd-java-agent/instrumentation/reactive-streams/src/main/java/datadog/trace/instrumentation/reactivestreams/PublisherInstrumentation.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -60,10 +60,10 @@ public Map<String, String> contextStore() {
6060
@Override
6161
public String[] helperClassNames() {
6262
return new String[] {
63-
this.packageName + ".ReactiveStreamsAsyncResultSupportExtension",
64-
this.packageName + ".ReactiveStreamsAsyncResultSupportExtension$WrappedPublisher",
65-
this.packageName + ".ReactiveStreamsAsyncResultSupportExtension$WrappedSubscriber",
66-
this.packageName + ".ReactiveStreamsAsyncResultSupportExtension$WrappedSubscription",
63+
this.packageName + ".ReactiveStreamsAsyncResultExtension",
64+
this.packageName + ".ReactiveStreamsAsyncResultExtension$WrappedPublisher",
65+
this.packageName + ".ReactiveStreamsAsyncResultExtension$WrappedSubscriber",
66+
this.packageName + ".ReactiveStreamsAsyncResultExtension$WrappedSubscription",
6767
};
6868
}
6969

@@ -82,7 +82,7 @@ public void methodAdvice(MethodTransformer transformer) {
8282
public static class PublisherAdvice {
8383
@Advice.OnMethodExit(suppress = Throwable.class)
8484
public static void init() {
85-
ReactiveStreamsAsyncResultSupportExtension.initialize();
85+
ReactiveStreamsAsyncResultExtension.initialize();
8686
}
8787
}
8888

‎dd-java-agent/instrumentation/reactive-streams/src/main/java/datadog/trace/instrumentation/reactivestreams/ReactiveStreamsAsyncResultSupportExtension.java‎ renamed to ‎dd-java-agent/instrumentation/reactive-streams/src/main/java/datadog/trace/instrumentation/reactivestreams/ReactiveStreamsAsyncResultExtension.java‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,23 @@
11
package datadog.trace.instrumentation.reactivestreams;
22

33
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
4-
import datadog.trace.bootstrap.instrumentation.decorator.AsyncResultDecorator;
5-
import datadog.trace.bootstrap.instrumentation.decorator.AsyncResultDecorator.AsyncResultSupportExtension;
4+
import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtension;
5+
import datadog.trace.bootstrap.instrumentation.java.concurrent.AsyncResultExtensions;
66
import org.reactivestreams.Publisher;
77
import org.reactivestreams.Subscriber;
88
import org.reactivestreams.Subscription;
99

10-
public class ReactiveStreamsAsyncResultSupportExtension implements AsyncResultSupportExtension {
10+
public class ReactiveStreamsAsyncResultExtension implements AsyncResultExtension {
1111
static {
12-
AsyncResultDecorator.registerExtension(new ReactiveStreamsAsyncResultSupportExtension());
12+
AsyncResultExtensions.register(new ReactiveStreamsAsyncResultExtension());
1313
}
1414

1515
/**
16-
* Register the extension as an {@link AsyncResultSupportExtension} using static class
17-
* initialization.<br>
16+
* Register the extension as an {@link AsyncResultExtension} using static class initialization.
17+
* <br>
1818
* It uses an empty static method call to ensure the class loading and the one-time-only static
19-
* class initialization. This will ensure this extension will only be registered once to the
20-
* {@link AsyncResultDecorator}.
19+
* class initialization. This will ensure this extension will only be registered once under {@link
20+
* AsyncResultExtensions}.
2121
*/
2222
public static void initialize() {}
2323

‎dd-java-agent/instrumentation/reactive-streams/src/test/groovy/ReactiveStreamsAsyncResultSupportExtensionTest.groovy‎ renamed to ‎dd-java-agent/instrumentation/reactive-streams/src/test/groovy/ReactiveStreamsAsyncResultExtensionTest.groovy‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import datadog.trace.bootstrap.instrumentation.api.Tags
44

55
import java.util.concurrent.CountDownLatch
66

7-
class ReactiveStreamsAsyncResultSupportExtensionTest extends AgentTestRunner {
7+
class ReactiveStreamsAsyncResultExtensionTest extends AgentTestRunner {
88
@Override
99
void configurePreAgent() {
1010
super.configurePreAgent()

0 commit comments

Comments
 (0)