Update as per review

Author: zsxwing

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

@@ -1060,7 +1060,7 @@ class DAGScheduler(
         if (runningStages.contains(failedStage)) {
           logInfo(s"Marking $failedStage (${failedStage.name}) as failed " +
             s"due to a fetch failure from $mapStage (${mapStage.name})")
-          markStageAsFinished(failedStage, Some(failureMessage))
+          markStageAsFinished(failedStage, Some("Fetch failure: " + failureMessage))
           runningStages -= failedStage
         }
 

core/src/main/scala/org/apache/spark/shuffle/FetchFailedException.scala

@@ -35,11 +35,14 @@ private[spark] class FetchFailedException(
     message: String)
   extends Exception {
 
+  def this(bmAddress: BlockManagerId, shuffleId: Int, mapId: Int, reduceId: Int, e: Throwable) {
+    this(bmAddress, shuffleId, mapId, reduceId, Utils.exceptionString(e))
+  }
+
   override def getMessage: String =
     "Fetch failed: %s %d %d %d".format(bmAddress, shuffleId, mapId, reduceId)
 
-  def toTaskEndReason: TaskEndReason =
-    FetchFailed(bmAddress, shuffleId, mapId, reduceId, message)
+  def toTaskEndReason: TaskEndReason = FetchFailed(bmAddress, shuffleId, mapId, reduceId, message)
 }
 
 /**

core/src/main/scala/org/apache/spark/shuffle/hash/BlockStoreShuffleFetcher.scala

@@ -25,7 +25,7 @@ import org.apache.spark._
 import org.apache.spark.serializer.Serializer
 import org.apache.spark.shuffle.FetchFailedException
 import org.apache.spark.storage.{BlockId, BlockManagerId, ShuffleBlockFetcherIterator, ShuffleBlockId}
-import org.apache.spark.util.{Utils, CompletionIterator}
+import org.apache.spark.util.CompletionIterator
 
 private[hash] object BlockStoreShuffleFetcher extends Logging {
   def fetch[T](
@@ -64,8 +64,7 @@ private[hash] object BlockStoreShuffleFetcher extends Logging {
           blockId match {
             case ShuffleBlockId(shufId, mapId, _) =>
               val address = statuses(mapId.toInt)._1
-              throw new FetchFailedException(address, shufId.toInt, mapId.toInt, reduceId,
-                Utils.exceptionString(e))
+              throw new FetchFailedException(address, shufId.toInt, mapId.toInt, reduceId, e)
             case _ =>
               throw new SparkException(
                 "Failed to get block " + blockId + ", which is not a shuffle block")

core/src/main/scala/org/apache/spark/storage/ShuffleBlockFetcherIterator.scala

@@ -321,7 +321,7 @@ object ShuffleBlockFetcherIterator {
   /**
    * Result of a fetch from a remote block.
    */
-  trait FetchResult {
+  sealed trait FetchResult {
     val blockId: BlockId
   }
 
@@ -332,7 +332,7 @@ object ShuffleBlockFetcherIterator {
    *             Note that this is NOT the exact bytes.
    * @param buf [[ManagedBuffer]] for the content.
    */
-  case class SuccessFetchResult(blockId: BlockId, size: Long, buf: ManagedBuffer)
+  sealed case class SuccessFetchResult(blockId: BlockId, size: Long, buf: ManagedBuffer)
     extends FetchResult {
     require(buf != null)
     require(size >= 0)
@@ -343,6 +343,5 @@ object ShuffleBlockFetcherIterator {
    * @param blockId block id
    * @param e the failure exception
    */
-  case class FailureFetchResult(blockId: BlockId, e: Throwable) extends FetchResult {
-  }
+  sealed case class FailureFetchResult(blockId: BlockId, e: Throwable) extends FetchResult
 }

core/src/main/scala/org/apache/spark/util/JsonProtocol.scala

@@ -630,7 +630,7 @@ private[spark] object JsonProtocol {
         val reduceId = (json \ "Reduce ID").extract[Int]
         val message = Utils.jsonOption(json \ "Message").map(_.extract[String])
         new FetchFailed(blockManagerAddress, shuffleId, mapId, reduceId,
-          message.getOrElse("Unknown"))
+          message.getOrElse("Unknown reason"))
       case `exceptionFailure` =>
         val className = (json \ "Class Name").extract[String]
         val description = (json \ "Description").extract[String]

core/src/test/scala/org/apache/spark/util/JsonProtocolSuite.scala

@@ -107,7 +107,8 @@ class JsonProtocolSuite extends FunSuite {
     testJobResult(jobFailed)
 
     // TaskEndReason
-    val fetchFailed = FetchFailed(BlockManagerId("With or", "without you", 15), 17, 18, 19, "Some exception")
+    val fetchFailed = FetchFailed(BlockManagerId("With or", "without you", 15), 17, 18, 19,
+      "Some exception")
     val exceptionFailure = ExceptionFailure("To be", "or not to be", stackTrace, None)
     testTaskEndReason(Success)
     testTaskEndReason(Resubmitted)