[K8s] Clean up terminated driver pods on informer add

kyuubi-server/src/main/scala/org/apache/kyuubi/engine/KubernetesApplicationOperation.scala

@@ -83,45 +83,9 @@ class KubernetesApplicationOperation extends ApplicationOperation with Logging {
 
   private var cleanupCanceledAppPodExecutor: ThreadPoolExecutor = _
 
-  private var kubernetesClientInitializeCleanupTerminatedPodExecutor: ThreadPoolExecutor = _
-
   private def getOrCreateKubernetesClient(kubernetesInfo: KubernetesInfo): KubernetesClient = {
     checkKubernetesInfo(kubernetesInfo)
-    kubernetesClients.computeIfAbsent(
-      kubernetesInfo,
-      kInfo => {
-        val kubernetesClient = buildKubernetesClient(kInfo)
-        cleanTerminatedAppPodsOnKubernetesClientInitialize(kInfo, kubernetesClient)
-        kubernetesClient
-      })
-  }
-
-  private def cleanTerminatedAppPodsOnKubernetesClientInitialize(
-      kubernetesInfo: KubernetesInfo,
-      kubernetesClient: KubernetesClient): Unit = {
-    if (kubernetesClientInitializeCleanupTerminatedPodExecutor != null) {
-      kubernetesClientInitializeCleanupTerminatedPodExecutor.submit(new Runnable {
-        override def run(): Unit = {
-          val existingPods =
-            kubernetesClient.pods().withLabel(LABEL_KYUUBI_UNIQUE_KEY).list().getItems
-          info(s"[$kubernetesInfo] Found ${existingPods.size()} existing pods with label " +
-            s"$LABEL_KYUUBI_UNIQUE_KEY")
-          val eventType = KubernetesResourceEventTypes.UPDATE
-          existingPods.asScala.filter(isSparkEnginePod).foreach { pod =>
-            val appState = toApplicationState(pod, appStateSource, appStateContainer, eventType)
-            if (isTerminated(appState)) {
-              val kyuubiUniqueKey = pod.getMetadata.getLabels.get(LABEL_KYUUBI_UNIQUE_KEY)
-              info(s"[$kubernetesInfo] Found existing pod ${pod.getMetadata.getName} with " +
-                s"${toLabel(kyuubiUniqueKey)} in app state $appState, marking it as terminated")
-              if (appInfoStore.get(kyuubiUniqueKey) == null) {
-                updateApplicationState(kubernetesInfo, pod, eventType)
-              }
-              markApplicationTerminated(kubernetesInfo, pod, eventType)
-            }
-          }
-        }
-      })
-    }
+    kubernetesClients.computeIfAbsent(kubernetesInfo, kInfo => buildKubernetesClient(kInfo))
   }
 
   private var metadataManager: Option[MetadataManager] = _
@@ -210,9 +174,6 @@ class KubernetesApplicationOperation extends ApplicationOperation with Logging {
       TimeUnit.MILLISECONDS)
     cleanupCanceledAppPodExecutor = ThreadUtils.newDaemonCachedThreadPool(
       "cleanup-canceled-app-pod-thread")
-    kubernetesClientInitializeCleanupTerminatedPodExecutor =
-      ThreadUtils.newDaemonCachedThreadPool(
-        "kubernetes-client-initialize-cleanup-terminated-pod-thread")
     initializeKubernetesClient(kyuubiConf)
   }
 
@@ -373,11 +334,6 @@ class KubernetesApplicationOperation extends ApplicationOperation with Logging {
       ThreadUtils.shutdown(cleanupCanceledAppPodExecutor)
       cleanupCanceledAppPodExecutor = null
     }
-
-    if (kubernetesClientInitializeCleanupTerminatedPodExecutor != null) {
-      ThreadUtils.shutdown(kubernetesClientInitializeCleanupTerminatedPodExecutor)
-      kubernetesClientInitializeCleanupTerminatedPodExecutor = null
-    }
   }
 
   private class SparkEnginePodEventHandler(kubernetesInfo: KubernetesInfo)
