Instrument Runtime.exit() to finish spark application spans

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");

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) {
+    }
+  }
+}

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");
+  }
+}