diff --git a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/KubernetesApplicationOperation.scala b/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/KubernetesApplicationOperation.scala index 0160f9077ba..8989672cdaf 100644 --- a/kyuubi-server/src/main/scala/org/apache/kyuubi/engine/KubernetesApplicationOperation.scala +++ b/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, diff --git a/kyuubi-server/src/test/scala/org/apache/kyuubi/engine/KubernetesApplicationOperationSuite.scala b/kyuubi-server/src/test/scala/org/apache/kyuubi/engine/KubernetesApplicationOperationSuite.scala index 61f2cd82715..278d8a22526 100644 --- a/kyuubi-server/src/test/scala/org/apache/kyuubi/engine/KubernetesApplicationOperationSuite.scala +++ b/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")