You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kyuubi.apache.org by fe...@apache.org on 2022/06/20 08:27:27 UTC

[incubator-kyuubi] branch master updated: [KYUUBI #2911] [SUB-TASK][KPIP-4] If the kyuubi instance unreachable, support to backfill state from resource manager and mark batch closed by remote kyuubi instance

This is an automated email from the ASF dual-hosted git repository.

feiwang pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-kyuubi.git


The following commit(s) were added to refs/heads/master by this push:
     new 089cf412f [KYUUBI #2911] [SUB-TASK][KPIP-4] If the kyuubi instance unreachable, support to backfill state from resource manager and mark batch closed by remote kyuubi instance
089cf412f is described below

commit 089cf412fd49b0635bf591fc15264846675acf1d
Author: Fei Wang <fw...@ebay.com>
AuthorDate: Mon Jun 20 16:27:19 2022 +0800

    [KYUUBI #2911] [SUB-TASK][KPIP-4] If the kyuubi instance unreachable, support to backfill state from resource manager and mark batch closed by remote kyuubi instance
    
    ### _Why are the changes needed?_
    
    Now for the batch state, we rely on the metadata store, which is only updated by the local kyuubi instance.
    
    But if the kyuubi instance is down or unreachable sometimes, the metadata state will not change and we can not get the correct state, it might cause the client side stuck if it is waiting the batch completion.
    
    So, if the kyuubi instance is not reachable, we need backfill the batch state according to resource manager batch application state it the batch was in RUNNING state.
    
    If the batch was stuck in PENDING state and the kyuubi instance can not get recovered in short-time, we need support to close the batch session and prevent the batch job re-submitted if it get recovered.
    
    In this pr, I introduce a new field `remoteClosed` to present whether it has been marked as CLOSED state by remote kyuubi instance.
    
    The `remoteClosed` should only be updated by remote kyuubi instance and set its value to true, so there is no race condition.
    
    - When opening the batch session in recovery mode, if `remoteClosed` is true, set the state to CANCELED directly
    - check the local batch sessions that is marked as remoteClosed periodically and close them.
    
    ### _How was this patch tested?_
    - [x] Add some test cases that check the changes thoroughly including negative and positive cases if possible
    
    - [ ] Add screenshots for manual tests if appropriate
    
    - [x] [Run test](https://kyuubi.apache.org/docs/latest/develop_tools/testing.html#running-tests) locally before make a pull request
    
    Closes #2911 from turboFei/across_kill_one.
    
    Closes #2911
    
    d582459c [Fei Wang] remoteClosed => peerInstanceClosed
    cd4207c2 [Fei Wang] unify the state
    ede40d0a [Fei Wang] Redirect the get batch request
    f4b86868 [Fei Wang] save
    5faee761 [Fei Wang] save
    
    Authored-by: Fei Wang <fw...@ebay.com>
    Signed-off-by: Fei Wang <fw...@ebay.com>
---
 .../org/apache/kyuubi/config/KyuubiConf.scala      |  12 ++
 .../sql/derby/metadata-store-schema-derby.sql      |   3 +-
 .../sql/mysql/metadata-store-schema-mysql.sql      |   1 +
 .../kyuubi/operation/BatchJobSubmission.scala      |  61 ++++---
 .../kyuubi/server/KyuubiRestFrontendService.scala  |  28 ++-
 .../kyuubi/server/api/v1/BatchesResource.scala     |  57 +++++-
 .../kyuubi/server/api/v1/InternalRestClient.scala  |   8 +-
 .../kyuubi/server/metadata/MetadataManager.scala   | 105 ++++++-----
 .../kyuubi/server/metadata/MetadataStore.scala     |  23 +--
 .../kyuubi/server/metadata/api/Metadata.scala      |   4 +-
 .../server/metadata/api/MetadataFilter.scala       |  33 ++++
 .../server/metadata/jdbc/JDBCMetadataStore.scala   |  47 ++---
 .../kyuubi/session/KyuubiSessionManager.scala      |  10 +
 .../metadata/jdbc/JDBCMetadataStoreSuite.scala     | 201 +++++++++------------
 14 files changed, 366 insertions(+), 227 deletions(-)

diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
index 0172c4cba..57b839fc9 100644
--- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
+++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala
@@ -984,6 +984,18 @@ object KyuubiConf {
       .timeConf
       .createWithDefault(Duration.ofSeconds(20).toMillis)
 
+  val BATCH_CHECK_INTERVAL: ConfigEntry[Long] =
+    buildConf("kyuubi.batch.check.interval")
+      .internal
+      .doc("The interval to check the batch session state. For batch session, it is not" +
+        " stateless, and some operations, such as close batch session, must be processed in the" +
+        " local kyuubi instance. But sometimes, the kyuubi instance might be unreachable, we" +
+        " need mark the batch session be CLOSED state in remote kyuubi instance. And the kyuubi" +
+        " instance should check whether there are local batch session are marked as CLOSED" +
+        " by remote kyuubi instance and close them periodically.")
+      .timeConf
+      .createWithDefault(Duration.ofSeconds(5).toMillis)
+
   val SERVER_EXEC_POOL_SIZE: ConfigEntry[Int] =
     buildConf("kyuubi.backend.server.exec.pool.size")
       .doc("Number of threads in the operation execution thread pool of Kyuubi server")
diff --git a/kyuubi-server/src/main/resources/sql/derby/metadata-store-schema-derby.sql b/kyuubi-server/src/main/resources/sql/derby/metadata-store-schema-derby.sql
index c1b84fda8..a4311fede 100644
--- a/kyuubi-server/src/main/resources/sql/derby/metadata-store-schema-derby.sql
+++ b/kyuubi-server/src/main/resources/sql/derby/metadata-store-schema-derby.sql
@@ -22,7 +22,8 @@ CREATE TABLE metadata(
     engine_url varchar(1024), -- the engine tracking url
     engine_state varchar(128), -- the engine application state
     engine_error clob, -- the engine application diagnose
-    end_time bigint  -- the metadata end time
+    end_time bigint,  -- the metadata end time
+    peer_instance_closed boolean default FALSE -- closed by peer kyuubi instance
 );
 
 CREATE INDEX metadata_kyuubi_instance_index ON metadata(kyuubi_instance);
diff --git a/kyuubi-server/src/main/resources/sql/mysql/metadata-store-schema-mysql.sql b/kyuubi-server/src/main/resources/sql/mysql/metadata-store-schema-mysql.sql
index aa8e2663c..43c0510d0 100644
--- a/kyuubi-server/src/main/resources/sql/mysql/metadata-store-schema-mysql.sql
+++ b/kyuubi-server/src/main/resources/sql/mysql/metadata-store-schema-mysql.sql
@@ -23,6 +23,7 @@ CREATE TABLE metadata(
     engine_state varchar(128) COMMENT 'the engine application state',
     engine_error mediumtext COMMENT 'the engine application diagnose',
     end_time bigint COMMENT 'the metadata end time',
+    peer_instance_closed boolean default '0' COMMENT 'closed by peer kyuubi instance',
     INDEX kyuubi_instance_index(kyuubi_instance),
     UNIQUE INDEX unique_identifier_index(identifier),
     INDEX user_name_index(user_name),
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala
index a81117d2e..9259fe04b 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/operation/BatchJobSubmission.scala
@@ -64,6 +64,7 @@ class BatchJobSubmission(
     batchArgs: Seq[String],
     recoveryMetadata: Option[Metadata])
   extends KyuubiOperation(OperationType.UNKNOWN_OPERATION, session) {
+  import BatchJobSubmission._
 
   override def statement: String = "BATCH_JOB_SUBMISSION"
 
@@ -153,25 +154,28 @@ class BatchJobSubmission(
 
   override protected def runInternal(): Unit = session.handleSessionException {
     val asyncOperation: Runnable = () => {
-      setStateIfNotCanceled(OperationState.RUNNING)
       try {
-        // If it is in recovery mode, only re-submit batch job if previous state is PENDING and
-        // fail to fetch the status including appId from resource manager. Otherwise, monitor the
-        // submitted batch application.
-        recoveryMetadata.map { metadata =>
-          if (metadata.state == OperationState.PENDING.toString) {
-            applicationStatus = currentApplicationState
-            applicationStatus.map(_.get(APP_ID_KEY)).map {
-              case Some(appId) => monitorBatchJob(appId)
-              case None => submitAndMonitorBatchJob()
+        if (recoveryMetadata.exists(_.peerInstanceClosed)) {
+          setState(OperationState.CANCELED)
+        } else {
+          // If it is in recovery mode, only re-submit batch job if previous state is PENDING and
+          // fail to fetch the status including appId from resource manager. Otherwise, monitor the
+          // submitted batch application.
+          recoveryMetadata.map { metadata =>
+            if (metadata.state == OperationState.PENDING.toString) {
+              applicationStatus = currentApplicationState
+              applicationStatus.map(_.get(APP_ID_KEY)).map {
+                case Some(appId) => monitorBatchJob(appId)
+                case None => submitAndMonitorBatchJob()
+              }
+            } else {
+              monitorBatchJob(metadata.engineId)
             }
-          } else {
-            monitorBatchJob(metadata.engineId)
+          }.getOrElse {
+            submitAndMonitorBatchJob()
           }
-        }.getOrElse {
-          submitAndMonitorBatchJob()
+          setStateIfNotCanceled(OperationState.FINISHED)
         }
-        setStateIfNotCanceled(OperationState.FINISHED)
       } catch {
         onError()
       } finally {
@@ -191,16 +195,6 @@ class BatchJobSubmission(
     }
   }
 
-  private def applicationFailed(applicationStatus: Option[Map[String, String]]): Boolean = {
-    applicationStatus.map(_.get(ApplicationOperation.APP_STATE_KEY)).exists(s =>
-      s.contains("KILLED") || s.contains("FAILED"))
-  }
-
-  private def applicationTerminated(applicationStatus: Option[Map[String, String]]): Boolean = {
-    applicationStatus.map(_.get(ApplicationOperation.APP_STATE_KEY)).exists(s =>
-      s.contains("KILLED") || s.contains("FAILED") || s.contains("FINISHED"))
-  }
-
   private def submitAndMonitorBatchJob(): Unit = {
     var appStatusFirstUpdated = false
     try {
@@ -209,6 +203,7 @@ class BatchJobSubmission(
       applicationStatus = currentApplicationState
       while (!applicationFailed(applicationStatus) && process.isAlive) {
         if (!appStatusFirstUpdated && applicationStatus.isDefined) {
+          setStateIfNotCanceled(OperationState.RUNNING)
           updateBatchMetadata()
           appStatusFirstUpdated = true
         }
@@ -240,12 +235,16 @@ class BatchJobSubmission(
     if (applicationStatus.isEmpty) {
       applicationStatus = currentApplicationState
     }
+    if (state == OperationState.PENDING) {
+      setStateIfNotCanceled(OperationState.RUNNING)
+    }
     if (applicationStatus.isEmpty) {
       info(s"The $batchType batch[$batchId] job: $appId not found, assume that it has finished.")
     } else if (applicationFailed(applicationStatus)) {
       throw new RuntimeException(s"$batchType batch[$batchId] job failed:" +
         applicationStatus.get.mkString(","))
     } else {
+      updateBatchMetadata()
       // TODO: add limit for max batch job submission lifetime
       while (applicationStatus.isDefined && !applicationTerminated(applicationStatus)) {
         Thread.sleep(applicationCheckInterval)
@@ -340,3 +339,15 @@ class BatchJobSubmission(
     throw new IllegalStateException("Use close instead.")
   }
 }
+
+object BatchJobSubmission {
+  def applicationFailed(applicationStatus: Option[Map[String, String]]): Boolean = {
+    applicationStatus.map(_.get(ApplicationOperation.APP_STATE_KEY)).exists(s =>
+      s.contains("KILLED") || s.contains("FAILED"))
+  }
+
+  def applicationTerminated(applicationStatus: Option[Map[String, String]]): Boolean = {
+    applicationStatus.map(_.get(ApplicationOperation.APP_STATE_KEY)).exists(s =>
+      s.contains("KILLED") || s.contains("FAILED") || s.contains("FINISHED"))
+  }
+}
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiRestFrontendService.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiRestFrontendService.scala
index 776097ca0..8da83a726 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiRestFrontendService.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/KyuubiRestFrontendService.scala
@@ -18,7 +18,7 @@
 package org.apache.kyuubi.server
 
 import java.util.EnumSet
-import java.util.concurrent.Future
+import java.util.concurrent.{Future, TimeUnit}
 import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger}
 import javax.servlet.DispatcherType
 
@@ -34,7 +34,7 @@ import org.apache.kyuubi.server.http.authentication.{AuthenticationFilter, Kyuub
 import org.apache.kyuubi.server.ui.JettyServer
 import org.apache.kyuubi.service.{AbstractFrontendService, Serverable, Service, ServiceUtils}
 import org.apache.kyuubi.service.authentication.KyuubiAuthenticationFactory
-import org.apache.kyuubi.session.KyuubiSessionManager
+import org.apache.kyuubi.session.{KyuubiSessionManager, SessionHandle}
 import org.apache.kyuubi.util.ThreadUtils
 
 /**
@@ -52,6 +52,8 @@ class KyuubiRestFrontendService(override val serverable: Serverable)
 
   private def sessionManager = be.sessionManager.asInstanceOf[KyuubiSessionManager]
 
+  private val batchChecker = ThreadUtils.newDaemonSingleThreadScheduledExecutor("batch-checker")
+
   override def initialize(conf: KyuubiConf): Unit = synchronized {
     val host = conf.get(FRONTEND_REST_BIND_HOST)
       .getOrElse(Utils.findLocalInetAddress.getHostAddress)
@@ -77,6 +79,26 @@ class KyuubiRestFrontendService(override val serverable: Serverable)
     server.addRedirectHandler("/docs", "/swagger")
   }
 
+  private def startBatchChecker(): Unit = {
+    val interval = conf.get(KyuubiConf.BATCH_CHECK_INTERVAL)
+    val task = new Runnable {
+      override def run(): Unit = {
+        try {
+          sessionManager.getPeerInstanceClosedBatchSessions(connectionUrl).foreach { batch =>
+            Utils.tryLogNonFatalError {
+              val sessionHandle = SessionHandle.fromUUID(batch.identifier)
+              Option(sessionManager.getBatchSessionImpl(sessionHandle)).foreach(_.close())
+            }
+          }
+        } catch {
+          case e: Throwable => error("Error checking batch sessions", e)
+        }
+      }
+    }
+
+    batchChecker.scheduleWithFixedDelay(task, interval, interval, TimeUnit.MILLISECONDS)
+  }
+
   @VisibleForTesting
   private[kyuubi] def recoverBatchSessions(): Unit = {
     val recoveryNumThreads = conf.get(METADATA_RECOVERY_THREADS)
@@ -123,6 +145,7 @@ class KyuubiRestFrontendService(override val serverable: Serverable)
         recoverBatchSessions()
         isStarted.set(true)
         info(s"$getName has started at ${server.getServerUri}")
+        startBatchChecker()
         startInternal()
       } catch {
         case e: Exception => throw new KyuubiException(s"Cannot start $getName", e)
@@ -133,6 +156,7 @@ class KyuubiRestFrontendService(override val serverable: Serverable)
   }
 
   override def stop(): Unit = synchronized {
+    ThreadUtils.shutdown(batchChecker)
     if (isStarted.getAndSet(false)) {
       server.stop()
     }
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala
index 2034be8d0..01c3f5616 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/BatchesResource.scala
@@ -33,10 +33,12 @@ import org.apache.kyuubi.{Logging, Utils}
 import org.apache.kyuubi.client.api.v1.dto._
 import org.apache.kyuubi.client.exception.KyuubiRestException
 import org.apache.kyuubi.config.KyuubiConf
-import org.apache.kyuubi.operation.{FetchOrientation, OperationState}
+import org.apache.kyuubi.operation.{BatchJobSubmission, FetchOrientation, OperationState}
 import org.apache.kyuubi.server.api.ApiRequestContext
 import org.apache.kyuubi.server.api.v1.BatchesResource._
 import org.apache.kyuubi.server.http.authentication.AuthenticationFilter
+import org.apache.kyuubi.server.metadata.MetadataManager
+import org.apache.kyuubi.server.metadata.api.Metadata
 import org.apache.kyuubi.service.authentication.KyuubiAuthenticationFactory
 import org.apache.kyuubi.session.{KyuubiBatchSessionImpl, KyuubiSessionManager, SessionHandle}
 
@@ -77,6 +79,34 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
       batchOpStatus.completed)
   }
 
+  private def buildBatch(
+      metadata: Metadata,
+      batchAppStatus: Option[Map[String, String]]): Batch = {
+    batchAppStatus.map { appStatus =>
+      val currentBatchState =
+        if (BatchJobSubmission.applicationFailed(batchAppStatus)) {
+          OperationState.ERROR.toString
+        } else if (BatchJobSubmission.applicationTerminated(batchAppStatus)) {
+          OperationState.FINISHED.toString
+        } else if (batchAppStatus.isDefined) {
+          OperationState.RUNNING.toString
+        } else {
+          metadata.state
+        }
+
+      new Batch(
+        metadata.identifier,
+        metadata.username,
+        metadata.engineType,
+        metadata.requestName,
+        appStatus.asJava,
+        metadata.kyuubiInstance,
+        currentBatchState,
+        metadata.createTime,
+        metadata.endTime)
+    }.getOrElse(MetadataManager.buildBatch(metadata))
+  }
+
   private def formatSessionHandle(sessionHandleStr: String): SessionHandle = {
     try {
       SessionHandle.fromUUID(sessionHandleStr)
@@ -122,11 +152,29 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
   @GET
   @Path("{batchId}")
   def batchInfo(@PathParam("batchId") batchId: String): Batch = {
+    val userName = fe.getUserName(Map.empty)
     val sessionHandle = formatSessionHandle(batchId)
     Option(sessionManager.getBatchSessionImpl(sessionHandle)).map { batchSession =>
       buildBatch(batchSession)
     }.getOrElse {
-      Option(sessionManager.getBatchFromMetadataStore(batchId)).getOrElse {
+      Option(sessionManager.getBatchMetadata(batchId)).map { metadata =>
+        if (OperationState.isTerminal(OperationState.withName(metadata.state)) ||
+          metadata.kyuubiInstance == fe.connectionUrl) {
+          MetadataManager.buildBatch(metadata)
+        } else {
+          val internalRestClient = getInternalRestClient(metadata.kyuubiInstance)
+          try {
+            internalRestClient.getBatch(userName, batchId)
+          } catch {
+            case e: KyuubiRestException =>
+              error(s"Error redirecting get batch[$batchId] to ${metadata.kyuubiInstance}", e)
+              val batchAppStatus = sessionManager.applicationManager.getApplicationInfo(
+                metadata.clusterManager,
+                batchId)
+              buildBatch(metadata, batchAppStatus)
+          }
+        }
+      }.getOrElse {
         error(s"Invalid batchId: $batchId")
         throw new NotFoundException(s"Invalid batchId: $batchId")
       }
@@ -264,6 +312,11 @@ private[v1] class BatchesResource extends ApiRequestContext with Logging {
               val appMgrKillResp = sessionManager.applicationManager.killApplication(
                 metadata.clusterManager,
                 batchId)
+              info(
+                s"Marking batch[$batchId/${metadata.kyuubiInstance}] closed by ${fe.connectionUrl}")
+              sessionManager.updateMetadata(Metadata(
+                identifier = batchId,
+                peerInstanceClosed = true))
               if (appMgrKillResp._1) {
                 new CloseBatchResponse(appMgrKillResp._1, appMgrKillResp._2)
               } else {
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/InternalRestClient.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/InternalRestClient.scala
index 433c020f4..8b8a61513 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/InternalRestClient.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/api/v1/InternalRestClient.scala
@@ -20,7 +20,7 @@ package org.apache.kyuubi.server.api.v1
 import java.util.Base64
 
 import org.apache.kyuubi.client.{BatchRestApi, KyuubiRestClient}
-import org.apache.kyuubi.client.api.v1.dto.{CloseBatchResponse, OperationLog}
+import org.apache.kyuubi.client.api.v1.dto.{Batch, CloseBatchResponse, OperationLog}
 import org.apache.kyuubi.client.auth.AuthHeaderGenerator
 import org.apache.kyuubi.server.http.authentication.AuthSchemes
 import org.apache.kyuubi.service.authentication.InternalSecurityAccessor
@@ -40,6 +40,12 @@ class InternalRestClient(kyuubiInstance: String, socketTimeout: Int, connectTime
 
   private val internalBatchRestApi = new BatchRestApi(initKyuubiRestClient())
 
+  def getBatch(user: String, batchId: String): Batch = {
+    withAuthUser(user) {
+      internalBatchRestApi.getBatchById(batchId)
+    }
+  }
+
   def getBatchLocalLog(user: String, batchId: String, from: Int, size: Int): OperationLog = {
     withAuthUser(user) {
       internalBatchRestApi.getBatchLocalLog(batchId, from, size)
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala
index dff9066a9..7e09ebf0a 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataManager.scala
@@ -27,12 +27,15 @@ import org.apache.kyuubi.client.api.v1.dto.Batch
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.config.KyuubiConf.METADATA_MAX_AGE
 import org.apache.kyuubi.engine.ApplicationOperation._
-import org.apache.kyuubi.server.metadata.api.Metadata
+import org.apache.kyuubi.operation.OperationState
+import org.apache.kyuubi.server.metadata.api.{Metadata, MetadataFilter}
 import org.apache.kyuubi.service.CompositeService
 import org.apache.kyuubi.session.SessionType
 import org.apache.kyuubi.util.{ClassUtils, ThreadUtils}
 
 class MetadataManager extends CompositeService("MetadataManager") {
+  import MetadataManager._
+
   private var _metadataStore: MetadataStore = _
 
   private val identifierRequestsRetryRefs =
@@ -103,17 +106,14 @@ class MetadataManager extends CompositeService("MetadataManager") {
       endTime: Long,
       from: Int,
       size: Int): Seq[Batch] = {
-    _metadataStore.getMetadataList(
-      SessionType.BATCH,
-      batchType,
-      batchUser,
-      batchState,
-      null,
-      createTime,
-      endTime,
-      from,
-      size,
-      true).map(buildBatch)
+    val filter = MetadataFilter(
+      sessionType = SessionType.BATCH,
+      engineType = batchType,
+      username = batchUser,
+      state = batchState,
+      createTime = createTime,
+      endTime = endTime)
+    _metadataStore.getMetadataList(filter, from, size, true).map(buildBatch)
   }
 
   def getBatchesRecoveryMetadata(
@@ -121,17 +121,24 @@ class MetadataManager extends CompositeService("MetadataManager") {
       kyuubiInstance: String,
       from: Int,
       size: Int): Seq[Metadata] = {
-    _metadataStore.getMetadataList(
-      SessionType.BATCH,
-      null,
-      null,
-      state,
-      kyuubiInstance,
-      0,
-      0,
-      from,
-      size,
-      false)
+    val filter = MetadataFilter(
+      sessionType = SessionType.BATCH,
+      state = state,
+      kyuubiInstance = kyuubiInstance)
+    _metadataStore.getMetadataList(filter, from, size, false)
+  }
+
+  def getPeerInstanceClosedBatchesMetadata(
+      state: String,
+      kyuubiInstance: String,
+      from: Int,
+      size: Int): Seq[Metadata] = {
+    val filter = MetadataFilter(
+      sessionType = SessionType.BATCH,
+      state = state,
+      kyuubiInstance = kyuubiInstance,
+      peerInstanceClosed = true)
+    _metadataStore.getMetadataList(filter, from, size, true)
   }
 
   def updateMetadata(metadata: Metadata, retryOnError: Boolean = true): Unit = {
@@ -148,28 +155,6 @@ class MetadataManager extends CompositeService("MetadataManager") {
     _metadataStore.cleanupMetadataByIdentifier(batchId)
   }
 
-  private def buildBatch(batchMetadata: Metadata): Batch = {
-    val batchAppInfo = Map(
-      APP_ID_KEY -> Option(batchMetadata.engineId),
-      APP_NAME_KEY -> Option(batchMetadata.engineName),
-      APP_STATE_KEY -> Option(batchMetadata.engineState),
-      APP_URL_KEY -> Option(batchMetadata.engineUrl),
-      APP_ERROR_KEY -> batchMetadata.engineError)
-      .filter(_._2.isDefined)
-      .map(info => (info._1, info._2.get))
-
-    new Batch(
-      batchMetadata.identifier,
-      batchMetadata.username,
-      batchMetadata.engineType,
-      batchMetadata.requestName,
-      batchAppInfo.asJava,
-      batchMetadata.kyuubiInstance,
-      batchMetadata.state,
-      batchMetadata.createTime,
-      batchMetadata.endTime)
-  }
-
   private def startMetadataCleaner(): Unit = {
     val cleanerEnabled = conf.get(KyuubiConf.METADATA_CLEANER_ENABLED)
     val stateMaxAge = conf.get(METADATA_MAX_AGE)
@@ -290,4 +275,34 @@ object MetadataManager extends Logging {
     }
     ClassUtils.createInstance(className, classOf[MetadataStore], conf)
   }
+
+  def buildBatch(batchMetadata: Metadata): Batch = {
+    val batchAppInfo = Map(
+      APP_ID_KEY -> Option(batchMetadata.engineId),
+      APP_NAME_KEY -> Option(batchMetadata.engineName),
+      APP_STATE_KEY -> Option(batchMetadata.engineState),
+      APP_URL_KEY -> Option(batchMetadata.engineUrl),
+      APP_ERROR_KEY -> batchMetadata.engineError)
+      .filter(_._2.isDefined)
+      .map(info => (info._1, info._2.get))
+
+    val batchState =
+      if (batchMetadata.peerInstanceClosed &&
+        !OperationState.isTerminal(OperationState.withName(batchMetadata.state))) {
+        OperationState.CANCELED.toString
+      } else {
+        batchMetadata.state
+      }
+
+    new Batch(
+      batchMetadata.identifier,
+      batchMetadata.username,
+      batchMetadata.engineType,
+      batchMetadata.requestName,
+      batchAppInfo.asJava,
+      batchMetadata.kyuubiInstance,
+      batchState,
+      batchMetadata.createTime,
+      batchMetadata.endTime)
+  }
 }
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataStore.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataStore.scala
index 07a549814..4416c4a6d 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataStore.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/MetadataStore.scala
@@ -19,8 +19,7 @@ package org.apache.kyuubi.server.metadata
 
 import java.io.Closeable
 
-import org.apache.kyuubi.server.metadata.api.Metadata
-import org.apache.kyuubi.session.SessionType.SessionType
+import org.apache.kyuubi.server.metadata.api.{Metadata, MetadataFilter}
 
 trait MetadataStore extends Closeable {
 
@@ -39,26 +38,14 @@ trait MetadataStore extends Closeable {
 
   /**
    * Get the metadata list with filter conditions, offset and size.
-   * @param sessionType the session type.
-   * @param engineType the engine type.
-   * @param userName the user name.
-   * @param state the state.
-   * @param kyuubiInstance the kyuubi instance.
-   * @param createTime the metadata create time.
-   * @param endTime the end time.
-   * @param from the batch offset.
-   * @param size the batch size to get.
+   * @param filter the metadata filter conditions.
+   * @param from the metadata offset.
+   * @param size the size to get.
    * @param stateOnly only return the state related column values.
    * @return selected metadata list.
    */
   def getMetadataList(
-      sessionType: SessionType,
-      engineType: String,
-      userName: String,
-      state: String,
-      kyuubiInstance: String,
-      createTime: Long,
-      endTime: Long,
+      filter: MetadataFilter,
       from: Int,
       size: Int,
       stateOnly: Boolean): Seq[Metadata]
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/Metadata.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/Metadata.scala
index c1003cb98..d800104d0 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/Metadata.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/Metadata.scala
@@ -47,6 +47,7 @@ import org.apache.kyuubi.session.SessionType.SessionType
  * @param engineState the engine state.
  * @param engineError the engine error diagnose.
  * @param endTime the end time.
+ * @param peerInstanceClosed closed by peer kyuubi instance.
  */
 case class Metadata(
     identifier: String,
@@ -69,4 +70,5 @@ case class Metadata(
     engineUrl: String = null,
     engineState: String = null,
     engineError: Option[String] = None,
-    endTime: Long = 0L)
+    endTime: Long = 0L,
+    peerInstanceClosed: Boolean = false)
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/MetadataFilter.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/MetadataFilter.scala
new file mode 100644
index 000000000..6213f8e64
--- /dev/null
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/api/MetadataFilter.scala
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.kyuubi.server.metadata.api
+
+import org.apache.kyuubi.session.SessionType.SessionType
+
+/**
+ * The conditions to filter the metadata.
+ */
+case class MetadataFilter(
+    sessionType: SessionType = null,
+    engineType: String = null,
+    username: String = null,
+    state: String = null,
+    kyuubiInstance: String = null,
+    createTime: Long = 0L,
+    endTime: Long = 0L,
+    peerInstanceClosed: Boolean = false)
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala
index 025d6316d..18b840ddb 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStore.scala
@@ -33,11 +33,10 @@ import org.apache.kyuubi.{KyuubiException, Logging, Utils}
 import org.apache.kyuubi.config.KyuubiConf
 import org.apache.kyuubi.operation.OperationState
 import org.apache.kyuubi.server.metadata.MetadataStore
-import org.apache.kyuubi.server.metadata.api.Metadata
+import org.apache.kyuubi.server.metadata.api.{Metadata, MetadataFilter}
 import org.apache.kyuubi.server.metadata.jdbc.DatabaseType._
 import org.apache.kyuubi.server.metadata.jdbc.JDBCMetadataStoreConf._
 import org.apache.kyuubi.session.SessionType
-import org.apache.kyuubi.session.SessionType.SessionType
 
 class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
   import JDBCMetadataStore._
@@ -169,13 +168,7 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
   }
 
   override def getMetadataList(
-      sessionType: SessionType,
-      engineType: String,
-      userName: String,
-      state: String,
-      kyuubiInstance: String,
-      createTime: Long,
-      endTime: Long,
+      filter: MetadataFilter,
       from: Int,
       size: Int,
       stateOnly: Boolean): Seq[Metadata] = {
@@ -187,34 +180,38 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
       queryBuilder.append(s"SELECT $METADATA_ALL_COLUMNS FROM $METADATA_TABLE")
     }
     val whereConditions = ListBuffer[String]()
-    Option(sessionType).foreach { _ =>
+    Option(filter.sessionType).foreach { sessionType =>
       whereConditions += " session_type = ?"
       params += sessionType.toString
     }
-    Option(engineType).filter(_.nonEmpty).foreach { _ =>
+    Option(filter.engineType).filter(_.nonEmpty).foreach { engineType =>
       whereConditions += " UPPER(engine_type) = ? "
       params += engineType.toUpperCase(Locale.ROOT)
     }
-    Option(userName).filter(_.nonEmpty).foreach { _ =>
+    Option(filter.username).filter(_.nonEmpty).foreach { username =>
       whereConditions += " user_name = ? "
-      params += userName
+      params += username
     }
-    Option(state).filter(_.nonEmpty).foreach { _ =>
+    Option(filter.state).filter(_.nonEmpty).foreach { state =>
       whereConditions += " state = ? "
       params += state.toUpperCase(Locale.ROOT)
     }
-    Option(kyuubiInstance).filter(_.nonEmpty).foreach { _ =>
+    Option(filter.kyuubiInstance).filter(_.nonEmpty).foreach { kyuubiInstance =>
       whereConditions += " kyuubi_instance = ? "
       params += kyuubiInstance
     }
-    if (createTime > 0) {
+    if (filter.createTime > 0) {
       whereConditions += " create_time >= ? "
-      params += createTime
+      params += filter.createTime
     }
-    if (endTime > 0) {
+    if (filter.endTime > 0) {
       whereConditions += " end_time > 0 "
       whereConditions += " end_time <= ? "
-      params += endTime
+      params += filter.endTime
+    }
+    if (filter.peerInstanceClosed) {
+      whereConditions += " peer_instance_closed = ? "
+      params += filter.peerInstanceClosed
     }
     if (whereConditions.nonEmpty) {
       queryBuilder.append(whereConditions.mkString(" WHERE ", " AND ", " "))
@@ -262,6 +259,10 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
       setClauses += " engine_error = ? "
       params += error
     }
+    if (metadata.peerInstanceClosed) {
+      setClauses += " peer_instance_closed = ? "
+      params += metadata.peerInstanceClosed
+    }
     if (setClauses.nonEmpty) {
       queryBuilder.append(setClauses.mkString(" SET ", " , ", " "))
     }
@@ -310,6 +311,7 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
         val engineState = resultSet.getString("engine_state")
         val engineError = Option(resultSet.getString("engine_error"))
         val endTime = resultSet.getLong("end_time")
+        val peerInstanceClosed = resultSet.getBoolean("peer_instance_closed")
 
         var resource: String = null
         var className: String = null
@@ -343,7 +345,8 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
           engineUrl = engineUrl,
           engineState = engineState,
           engineError = engineError,
-          endTime = endTime)
+          endTime = endTime,
+          peerInstanceClosed = peerInstanceClosed)
         metadataList += metadata
       }
       metadataList
@@ -403,6 +406,7 @@ class JDBCMetadataStore(conf: KyuubiConf) extends MetadataStore with Logging {
         case l: Long => statement.setLong(index + 1, l)
         case d: Double => statement.setDouble(index + 1, d)
         case f: Float => statement.setFloat(index + 1, f)
+        case b: Boolean => statement.setBoolean(index + 1, b)
         case _ => throw new KyuubiException(s"Unsupported param type ${param.getClass.getName}")
       }
     }
@@ -464,7 +468,8 @@ object JDBCMetadataStore {
     "engine_url",
     "engine_state",
     "engine_error",
-    "end_time").mkString(",")
+    "end_time",
+    "peer_instance_closed").mkString(",")
   private val METADATA_ALL_COLUMNS = Seq(
     METADATA_STATE_ONLY_COLUMNS,
     "resource",
diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
index ce11a0ad7..99265c6bc 100644
--- a/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
+++ b/kyuubi-server/src/main/scala/org/apache/kyuubi/session/KyuubiSessionManager.scala
@@ -242,6 +242,16 @@ class KyuubiSessionManager private (name: String) extends SessionManager(name) {
     }
   }
 
+  def getPeerInstanceClosedBatchSessions(kyuubiInstance: String): Seq[Metadata] = {
+    Seq(OperationState.PENDING, OperationState.RUNNING).flatMap { stateToKill =>
+      metadataManager.getPeerInstanceClosedBatchesMetadata(
+        stateToKill.toString,
+        kyuubiInstance,
+        0,
+        Int.MaxValue)
+    }
+  }
+
   override protected def isServer: Boolean = true
 
   private def initSessionLimiter(conf: KyuubiConf): Unit = {
diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala
index 7e81931f2..c7adf4a2c 100644
--- a/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala
+++ b/kyuubi-server/src/test/scala/org/apache/kyuubi/server/metadata/jdbc/JDBCMetadataStoreSuite.scala
@@ -24,7 +24,7 @@ import org.scalatest.time.SpanSugar._
 
 import org.apache.kyuubi.KyuubiFunSuite
 import org.apache.kyuubi.config.KyuubiConf
-import org.apache.kyuubi.server.metadata.api.Metadata
+import org.apache.kyuubi.server.metadata.api.{Metadata, MetadataFilter}
 import org.apache.kyuubi.server.metadata.jdbc.JDBCMetadataStoreConf._
 import org.apache.kyuubi.session.SessionType
 
@@ -40,13 +40,7 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite {
   override def afterAll(): Unit = {
     super.afterAll()
     jdbcMetadataStore.getMetadataList(
-      null,
-      null,
-      null,
-      null,
-      null,
-      0,
-      0,
+      MetadataFilter(),
       0,
       Int.MaxValue,
       true).foreach {
@@ -108,26 +102,19 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite {
 
     var batches =
       jdbcMetadataStore.getMetadataList(
-        SessionType.BATCH,
-        "Spark",
-        null,
-        null,
-        null,
-        0,
-        0,
+        MetadataFilter(
+          sessionType = SessionType.BATCH,
+          engineType = "Spark"),
         0,
         1,
         true)
     assert(batches == Seq(batchStateOnlyMetadata))
 
     batches = jdbcMetadataStore.getMetadataList(
-      SessionType.BATCH,
-      "spark",
-      "kyuubi",
-      null,
-      null,
-      0,
-      0,
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        engineType = "Spark",
+        username = "kyuubi"),
       0,
       Int.MaxValue,
       true)
@@ -135,109 +122,107 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite {
 
     jdbcMetadataStore.cleanupMetadataByIdentifier(batchState2.identifier)
 
-    batches =
-      jdbcMetadataStore.getMetadataList(
-        SessionType.SQL,
-        "SPARK",
-        "kyuubi",
-        "PENDING",
-        null,
-        0,
-        0,
-        0,
-        Int.MaxValue,
-        true)
+    batches = jdbcMetadataStore.getMetadataList(
+      MetadataFilter(
+        sessionType = SessionType.SQL,
+        engineType = "Spark",
+        username = "kyuubi",
+        state = "PENDING"),
+      0,
+      Int.MaxValue,
+      true)
     assert(batches.isEmpty)
 
-    batches =
-      jdbcMetadataStore.getMetadataList(
-        SessionType.BATCH,
-        "SPARK",
-        "kyuubi",
-        "PENDING",
-        null,
-        0,
-        0,
-        0,
-        Int.MaxValue,
-        true)
+    batches = jdbcMetadataStore.getMetadataList(
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        engineType = "Spark",
+        username = "kyuubi",
+        state = "PENDING"),
+      0,
+      Int.MaxValue,
+      true)
     assert(batches == Seq(batchStateOnlyMetadata))
 
-    batches =
-      jdbcMetadataStore.getMetadataList(
-        SessionType.BATCH,
-        "SPARK",
-        "kyuubi",
-        "RUNNING",
-        null,
-        0,
-        0,
-        0,
-        Int.MaxValue,
-        true)
+    batches = jdbcMetadataStore.getMetadataList(
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        engineType = "Spark",
+        username = "kyuubi",
+        state = "RUNNING"),
+      0,
+      Int.MaxValue,
+      true)
     assert(batches.isEmpty)
 
-    batches =
-      jdbcMetadataStore.getMetadataList(
-        SessionType.BATCH,
-        "SPARK",
-        "no_kyuubi",
-        "PENDING",
-        null,
-        0,
-        0,
-        0,
-        Int.MaxValue,
-        true)
+    batches = jdbcMetadataStore.getMetadataList(
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        engineType = "Spark",
+        username = "no_kyuubi",
+        state = "PENDING"),
+      0,
+      Int.MaxValue,
+      true)
     assert(batches.isEmpty)
 
     batches = jdbcMetadataStore.getMetadataList(
-      SessionType.BATCH,
-      "SPARK",
-      null,
-      "PENDING",
-      null,
-      0,
-      0,
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        engineType = "SPARK",
+        state = "PENDING"),
       0,
       Int.MaxValue,
       true)
     assert(batches == Seq(batchStateOnlyMetadata))
 
     batches = jdbcMetadataStore.getMetadataList(
-      SessionType.BATCH,
-      null,
-      null,
-      null,
-      null,
-      0,
-      0,
+      MetadataFilter(sessionType = SessionType.BATCH),
       0,
       Int.MaxValue,
       true)
     assert(batches == Seq(batchStateOnlyMetadata))
 
-    var batchesToRecover = jdbcMetadataStore.getMetadataList(
-      SessionType.BATCH,
-      null,
-      null,
-      "PENDING",
-      kyuubiInstance,
+    batches = jdbcMetadataStore.getMetadataList(
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        peerInstanceClosed = true),
       0,
+      Int.MaxValue,
+      true)
+    assert(batches.isEmpty)
+
+    jdbcMetadataStore.updateMetadata(Metadata(
+      identifier = batchStateOnlyMetadata.identifier,
+      peerInstanceClosed = true))
+
+    batchStateOnlyMetadata = batchStateOnlyMetadata.copy(peerInstanceClosed = true)
+    batchMetadata = batchMetadata.copy(peerInstanceClosed = true)
+
+    batches = jdbcMetadataStore.getMetadataList(
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        peerInstanceClosed = true),
       0,
+      Int.MaxValue,
+      true)
+    assert(batches === Seq(batchStateOnlyMetadata))
+
+    var batchesToRecover = jdbcMetadataStore.getMetadataList(
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        state = "PENDING",
+        kyuubiInstance = kyuubiInstance),
       0,
       Int.MaxValue,
       false)
     assert(batchesToRecover == Seq(batchMetadata))
 
     batchesToRecover = jdbcMetadataStore.getMetadataList(
-      SessionType.BATCH,
-      null,
-      null,
-      "RUNNING",
-      kyuubiInstance,
-      0,
-      0,
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        state = "RUNNING",
+        kyuubiInstance = kyuubiInstance),
       0,
       Int.MaxValue,
       false)
@@ -259,25 +244,19 @@ class JDBCMetadataStoreSuite extends KyuubiFunSuite {
     assert(jdbcMetadataStore.getMetadata(batchId, true) == newBatchState)
 
     assert(jdbcMetadataStore.getMetadataList(
-      SessionType.BATCH,
-      null,
-      null,
-      "PENDING",
-      kyuubiInstance,
-      0,
-      0,
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        state = "PENDING",
+        kyuubiInstance = kyuubiInstance),
       0,
       Int.MaxValue,
       false).isEmpty)
 
     assert(jdbcMetadataStore.getMetadataList(
-      SessionType.BATCH,
-      null,
-      null,
-      "RUNNING",
-      kyuubiInstance,
-      0,
-      0,
+      MetadataFilter(
+        sessionType = SessionType.BATCH,
+        state = "RUNNING",
+        kyuubiInstance = kyuubiInstance),
       0,
       Int.MaxValue,
       false).isEmpty)