From e6f10d2488efd0abd0ea9ae4564e93d1e724fbb1 Mon Sep 17 00:00:00 2001
From: Paul Laffon <paul.laffon@datadoghq.com>
Date: Mon, 17 Mar 2025 11:22:30 +0100
Subject: [PATCH 1/2] Instrument Runtime.exit() to finish spark application
 spans

---
 .../spark/SparkExitAdvice.java                | 28 +++++++++++++++
 .../spark/SparkExitInstrumentation.java       | 35 +++++++++++++++++++
 2 files changed, 63 insertions(+)
 create mode 100644 dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/SparkExitAdvice.java
 create mode 100644 dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/SparkExitInstrumentation.java

diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/SparkExitAdvice.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/SparkExitAdvice.java
new file mode 100644
index 00000000000..1306da2b373
--- /dev/null
+++ b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/SparkExitAdvice.java
@@ -0,0 +1,28 @@
+package datadog.trace.instrumentation.spark;
+
+import java.lang.reflect.Method;
+import net.bytebuddy.asm.Advice;
+
+class SparkExitAdvice {
+  @Advice.OnMethodEnter(suppress = Throwable.class)
+  public static void enter(@Advice.Argument(0) int exitCode) {
+    try {
+      // Using reflection as java.lang.* instrumentation have to be done at boostrap,
+      // when spark classes have not been injected yet
+      Class<?> klass =
+          Thread.currentThread()
+              .getContextClassLoader()
+              .loadClass("datadog.trace.instrumentation.spark.AbstractDatadogSparkListener");
+      Object datadogListener = klass.getDeclaredField("listener").get(null);
+      if (datadogListener != null) {
+        Method method =
+            datadogListener
+                .getClass()
+                .getMethod(
+                    "finishApplication", long.class, Throwable.class, int.class, String.class);
+        method.invoke(datadogListener, System.currentTimeMillis(), null, exitCode, null);
+      }
+    } catch (Exception ignored) {
+    }
+  }
+}
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/SparkExitInstrumentation.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/SparkExitInstrumentation.java
new file mode 100644
index 00000000000..925a92fc2e4
--- /dev/null
+++ b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/SparkExitInstrumentation.java
@@ -0,0 +1,35 @@
+package datadog.trace.instrumentation.spark;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.agent.tooling.InstrumenterModule;
+import datadog.trace.api.Config;
+
+@AutoService(InstrumenterModule.class)
+public class SparkExitInstrumentation extends InstrumenterModule.Tracing
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice, Instrumenter.ForBootstrap {
+
+  public SparkExitInstrumentation() {
+    super("spark-exit");
+  }
+
+  @Override
+  protected boolean defaultEnabled() {
+    return Config.get().isDataJobsEnabled();
+  }
+
+  @Override
+  public String instrumentedType() {
+    return "java.lang.Runtime";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        named("exit").and(isDeclaredBy(named("java.lang.Runtime"))),
+        packageName + ".SparkExitAdvice");
+  }
+}

From 62109bfe232d6123a5030c3b028d001ecb640ae5 Mon Sep 17 00:00:00 2001
From: Paul Laffon <paul.laffon@datadoghq.com>
Date: Thu, 27 Mar 2025 13:25:38 +0100
Subject: [PATCH 2/2] Add comment on finishApplication called using reflection

---
 .../instrumentation/spark/AbstractDatadogSparkListener.java     | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
index bbd03d5b897..8da8865c982 100644
--- a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
+++ b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -239,6 +239,8 @@ public void onApplicationEnd(SparkListenerApplicationEnd applicationEnd) {
     }
   }
 
+  // This function is called using reflection in SparkExitInstrumentation, make sure to update if
+  // the signature of this function is changed
   public synchronized void finishApplication(
       long time, Throwable throwable, int exitCode, String msg) {
     log.info("Finishing spark application trace");