[SPARK-55240][CORE] Refactor LazyTry stacktrace handling to use wrapper instead of suppressed exception

cloud-fan · cloud-fan · commit 8719c6c8dfd3 · 2026-01-28T22:37:17.000+08:00
### What changes were proposed in this pull request? This PR refactors how `doTryWithCallerStacktrace` and `getTryWithCallerStacktrace` handle stacktrace stitching, with improved behavior for `LazyTry`: **Key changes:** 1. **New internal wrapper class `TryStackTraceWrapper`**: Instead of modifying the original exception's suppressed list in `doTryWithCallerStacktrace`, we now wrap the exception in a `TryStackTraceWrapper` that carries: - The original exception - The depth of stacktrace frames to preserve (optimization: just an int, not a copy of the array) - An `OriginalTryStackTraceException` holding the original stacktrace for later use 2. **Conditional suppressed exception in `LazyTry`**: - `LazyTry` now tracks whether the lazy value has been accessed (`materialized` flag) - On **first access**: No suppressed exception is added (clean output since the caller context is the same as initialization) - On **subsequent accesses**: The original stacktrace is added as a suppressed exception to help with debugging 3. **`getTryWithCallerStacktrace` now accepts `isFirstAccess` parameter**: - Default is `true` (no suppressed exception added) - When `false`, adds the original stacktrace as suppressed if not already present (idempotent check) ### Why are the changes needed? 1. **Cleaner first-access output**: On the first access to a `LazyTry` value, the suppressed exception showing the "original" stacktrace is redundant because it's essentially the same context. Removing it makes error output cleaner. 2. **Preserved debugging info for subsequent accesses**: When accessing a failed `LazyTry` from a different call site, the suppressed exception showing where the error originally occurred is valuable for debugging. 3. **Better encapsulation**: The wrapper is internal-only and never exposed to users. The original exception type is preserved when thrown. ### Does this PR introduce _any_ user-facing change? Yes: - On first access to a failed `LazyTry`, exceptions will no longer have a suppressed `OriginalTryStackTraceException` (cleaner output) - On subsequent accesses, the behavior is unchanged (suppressed exception shows original stacktrace) ### How was this patch tested? Existing unit tests in `UtilsSuite` and `LazyTrySuite` updated and passing. ### Was this patch authored or co-authored using generative AI tooling? Yes. cursor 2.4 Closes apache#54007 from cloud-fan/error. Authored-by: Wenchen Fan <wenchen@databricks.com> Signed-off-by: Wenchen Fan <wenchen@databricks.com>
diff --git a/core/src/main/scala/org/apache/spark/util/LazyTry.scala b/core/src/main/scala/org/apache/spark/util/LazyTry.scala
@@ -45,17 +45,21 @@ import scala.util.Try
 private[spark] class LazyTry[T](initialize: => T) extends Serializable {
   private lazy val tryT: Try[T] = Utils.doTryWithCallerStacktrace { initialize }
 
+  // Track whether the lazy val has been materialized or not.
+  @volatile private var materialized: Boolean = false
+
   /**
    * Get the lazy value. If the initialization block threw an exception, it will be re-thrown here.
    * The exception will be re-thrown with the current caller's stacktrace.
-   * An exception with stack trace from when the exception was first thrown can be accessed with
-   * ```
-   * ex.getSuppressed.find { e =>
-   *   e.getMessage == org.apache.spark.util.Utils.TRY_WITH_CALLER_STACKTRACE_FULL_STACKTRACE
-   * }
-   * ```
+   *
+   * On the first access, no suppressed exception is added. On subsequent accesses, the original
+   * stacktrace is added as a suppressed exception to help with debugging.
    */
-  def get: T = Utils.getTryWithCallerStacktrace(tryT)
+  def get: T = {
+    val isFirstAccess = !materialized
+    materialized = true
+    Utils.getTryWithCallerStacktrace(tryT, isFirstAccess)
+  }
 }
 
 private[spark] object LazyTry {
diff --git a/core/src/main/scala/org/apache/spark/util/Utils.scala b/core/src/main/scala/org/apache/spark/util/Utils.scala
@@ -1368,20 +1368,35 @@ private[spark] object Utils
   val TRY_WITH_CALLER_STACKTRACE_FULL_STACKTRACE =
     "Full stacktrace of original doTryWithCallerStacktrace caller"
 
+  /**
+   * Exception used to preserve the original stacktrace before stitching. On subsequent accesses
+   * (not the first), this is added as a suppressed exception to show where the error originally
+   * occurred.
+   */
   class OriginalTryStackTraceException()
-    extends Exception(TRY_WITH_CALLER_STACKTRACE_FULL_STACKTRACE) {
-    var doTryWithCallerStacktraceDepth: Int = 0
-  }
+    extends Exception(TRY_WITH_CALLER_STACKTRACE_FULL_STACKTRACE)
+
+  /**
+   * Internal wrapper used to carry the original exception along with metadata needed for
+   * stacktrace stitching. This wrapper is never exposed to users.
+   *
+   * @param originalException The original exception
+   * @param depth The number of stacktrace frames below doTryWithCallerStacktrace to preserve
+   * @param originalStacktraceEx Exception holding the original stacktrace (for suppressed display)
+   */
+  class TryStackTraceWrapper(
+      val originalException: Throwable,
+      val depth: Int,
+      val originalStacktraceEx: OriginalTryStackTraceException) extends Exception(originalException)
 
   /**
    * Use Try with stacktrace substitution for the caller retrieving the error.
    *
    * Normally in case of failure, the exception would have the stacktrace of the caller that
    * originally called doTryWithCallerStacktrace. However, we want to replace the part above
    * this function with the stacktrace of the caller who calls getTryWithCallerStacktrace.
-   * So here we save the part of the stacktrace below doTryWithCallerStacktrace, and
-   * getTryWithCallerStacktrace will stitch it with the new stack trace of the caller.
-   * The full original stack trace is kept in ex.getSuppressed.
+   * So here we wrap the exception with metadata, and getTryWithCallerStacktrace will
+   * unwrap it and stitch the stacktrace with the new caller's stack trace.
    *
    * @param f Code block to be wrapped in Try
    * @return Try with Success or Failure of the code block. Use with getTryWithCallerStacktrace.
@@ -1392,29 +1407,27 @@ private[spark] object Utils
     }
     t match {
       case Failure(ex) =>
+        // If already wrapped, return as-is (nested call)
+        if (ex.isInstanceOf[TryStackTraceWrapper]) {
+          return t
+        }
         // Note: we remove the common suffix instead of e.g. finding the call to this function, to
         // account for recursive calls with multiple doTryWithCallerStacktrace on the stack trace.
         val origStackTrace = ex.getStackTrace
         val currentStackTrace = Thread.currentThread().getStackTrace
         val commonSuffixLen = origStackTrace.reverse.zip(currentStackTrace.reverse).takeWhile {
           case (exElem, currentElem) => exElem == currentElem
         }.length
-        // Add the full stack trace of the original caller as the suppressed exception.
-        // It may already be there if it's a nested call to doTryWithCallerStacktrace.
-        val origEx = ex.getSuppressed.find { e =>
-          e.isInstanceOf[OriginalTryStackTraceException]
-        }.getOrElse {
-          val fullEx = new OriginalTryStackTraceException()
-          fullEx.setStackTrace(origStackTrace)
-          ex.addSuppressed(fullEx)
-          fullEx
-        }.asInstanceOf[OriginalTryStackTraceException]
-        // Update the depth of the stack of the current doTryWithCallerStacktrace, for stitching
-        // it with the stack of getTryWithCallerStacktrace.
-        origEx.doTryWithCallerStacktraceDepth = origStackTrace.size - commonSuffixLen
-      case Success(_) => // nothing
+        // Compute the depth of the "below" portion to preserve during stitching
+        val depth = origStackTrace.size - commonSuffixLen
+        // Create exception to hold original stacktrace (for suppressed on subsequent access)
+        val originalStacktraceEx = new OriginalTryStackTraceException()
+        originalStacktraceEx.setStackTrace(origStackTrace)
+        // Wrap everything
+        val wrapper = new TryStackTraceWrapper(ex, depth, originalStacktraceEx)
+        Failure(wrapper)
+      case Success(_) => t
     }
-    t
   }
 
   /**
@@ -1424,32 +1437,34 @@ private[spark] object Utils
    * below the original doTryWithCallerStacktrace which triggered it, with the caller stack trace
    * of the current caller of getTryWithCallerStacktrace.
    *
-   * Full stack trace of the original doTryWithCallerStacktrace caller can be retrieved with
-   * ```
-   * ex.getSuppressed.find { e =>
-   *   e.isInstanceOf[Utils.OriginalTryStackTraceException]
-   * }
-   * ```
-   *
+   * On subsequent accesses (not the first), the original stacktrace is added as a suppressed
+   * exception to help with debugging.
    *
    * @param t Try from doTryWithCallerStacktrace
+   * @param isFirstAccess Whether this is the first access to the Try value
    * @return Result of the Try or rethrows the failure exception with modified stacktrace.
    */
-  def getTryWithCallerStacktrace[T](t: Try[T]): T = t match {
-    case Failure(ex) =>
-      val originalStacktraceEx = ex.getSuppressed.find { e =>
-        // added in doTryWithCallerStacktrace
-        e.isInstanceOf[OriginalTryStackTraceException]
-      }.getOrElse {
-        // If we don't have the expected stacktrace information, just rethrow
-        throw ex
-      }.asInstanceOf[OriginalTryStackTraceException]
-      val belowStacktrace = originalStacktraceEx.getStackTrace
-        .take(originalStacktraceEx.doTryWithCallerStacktraceDepth)
+  def getTryWithCallerStacktrace[T](t: Try[T], isFirstAccess: Boolean = true): T = t match {
+    case Failure(wrapper: TryStackTraceWrapper) =>
+      val originalEx = wrapper.originalException
+      // Stitch the stacktrace: keep the first `depth` frames, replace the rest with current caller
+      val belowStacktrace = originalEx.getStackTrace.take(wrapper.depth)
       // We are modifying and throwing the original exception. It would be better if we could
       // return a copy, but we can't easily clone it and preserve. If this is accessed from
       // multiple threads that then look at the stack trace, this could break.
-      ex.setStackTrace(belowStacktrace ++ Thread.currentThread().getStackTrace.drop(1))
+      originalEx.setStackTrace(
+        belowStacktrace ++ Thread.currentThread().getStackTrace.drop(1))
+      // On subsequent accesses, add suppressed exception showing original stacktrace
+      if (!isFirstAccess) {
+        val alreadyAdded = originalEx.getSuppressed.exists(
+          _.isInstanceOf[OriginalTryStackTraceException])
+        if (!alreadyAdded) {
+          originalEx.addSuppressed(wrapper.originalStacktraceEx)
+        }
+      }
+      throw originalEx
+    case Failure(ex) =>
+      // Not wrapped (shouldn't happen if used correctly), just rethrow
       throw ex
     case Success(s) => s
   }
diff --git a/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala b/core/src/test/scala/org/apache/spark/util/UtilsSuite.scala
@@ -1640,7 +1640,7 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties {
   }
 
   private def callGetTryAgain(t: Try[String]): String = {
-    Utils.getTryWithCallerStacktrace(t)
+    Utils.getTryWithCallerStacktrace(t, isFirstAccess = false)
   }
 
   test("doTryWithCallerStacktrace and getTryWithCallerStacktrace") {
@@ -1669,26 +1669,8 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties {
     assert(!st1.exists(_.getMethodName == "callDoTry"))
     assert(st1.exists(_.getMethodName == "callGetTry"))
 
-    // The original stack trace with callDoTry should be in the suppressed exceptions.
-    // Example:
-    // scalastyle:off line.size.limit
-    // Suppressed: java.lang.Exception: Full stacktrace of original doTryWithCallerStacktrace caller
-    //   at org.apache.spark.util.UtilsSuite.throwException(UtilsSuite.scala:1640)
-    //   at org.apache.spark.util.UtilsSuite.$anonfun$callDoTry$1(UtilsSuite.scala:1645)
-    //   at scala.util.Try$.apply(Try.scala:213)
-    //   at org.apache.spark.util.Utils$.doTryWithCallerStacktrace(Utils.scala:1586)
-    //   at org.apache.spark.util.UtilsSuite.callDoTry(UtilsSuite.scala:1645)
-    //   at org.apache.spark.util.UtilsSuite.$anonfun$new$165(UtilsSuite.scala:1658)
-    //   ... 56 more
-    // scalastyle:on line.size.limit
-    val origSt = e1.getSuppressed.find(_.isInstanceOf[Utils.OriginalTryStackTraceException])
-    assert(origSt.isDefined)
-    assert(origSt.get.getStackTrace.exists(_.getMethodName == "throwException"))
-    assert(origSt.get.getStackTrace.exists(_.getMethodName == "callDoTry"))
-
-    // Should save the depth of the stack trace under doTryWithCallerStacktrace.
-    assert(origSt.get.asInstanceOf[Utils.OriginalTryStackTraceException]
-      .doTryWithCallerStacktraceDepth == 4)
+    // First access (isFirstAccess = true by default) - no suppressed exception
+    assert(!e1.getSuppressed.exists(_.isInstanceOf[Utils.OriginalTryStackTraceException]))
 
     val e2 = intercept[Exception] {
       callGetTryAgain(t)
@@ -1715,6 +1697,24 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties {
     // callGetTry that we called before shouldn't be on the stack trace.
     assert(!st2.exists(_.getMethodName == "callGetTry"))
 
+    // Second access (isFirstAccess = false) - suppressed exception should be added.
+    // The original stack trace with callDoTry should be in the suppressed exceptions.
+    // Example:
+    // scalastyle:off line.size.limit
+    // Suppressed: org.apache.spark.util.Utils$OriginalTryStackTraceException: Full stacktrace of original doTryWithCallerStacktrace caller
+    //   at org.apache.spark.util.UtilsSuite.throwException(UtilsSuite.scala:1640)
+    //   at org.apache.spark.util.UtilsSuite.$anonfun$callDoTry$1(UtilsSuite.scala:1645)
+    //   at scala.util.Try$.apply(Try.scala:213)
+    //   at org.apache.spark.util.Utils$.doTryWithCallerStacktrace(Utils.scala:1586)
+    //   at org.apache.spark.util.UtilsSuite.callDoTry(UtilsSuite.scala:1645)
+    //   at org.apache.spark.util.UtilsSuite.$anonfun$new$165(UtilsSuite.scala:1658)
+    //   ... 56 more
+    // scalastyle:on line.size.limit
+    val origSt = e2.getSuppressed.find(_.isInstanceOf[Utils.OriginalTryStackTraceException])
+    assert(origSt.isDefined, "Suppressed exception should be added on subsequent access")
+    assert(origSt.get.getStackTrace.exists(_.getMethodName == "throwException"))
+    assert(origSt.get.getStackTrace.exists(_.getMethodName == "callDoTry"))
+
     // Unfortunately, this utility is not able to clone the exception, but modifies it in place,
     // so now e1 is also pointing to "callGetTryAgain" instead of "callGetTry".
     val st1Again = e1.getStackTrace
@@ -1760,20 +1760,13 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties {
       //  at org.apache.spark.util.UtilsSuite.callDoTryNestedNested(UtilsSuite.scala:1654)
       //  at org.apache.spark.util.UtilsSuite.$anonfun$new$172(UtilsSuite.scala:1674)
       // ...
-      // Suppressed: org.apache.spark.util.Utils$OriginalTryStackTraceException: Full stacktrace of original doTryWithCallerStacktrace caller
-      //  at org.apache.spark.util.UtilsSuite.throwException(UtilsSuite.scala:1529)
-      //  at org.apache.spark.util.UtilsSuite.$anonfun$callDoTry$1(UtilsSuite.scala:1534)
-      //  at scala.util.Try$.apply(Try.scala:217)
-      //  at org.apache.spark.util.Utils$.doTryWithCallerStacktrace(Utils.scala:1377)
-      //  at org.apache.spark.util.UtilsSuite.callDoTry(UtilsSuite.scala:1534)
-      //  at org.apache.spark.util.UtilsSuite.$anonfun$callDoTryNested$1(UtilsSuite.scala:1631)
-      //  ...
       // scalastyle:on line.size.limit
 
       assert(e.getStackTrace.exists(_.getMethodName == "callGetTryFromNested"))
       assert(!e.getStackTrace.exists(_.getMethodName == "callGetTryFromNestedNested"))
       assert(!e.getStackTrace.exists(_.getMethodName == "callGetTry"))
-      assert(e.getSuppressed.length == 1)
+      // No suppressed exception - the wrapper is used internally only
+      assert(e.getSuppressed.length == 0)
 
       Utils.getTryWithCallerStacktrace(t)
     }
@@ -1821,7 +1814,8 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties {
       assert(e.getStackTrace.exists(_.getMethodName == "callGetTryFromNestedNested"))
       assert(!e.getStackTrace.exists(_.getMethodName == "callGetTryFromNested"))
       assert(!e.getStackTrace.exists(_.getMethodName == "callGetTry"))
-      assert(e.getSuppressed.length == 1)
+      // No suppressed exception - the wrapper is used internally only
+      assert(e.getSuppressed.length == 0)
 
       Utils.getTryWithCallerStacktrace(t)
     }
@@ -1865,7 +1859,8 @@ class UtilsSuite extends SparkFunSuite with ResetSystemProperties {
     assert(e.getStackTrace.exists(_.getMethodName == "callGetTry"))
     assert(!e.getStackTrace.exists(_.getMethodName == "callGetTryFromNested"))
     assert(!e.getStackTrace.exists(_.getMethodName == "callGetTryFromNestedNested"))
-    assert(e.getSuppressed.length == 1)
+    // No suppressed exception - the wrapper is used internally only
+    assert(e.getSuppressed.length == 0)
   }
 }
 
