Skip to content

Commit 0a1776a

Browse files
Instrument Runtime.exit() to finish spark application spans (#8572)
Explicitly exiting spark using System.exit(code) will bypass the regular code to finish spark application spans We are currently using a shutdown hook https://github.com/DataDog/dd-trace-java/blob/v1.47.2/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java#L139-L153, but it is only best effort: - no guarantee it will be executed before the tracer shuts down - no access to the exit code Instrumenting the Runtime.exit() allow to prevent those drawbacks
1 parent b56bb15 commit 0a1776a

File tree

3 files changed

+65
-0
lines changed

3 files changed

+65
-0
lines changed

dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -239,6 +239,8 @@ public void onApplicationEnd(SparkListenerApplicationEnd applicationEnd) {
239239
}
240240
}
241241

242+
// This function is called using reflection in SparkExitInstrumentation, make sure to update if
243+
// the signature of this function is changed
242244
public synchronized void finishApplication(
243245
long time, Throwable throwable, int exitCode, String msg) {
244246
log.info("Finishing spark application trace");
Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
package datadog.trace.instrumentation.spark;
2+
3+
import java.lang.reflect.Method;
4+
import net.bytebuddy.asm.Advice;
5+
6+
class SparkExitAdvice {
7+
@Advice.OnMethodEnter(suppress = Throwable.class)
8+
public static void enter(@Advice.Argument(0) int exitCode) {
9+
try {
10+
// Using reflection as java.lang.* instrumentation have to be done at boostrap,
11+
// when spark classes have not been injected yet
12+
Class<?> klass =
13+
Thread.currentThread()
14+
.getContextClassLoader()
15+
.loadClass("datadog.trace.instrumentation.spark.AbstractDatadogSparkListener");
16+
Object datadogListener = klass.getDeclaredField("listener").get(null);
17+
if (datadogListener != null) {
18+
Method method =
19+
datadogListener
20+
.getClass()
21+
.getMethod(
22+
"finishApplication", long.class, Throwable.class, int.class, String.class);
23+
method.invoke(datadogListener, System.currentTimeMillis(), null, exitCode, null);
24+
}
25+
} catch (Exception ignored) {
26+
}
27+
}
28+
}
Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
package datadog.trace.instrumentation.spark;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
4+
import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
5+
6+
import com.google.auto.service.AutoService;
7+
import datadog.trace.agent.tooling.Instrumenter;
8+
import datadog.trace.agent.tooling.InstrumenterModule;
9+
import datadog.trace.api.Config;
10+
11+
@AutoService(InstrumenterModule.class)
12+
public class SparkExitInstrumentation extends InstrumenterModule.Tracing
13+
implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice, Instrumenter.ForBootstrap {
14+
15+
public SparkExitInstrumentation() {
16+
super("spark-exit");
17+
}
18+
19+
@Override
20+
protected boolean defaultEnabled() {
21+
return Config.get().isDataJobsEnabled();
22+
}
23+
24+
@Override
25+
public String instrumentedType() {
26+
return "java.lang.Runtime";
27+
}
28+
29+
@Override
30+
public void methodAdvice(MethodTransformer transformer) {
31+
transformer.applyAdvice(
32+
named("exit").and(isDeclaredBy(named("java.lang.Runtime"))),
33+
packageName + ".SparkExitAdvice");
34+
}
35+
}

0 commit comments

Comments
 (0)