Fix the error masking problem

Author: JoshRosen

core/src/main/scala/org/apache/spark/executor/Executor.scala

@@ -209,18 +209,21 @@ private[spark] class Executor(
 
         // Run the actual task and measure its runtime.
         taskStart = System.currentTimeMillis()
+        var threwException = true
         val (value, accumUpdates) = try {
-          task.run(
+          val res = task.run(
             taskAttemptId = taskId,
             attemptNumber = attemptNumber,
             metricsSystem = env.metricsSystem)
+          threwException = false
+          res
         } finally {
           // Note: this memory freeing logic is duplicated in DAGScheduler.runLocallyWithinThread;
           // when changing this, make sure to update both copies.
           val freedMemory = taskMemoryManager.cleanUpAllAllocatedMemory()
           if (freedMemory > 0) {
             val errMsg = s"Managed memory leak detected; size = $freedMemory bytes, TID = $taskId"
-            if (conf.getBoolean("spark.unsafe.exceptionOnMemoryLeak", false)) {
+            if (conf.getBoolean("spark.unsafe.exceptionOnMemoryLeak", false) && !threwException) {
               throw new SparkException(errMsg)
             } else {
               logError(errMsg)

core/src/main/scala/org/apache/spark/scheduler/DAGScheduler.scala

@@ -685,17 +685,20 @@ class DAGScheduler(
           metricsSystem = env.metricsSystem,
           runningLocally = true)
       TaskContext.setTaskContext(taskContext)
+      var threwException = true
       try {
         val result = job.func(taskContext, rdd.iterator(split, taskContext))
         job.listener.taskSucceeded(0, result)
+        threwException = false
       } finally {
         taskContext.markTaskCompleted()
         TaskContext.unset()
         // Note: this memory freeing logic is duplicated in Executor.run(); when changing this,
         // make sure to update both copies.
         val freedMemory = taskMemoryManager.cleanUpAllAllocatedMemory()
         if (freedMemory > 0) {
-          if (sc.getConf.getBoolean("spark.unsafe.exceptionOnMemoryLeak", false)) {
+          if (sc.getConf.getBoolean("spark.unsafe.exceptionOnMemoryLeak", false)
+              && !threwException) {
             throw new SparkException(s"Managed memory leak detected; size = $freedMemory bytes")
           } else {
             logError(s"Managed memory leak detected; size = $freedMemory bytes")