[CARMEL-6632] Fix the running time of download statement in query log (#1275)

Author: fenzhu

sql/hive-thriftserver/src/main/scala/org/apache/spark/sql/hive/thriftserver/SparkDownloadDataOperation.scala

@@ -143,15 +143,6 @@ private[hive] class SparkDownloadDataOperation(
     if (getStatus.getState eq OperationState.FINISHED) {
       HiveThriftServer2.eventManager.onStatementFinish(statementId)
     }
-    HiveThriftServer2.listener.onQueryExit(
-      SparkListenerThriftServerQueryExit(
-        statementId,
-        QueryLogObjectList(Option(result).map(_.queryExecution)),
-        QueryLogExtInfo(
-          finalQueryStateId,
-          failDetails,
-          System.currentTimeMillis(),
-          false, totalDataSize)))
     logInfo(s"CLOSING $statementId")
     cleanup(OperationState.CLOSED)
     sqlContext.sparkContext.clearJobGroup()
@@ -324,6 +315,16 @@ private[hive] class SparkDownloadDataOperation(
         setOperationException(exception)
     } finally {
       TaskSchedulerImpl.activeSqls.decrementAndGet()
+      logInfo("Post data to query audit log for download operation")
+      HiveThriftServer2.listener.onQueryExit(
+        SparkListenerThriftServerQueryExit(
+          statementId,
+          QueryLogObjectList(Option(result).map(_.queryExecution)),
+          QueryLogExtInfo(
+            finalQueryStateId,
+            failDetails,
+            System.currentTimeMillis(),
+            false, totalDataSize)))
     }
   }
 
@@ -609,21 +610,17 @@ private[hive] class SparkDownloadDataOperation(
   }
 
   override def cancel(): Unit = {
-    failDetails = s"Canceling operation, stack trace:\n" +
-      Thread.currentThread().getStackTrace.mkString("\n")
-    if (statementId != null) {
-      HiveThriftServer2.eventManager.onStatementCanceled(statementId, failDetails)
+    synchronized {
+      if (!getStatus.getState.isTerminal) {
+        logInfo(s"Cancel download with $statementId")
+        failDetails = s"Canceling operation, stack trace:\n" +
+          Thread.currentThread().getStackTrace.mkString("\n")
+        if (statementId != null) {
+          HiveThriftServer2.eventManager.onStatementCanceled(statementId, failDetails)
+        }
+        cleanup(OperationState.CANCELED)
+      }
     }
-    HiveThriftServer2.listener.onQueryExit(
-      SparkListenerThriftServerQueryExit(
-        statementId,
-        QueryLogObjectList(Option(result).map(_.queryExecution)),
-        QueryLogExtInfo(
-          ExecutionState.CANCELED.id,
-          failDetails,
-          System.currentTimeMillis(),
-          false, totalDataSize)))
-    cleanup(OperationState.CANCELED)
   }
 
   private def withRetry[T](f: => T): T = {