Finish pyspark-shell trace using SparkListenerApplicationEnd (#6956)

Finish pyspark-shell application span when the spark listener is receiving the SparkListenerApplicationEnd event

The application span is currently finished when exiting the runMain method to capture JVM errors, however this function is not exited as expected (System.exit(0)) when executing as pyspark-shell

Pyspark application can be launched using multiple method, mainly:
- spark-submit script.py
- python script.py

The spark instrumentation is working as expected in the spark-submit case, because this is java/scala code launched behind the scene

However, the python case, the listener is being injected as expected, but the spark.application span is not being closed properly because the JVM is launched as PythonGatewayServer which is exited using System.exit(0), leading to the exit advice not being called

Author: paul-laffon-dd

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

@@ -55,6 +55,7 @@
 public abstract class AbstractDatadogSparkListener extends SparkListener {
   public static volatile AbstractDatadogSparkListener listener = null;
   public static volatile boolean finishTraceOnApplicationEnd = true;
+  public static volatile boolean isPysparkShell = false;
 
   private final int MAX_COLLECTION_SIZE = 1000;
   private final int MAX_ACCUMULATOR_SIZE = 10000;

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

@@ -4,10 +4,12 @@
 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.takesArgument;
 
 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 +34,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,10 +57,29 @@ 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 pyspark `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.isPysparkShell = true;
+
+        // prepareSubmitEnvironment might be called before/after runMain depending on spark version
+        AbstractDatadogSparkListener.finishTraceOnApplicationEnd = true;
+      }
+    }
+  }
+
   public static class RunMainAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter() {
-      AbstractDatadogSparkListener.finishTraceOnApplicationEnd = false;
+      if (!AbstractDatadogSparkListener.isPysparkShell) {
+        AbstractDatadogSparkListener.finishTraceOnApplicationEnd = false;
+      }
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)

dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy

@@ -221,6 +221,41 @@ abstract class AbstractSparkTest extends AgentTestRunner {
     AbstractDatadogSparkListener.finishTraceOnApplicationEnd = true
   }
 
+  def "finish pyspark span launched with python onApplicationEnd"() {
+    setup:
+    def sparkSession = SparkSession.builder()
+      .config("spark.master", "local[2]")
+      .getOrCreate()
+
+    try {
+      // Generating a fake submit of pyspark-shell
+      def sparkSubmit = new SparkSubmit()
+      sparkSubmit.doSubmit(["--verbose", "pyspark-shell"] as String[])
+    }
+    catch (Exception ignored) {}
+    sparkSession.stop()
+
+    expect:
+    assert AbstractDatadogSparkListener.isPysparkShell
+    assert AbstractDatadogSparkListener.finishTraceOnApplicationEnd
+
+    assertTraces(1) {
+      trace(1) {
+        span {
+          operationName "spark.application"
+          resourceName "spark.application"
+          spanType "spark"
+          errored true
+          parent()
+        }
+      }
+    }
+
+    cleanup:
+    AbstractDatadogSparkListener.isPysparkShell = false
+    AbstractDatadogSparkListener.finishTraceOnApplicationEnd = true
+  }
+
   def "generate databricks spans"() {
     setup:
     def sparkSession = SparkSession.builder()