Fix a race.

Author: Marcelo Vanzin

sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLAppStatusListener.scala

@@ -131,15 +131,13 @@ private[sql] class SQLAppStatusListener(
       info.successful)
   }
 
-  def executionMetrics(executionId: Long): Map[Long, String] = {
+  def liveExecutionMetrics(executionId: Long): Option[Map[Long, String]] = {
     Option(liveExecutions.get(executionId)).map { exec =>
       if (exec.metricsValues != null) {
         exec.metricsValues
       } else {
         aggregateMetrics(exec)
       }
-    }.getOrElse {
-      throw new NoSuchElementException(s"execution $executionId not found")
     }
   }
 
@@ -207,7 +205,8 @@ private[sql] class SQLAppStatusListener(
         }
       }
 
-      // TODO: storing metrics by task ID can lead to innacurate metrics when speculation is on.
+      // TODO: storing metrics by task ID can cause metrics for the same task index to be
+      // counted multiple times, for example due to speculation or re-attempts.
       metrics.taskMetrics.put(taskId, new LiveTaskMetrics(ids, values, succeeded))
     }
   }

sql/core/src/main/scala/org/apache/spark/sql/execution/ui/SQLAppStatusStore.scala

@@ -58,9 +58,16 @@ private[sql] class SQLAppStatusStore(
   }
 
   def executionMetrics(executionId: Long): Map[Long, String] = {
-    val exec = store.read(classOf[SQLExecutionUIData], executionId)
-    Option(exec.metricValues)
-      .orElse(listener.map(_.executionMetrics(executionId)))
+    def metricsFromStore(): Option[Map[Long, String]] = {
+      val exec = store.read(classOf[SQLExecutionUIData], executionId)
+      Option(exec.metricValues)
+    }
+
+    metricsFromStore()
+      .orElse(listener.flatMap(_.liveExecutionMetrics(executionId)))
+      // Try a second time in case the execution finished while this method is trying to
+      // get the metrics.
+      .orElse(metricsFromStore())
       .getOrElse(Map())
   }