@@ -387,6 +343,10 @@ class KubernetesApplicationOperation extends ApplicationOperation with Logging {
       if (isSparkEnginePod(pod)) {
         val eventType = KubernetesResourceEventTypes.ADD
         updateApplicationState(kubernetesInfo, pod, eventType)
+        val appState = toApplicationState(pod, appStateSource, appStateContainer, eventType)
+        if (isTerminated(appState)) {
+          markApplicationTerminated(kubernetesInfo, pod, eventType)
+        }
         KubernetesApplicationAuditLogger.audit(
           eventType,
           kubernetesInfo,

kyuubi-server/src/test/scala/org/apache/kyuubi/engine/KubernetesApplicationOperationSuite.scala

@@ -17,14 +17,63 @@
 
 package org.apache.kyuubi.engine
 
-import io.fabric8.kubernetes.api.model.{ContainerState, ContainerStateWaiting}
+import com.google.common.cache.Cache
+import io.fabric8.kubernetes.api.model.{ContainerState, ContainerStateWaiting, Pod, PodBuilder}
+import io.fabric8.kubernetes.client.informers.ResourceEventHandler
 
 import org.apache.kyuubi.{KyuubiException, KyuubiFunSuite}
 import org.apache.kyuubi.config.KyuubiConf
-import org.apache.kyuubi.engine.ApplicationState.{FAILED, PENDING}
+import org.apache.kyuubi.engine.ApplicationState.{ApplicationState, FAILED, FINISHED, PENDING}
+import org.apache.kyuubi.engine.KubernetesApplicationOperation.LABEL_KYUUBI_UNIQUE_KEY
 
 class KubernetesApplicationOperationSuite extends KyuubiFunSuite {
 
+  private def podEventHandler(
+      operation: KubernetesApplicationOperation,
+      kubernetesInfo: KubernetesInfo): ResourceEventHandler[Pod] = {
+    val handlerClass = classOf[KubernetesApplicationOperation].getDeclaredClasses
+      .find(_.getSimpleName == "SparkEnginePodEventHandler")
+      .getOrElse(fail("SparkEnginePodEventHandler not found"))
+    val constructor = handlerClass.getDeclaredConstructors.head
+    constructor.setAccessible(true)
+    constructor
+      .newInstance(operation, kubernetesInfo)
+      .asInstanceOf[ResourceEventHandler[Pod]]
+  }
+
+  private def cleanupTrigger(
+      operation: KubernetesApplicationOperation): Cache[String, ApplicationState] = {
+    val field = classOf[KubernetesApplicationOperation].getDeclaredFields
+      .find(_.getName.endsWith("cleanupTerminatedAppInfoTrigger"))
+      .getOrElse(fail("cleanupTerminatedAppInfoTrigger not found"))
+    field.setAccessible(true)
+    field.get(operation).asInstanceOf[Cache[String, ApplicationState]]
+  }
+
+  test("mark terminated application received from pod add event") {
+    val operation = new KubernetesApplicationOperation()
+    operation.initialize(KyuubiConf(), None)
+    val tag = "terminated-app"
+    val pod = new PodBuilder()
+      .withNewMetadata()
+      .withName("terminated-driver")
+      .addToLabels(LABEL_KYUUBI_UNIQUE_KEY, tag)
+      .addToLabels("spark-app-selector", "spark-application")
+      .endMetadata()
+      .withNewStatus()
+      .withPhase("Succeeded")
+      .withContainerStatuses()
+      .endStatus()
+      .build()
+
+    try {
+      podEventHandler(operation, KubernetesInfo()).onAdd(pod)
+      assert(cleanupTrigger(operation).getIfPresent(tag) === FINISHED)
+    } finally {
+      operation.stop()
+    }
+  }
+
   test("test check kubernetes info") {
     val kyuubiConf = KyuubiConf()
     kyuubiConf.set(KyuubiConf.KUBERNETES_CONTEXT_ALLOW_LIST.key, "1,2")