Instrument QueryExecution.assertAnalyzed() to catch DataFrame analysis failures

aboitreaud · aboitreaud · commit 66cec2932764 · 2026-04-21T18:57:07.000+02:00
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleCheck.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/muzzle/MuzzleCheck.java
@@ -31,8 +31,32 @@ public boolean matches(ClassLoader classLoader) {
     boolean muzzleMatches = muzzle().matches(classLoader);
     if (muzzleMatches) {
       InstrumenterState.applyInstrumentation(classLoader, instrumentationId);
+      if (instrumentationClass.contains("spark") || instrumentationClass.contains("Spark")) {
+        System.err.println(
+            "[DD-SPARK-DEBUG] MuzzleCheck PASSED: "
+                + InstrumenterState.describe(instrumentationId)
+                + " classloader="
+                + classLoader);
+      }
     } else {
       InstrumenterState.blockInstrumentation(classLoader, instrumentationId);
+      if (instrumentationClass.contains("spark") || instrumentationClass.contains("Spark")) {
+        final List<Reference.Mismatch> mismatches =
+            muzzle.getMismatchedReferenceSources(classLoader);
+        System.err.println(
+            "[DD-SPARK-DEBUG] MuzzleCheck FAILED: "
+                + InstrumenterState.describe(instrumentationId)
+                + " classloader="
+                + classLoader);
+        for (final Reference.Mismatch mismatch : mismatches) {
+          System.err.println(
+              "[DD-SPARK-DEBUG] MuzzleCheck mismatch: "
+                  + InstrumenterState.describe(instrumentationId)
+                  + " muzzle.mismatch=\""
+                  + mismatch
+                  + "\"");
+        }
+      }
       if (log.isDebugEnabled()) {
         final List<Reference.Mismatch> mismatches =
             muzzle.getMismatchedReferenceSources(classLoader);
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -173,8 +173,17 @@ public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sp
             AgentThreadFactory.newAgentThread(
                 AgentThreadFactory.AgentThread.DATA_JOBS_MONITORING_SHUTDOWN_HOOK,
                 () -> {
+                  System.err.println(
+                      "[DD-SPARK-DEBUG] shutdownHook: applicationEnded="
+                          + applicationEnded
+                          + ", lastJobFailed="
+                          + lastJobFailed
+                          + ", lastSqlFailed="
+                          + lastSqlFailed
+                          + ", thread="
+                          + Thread.currentThread().getName());
                   if (!applicationEnded) {
-                    log.info("Finishing application trace from shutdown hook");
+                    System.err.println("[DD-SPARK-DEBUG] shutdownHook: calling finishApplication");
                     finishApplication(System.currentTimeMillis(), null, 0, null);
                   }
                 }));
@@ -319,7 +328,18 @@ private void captureOpenlineageContextIfPresent(
    * has an error signal even when no job/stage/task events fire.
    */
   public synchronized void onSqlFailure(Throwable throwable) {
+    System.err.println(
+        "[DD-SPARK-DEBUG] onSqlFailure: applicationEnded="
+            + applicationEnded
+            + ", throwable="
+            + throwable.getClass().getName()
+            + ": "
+            + throwable.getMessage()
+            + ", thread="
+            + Thread.currentThread().getName());
+
     if (applicationEnded) {
+      System.err.println("[DD-SPARK-DEBUG] onSqlFailure: skipping because applicationEnded=true");
       return;
     }
     lastSqlFailed = true;
@@ -328,13 +348,19 @@ public synchronized void onSqlFailure(Throwable throwable) {
     StringWriter sw = new StringWriter();
     throwable.printStackTrace(new PrintWriter(sw));
     lastSqlFailedStackTrace = sw.toString();
+
+    System.err.println("[DD-SPARK-DEBUG] onSqlFailure: lastSqlFailed set to true");
   }
 
   @Override
   public void onApplicationEnd(SparkListenerApplicationEnd applicationEnd) {
-    log.info(
-        "Received spark application end event, finish trace on this event: {}",
-        finishTraceOnApplicationEnd);
+    System.err.println(
+        "[DD-SPARK-DEBUG] onApplicationEnd: finishTraceOnApplicationEnd="
+            + finishTraceOnApplicationEnd
+            + ", applicationEnded="
+            + applicationEnded
+            + ", thread="
+            + Thread.currentThread().getName());
     notifyOl(x -> openLineageSparkListener.onApplicationEnd(x), applicationEnd);
 
     if (finishTraceOnApplicationEnd) {
@@ -346,9 +372,29 @@ public void onApplicationEnd(SparkListenerApplicationEnd applicationEnd) {
   // 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");
+    System.err.println(
+        "[DD-SPARK-DEBUG] finishApplication: thread="
+            + Thread.currentThread().getName()
+            + ", applicationEnded="
+            + applicationEnded
+            + ", throwable="
+            + (throwable != null
+                ? throwable.getClass().getName() + ": " + throwable.getMessage()
+                : "null")
+            + ", exitCode="
+            + exitCode
+            + ", msg="
+            + msg
+            + ", lastJobFailed="
+            + lastJobFailed
+            + ", lastSqlFailed="
+            + lastSqlFailed
+            + ", jobCount="
+            + jobCount);
 
     if (applicationEnded) {
+      System.err.println(
+          "[DD-SPARK-DEBUG] finishApplication: skipping because applicationEnded=true (duplicate call)");
       return;
     }
     applicationEnded = true;
@@ -361,9 +407,12 @@ public synchronized void finishApplication(
 
     initApplicationSpanIfNotInitialized();
 
+    String errorDecision;
     if (throwable != null) {
+      errorDecision = "throwable";
       applicationSpan.addThrowable(throwable);
     } else if (exitCode != 0) {
+      errorDecision = "exitCode=" + exitCode;
       applicationSpan.setError(true);
       applicationSpan.setTag(
           DDTags.ERROR_TYPE, "Spark Application Failed with exit code " + exitCode);
@@ -372,16 +421,21 @@ public synchronized void finishApplication(
       applicationSpan.setTag(DDTags.ERROR_MSG, errorMessage);
       applicationSpan.setTag(DDTags.ERROR_STACK, msg);
     } else if (lastJobFailed) {
+      errorDecision = "lastJobFailed";
       applicationSpan.setError(true);
       applicationSpan.setTag(DDTags.ERROR_TYPE, "Spark Application Failed");
       applicationSpan.setTag(DDTags.ERROR_MSG, lastJobFailedMessage);
       applicationSpan.setTag(DDTags.ERROR_STACK, lastJobFailedStackTrace);
     } else if (lastSqlFailed) {
+      errorDecision = "lastSqlFailed";
       applicationSpan.setError(true);
       applicationSpan.setTag(DDTags.ERROR_TYPE, "Spark SQL Failed");
       applicationSpan.setTag(DDTags.ERROR_MSG, lastSqlFailedMessage);
       applicationSpan.setTag(DDTags.ERROR_STACK, lastSqlFailedStackTrace);
+    } else {
+      errorDecision = "none (SUCCESS)";
     }
+    System.err.println("[DD-SPARK-DEBUG] finishApplication: errorDecision=" + errorDecision);
 
     applicationMetrics.setSpanMetrics(applicationSpan);
     applicationSpan.setMetric("spark.max_executor_count", maxExecutorCount);
@@ -486,6 +540,13 @@ private AgentSpan getOrCreateSqlSpan(
 
   @Override
   public synchronized void onJobStart(SparkListenerJobStart jobStart) {
+    System.err.println(
+        "[DD-SPARK-DEBUG] onJobStart: jobId="
+            + jobStart.jobId()
+            + ", stageCount="
+            + getStageCount(jobStart)
+            + ", thread="
+            + Thread.currentThread().getName());
     jobCount++;
     if (jobSpans.size() > MAX_COLLECTION_SIZE) {
       return;
@@ -544,8 +605,22 @@ public synchronized void onJobStart(SparkListenerJobStart jobStart) {
 
   @Override
   public synchronized void onJobEnd(SparkListenerJobEnd jobEnd) {
+    System.err.println(
+        "[DD-SPARK-DEBUG] onJobEnd: jobId="
+            + jobEnd.jobId()
+            + ", result="
+            + jobEnd.jobResult().getClass().getSimpleName()
+            + ", lastJobFailed="
+            + lastJobFailed
+            + ", lastSqlFailed="
+            + lastSqlFailed
+            + ", thread="
+            + Thread.currentThread().getName());
+
     AgentSpan jobSpan = jobSpans.remove(jobEnd.jobId());
     if (jobSpan == null) {
+      System.err.println(
+          "[DD-SPARK-DEBUG] onJobEnd: no span found for jobId=" + jobEnd.jobId() + ", skipping");
       return;
     }
 
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkInstrumentation.java b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkInstrumentation.java
@@ -40,7 +40,8 @@ public String[] knownMatchingTypes() {
       "org.apache.spark.util.SparkClassUtils",
       "org.apache.spark.scheduler.LiveListenerBus",
       "org.apache.spark.sql.execution.SparkPlanInfo$",
-      "org.apache.spark.sql.SparkSession"
+      "org.apache.spark.sql.SparkSession",
+      "org.apache.spark.sql.execution.QueryExecution"
     };
   }
 
@@ -77,6 +78,14 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(isDeclaredBy(named("org.apache.spark.sql.SparkSession"))),
         AbstractSparkInstrumentation.class.getName() + "$SparkSqlFailureAdvice");
 
+    // QueryExecution.assertAnalyzed() — catch all Catalyst analysis failures regardless of
+    // entry point (SparkSession.sql, Dataset.select, Dataset.filter, etc.)
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("assertAnalyzed"))
+            .and(isDeclaredBy(named("org.apache.spark.sql.execution.QueryExecution"))),
+        AbstractSparkInstrumentation.class.getName() + "$QueryExecutionFailureAdvice");
+
     // LiveListenerBus class is used to manage spark listeners
     transformer.applyAdvice(
         isMethod()
@@ -135,8 +144,39 @@ public static void enter(@Advice.Argument(1) int exitCode, @Advice.Argument(2) S
   public static class SparkSqlFailureAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
     public static void exit(@Advice.Thrown Throwable throwable) {
-      if (throwable != null && AbstractDatadogSparkListener.listener != null) {
-        AbstractDatadogSparkListener.listener.onSqlFailure(throwable);
+      if (throwable != null) {
+        System.err.println(
+            "[DD-SPARK-DEBUG] SparkSqlFailureAdvice.exit: thread="
+                + Thread.currentThread().getName()
+                + ", throwable="
+                + throwable.getClass().getName()
+                + ": "
+                + throwable.getMessage()
+                + ", listenerNull="
+                + (AbstractDatadogSparkListener.listener == null));
+        if (AbstractDatadogSparkListener.listener != null) {
+          AbstractDatadogSparkListener.listener.onSqlFailure(throwable);
+        }
+      }
+    }
+  }
+
+  public static class QueryExecutionFailureAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(@Advice.Thrown Throwable throwable) {
+      if (throwable != null) {
+        System.err.println(
+            "[DD-SPARK-DEBUG] QueryExecutionFailureAdvice.exit: thread="
+                + Thread.currentThread().getName()
+                + ", throwable="
+                + throwable.getClass().getName()
+                + ": "
+                + throwable.getMessage()
+                + ", listenerNull="
+                + (AbstractDatadogSparkListener.listener == null));
+        if (AbstractDatadogSparkListener.listener != null) {
+          AbstractDatadogSparkListener.listener.onSqlFailure(throwable);
+        }
       }
     }
   }
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/SparkExitAdvice.java b/dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/SparkExitAdvice.java
@@ -14,6 +14,13 @@ public static void enter(@Advice.Argument(0) int exitCode) {
               .getContextClassLoader()
               .loadClass("datadog.trace.instrumentation.spark.AbstractDatadogSparkListener");
       Object datadogListener = klass.getDeclaredField("listener").get(null);
+      System.err.println(
+          "[DD-SPARK-DEBUG] SparkExitAdvice.enter: exitCode="
+              + exitCode
+              + ", listenerNull="
+              + (datadogListener == null)
+              + ", thread="
+              + Thread.currentThread().getName());
       if (datadogListener != null) {
         Method method =
             datadogListener
diff --git a/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy b/dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy
@@ -223,6 +223,38 @@ abstract class AbstractSparkTest extends InstrumentationSpecification {
     }
   }
 
+  def "DataFrame analysis failure on unresolved column marks application span as error"() {
+    setup:
+    def sparkSession = SparkSession.builder()
+      .config("spark.master", "local[2]")
+      .getOrCreate()
+
+    try {
+      // Triggers AnalysisException via Dataset.select() -> QueryExecution.assertAnalyzed(),
+      // NOT through SparkSession.sql(). This exercises the QueryExecutionFailureAdvice.
+      sparkSession.range(1).toDF("id").select("nonexistent_column")
+    } catch (Exception ignored) {
+      // Expected: AnalysisException thrown by Catalyst analysis
+    }
+    sparkSession.stop()
+
+    expect:
+    assertTraces(1) {
+      trace(1) {
+        span {
+          operationName "spark.application"
+          resourceName "spark.application"
+          spanType "spark"
+          errored true
+          parent()
+          assert span.tags["error.type"] == "Spark SQL Failed"
+          assert span.tags["error.message"] =~ /(?i).*nonexistent_column.*/
+          assert span.tags["error.stack"] =~ /(?s).*AnalysisException.*/
+        }
+      }
+    }
+  }
+
   def "capture SparkSubmit.runMain() errors"() {
     setup:
     def sparkSession = SparkSession.builder()
diff --git a/dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java b/dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java
@@ -94,6 +94,13 @@ public static void enter(@Advice.This SparkContext sparkContext) {
       AbstractDatadogSparkListener.listener =
           new DatadogSpark212Listener(
               sparkContext.getConf(), sparkContext.applicationId(), sparkContext.version());
+      System.err.println(
+          "[DD-SPARK-DEBUG] listener registered: class="
+              + AbstractDatadogSparkListener.listener.getClass().getName()
+              + ", appId="
+              + sparkContext.applicationId()
+              + ", thread="
+              + Thread.currentThread().getName());
       sparkContext.listenerBus().addToSharedQueue(AbstractDatadogSparkListener.listener);
     }
   }
diff --git a/dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java b/dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java
@@ -95,6 +95,13 @@ public static void enter(@Advice.This SparkContext sparkContext) {
       AbstractDatadogSparkListener.listener =
           new DatadogSpark213Listener(
               sparkContext.getConf(), sparkContext.applicationId(), sparkContext.version());
+      System.err.println(
+          "[DD-SPARK-DEBUG] listener registered: class="
+              + AbstractDatadogSparkListener.listener.getClass().getName()
+              + ", appId="
+              + sparkContext.applicationId()
+              + ", thread="
+              + Thread.currentThread().getName());
       sparkContext.listenerBus().addToSharedQueue(AbstractDatadogSparkListener.listener);
     }
   }

Original file line number	Diff line number	Diff line change
`@@ -94,6 +94,13 @@ public static void enter(@Advice.This SparkContext sparkContext) {`
`94`	`94`	`AbstractDatadogSparkListener.listener =`
`95`	`95`	`new DatadogSpark212Listener(`
`96`	`96`	`sparkContext.getConf(), sparkContext.applicationId(), sparkContext.version());`
	`97`	`+ System.err.println(`
	`98`	`+ "[DD-SPARK-DEBUG] listener registered: class="`
	`99`	`+ + AbstractDatadogSparkListener.listener.getClass().getName()`
	`100`	`+ + ", appId="`
	`101`	`+ + sparkContext.applicationId()`
	`102`	`+ + ", thread="`
	`103`	`+ + Thread.currentThread().getName());`
`97`	`104`	`sparkContext.listenerBus().addToSharedQueue(AbstractDatadogSparkListener.listener);`
`98`	`105`	`}`
`99`	`106`	`}`
Original file line number	Diff line number	Diff line change
`@@ -95,6 +95,13 @@ public static void enter(@Advice.This SparkContext sparkContext) {`
`95`	`95`	`AbstractDatadogSparkListener.listener =`
`96`	`96`	`new DatadogSpark213Listener(`
`97`	`97`	`sparkContext.getConf(), sparkContext.applicationId(), sparkContext.version());`
	`98`	`+ System.err.println(`
	`99`	`+ "[DD-SPARK-DEBUG] listener registered: class="`
	`100`	`+ + AbstractDatadogSparkListener.listener.getClass().getName()`
	`101`	`+ + ", appId="`
	`102`	`+ + sparkContext.applicationId()`
	`103`	`+ + ", thread="`
	`104`	`+ + Thread.currentThread().getName());`
`98`	`105`	`sparkContext.listenerBus().addToSharedQueue(AbstractDatadogSparkListener.listener);`
`99`	`106`	`}`
`100`	`107`	`}`