Revert "[SPARK-31387] Handle unknown operation/session ID in HiveThriftServer2Listener"

This reverts commit 6994c64.

Author: dongjoon-hyun

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

@@ -25,7 +25,6 @@ import scala.collection.mutable.ArrayBuffer
 import org.apache.hive.service.server.HiveServer2
 
 import org.apache.spark.{SparkConf, SparkContext}
-import org.apache.spark.internal.Logging
 import org.apache.spark.internal.config.Status.LIVE_ENTITY_UPDATE_PERIOD
 import org.apache.spark.scheduler._
 import org.apache.spark.sql.hive.thriftserver.HiveThriftServer2.ExecutionState
@@ -39,7 +38,7 @@ private[thriftserver] class HiveThriftServer2Listener(
     kvstore: ElementTrackingStore,
     sparkConf: SparkConf,
     server: Option[HiveServer2],
-    live: Boolean = true) extends SparkListener with Logging {
+    live: Boolean = true) extends SparkListener {
 
   private val sessionList = new ConcurrentHashMap[String, LiveSessionData]()
   private val executionList = new ConcurrentHashMap[String, LiveExecutionData]()
@@ -132,81 +131,60 @@ private[thriftserver] class HiveThriftServer2Listener(
     updateLiveStore(session)
   }
 
-  private def onSessionClosed(e: SparkListenerThriftServerSessionClosed): Unit =
-    Option(sessionList.get(e.sessionId)) match {
-      case None => logWarning(s"onSessionClosed called with unknown session id: ${e.sessionId}")
-      case Some(sessionData) =>
-        val session = sessionData
-        session.finishTimestamp = e.finishTime
-        updateStoreWithTriggerEnabled(session)
-        sessionList.remove(e.sessionId)
-    }
+  private def onSessionClosed(e: SparkListenerThriftServerSessionClosed): Unit = {
+    val session = sessionList.get(e.sessionId)
+    session.finishTimestamp = e.finishTime
+    updateStoreWithTriggerEnabled(session)
+    sessionList.remove(e.sessionId)
+  }
 
-  private def onOperationStart(e: SparkListenerThriftServerOperationStart): Unit =
-    Option(sessionList.get(e.sessionId)) match {
-      case None => logWarning(s"onOperationStart called with unknown session id: ${e.sessionId}")
-      case Some(sessionData) =>
-        val info = getOrCreateExecution(
-          e.id,
-          e.statement,
-          e.sessionId,
-          e.startTime,
-          e.userName)
-
-        info.state = ExecutionState.STARTED
-        executionList.put(e.id, info)
-        sessionData.totalExecution += 1
-        executionList.get(e.id).groupId = e.groupId
-        updateLiveStore(executionList.get(e.id))
-        updateLiveStore(sessionData)
-    }
+  private def onOperationStart(e: SparkListenerThriftServerOperationStart): Unit = {
+    val info = getOrCreateExecution(
+      e.id,
+      e.statement,
+      e.sessionId,
+      e.startTime,
+      e.userName)
+
+    info.state = ExecutionState.STARTED
+    executionList.put(e.id, info)
+    sessionList.get(e.sessionId).totalExecution += 1
+    executionList.get(e.id).groupId = e.groupId
+    updateLiveStore(executionList.get(e.id))
+    updateLiveStore(sessionList.get(e.sessionId))
+  }
 
-  private def onOperationParsed(e: SparkListenerThriftServerOperationParsed): Unit =
-    Option(executionList.get(e.id)) match {
-      case None => logWarning(s"onOperationParsed called with unknown operation id: ${e.id}")
-      case Some(executionData) =>
-        executionData.executePlan = e.executionPlan
-        executionData.state = ExecutionState.COMPILED
-        updateLiveStore(executionData)
-    }
+  private def onOperationParsed(e: SparkListenerThriftServerOperationParsed): Unit = {
+    executionList.get(e.id).executePlan = e.executionPlan
+    executionList.get(e.id).state = ExecutionState.COMPILED
+    updateLiveStore(executionList.get(e.id))
+  }
 
-  private def onOperationCanceled(e: SparkListenerThriftServerOperationCanceled): Unit =
-    Option(executionList.get(e.id)) match {
-      case None => logWarning(s"onOperationCanceled called with unknown operation id: ${e.id}")
-      case Some(executionData) =>
-        executionData.finishTimestamp = e.finishTime
-        executionData.state = ExecutionState.CANCELED
-        updateLiveStore(executionData)
-    }
+  private def onOperationCanceled(e: SparkListenerThriftServerOperationCanceled): Unit = {
+    executionList.get(e.id).finishTimestamp = e.finishTime
+    executionList.get(e.id).state = ExecutionState.CANCELED
+    updateLiveStore(executionList.get(e.id))
+  }
 
-  private def onOperationError(e: SparkListenerThriftServerOperationError): Unit =
-    Option(executionList.get(e.id)) match {
-      case None => logWarning(s"onOperationError called with unknown operation id: ${e.id}")
-      case Some(executionData) =>
-        executionData.finishTimestamp = e.finishTime
-        executionData.detail = e.errorMsg
-        executionData.state = ExecutionState.FAILED
-        updateLiveStore(executionData)
-    }
+  private def onOperationError(e: SparkListenerThriftServerOperationError): Unit = {
+    executionList.get(e.id).finishTimestamp = e.finishTime
+    executionList.get(e.id).detail = e.errorMsg
+    executionList.get(e.id).state = ExecutionState.FAILED
+    updateLiveStore(executionList.get(e.id))
+  }
 
-  private def onOperationFinished(e: SparkListenerThriftServerOperationFinish): Unit =
-    Option(executionList.get(e.id)) match {
-      case None => logWarning(s"onOperationFinished called with unknown operation id: ${e.id}")
-      case Some(executionData) =>
-        executionData.finishTimestamp = e.finishTime
-        executionData.state = ExecutionState.FINISHED
-        updateLiveStore(executionData)
-    }
+  private def onOperationFinished(e: SparkListenerThriftServerOperationFinish): Unit = {
+    executionList.get(e.id).finishTimestamp = e.finishTime
+    executionList.get(e.id).state = ExecutionState.FINISHED
+    updateLiveStore(executionList.get(e.id))
+  }
 
-  private def onOperationClosed(e: SparkListenerThriftServerOperationClosed): Unit =
-    Option(executionList.get(e.id)) match {
-      case None => logWarning(s"onOperationClosed called with unknown operation id: ${e.id}")
-      case Some(executionData) =>
-        executionData.closeTimestamp = e.closeTime
-        executionData.state = ExecutionState.CLOSED
-        updateStoreWithTriggerEnabled(executionData)
-        executionList.remove(e.id)
-    }
+  private def onOperationClosed(e: SparkListenerThriftServerOperationClosed): Unit = {
+    executionList.get(e.id).closeTimestamp = e.closeTime
+    executionList.get(e.id).state = ExecutionState.CLOSED
+    updateStoreWithTriggerEnabled(executionList.get(e.id))
+    executionList.remove(e.id)
+  }
 
   // Update both live and history stores. Trigger is enabled by default, hence
   // it will cleanup the entity which exceeds the threshold.

sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/HiveSessionImplSuite.scala

@@ -140,22 +140,6 @@ class HiveThriftServer2ListenerSuite extends SparkFunSuite with BeforeAndAfter {
     assert(listener.noLiveData())
   }
 
-  test("SPARK-31387 - listener update methods should not throw exception with unknown input") {
-    val (statusStore: HiveThriftServer2AppStatusStore,
-    listener: HiveThriftServer2Listener) = createAppStatusStore(true)
-    val unknownSession = "unknown_session"
-    val unknownOperation = "unknown_operation"
-    listener.onOtherEvent(SparkListenerThriftServerSessionClosed(unknownSession, 0))
-    listener.onOtherEvent(SparkListenerThriftServerOperationStart("id", unknownSession,
-      "stmt", "groupId", 0))
-    listener.onOtherEvent(SparkListenerThriftServerOperationParsed(unknownOperation, "query"))
-    listener.onOtherEvent(SparkListenerThriftServerOperationCanceled(unknownOperation, 0))
-    listener.onOtherEvent(SparkListenerThriftServerOperationError(unknownOperation,
-      "msg", "trace", 0))
-    listener.onOtherEvent(SparkListenerThriftServerOperationFinish(unknownOperation, 0))
-    listener.onOtherEvent(SparkListenerThriftServerOperationClosed(unknownOperation, 0))
-  }
-
   private def createProperties: Properties = {
     val properties = new Properties()
     properties.setProperty(SparkContext.SPARK_JOB_GROUP_ID, "groupId")

sql/hive-thriftserver/src/test/scala/org/apache/spark/sql/hive/thriftserver/ui/HiveThriftServer2ListenerSuite.scala

@@ -636,11 +636,7 @@ public void close() throws HiveSQLException {
       acquire(true);
       // Iterate through the opHandles and close their operations
       for (OperationHandle opHandle : opHandleSet) {
-        try {
-          operationManager.closeOperation(opHandle);
-        } catch (Exception e) {
-          LOG.warn("Exception is thrown closing operation " + opHandle, e);
-        }
+        operationManager.closeOperation(opHandle);
       }
       opHandleSet.clear();
       // Cleanup session log directory.

sql/hive-thriftserver/v1.2/src/main/java/org/apache/hive/service/cli/session/HiveSessionImpl.java

@@ -650,11 +650,7 @@ public void close() throws HiveSQLException {
       acquire(true);
       // Iterate through the opHandles and close their operations
       for (OperationHandle opHandle : opHandleSet) {
-        try {
-          operationManager.closeOperation(opHandle);
-        } catch (Exception e) {
-          LOG.warn("Exception is thrown closing operation " + opHandle, e);
-        }
+        operationManager.closeOperation(opHandle);
       }
       opHandleSet.clear();
       // Cleanup session log directory.