Finish pyspark-shell trace using SparkListenerApplicationEnd

Author: paul-laffon-dd

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

@@ -1,13 +1,12 @@
 package datadog.trace.instrumentation.spark;
 
 import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
-import static net.bytebuddy.matcher.ElementMatchers.isDeclaredBy;
-import static net.bytebuddy.matcher.ElementMatchers.isMethod;
-import static net.bytebuddy.matcher.ElementMatchers.nameEndsWith;
+import static net.bytebuddy.matcher.ElementMatchers.*;
 
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import net.bytebuddy.asm.Advice;
+import org.apache.spark.deploy.SparkSubmitArguments;
 
 public abstract class AbstractSparkInstrumentation extends InstrumenterModule.Tracing
     implements Instrumenter.ForKnownTypes {
@@ -32,6 +31,14 @@ public String[] knownMatchingTypes() {
 
   @Override
   public void methodAdvice(MethodTransformer transformer) {
+    // Capture spark submit arguments
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("prepareSubmitEnvironment"))
+            .and(takesArgument(0, named("org.apache.spark.deploy.SparkSubmitArguments")))
+            .and(isDeclaredBy(named("org.apache.spark.deploy.SparkSubmit"))),
+        AbstractSparkInstrumentation.class.getName() + "$PrepareSubmitEnvAdvice");
+
     // SparkSubmit class used for non YARN/Mesos environment
     transformer.applyAdvice(
         isMethod()
@@ -47,6 +54,20 @@ public void methodAdvice(MethodTransformer transformer) {
         AbstractSparkInstrumentation.class.getName() + "$YarnFinishAdvice");
   }
 
+  public static class PrepareSubmitEnvAdvice {
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void enter(@Advice.Argument(0) SparkSubmitArguments submitArgs) {
+
+      // Using `python script.py`, spark JVM is launched as PythonGatewayServer, which is exited
+      // using System.exit(0), leading to the exit advice not being called
+      // https://github.com/apache/spark/blob/v3.5.1/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala#L540-L542
+      // https://github.com/apache/spark/blob/v3.5.1/core/src/main/scala/org/apache/spark/api/python/PythonGatewayServer.scala#L74
+      if ("pyspark-shell".equals(submitArgs.primaryResource())) {
+        AbstractDatadogSparkListener.finishTraceOnApplicationEnd = true;
+      }
+    }
+  }
+
   public static class RunMainAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter() {