diff --git a/build.sbt b/build.sbt index ee6882600d8..940b0bcb101 100644 --- a/build.sbt +++ b/build.sbt @@ -161,7 +161,12 @@ lazy val WorkflowCore = (project in file("common/workflow-core")) lazy val ComputingUnitManagingService = (project in file("computing-unit-managing-service")) .dependsOn(WorkflowCore, Auth, Config, Resource) .settings(commonModuleSettings) + .configs(Test) + .dependsOn(DAO % "test->test", Auth % "test->test") // reuse MockTexeraDB embedded Postgres in tests .settings( + // Tests mutate JVM-wide singletons (SqlServer via MockTexeraDB, the + // KubernetesClient fabric8 client), 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/ComputingUnitManagingService.scala b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/ComputingUnitManagingService.scala index db63bbf2eb2..f0dffc89e11 100644 --- a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/ComputingUnitManagingService.scala +++ b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/ComputingUnitManagingService.scala @@ -27,6 +27,7 @@ import org.apache.texera.common.config.StorageConfig import org.apache.texera.auth.{AuthFeatures, RequestLoggingFilter, RoleAnnotationEnforcer} import org.apache.texera.dao.SqlServer import org.apache.texera.service.resource.{ + AdminComputingUnitResource, ComputingUnitAccessResource, ComputingUnitManagingResource, HealthCheckResource @@ -66,6 +67,7 @@ class ComputingUnitManagingService extends Application[ComputingUnitManagingServ environment.jersey().register(new ComputingUnitManagingResource) environment.jersey().register(new ComputingUnitAccessResource) + environment.jersey().register(new AdminComputingUnitResource) RoleAnnotationEnforcer.enforce( environment.jersey.getResourceConfig, diff --git a/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/AdminComputingUnitResource.scala b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/AdminComputingUnitResource.scala new file mode 100644 index 00000000000..e08cd78b004 --- /dev/null +++ b/computing-unit-managing-service/src/main/scala/org/apache/texera/service/resource/AdminComputingUnitResource.scala @@ -0,0 +1,113 @@ +/* + * 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 io.dropwizard.auth.Auth +import jakarta.annotation.security.RolesAllowed +import jakarta.ws.rs.{GET, Path, Produces} +import jakarta.ws.rs.core.MediaType +import org.apache.texera.auth.SessionUser +import org.apache.texera.dao.SqlServer +import org.apache.texera.dao.jooq.generated.Tables.WORKFLOW_COMPUTING_UNIT +import org.apache.texera.dao.jooq.generated.enums.PrivilegeEnum +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.DashboardWorkflowComputingUnit +import org.apache.texera.service.util.ComputingUnitHelpers +import org.jooq.DSLContext + +import scala.jdk.CollectionConverters.CollectionHasAsScala + +object AdminComputingUnitResource { + private def context: DSLContext = + SqlServer + .getInstance() + .createDSLContext() + + /** + * Render already-active `units` as admin dashboard rows (pure — no filtering). Every row gets + * WRITE access (an admin controls every unit it sees); `isOwner` reflects `callerUid`. + */ + def buildDashboardUnits( + units: List[WorkflowComputingUnit], + ownerInfo: Map[Integer, (String, String)], + callerUid: Integer, + podPhases: Map[String, String], + podMetrics: Map[String, Map[String, String]] + ): List[DashboardWorkflowComputingUnit] = + units.map { unit => + ComputingUnitHelpers.buildDashboardUnit( + unit, + isOwner = unit.getUid.equals(callerUid), + accessPrivilege = PrivilegeEnum.WRITE, + ownerInfo = ownerInfo, + podPhases = podPhases, + podMetrics = podMetrics + ) + } +} + +@Produces(Array(MediaType.APPLICATION_JSON)) +@Path("/computing-unit/admin") +@RolesAllowed(Array("ADMIN")) +class AdminComputingUnitResource { + + import AdminComputingUnitResource._ + + /** + * List every non-terminated computing unit across all users (ADMIN-only). Like the per-user + * endpoint, a Kubernetes unit whose pod has vanished is marked terminated and dropped, so ghost + * units don't accumulate. Status/metrics use one namespace-wide call each (O(1) round trips). + */ + @GET + @Produces(Array(MediaType.APPLICATION_JSON)) + @Path("/list") + def listAllComputingUnits( + @Auth user: SessionUser + ): List[DashboardWorkflowComputingUnit] = { + val ctx = context + + // Filter to active units in SQL so terminated rows are never loaded. + val activeUnits: List[WorkflowComputingUnit] = + ctx + .selectFrom(WORKFLOW_COMPUTING_UNIT) + .where(WORKFLOW_COMPUTING_UNIT.TERMINATE_TIME.isNull) + .fetchInto(classOf[WorkflowComputingUnit]) + .asScala + .toList + + // Pod phases decide which Kubernetes units are still alive. + val podPhases = ComputingUnitHelpers.podPhasesFor(activeUnits) + + val liveUnits = ComputingUnitHelpers.reconcileVanishedKubernetesUnits( + new WorkflowComputingUnitDao(ctx.configuration()), + activeUnits, + podPhases + ) + + // Metrics only for survivors, so fetch after reconciliation. + val podMetrics = ComputingUnitHelpers.podMetricsFor(liveUnits) + + val userDao = new UserDao(ctx.configuration()) + val ownerInfo = ComputingUnitHelpers.resolveOwnerInfo(userDao, liveUnits.map(_.getUid).distinct) + + buildDashboardUnits(liveUnits, ownerInfo, user.getUid, podPhases, podMetrics) + } +} 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..e877d0209ff 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,11 +19,41 @@ 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 { + + /** + * 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 = { unit.getType match { // Local CUs are always “running” @@ -58,4 +88,113 @@ 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 for the namespace; skipped (empty) when no Kubernetes unit is present. */ + def podPhasesFor(units: Seq[WorkflowComputingUnit]): Map[String, String] = + if (units.exists(isKubernetes)) KubernetesClient.getAllPodPhases else Map.empty + + /** Pod metrics for the namespace; skipped (empty) when no Kubernetes unit is present. */ + def podMetricsFor(units: Seq[WorkflowComputingUnit]): Map[String, Map[String, String]] = + if (units.exists(isKubernetes)) KubernetesClient.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..2f75c3f1562 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,7 +20,7 @@ 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 @@ -28,10 +28,15 @@ import scala.jdk.CollectionConverters._ object KubernetesClient { - // Initialize the Kubernetes client - private val client: io.fabric8.kubernetes.client.KubernetesClient = + private var client: io.fabric8.kubernetes.client.KubernetesClient = new KubernetesClientBuilder().build() private val namespace: String = KubernetesConfig.computeUnitPoolNamespace + + /** Test-only seam to swap in a stubbed client (exercises the wrappers without a live cluster). */ + private[util] def setClientForTesting( + testClient: io.fabric8.kubernetes.client.KubernetesClient + ): Unit = + client = testClient private val podNamePrefix = "computing-unit" def generatePodURI(cuid: Int): String = { @@ -48,19 +53,49 @@ 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] = { + client + .pods() + .inNamespace(namespace) + .list() + .getItems + .asScala + .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 + + // 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]] = + fetchPodMetricsItems() + .map(podMetrics => podMetrics.getMetadata.getName -> containerUsage(podMetrics)) + .toMap + 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]) } diff --git a/computing-unit-managing-service/src/test/scala/org/apache/texera/service/ComputingUnitManagingServiceRunSpec.scala b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/ComputingUnitManagingServiceRunSpec.scala index d2162d48c77..c6712e9a815 100644 --- a/computing-unit-managing-service/src/test/scala/org/apache/texera/service/ComputingUnitManagingServiceRunSpec.scala +++ b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/ComputingUnitManagingServiceRunSpec.scala @@ -19,25 +19,76 @@ package org.apache.texera.service +import io.dropwizard.core.setup.Environment +import io.dropwizard.jersey.DropwizardResourceConfig +import io.dropwizard.jersey.setup.JerseyEnvironment +import io.dropwizard.jetty.MutableServletContextHandler +import io.dropwizard.jetty.setup.ServletEnvironment import org.apache.texera.auth.RoleAnnotationEnforcer +import org.apache.texera.common.config.StorageConfig import org.apache.texera.service.resource.{ + AdminComputingUnitResource, ComputingUnitAccessResource, ComputingUnitManagingResource, HealthCheckResource } +import org.mockito.ArgumentMatchers.isA +import org.mockito.Mockito.{mock, verify, when} import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers +import java.sql.DriverManager + class ComputingUnitManagingServiceRunSpec extends AnyFlatSpec with Matchers { + // Booting the service opens a real connection pool, so this path is only exercisable where + // Postgres is provisioned at the configured JDBC URL (as in CI). + private def databaseReachable: Boolean = + try { + DriverManager + .getConnection( + StorageConfig.jdbcUrl, + StorageConfig.jdbcUsername, + StorageConfig.jdbcPassword + ) + .close() + true + } catch { + case _: Throwable => false + } + // Every endpoint this service registers declares @RolesAllowed/@PermitAll/@DenyAll. "ComputingUnitManagingService's registered resources" should "all declare access control" in { RoleAnnotationEnforcer.findUnannotatedEndpoints( Seq( classOf[ComputingUnitManagingResource], classOf[ComputingUnitAccessResource], + classOf[AdminComputingUnitResource], classOf[HealthCheckResource] ) ) shouldBe empty } + + "ComputingUnitManagingService.run" should "register the admin resource on the Jersey environment" in { + val jersey = mock(classOf[JerseyEnvironment]) + val servlets = mock(classOf[ServletEnvironment]) + val context = mock(classOf[MutableServletContextHandler]) + val env = mock(classOf[Environment]) + when(env.jersey).thenReturn(jersey) + when(env.servlets).thenReturn(servlets) + when(env.getApplicationContext).thenReturn(context) + when(jersey.getResourceConfig).thenReturn(DropwizardResourceConfig.forTesting()) + + assume( + databaseReachable, + "run() requires a reachable Postgres at the configured JDBC URL (provided in CI)" + ) + new ComputingUnitManagingService() + .run(mock(classOf[ComputingUnitManagingServiceConfiguration]), env) + + verify(jersey).register(isA(classOf[ComputingUnitManagingResource])) + verify(jersey).register(isA(classOf[ComputingUnitAccessResource])) + verify(jersey).register(isA(classOf[AdminComputingUnitResource])) + verify(jersey).setUrlPattern("/api/*") + } } diff --git a/computing-unit-managing-service/src/test/scala/org/apache/texera/service/resource/AdminComputingUnitResourceSpec.scala b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/resource/AdminComputingUnitResourceSpec.scala new file mode 100644 index 00000000000..0dfa1b4da94 --- /dev/null +++ b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/resource/AdminComputingUnitResourceSpec.scala @@ -0,0 +1,217 @@ +/* + * 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 jakarta.annotation.security.RolesAllowed +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.util.KubernetesClient +import org.scalatest.BeforeAndAfterAll +import org.scalatest.flatspec.AnyFlatSpec +import org.scalatest.matchers.should.Matchers + +import java.sql.Timestamp + +class AdminComputingUnitResourceSpec + extends AnyFlatSpec + with Matchers + with BeforeAndAfterAll + with MockTexeraDB { + + override protected def beforeAll(): Unit = initializeDBAndReplaceDSLContext() + + override protected def afterAll(): Unit = shutdownDB() + + private def makeUser(uid: Int, name: String): User = { + val u = new User() + u.setUid(uid) + u.setName(name) + u.setEmail(s"user$uid@example.com") + u.setRole(UserRoleEnum.ADMIN) + u.setPassword("password") + u.setGoogleAvatar(s"avatar-$uid") + u + } + + private def makeUnit( + cuid: Int, + uid: Int, + name: String, + tpe: WorkflowComputingUnitTypeEnum + ): WorkflowComputingUnit = { + val u = new WorkflowComputingUnit() + u.setCuid(cuid) + u.setUid(uid) + u.setName(name) + u.setType(tpe) + u + } + + private def localUnit(cuid: Int, uid: Int, name: String): WorkflowComputingUnit = + makeUnit(cuid, uid, name, WorkflowComputingUnitTypeEnum.local) + + private def kubernetesUnit(cuid: Int, uid: Int, name: String): WorkflowComputingUnit = + makeUnit(cuid, uid, name, WorkflowComputingUnitTypeEnum.kubernetes) + + // The class-level @RolesAllowed(ADMIN) is what makes Jersey reject non-ADMIN callers. + "AdminComputingUnitResource" should "only permit the ADMIN role (non-ADMIN callers are rejected)" in { + val annotation = classOf[AdminComputingUnitResource].getAnnotation(classOf[RolesAllowed]) + annotation should not be null + annotation.value.toSeq shouldBe Seq("ADMIN") + } + + "buildDashboardUnits" should "render active units owned by multiple users with their owner info" in { + val units = List( + localUnit(cuid = 1, uid = 100, name = "alice-cu"), + localUnit(cuid = 2, uid = 200, name = "bob-cu") + ) + val ownerInfo: Map[Integer, (String, String)] = Map( + (100: Integer) -> ("alice-avatar", "alice"), + (200: Integer) -> (null, "bob") + ) + + // caller is an admin who owns none of these units + val result = AdminComputingUnitResource.buildDashboardUnits( + units, + ownerInfo, + callerUid = 999, + podPhases = Map.empty, + podMetrics = Map.empty + ) + + result.map(_.computingUnit.getCuid) should contain theSameElementsAs Seq(1, 2) + result.map(_.computingUnit.getUid).distinct should contain theSameElementsAs Seq(100, 200) + + val byCuid = result.map(u => u.computingUnit.getCuid.intValue() -> u).toMap + byCuid(1).ownerName shouldBe "alice" + byCuid(1).ownerGoogleAvatar shouldBe "alice-avatar" + byCuid(2).ownerName shouldBe "bob" + byCuid(2).ownerGoogleAvatar shouldBe null + + // local units report Running/NaN and the admin caller does not own them + all(result.map(_.status)) shouldBe "Running" + all(result.map(_.isOwner)) shouldBe false + all(result.map(_.accessPrivilege)) shouldBe PrivilegeEnum.WRITE + } + + it should "set isOwner only for units owned by the caller" in { + val units = List( + localUnit(cuid = 1, uid = 100, name = "owned-by-caller"), + localUnit(cuid = 2, uid = 200, name = "owned-by-other") + ) + + val result = AdminComputingUnitResource + .buildDashboardUnits( + units, + ownerInfo = Map.empty, + callerUid = 100, + podPhases = Map.empty, + podMetrics = Map.empty + ) + .map(u => u.computingUnit.getCuid.intValue() -> u.isOwner) + .toMap + + result(1) shouldBe true + result(2) shouldBe false + } + + it should "fall back to null owner info when the owner is missing from the map" in { + val units = List(localUnit(cuid = 1, uid = 100, name = "orphan-cu")) + + val result = AdminComputingUnitResource.buildDashboardUnits( + units, + ownerInfo = Map.empty, + callerUid = 999, + podPhases = Map.empty, + podMetrics = Map.empty + ) + + result.head.ownerName shouldBe null + result.head.ownerGoogleAvatar shouldBe null + } + + it should "resolve kubernetes status and metrics from the pre-fetched maps" in { + val unit = kubernetesUnit(cuid = 5, uid = 300, name = "k8s-cu") + val podName = KubernetesClient.generatePodName(5) + + val podPhases = Map(podName -> "Running") + val podMetrics = Map(podName -> Map("cpu" -> "250m", "memory" -> "128Mi")) + + val result = AdminComputingUnitResource.buildDashboardUnits( + List(unit), + ownerInfo = Map((300: Integer) -> ("k8s-avatar", "k8s-owner")), + callerUid = 999, + podPhases = podPhases, + podMetrics = podMetrics + ) + + result.head.status shouldBe "Running" + result.head.metrics.cpuUsage shouldBe "250m" + result.head.metrics.memoryUsage shouldBe "128Mi" + } + + it should "report a kubernetes unit as Pending with empty metrics when its pod is absent from the maps" in { + val unit = kubernetesUnit(cuid = 6, uid = 300, name = "k8s-no-pod") + + val result = AdminComputingUnitResource.buildDashboardUnits( + List(unit), + ownerInfo = Map.empty, + callerUid = 999, + podPhases = Map.empty, + podMetrics = Map.empty + ) + + result.head.status shouldBe "Pending" + result.head.metrics.cpuUsage shouldBe "" + result.head.metrics.memoryUsage shouldBe "" + } + + "listAllComputingUnits" should "return every non-terminated unit across users, marked WRITE" in { + val userDao = new UserDao(getDSLContext.configuration()) + val unitDao = new WorkflowComputingUnitDao(getDSLContext.configuration()) + val admin = makeUser(700, "admin") + userDao.insert(admin) + userDao.insert(makeUser(701, "other")) + unitDao.insert(localUnit(cuid = 700, uid = 700, name = "admin-cu")) + unitDao.insert(localUnit(cuid = 701, uid = 701, name = "other-cu")) + // A terminated unit must be excluded by the SQL filter. + val terminated = localUnit(cuid = 702, uid = 701, name = "terminated-cu") + terminated.setTerminateTime(new Timestamp(0L)) + unitDao.insert(terminated) + + val result = new AdminComputingUnitResource().listAllComputingUnits(new SessionUser(admin)) + + result.map(_.computingUnit.getCuid.intValue()) should contain theSameElementsAs Seq(700, 701) + all(result.map(_.accessPrivilege)) shouldBe PrivilegeEnum.WRITE + all(result.map(_.status)) shouldBe "Running" // local units + val byCuid = result.map(r => r.computingUnit.getCuid.intValue() -> r).toMap + byCuid(700).isOwner shouldBe true // owned by the requesting admin + byCuid(700).ownerName shouldBe "admin" + byCuid(701).isOwner shouldBe false + byCuid(701).ownerName shouldBe "other" + } +} 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..0fb6e214afa --- /dev/null +++ b/computing-unit-managing-service/src/test/scala/org/apache/texera/service/resource/ComputingUnitManagingResourceSpec.scala @@ -0,0 +1,101 @@ +/* + * 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 = { + 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 = shutdownDB() + + 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..73deffb4eb2 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,23 +19,63 @@ 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 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.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 = initializeDBAndReplaceDSLContext() + + override protected def afterAll(): Unit = shutdownDB() + + 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() "getComputingUnitStatus" should "return Running for a local unit" in { @@ -55,4 +95,170 @@ 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") + } + + // ── 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..fc98727e7bc 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,10 +19,109 @@ package org.apache.texera.service.util +import io.fabric8.kubernetes.api.model.metrics.v1beta1.{ + ContainerMetricsBuilder, + 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.{KubernetesClientBuilder, KubernetesClient => Fabric8Client} +import org.apache.texera.common.config.KubernetesConfig +import org.apache.texera.dao.jooq.generated.enums.WorkflowComputingUnitTypeEnum +import org.apache.texera.dao.jooq.generated.tables.pojos.WorkflowComputingUnit +import org.mockito.Mockito.{mock, when} +import org.scalatest.BeforeAndAfterAll import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers -class KubernetesClientSpec extends AnyFlatSpec with Matchers { +import scala.jdk.CollectionConverters._ + +// Exercises the fabric8 wrappers without a cluster by stubbing the client with Mockito. +class KubernetesClientSpec extends AnyFlatSpec with Matchers with BeforeAndAfterAll { + + private val namespace: String = KubernetesConfig.computeUnitPoolNamespace + + private def unitOfType(cuid: Int, tpe: WorkflowComputingUnitTypeEnum): WorkflowComputingUnit = { + val u = new WorkflowComputingUnit() + u.setCuid(cuid) + u.setType(tpe) + u + } + + private def pod(cuid: Int, phase: String): Pod = + new PodBuilder() + .withNewMetadata() + .withName(KubernetesClient.generatePodName(cuid)) + .endMetadata() + .withNewStatus() + .withPhase(phase) + .endStatus() + .build() + + // cuid 1 -> Running, cuid 2 -> Pending; both returned by the namespace-wide listing. + private val podList: PodList = + new PodListBuilder().addToItems(pod(1, "Running"), pod(2, "Pending")).build() + + // Only cuid 1 reports metrics. + private val metricsList: PodMetricsList = + new PodMetricsListBuilder() + .addToItems( + new PodMetricsBuilder() + .withNewMetadata() + .withName(KubernetesClient.generatePodName(1)) + .endMetadata() + .addToContainers( + new ContainerMetricsBuilder() + .withName("main") + .withUsage( + Map("cpu" -> new Quantity("250m"), "memory" -> new Quantity("128Mi")).asJava + ) + .build() + ) + .build() + ) + .build() + + override protected def beforeAll(): Unit = { + 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(podList) + + // cuid 1 -> a present pod; cuid 2 -> no pod. + val presentResource = mock(classOf[PodResource]) + when(presentResource.get()).thenReturn(pod(1, "Running")) + val absentResource = mock(classOf[PodResource]) + when(absentResource.get()).thenReturn(null) + when(podsInNamespace.withName(KubernetesClient.generatePodName(1))).thenReturn(presentResource) + when(podsInNamespace.withName(KubernetesClient.generatePodName(2))).thenReturn(absentResource) + + val top = mock(classOf[MetricAPIGroupDSL]) + val podMetricOp = mock(classOf[PodMetricOperation]) + when(client.top()).thenReturn(top) + when(top.pods()).thenReturn(podMetricOp) + when(podMetricOp.metrics(namespace)).thenReturn(metricsList) + + KubernetesClient.setClientForTesting(client) + } + + override protected def afterAll(): Unit = { + // Restore a real client so subsequent suites don't reuse the mock. + KubernetesClient.setClientForTesting(new KubernetesClientBuilder().build()) + } "generatePodName" should "prefix the cuid with computing-unit" in { KubernetesClient.generatePodName(42) shouldBe "computing-unit-42" @@ -31,4 +130,53 @@ class KubernetesClientSpec extends AnyFlatSpec with Matchers { it should "handle a cuid of 0" in { KubernetesClient.generatePodName(0) shouldBe "computing-unit-0" } + + "getPodByName" should "return the pod when present and None when absent" in { + KubernetesClient.getPodByName(KubernetesClient.generatePodName(1)) shouldBe defined + KubernetesClient.getPodByName(KubernetesClient.generatePodName(2)) shouldBe None + } + + "podExists" should "reflect whether the pod is present" in { + KubernetesClient.podExists(1) shouldBe true + KubernetesClient.podExists(2) shouldBe false + } + + "getAllPodPhases" should "map every pod name in the namespace to its phase" in { + val phases = KubernetesClient.getAllPodPhases + phases(KubernetesClient.generatePodName(1)) shouldBe "Running" + phases(KubernetesClient.generatePodName(2)) shouldBe "Pending" + } + + "getAllPodMetrics" should "map pod names to their cpu/memory usage" in { + val metrics = KubernetesClient.getAllPodMetrics + metrics(KubernetesClient.generatePodName(1)) shouldBe Map("cpu" -> "250m", "memory" -> "128Mi") + } + + "getPodMetrics" should "return the usage for the pod matching the cuid" in { + KubernetesClient.getPodMetrics(1) shouldBe Map("cpu" -> "250m", "memory" -> "128Mi") + } + + it should "return an empty map when no pod matches the cuid" in { + KubernetesClient.getPodMetrics(999) shouldBe empty + } + + "ComputingUnitHelpers.podPhasesFor" should "query the cluster iff a kubernetes unit is present" in { + ComputingUnitHelpers.podPhasesFor( + List(unitOfType(5, WorkflowComputingUnitTypeEnum.local)) + ) shouldBe empty + + ComputingUnitHelpers.podPhasesFor( + List(unitOfType(1, WorkflowComputingUnitTypeEnum.kubernetes)) + )(KubernetesClient.generatePodName(1)) shouldBe "Running" + } + + "ComputingUnitHelpers.podMetricsFor" should "query the cluster iff a kubernetes unit is present" in { + ComputingUnitHelpers.podMetricsFor( + List(unitOfType(5, WorkflowComputingUnitTypeEnum.local)) + ) shouldBe empty + + ComputingUnitHelpers.podMetricsFor( + List(unitOfType(1, WorkflowComputingUnitTypeEnum.kubernetes)) + )(KubernetesClient.generatePodName(1)) shouldBe Map("cpu" -> "250m", "memory" -> "128Mi") + } }