diff --git a/build.sbt b/build.sbt index b1ab4b6db31..52b69c77977 100644 --- a/build.sbt +++ b/build.sbt @@ -164,7 +164,12 @@ lazy val ComputingUnitManagingService = (project in file("computing-unit-managin .configs(Test) .dependsOn(DAO % "test->test") // reuse MockTexeraDB embedded Postgres in tests .settings(commonModuleSettings) + .configs(Test) + .dependsOn(DAO % "test->test", Auth % "test->test") // reuse MockTexeraDB embedded Postgres in tests .settings( + // MockTexeraDB swaps a JVM-wide singleton (SqlServer's embedded Postgres), + // so run suites serially to avoid cross-suite races. + Test / parallelExecution := false, dependencyOverrides ++= Seq( // override it as io.dropwizard 4 require 2.16.1 or higher "com.fasterxml.jackson.module" %% "jackson-module-scala" % jacksonVersion, diff --git a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/ComputingUnitManagingResource.scala b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/ComputingUnitManagingResource.scala index aa02f73387e..df67fb3bfe0 100644 --- a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/ComputingUnitManagingResource.scala +++ b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/ComputingUnitManagingResource.scala @@ -50,8 +50,8 @@ import org.apache.texera.dao.jooq.generated.tables.daos.{ } import org.apache.texera.dao.jooq.generated.tables.pojos.WorkflowComputingUnit import org.apache.texera.service.resource.ComputingUnitManagingResource._ -import org.apache.texera.service.resource.ComputingUnitState._ import org.apache.texera.service.util.{ + ComputingUnitHelpers, ComputingUnitManagingServiceException, InsufficientComputingUnitQuota, KubernetesClient @@ -197,41 +197,6 @@ class ComputingUnitManagingResource { } } - private def getComputingUnitStatus(unit: WorkflowComputingUnit): ComputingUnitState = { - unit.getType match { - // ── Local CUs are always "running" ────────────────────────────── - case WorkflowComputingUnitTypeEnum.local => - Running - - // ── Kubernetes CUs – only explicit "Running" counts as running ─ - case WorkflowComputingUnitTypeEnum.kubernetes => - val phaseOpt = KubernetesClient - .getPodByName(KubernetesClient.generatePodName(unit.getCuid)) - .map(_.getStatus.getPhase) - - if (phaseOpt.contains("Running")) Running else Pending - - // ── Any other (unknown) type is treated as pending ────────────── - case _ => - Pending - } - } - - private def getComputingUnitMetrics(unit: WorkflowComputingUnit): WorkflowComputingUnitMetrics = { - unit.getType match { - case WorkflowComputingUnitTypeEnum.local => - WorkflowComputingUnitMetrics("NaN", "NaN") - case WorkflowComputingUnitTypeEnum.kubernetes => - val metrics = KubernetesClient.getPodMetrics(unit.getCuid) - WorkflowComputingUnitMetrics( - metrics.getOrElse("cpu", ""), - metrics.getOrElse("memory", "") - ) - case _ => - WorkflowComputingUnitMetrics("NaN", "NaN") - } - } - private def getComputingUnitResourceLimit( unit: WorkflowComputingUnit ): WorkflowComputingUnitResourceLimit = { @@ -476,8 +441,8 @@ class ComputingUnitManagingResource { DashboardWorkflowComputingUnit( insertedUnit, - getComputingUnitStatus(insertedUnit).toString, - getComputingUnitMetrics(insertedUnit), + ComputingUnitHelpers.getComputingUnitStatus(insertedUnit).toString, + ComputingUnitHelpers.getComputingUnitMetrics(insertedUnit), isOwner = true, accessPrivilege = PrivilegeEnum.WRITE, ownerGoogleAvatar, @@ -528,63 +493,46 @@ class ComputingUnitManagingResource { (List.empty[WorkflowComputingUnit], Map.empty[Integer, PrivilegeEnum]) } - val allUnits = ownedUnits ++ sharedUnits - val ownerUids: List[Integer] = allUnits.map(_.getUid).distinct val userDao = new UserDao(ctx.configuration()) - val ownerInfoMap: Map[Integer, (String, String)] = - userDao - .fetchByUid(ownerUids: _*) - .asScala - .map { u => - val avatar = Option(u.getGoogleAvatar).filter(_.nonEmpty).orNull - val name = Option(u.getName).filter(_.nonEmpty).orNull - u.getUid -> (avatar, name) - } - .toMap - - // If a Kubernetes pod has already disappeared (e.g., manually deleted or TTL - // GC-ed by the cluster), we treat the corresponding computing unit as - // terminated from the system's point of view. Here we eagerly update its - // terminateTime in the database **before** we build the response list so - // that subsequent API calls will no longer return this unit. - allUnits.foreach { unit => - if ( - unit.getType == WorkflowComputingUnitTypeEnum.kubernetes && - !KubernetesClient.podExists(unit.getCuid) - ) { - unit.setTerminateTime(new Timestamp(System.currentTimeMillis())) - computingUnitDao.update(unit) - } - } - // For shared units, we need to check the access privilege which are saved in different table - // to streamline the process, we combine owned units with default WRITE privilege and use sharedUnitInfo - // to get the privilege for shared units. - (ownedUnits.map(u => (u, PrivilegeEnum.WRITE)) ++ sharedUnits.map(u => - (u, sharedUnitInfo(u.getCuid)) - )) - .distinctBy { case (unit, _) => unit.getCuid } - .filter { case (unit, _) => unit.getTerminateTime == null } - .filter { - case (unit, _) => - unit.getType match { - case WorkflowComputingUnitTypeEnum.kubernetes => - KubernetesClient.podExists(unit.getCuid) - case _ => true - } - } - .map { - case (unit, privilege) => - DashboardWorkflowComputingUnit( - computingUnit = unit, - isOwner = unit.getUid.equals(uid), - accessPrivilege = privilege, - status = getComputingUnitStatus(unit).toString, - metrics = getComputingUnitMetrics(unit), - ownerGoogleAvatar = ownerInfoMap.getOrElse(unit.getUid, (null, null))._1, - ownerName = ownerInfoMap.getOrElse(unit.getUid, (null, null))._2 - ) - } + // Pair each unit with the caller's privilege (owned default to WRITE), one row per cuid, so + // a unit that is both owned and shared is reconciled/rendered exactly once. + val unitsWithPrivilege = + (ownedUnits.map(u => (u, PrivilegeEnum.WRITE)) ++ + sharedUnits.map(u => (u, sharedUnitInfo(u.getCuid)))) + .distinctBy { case (unit, _) => unit.getCuid } + .filter { case (unit, _) => unit.getTerminateTime == null } + val privilegeByCuid = unitsWithPrivilege.map { + case (unit, privilege) => unit.getCuid -> privilege + }.toMap + val candidateUnits = unitsWithPrivilege.map { case (unit, _) => unit } + + // Pod phases decide which Kubernetes units are still alive. + val podPhases = ComputingUnitHelpers.podPhasesFor(candidateUnits) + + val liveUnits = + ComputingUnitHelpers.reconcileVanishedKubernetesUnits( + computingUnitDao, + candidateUnits, + podPhases + ) + + // Metrics only for survivors, so fetch after reconciliation. + val podMetrics = ComputingUnitHelpers.podMetricsFor(liveUnits) + + val ownerInfoMap = + ComputingUnitHelpers.resolveOwnerInfo(userDao, liveUnits.map(_.getUid).distinct) + + liveUnits.map { unit => + ComputingUnitHelpers.buildDashboardUnit( + unit, + isOwner = unit.getUid.equals(uid), + accessPrivilege = privilegeByCuid(unit.getCuid), + ownerInfo = ownerInfoMap, + podPhases = podPhases, + podMetrics = podMetrics + ) + } } } @@ -613,8 +561,8 @@ class ComputingUnitManagingResource { DashboardWorkflowComputingUnit( computingUnit = unit, - status = getComputingUnitStatus(unit).toString, - metrics = getComputingUnitMetrics(unit), + status = ComputingUnitHelpers.getComputingUnitStatus(unit).toString, + metrics = ComputingUnitHelpers.getComputingUnitMetrics(unit), isOwner = unit.getUid.equals(user.getUid), accessPrivilege = { val cuAccessDao = new ComputingUnitUserAccessDao(context.configuration()) @@ -748,7 +696,7 @@ class ComputingUnitManagingResource { throw new BadRequestException("User has no access to the computing unit") } val computingUnit = getComputingUnitByCuid(context, cuid.toInt) - getComputingUnitMetrics(computingUnit) + ComputingUnitHelpers.getComputingUnitMetrics(computingUnit) } @GET diff --git a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/ComputingUnitHelpers.scala b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/ComputingUnitHelpers.scala index 8ae31bba7fb..ea7a65ef517 100644 --- a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/ComputingUnitHelpers.scala +++ b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/ComputingUnitHelpers.scala @@ -19,12 +19,56 @@ package org.apache.texera.service.util import org.apache.texera.dao.jooq.generated.enums.WorkflowComputingUnitTypeEnum +import org.apache.texera.dao.jooq.generated.tables.daos.{UserDao, WorkflowComputingUnitDao} import org.apache.texera.dao.jooq.generated.tables.pojos.WorkflowComputingUnit -import org.apache.texera.service.resource.ComputingUnitManagingResource.WorkflowComputingUnitMetrics +import org.apache.texera.service.resource.ComputingUnitManagingResource.{ + DashboardWorkflowComputingUnit, + WorkflowComputingUnitMetrics +} import org.apache.texera.service.resource.ComputingUnitState.{ComputingUnitState, Pending, Running} +import org.jooq.EnumType + +import java.sql.Timestamp +import scala.jdk.CollectionConverters.{CollectionHasAsScala, SeqHasAsJava} object ComputingUnitHelpers { - def getComputingUnitStatus(unit: WorkflowComputingUnit): ComputingUnitState = { + + /** + * Owner (avatar, name) keyed by uid, resolved in one query. Blank values collapse to `null`; + * empty `uids` returns empty without querying. + */ + def resolveOwnerInfo( + userDao: UserDao, + uids: Seq[Integer] + ): Map[Integer, (String, String)] = { + if (uids.isEmpty) Map.empty + else + userDao + .fetchByUid(uids: _*) + .asScala + .map { u => + val avatar = Option(u.getGoogleAvatar).filter(_.nonEmpty).orNull + val name = Option(u.getName).filter(_.nonEmpty).orNull + u.getUid -> (avatar, name) + } + .toMap + } + + def getComputingUnitStatus(unit: WorkflowComputingUnit): ComputingUnitState = + singleUnitStatus(unit, KubernetesClient) + + /** + * Single-unit status via a per-unit pod lookup (a targeted GET, cheaper than listing the whole + * namespace). The client is a by-name parameter — not the global singleton — so the kubernetes + * branch is unit-testable with a stub and the local/unknown branches never force the singleton; + * the public overload binds the production [[KubernetesClient]]. (Metrics has no analogous seam: + * its per-unit lookup already fans out to the whole namespace and the bulk (unit, podMetrics) + * overload already covers the cpu/memory resolution, so nothing there is worth pinning.) + */ + private[util] def singleUnitStatus( + unit: WorkflowComputingUnit, + k8s: => KubernetesClient + ): ComputingUnitState = { unit.getType match { // Local CUs are always “running” case WorkflowComputingUnitTypeEnum.local => @@ -32,9 +76,12 @@ object ComputingUnitHelpers { // Kubernetes CUs – only explicit “Running” counts as running case WorkflowComputingUnitTypeEnum.kubernetes => - val phaseOpt = KubernetesClient - .getPodByName(KubernetesClient.generatePodName(unit.getCuid)) - .map(_.getStatus.getPhase) + // Guard the pod status the same way the bulk getAllPodPhases does: a pod with no + // status yet has a null getStatus, so map through Option to avoid an NPE. + val client = k8s + val phaseOpt = client + .getPodByName(client.generatePodName(unit.getCuid)) + .flatMap(pod => Option(pod.getStatus).map(_.getPhase)) if (phaseOpt.contains("Running")) Running else Pending @@ -58,4 +105,128 @@ object ComputingUnitHelpers { WorkflowComputingUnitMetrics("NaN", "NaN") } } + + /** + * Resolves status from a pre-fetched pod-phase map instead of a per-unit cluster call, so a + * listing costs O(1) round trips rather than one per unit. + */ + def getComputingUnitStatus( + unit: WorkflowComputingUnit, + podPhases: Map[String, String] + ): ComputingUnitState = { + unit.getType match { + case WorkflowComputingUnitTypeEnum.local => + Running + case WorkflowComputingUnitTypeEnum.kubernetes => + // A missing entry or null phase both count as not-Running. + if (podPhases.get(KubernetesClient.generatePodName(unit.getCuid)).contains("Running")) + Running + else Pending + case _ => + Pending + } + } + + /** Resolves metrics from a pre-fetched pod-metrics map instead of a per-unit cluster call. */ + def getComputingUnitMetrics( + unit: WorkflowComputingUnit, + podMetrics: Map[String, Map[String, String]] + ): WorkflowComputingUnitMetrics = { + unit.getType match { + case WorkflowComputingUnitTypeEnum.local => + WorkflowComputingUnitMetrics("NaN", "NaN") + case WorkflowComputingUnitTypeEnum.kubernetes => + val metrics = podMetrics + .getOrElse(KubernetesClient.generatePodName(unit.getCuid), Map.empty[String, String]) + WorkflowComputingUnitMetrics( + metrics.getOrElse("cpu", ""), + metrics.getOrElse("memory", "") + ) + case _ => + WorkflowComputingUnitMetrics("NaN", "NaN") + } + } + + private def isKubernetes(unit: WorkflowComputingUnit): Boolean = + unit.getType match { + case WorkflowComputingUnitTypeEnum.kubernetes => true + case _ => false + } + + // Pod phases/metrics for the namespace; skipped (empty) when no Kubernetes unit is present, so a + // cluster-free listing issues no round trip. Same seam as singleUnitStatus: the public overload + // binds the production singleton, passing it by-name to the private[util] overload, which forces + // it only inside the guard's true branch — so the empty path never touches the client, and tests + // drive the private overload with a stub. + def podPhasesFor(units: Seq[WorkflowComputingUnit]): Map[String, String] = + podPhasesFor(units, KubernetesClient) + + private[util] def podPhasesFor( + units: Seq[WorkflowComputingUnit], + k8s: => KubernetesClient + ): Map[String, String] = + if (units.exists(isKubernetes)) k8s.getAllPodPhases else Map.empty + + def podMetricsFor(units: Seq[WorkflowComputingUnit]): Map[String, Map[String, String]] = + podMetricsFor(units, KubernetesClient) + + private[util] def podMetricsFor( + units: Seq[WorkflowComputingUnit], + k8s: => KubernetesClient + ): Map[String, Map[String, String]] = + if (units.exists(isKubernetes)) k8s.getAllPodMetrics else Map.empty + + /** A Kubernetes unit whose pod is absent from `podPhases` (deleted or TTL GC-ed). */ + private def isVanished(unit: WorkflowComputingUnit, podPhases: Map[String, String]): Boolean = + isKubernetes(unit) && !podPhases.contains(KubernetesClient.generatePodName(unit.getCuid)) + + /** Partition into `(live, vanished)` by `podPhases`. Pure (no I/O), so it is unit-testable. */ + def partitionLiveUnits( + units: List[WorkflowComputingUnit], + podPhases: Map[String, String] + ): (List[WorkflowComputingUnit], List[WorkflowComputingUnit]) = + units.partition(unit => !isVanished(unit, podPhases)) + + /** + * Stamp `terminateTime` on vanished Kubernetes units (one batched update) and return the live + * ones. Shared by both listing endpoints so they agree on when a unit is terminated. + */ + def reconcileVanishedKubernetesUnits( + dao: WorkflowComputingUnitDao, + units: List[WorkflowComputingUnit], + podPhases: Map[String, String] + ): List[WorkflowComputingUnit] = { + val partitioned = partitionLiveUnits(units, podPhases) + val vanished = partitioned._2 + if (vanished.nonEmpty) { + val now = new Timestamp(System.currentTimeMillis()) + vanished.foreach(_.setTerminateTime(now)) + dao.update(vanished.asJava) + } + partitioned._1 + } + + /** + * Build one dashboard row; status/metrics come from the pre-fetched maps (no per-unit K8s + * call). Shared by both listing endpoints so row shape and owner-info fallback stay identical. + */ + def buildDashboardUnit( + unit: WorkflowComputingUnit, + isOwner: Boolean, + accessPrivilege: EnumType, + ownerInfo: Map[Integer, (String, String)], + podPhases: Map[String, String], + podMetrics: Map[String, Map[String, String]] + ): DashboardWorkflowComputingUnit = { + val owner = ownerInfo.getOrElse(unit.getUid, (null, null)) + DashboardWorkflowComputingUnit( + computingUnit = unit, + status = getComputingUnitStatus(unit, podPhases).toString, + metrics = getComputingUnitMetrics(unit, podMetrics), + isOwner = isOwner, + accessPrivilege = accessPrivilege, + ownerGoogleAvatar = owner._1, + ownerName = owner._2 + ) + } } diff --git a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/KubernetesClient.scala b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/KubernetesClient.scala index 4f1d391cb30..6f97a5cf6b1 100644 --- a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/KubernetesClient.scala +++ b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/util/KubernetesClient.scala @@ -20,17 +20,20 @@ package org.apache.texera.service.util import io.fabric8.kubernetes.api.model._ -import io.fabric8.kubernetes.api.model.metrics.v1beta1.PodMetricsList +import io.fabric8.kubernetes.api.model.metrics.v1beta1.PodMetrics import io.fabric8.kubernetes.client.KubernetesClientBuilder import org.apache.texera.common.config.KubernetesConfig import scala.jdk.CollectionConverters._ -object KubernetesClient { +/** + * Thin wrapper over the fabric8 Kubernetes client. The production singleton is the companion + * object below, bound to a real in-cluster client. The fabric8 client is a constructor + * parameter (not a mutable global) so tests can construct an instance backed by a stubbed + * client and exercise the passthrough wrappers without a live cluster. + */ +class KubernetesClient(client: io.fabric8.kubernetes.client.KubernetesClient) { - // Initialize the Kubernetes client - private val client: io.fabric8.kubernetes.client.KubernetesClient = - new KubernetesClientBuilder().build() private val namespace: String = KubernetesConfig.computeUnitPoolNamespace private val podNamePrefix = "computing-unit" @@ -48,19 +51,51 @@ object KubernetesClient { Option(client.pods().inNamespace(namespace).withName(podName).get()) } + /** + * Phase of every pod in the namespace, keyed by pod name, in one call — so a bulk listing + * avoids a per-unit lookup. Unfiltered so callers can test a unit's presence by its pod-name + * key; a pod with no status yet maps to a `null` phase but still appears. + */ + def getAllPodPhases: Map[String, String] = + phasesByPodName(client.pods().inNamespace(namespace).list().getItems.asScala) + + /** Pure fabric8 -> map transform: a pod with no status yet maps to a `null` phase. */ + private[util] def phasesByPodName(pods: Iterable[Pod]): Map[String, String] = + pods + .map(pod => pod.getMetadata.getName -> Option(pod.getStatus).map(_.getPhase).orNull) + .toMap + + // Flatten a pod's per-container resource usage into a single metric -> value map. + private def containerUsage(podMetrics: PodMetrics): Map[String, String] = + podMetrics.getContainers.asScala.flatMap { container => + container.getUsage.asScala.map { + case (metric, value) => metric -> value.toString + } + }.toMap + + /** Pure fabric8 -> map transform over the raw per-pod metrics items. */ + private[util] def metricsByPodName( + items: Iterable[PodMetrics] + ): Map[String, Map[String, String]] = + items.map(podMetrics => podMetrics.getMetadata.getName -> containerUsage(podMetrics)).toMap + + // One namespace-wide metrics call, returning the raw per-pod items. + private def fetchPodMetricsItems(): Iterable[PodMetrics] = + client.top().pods().metrics(namespace).getItems.asScala + + /** + * CPU/memory of every pod in the namespace, keyed by pod name, in one call — the bulk + * counterpart to the single-unit lookup. + */ + def getAllPodMetrics: Map[String, Map[String, String]] = + metricsByPodName(fetchPodMetricsItems()) + def getPodMetrics(cuid: Int): Map[String, String] = { - val podMetricsList: PodMetricsList = client.top().pods().metrics(namespace) val targetPodName = generatePodName(cuid) - - podMetricsList.getItems.asScala + fetchPodMetricsItems() .collectFirst { case podMetrics if podMetrics.getMetadata.getName == targetPodName => - podMetrics.getContainers.asScala.flatMap { container => - container.getUsage.asScala.map { - case (metric, value) => - metric -> value.toString - } - }.toMap + containerUsage(podMetrics) } .getOrElse(Map.empty[String, String]) } @@ -182,3 +217,6 @@ object KubernetesClient { client.pods().inNamespace(namespace).withName(generatePodName(cuid)).delete() } } + +/** Production singleton bound to a real in-cluster fabric8 client. */ +object KubernetesClient extends KubernetesClient(new KubernetesClientBuilder().build()) diff --git a/computing-unit-managing-service/src/test/scala/org/apache/texera/service/resource/ComputingUnitManagingResourceSpec.scala b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/resource/ComputingUnitManagingResourceSpec.scala new file mode 100644 index 00000000000..54e495473d0 --- /dev/null +++ b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/resource/ComputingUnitManagingResourceSpec.scala @@ -0,0 +1,104 @@ +/* + * 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.texera.service.resource + +import org.apache.texera.auth.SessionUser +import org.apache.texera.dao.MockTexeraDB +import org.apache.texera.dao.jooq.generated.enums.{ + PrivilegeEnum, + UserRoleEnum, + WorkflowComputingUnitTypeEnum +} +import org.apache.texera.dao.jooq.generated.tables.daos.{UserDao, WorkflowComputingUnitDao} +import org.apache.texera.dao.jooq.generated.tables.pojos.{User, WorkflowComputingUnit} +import org.apache.texera.service.resource.ComputingUnitManagingResource.WorkflowComputingUnitMetrics +import org.scalatest.BeforeAndAfterAll +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +// Drives the per-user computing-unit endpoints against the embedded database using +// local units (so no Kubernetes calls are made). +class ComputingUnitManagingResourceSpec + extends AnyFlatSpec + with Matchers + with BeforeAndAfterAll + with MockTexeraDB { + + private val uid = 800 + private lazy val user: SessionUser = { + val u = new User() + u.setUid(uid) + u.setName("owner") + u.setEmail("owner@example.com") + u.setRole(UserRoleEnum.REGULAR) + u.setPassword("password") + u.setGoogleAvatar("owner-avatar") + new SessionUser(u) + } + + private def localUnit(cuid: Int, name: String): WorkflowComputingUnit = { + val unit = new WorkflowComputingUnit() + unit.setCuid(cuid) + unit.setUid(uid) + unit.setName(name) + unit.setType(WorkflowComputingUnitTypeEnum.local) + unit + } + + override protected def beforeAll(): Unit = { + super.beforeAll() + initializeDBAndReplaceDSLContext() + new UserDao(getDSLContext.configuration()).insert(user.getUser) + val unitDao = new WorkflowComputingUnitDao(getDSLContext.configuration()) + unitDao.insert(localUnit(800, "cu-a")) + unitDao.insert(localUnit(801, "cu-b")) + } + + override protected def afterAll(): Unit = + try shutdownDB() + finally super.afterAll() + + private val resource = new ComputingUnitManagingResource + + "getComputingUnitInfo" should "return the owner's local unit with WRITE access and Running status" in { + val info = resource.getComputingUnitInfo(800, user) + + info.computingUnit.getCuid shouldBe 800 + info.status shouldBe "Running" + info.metrics shouldBe WorkflowComputingUnitMetrics("NaN", "NaN") + info.isOwner shouldBe true + info.accessPrivilege shouldBe PrivilegeEnum.WRITE + info.ownerName shouldBe "owner" + } + + "getComputingUnitMetricsEndpoint" should "return NaN metrics for an owned local unit" in { + resource.getComputingUnitMetricsEndpoint("800", user) shouldBe + WorkflowComputingUnitMetrics("NaN", "NaN") + } + + "listComputingUnits" should "return the caller's owned, non-terminated units" in { + val result = resource.listComputingUnits(user) + + result.map(_.computingUnit.getCuid.intValue()) should contain theSameElementsAs Seq(800, 801) + all(result.map(_.isOwner)) shouldBe true + all(result.map(_.accessPrivilege)) shouldBe PrivilegeEnum.WRITE + all(result.map(_.status)) shouldBe "Running" + } +} diff --git a/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/ComputingUnitHelpersSpec.scala b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/ComputingUnitHelpersSpec.scala index 39a11ea2a54..cf54089a132 100644 --- a/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/ComputingUnitHelpersSpec.scala +++ b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/ComputingUnitHelpersSpec.scala @@ -19,25 +19,78 @@ package org.apache.texera.service.util -import org.apache.texera.dao.jooq.generated.enums.WorkflowComputingUnitTypeEnum -import org.apache.texera.dao.jooq.generated.tables.pojos.WorkflowComputingUnit +import io.fabric8.kubernetes.api.model.{Pod, PodBuilder} +import org.apache.texera.dao.MockTexeraDB +import org.apache.texera.dao.jooq.generated.enums.{ + PrivilegeEnum, + UserRoleEnum, + WorkflowComputingUnitTypeEnum +} +import org.apache.texera.dao.jooq.generated.tables.daos.{UserDao, WorkflowComputingUnitDao} +import org.apache.texera.dao.jooq.generated.tables.pojos.{User, WorkflowComputingUnit} import org.apache.texera.service.resource.ComputingUnitManagingResource.WorkflowComputingUnitMetrics import org.apache.texera.service.resource.ComputingUnitState.{Pending, Running} +import org.mockito.Mockito.{mock, when} +import org.scalatest.BeforeAndAfterAll import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers -class ComputingUnitHelpersSpec extends AnyFlatSpec with Matchers { +class ComputingUnitHelpersSpec + extends AnyFlatSpec + with Matchers + with BeforeAndAfterAll + with MockTexeraDB { + + override protected def beforeAll(): Unit = { + super.beforeAll() + initializeDBAndReplaceDSLContext() + } + + override protected def afterAll(): Unit = + try shutdownDB() + finally super.afterAll() + + private lazy val userDao = new UserDao(getDSLContext.configuration()) + private lazy val computingUnitDao = new WorkflowComputingUnitDao(getDSLContext.configuration()) - private def localUnit(): WorkflowComputingUnit = { + private def makeUser(uid: Int, name: String, email: String, avatar: String): User = { + val u = new User() + u.setUid(uid) + u.setName(name) + u.setEmail(email) + u.setRole(UserRoleEnum.REGULAR) + u.setPassword("password") + u.setGoogleAvatar(avatar) + u + } + + private def makeUnit( + cuid: Int, + uid: Int, + tpe: WorkflowComputingUnitTypeEnum + ): WorkflowComputingUnit = { val unit = new WorkflowComputingUnit() - unit.setType(WorkflowComputingUnitTypeEnum.local) + unit.setCuid(cuid) + unit.setUid(uid) + unit.setType(tpe) unit } - // WorkflowComputingUnitTypeEnum only defines `local` and `kubernetes`, so an - // untyped unit (getType == null) is what exercises the pure "unknown" branch. + private def localUnit(cuid: Int = 0, uid: Int = 0): WorkflowComputingUnit = + makeUnit(cuid, uid, WorkflowComputingUnitTypeEnum.local) + + private def kubernetesUnit(cuid: Int, uid: Int = 0): WorkflowComputingUnit = + makeUnit(cuid, uid, WorkflowComputingUnitTypeEnum.kubernetes) + + // A null-type unit (the enum has only local/kubernetes) exercises the "unknown" branch. private def untypedUnit(): WorkflowComputingUnit = new WorkflowComputingUnit() + private def podWithPhase(phase: String): Pod = + new PodBuilder().withNewStatus().withPhase(phase).endStatus().build() + + // A pod whose status has not been populated yet (getStatus == null). + private def statuslessPod(): Pod = new PodBuilder().build() + "getComputingUnitStatus" should "return Running for a local unit" in { ComputingUnitHelpers.getComputingUnitStatus(localUnit()) shouldBe Running } @@ -46,6 +99,29 @@ class ComputingUnitHelpersSpec extends AnyFlatSpec with Matchers { ComputingUnitHelpers.getComputingUnitStatus(untypedUnit()) shouldBe Pending } + // The kubernetes branch does a per-unit pod lookup, so singleUnitStatus is driven through a + // stubbed client (the public getComputingUnitStatus binds the production singleton). + "singleUnitStatus" should "return Running for a kubernetes unit whose pod phase is Running" in { + val k8s = mock(classOf[KubernetesClient]) + when(k8s.generatePodName(40)).thenReturn("computing-unit-40") + when(k8s.getPodByName("computing-unit-40")).thenReturn(Some(podWithPhase("Running"))) + ComputingUnitHelpers.singleUnitStatus(kubernetesUnit(40), k8s) shouldBe Running + } + + it should "return Pending for a kubernetes unit whose pod has no status yet" in { + val k8s = mock(classOf[KubernetesClient]) + when(k8s.generatePodName(41)).thenReturn("computing-unit-41") + when(k8s.getPodByName("computing-unit-41")).thenReturn(Some(statuslessPod())) + ComputingUnitHelpers.singleUnitStatus(kubernetesUnit(41), k8s) shouldBe Pending + } + + it should "return Pending for a kubernetes unit whose pod is absent" in { + val k8s = mock(classOf[KubernetesClient]) + when(k8s.generatePodName(42)).thenReturn("computing-unit-42") + when(k8s.getPodByName("computing-unit-42")).thenReturn(None) + ComputingUnitHelpers.singleUnitStatus(kubernetesUnit(42), k8s) shouldBe Pending + } + "getComputingUnitMetrics" should "return NaN metrics for a local unit" in { ComputingUnitHelpers.getComputingUnitMetrics(localUnit()) shouldBe WorkflowComputingUnitMetrics("NaN", "NaN") @@ -55,4 +131,194 @@ class ComputingUnitHelpersSpec extends AnyFlatSpec with Matchers { ComputingUnitHelpers.getComputingUnitMetrics(untypedUnit()) shouldBe WorkflowComputingUnitMetrics("NaN", "NaN") } + + // ── Bulk variants resolving from pre-fetched pod maps ──────────────── + + "getComputingUnitStatus(unit, podPhases)" should "return Running for a local unit" in { + ComputingUnitHelpers.getComputingUnitStatus(localUnit(), Map.empty) shouldBe Running + } + + it should "return Running for a kubernetes unit whose pod phase is Running" in { + val unit = kubernetesUnit(7) + val podPhases = Map(KubernetesClient.generatePodName(7) -> "Running") + ComputingUnitHelpers.getComputingUnitStatus(unit, podPhases) shouldBe Running + } + + it should "return Pending for a kubernetes unit whose pod is absent or not Running" in { + val unit = kubernetesUnit(8) + ComputingUnitHelpers.getComputingUnitStatus(unit, Map.empty) shouldBe Pending + ComputingUnitHelpers.getComputingUnitStatus( + unit, + Map(KubernetesClient.generatePodName(8) -> "Pending") + ) shouldBe Pending + } + + it should "treat a null phase as not Running" in { + val unit = kubernetesUnit(9) + val podPhases = Map(KubernetesClient.generatePodName(9) -> (null: String)) + ComputingUnitHelpers.getComputingUnitStatus(unit, podPhases) shouldBe Pending + } + + "getComputingUnitMetrics(unit, podMetrics)" should "return NaN metrics for a local unit" in { + ComputingUnitHelpers.getComputingUnitMetrics(localUnit(), Map.empty) shouldBe + WorkflowComputingUnitMetrics("NaN", "NaN") + } + + it should "resolve cpu/memory for a kubernetes unit from the map" in { + val unit = kubernetesUnit(10) + val podMetrics = Map( + KubernetesClient.generatePodName(10) -> Map("cpu" -> "500m", "memory" -> "256Mi") + ) + ComputingUnitHelpers.getComputingUnitMetrics(unit, podMetrics) shouldBe + WorkflowComputingUnitMetrics("500m", "256Mi") + } + + it should "return empty cpu/memory for a kubernetes unit absent from the map" in { + ComputingUnitHelpers.getComputingUnitMetrics(kubernetesUnit(11), Map.empty) shouldBe + WorkflowComputingUnitMetrics("", "") + } + + // ── partitionLiveUnits ─────────────────────────────────────────────── + + "partitionLiveUnits" should "treat local units as always live" in { + val units = List(localUnit(cuid = 1), localUnit(cuid = 2)) + val (live, vanished) = ComputingUnitHelpers.partitionLiveUnits(units, Map.empty) + live.map(_.getCuid) shouldBe List(1, 2) + vanished shouldBe empty + } + + it should "classify a kubernetes unit as live iff its pod is present in the map" in { + val present = kubernetesUnit(20) + val gone = kubernetesUnit(21) + val podPhases = Map(KubernetesClient.generatePodName(20) -> "Running") + + val (live, vanished) = ComputingUnitHelpers.partitionLiveUnits(List(present, gone), podPhases) + + live.map(_.getCuid) shouldBe List(20) + vanished.map(_.getCuid) shouldBe List(21) + } + + it should "treat an untyped (null-type) unit as live (never kubernetes)" in { + val (live, vanished) = ComputingUnitHelpers.partitionLiveUnits(List(untypedUnit()), Map.empty) + live should have size 1 + vanished shouldBe empty + } + + // ── buildDashboardUnit ─────────────────────────────────────────────── + + "buildDashboardUnit" should "populate the row from the caller flags and pre-fetched maps" in { + val unit = kubernetesUnit(cuid = 30, uid = 100) + val podName = KubernetesClient.generatePodName(30) + + val row = ComputingUnitHelpers.buildDashboardUnit( + unit, + isOwner = true, + accessPrivilege = PrivilegeEnum.READ, + ownerInfo = Map((100: Integer) -> ("avatar", "owner")), + podPhases = Map(podName -> "Running"), + podMetrics = Map(podName -> Map("cpu" -> "100m", "memory" -> "64Mi")) + ) + + row.computingUnit.getCuid shouldBe 30 + row.isOwner shouldBe true + row.accessPrivilege shouldBe PrivilegeEnum.READ + row.status shouldBe "Running" + row.metrics shouldBe WorkflowComputingUnitMetrics("100m", "64Mi") + row.ownerGoogleAvatar shouldBe "avatar" + row.ownerName shouldBe "owner" + } + + it should "fall back to null owner info when the owner is missing from the map" in { + val row = ComputingUnitHelpers.buildDashboardUnit( + localUnit(cuid = 31, uid = 200), + isOwner = false, + accessPrivilege = PrivilegeEnum.WRITE, + ownerInfo = Map.empty, + podPhases = Map.empty, + podMetrics = Map.empty + ) + + row.ownerGoogleAvatar shouldBe null + row.ownerName shouldBe null + row.status shouldBe "Running" + row.metrics shouldBe WorkflowComputingUnitMetrics("NaN", "NaN") + } + + // ── Bulk variants: unknown (untyped) branch ────────────────────────── + + "getComputingUnitStatus(unit, podPhases)" should "return Pending for an unknown (untyped) unit" in { + ComputingUnitHelpers.getComputingUnitStatus(untypedUnit(), Map.empty) shouldBe Pending + } + + "getComputingUnitMetrics(unit, podMetrics)" should "return NaN for an unknown (untyped) unit" in { + ComputingUnitHelpers.getComputingUnitMetrics(untypedUnit(), Map.empty) shouldBe + WorkflowComputingUnitMetrics("NaN", "NaN") + } + + // ── podPhasesFor / podMetricsFor guards ────────────────────────────── + + "podPhasesFor" should "return empty (issuing no cluster call) when no kubernetes unit is present" in { + ComputingUnitHelpers.podPhasesFor(List(localUnit(), untypedUnit())) shouldBe empty + } + + it should "fetch all pod phases once when a kubernetes unit is present" in { + val k8s = mock(classOf[KubernetesClient]) + val phases = Map("computing-unit-50" -> "Running") + when(k8s.getAllPodPhases).thenReturn(phases) + ComputingUnitHelpers.podPhasesFor(List(kubernetesUnit(50)), k8s) shouldBe phases + } + + "podMetricsFor" should "return empty (issuing no cluster call) when no kubernetes unit is present" in { + ComputingUnitHelpers.podMetricsFor(List(localUnit(), untypedUnit())) shouldBe empty + } + + it should "fetch all pod metrics once when a kubernetes unit is present" in { + val k8s = mock(classOf[KubernetesClient]) + val metrics = Map("computing-unit-51" -> Map("cpu" -> "100m", "memory" -> "64Mi")) + when(k8s.getAllPodMetrics).thenReturn(metrics) + ComputingUnitHelpers.podMetricsFor(List(kubernetesUnit(51)), k8s) shouldBe metrics + } + + // ── resolveOwnerInfo (backed by the embedded database) ─────────────── + + "resolveOwnerInfo" should "resolve avatar/name and collapse blank values to null" in { + userDao.insert(makeUser(500, "alice", "alice@example.com", "alice-avatar")) + userDao.insert(makeUser(501, "", "bob@example.com", "")) + + val info = ComputingUnitHelpers.resolveOwnerInfo(userDao, Seq[Integer](500, 501)) + info(500) shouldBe (("alice-avatar", "alice")) + info(501) shouldBe ((null, null)) + } + + it should "return an empty map (and issue no query) for no uids" in { + ComputingUnitHelpers.resolveOwnerInfo(userDao, Seq.empty) shouldBe empty + } + + // ── reconcileVanishedKubernetesUnits (backed by the embedded database) ─ + + "reconcileVanishedKubernetesUnits" should "terminate vanished kubernetes units and return the live ones" in { + userDao.insert(makeUser(600, "carol", "carol@example.com", null)) + + val present = kubernetesUnit(600, 600) + present.setName("present") + val gone = kubernetesUnit(601, 600) + gone.setName("gone") + val local = localUnit(602, 600) + local.setName("local") + Seq(present, gone, local).foreach(computingUnitDao.insert(_)) + + // Only the pod for cuid 600 exists; cuid 601's pod has vanished. + val podPhases = Map(KubernetesClient.generatePodName(600) -> "Running") + val live = + ComputingUnitHelpers.reconcileVanishedKubernetesUnits( + computingUnitDao, + List(present, gone, local), + podPhases + ) + + live.map(_.getCuid) should contain theSameElementsAs Seq(600, 602) + computingUnitDao.fetchOneByCuid(601).getTerminateTime should not be null + computingUnitDao.fetchOneByCuid(600).getTerminateTime shouldBe null + computingUnitDao.fetchOneByCuid(602).getTerminateTime shouldBe null + } } diff --git a/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/KubernetesClientSpec.scala b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/KubernetesClientSpec.scala index 8ce3b107ae4..f0d96347a91 100644 --- a/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/KubernetesClientSpec.scala +++ b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/util/KubernetesClientSpec.scala @@ -19,11 +19,97 @@ package org.apache.texera.service.util +import io.fabric8.kubernetes.api.model.metrics.v1beta1.{ + ContainerMetricsBuilder, + PodMetrics, + PodMetricsBuilder, + PodMetricsList, + PodMetricsListBuilder +} +import io.fabric8.kubernetes.api.model.{Pod, PodBuilder, PodList, PodListBuilder, Quantity} +import io.fabric8.kubernetes.client.dsl.{ + MetricAPIGroupDSL, + MixedOperation, + NonNamespaceOperation, + PodMetricOperation, + PodResource +} +import io.fabric8.kubernetes.client.{KubernetesClient => Fabric8Client} +import org.apache.texera.common.config.KubernetesConfig +import org.mockito.Mockito.{mock, when} import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers +import scala.jdk.CollectionConverters._ + +// Two layers are exercised here: +// * the pure fabric8 -> map transforms (phasesByPodName / metricsByPodName) with +// builder-constructed model objects, so the transform logic needs no client, and +// * the thin namespace-wide wrappers (getAllPodPhases / getAllPodMetrics / getPodMetrics), +// which are driven through a freshly constructed KubernetesClient whose fabric8 client is a +// Mockito stub — no live cluster and no mutable global. +// The status/metrics *decision* logic that consumes these maps (Running vs Pending, cpu/memory +// resolution) is covered by ComputingUnitHelpersSpec. class KubernetesClientSpec extends AnyFlatSpec with Matchers { + private val namespace: String = KubernetesConfig.computeUnitPoolNamespace + + // A fabric8 client stubbed just enough to answer the namespace-wide pod-list and pod-metrics + // calls the wrappers make. RETURNS_DEEP_STUBS can't be used: fabric8's fluent API returns type + // variables, so each step of the chain is mocked explicitly. + private def stubbedClient(pods: Seq[Pod], metrics: Seq[PodMetrics]): Fabric8Client = { + val client = mock(classOf[Fabric8Client]) + + val podsMixed = mock(classOf[MixedOperation[_, _, _]]) + .asInstanceOf[MixedOperation[Pod, PodList, PodResource]] + val podsInNamespace = mock(classOf[NonNamespaceOperation[_, _, _]]) + .asInstanceOf[NonNamespaceOperation[Pod, PodList, PodResource]] + when(client.pods()).thenReturn(podsMixed) + when(podsMixed.inNamespace(namespace)).thenReturn(podsInNamespace) + when(podsInNamespace.list()).thenReturn(new PodListBuilder().addAllToItems(pods.asJava).build()) + + val top = mock(classOf[MetricAPIGroupDSL]) + val podMetricOp = mock(classOf[PodMetricOperation]) + val metricsList: PodMetricsList = + new PodMetricsListBuilder().addAllToItems(metrics.asJava).build() + when(client.top()).thenReturn(top) + when(top.pods()).thenReturn(podMetricOp) + when(podMetricOp.metrics(namespace)).thenReturn(metricsList) + + client + } + + private def pod(cuid: Int, phase: String): Pod = + new PodBuilder() + .withNewMetadata() + .withName(KubernetesClient.generatePodName(cuid)) + .endMetadata() + .withNewStatus() + .withPhase(phase) + .endStatus() + .build() + + // A pod whose status has not been populated yet (getStatus == null). + private def statuslessPod(cuid: Int): Pod = + new PodBuilder() + .withNewMetadata() + .withName(KubernetesClient.generatePodName(cuid)) + .endMetadata() + .build() + + private def podMetrics(cuid: Int, cpu: String, memory: String): PodMetrics = + new PodMetricsBuilder() + .withNewMetadata() + .withName(KubernetesClient.generatePodName(cuid)) + .endMetadata() + .addToContainers( + new ContainerMetricsBuilder() + .withName("main") + .withUsage(Map("cpu" -> new Quantity(cpu), "memory" -> new Quantity(memory)).asJava) + .build() + ) + .build() + "generatePodName" should "prefix the cuid with computing-unit" in { KubernetesClient.generatePodName(42) shouldBe "computing-unit-42" } @@ -31,4 +117,47 @@ class KubernetesClientSpec extends AnyFlatSpec with Matchers { it should "handle a cuid of 0" in { KubernetesClient.generatePodName(0) shouldBe "computing-unit-0" } + + "phasesByPodName" should "map every pod name to its phase" in { + val phases = KubernetesClient.phasesByPodName(Seq(pod(1, "Running"), pod(2, "Pending"))) + phases(KubernetesClient.generatePodName(1)) shouldBe "Running" + phases(KubernetesClient.generatePodName(2)) shouldBe "Pending" + } + + it should "map a pod with no status to a null phase but still include it" in { + val phases = KubernetesClient.phasesByPodName(Seq(statuslessPod(3))) + phases should contain key KubernetesClient.generatePodName(3) + phases(KubernetesClient.generatePodName(3)) shouldBe null + } + + "metricsByPodName" should "flatten each pod's container usage into a cpu/memory map" in { + val metrics = + KubernetesClient.metricsByPodName(Seq(podMetrics(1, "250m", "128Mi"))) + metrics(KubernetesClient.generatePodName(1)) shouldBe Map("cpu" -> "250m", "memory" -> "128Mi") + } + + // ── namespace-wide wrappers, driven through a stubbed fabric8 client ── + // These pin the fabric8 fluent-chain plumbing (list() / top().pods().metrics()); the value + // transform they delegate to is already pinned by the phasesByPodName / metricsByPodName tests, + // so they only assert that the namespace items flow through keyed by pod name. + + "getAllPodPhases" should "list the namespace pods and key them by pod name" in { + val k8s = + new KubernetesClient(stubbedClient(Seq(pod(1, "Running"), pod(2, "Pending")), Seq.empty)) + k8s.getAllPodPhases.keySet shouldBe + Set(KubernetesClient.generatePodName(1), KubernetesClient.generatePodName(2)) + } + + "getAllPodMetrics" should "fetch the namespace metrics and key them by pod name" in { + val k8s = new KubernetesClient(stubbedClient(Seq.empty, Seq(podMetrics(1, "250m", "128Mi")))) + k8s.getAllPodMetrics.keySet shouldBe Set(KubernetesClient.generatePodName(1)) + } + + // getPodMetrics adds its own collectFirst-by-name lookup on top of the transform, so it asserts + // both the matched pod's usage and the no-match fallback. + "getPodMetrics" should "return the matching pod's usage and an empty map when none matches" in { + val k8s = new KubernetesClient(stubbedClient(Seq.empty, Seq(podMetrics(1, "250m", "128Mi")))) + k8s.getPodMetrics(1) shouldBe Map("cpu" -> "250m", "memory" -> "128Mi") + k8s.getPodMetrics(999) shouldBe empty + } }