From 4f24bb6cb8e79fd4edd960ac458b9a4ec086055f Mon Sep 17 00:00:00 2001 From: Xinyuan Lin Date: Thu, 30 Jul 2026 19:44:30 -0700 Subject: [PATCH 1/2] test(amber): cover workflow lifecycle cleanup --- .../web/WorkflowLifecycleManagerSpec.scala | 130 ++++++++++++++++++ 1 file changed, 130 insertions(+) create mode 100644 amber/src/test/scala/org/apache/texera/web/WorkflowLifecycleManagerSpec.scala diff --git a/amber/src/test/scala/org/apache/texera/web/WorkflowLifecycleManagerSpec.scala b/amber/src/test/scala/org/apache/texera/web/WorkflowLifecycleManagerSpec.scala new file mode 100644 index 00000000000..1b7c219d9a2 --- /dev/null +++ b/amber/src/test/scala/org/apache/texera/web/WorkflowLifecycleManagerSpec.scala @@ -0,0 +1,130 @@ +/* + * 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.web + +import org.apache.pekko.actor.ActorSystem +import org.apache.pekko.testkit.TestKit +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState.{ + COMPLETED, + RUNNING +} +import org.apache.texera.amber.engine.common.AmberRuntime +import org.apache.texera.web.storage.ExecutionStateStore +import org.scalatest.BeforeAndAfterAll +import org.scalatest.flatspec.AnyFlatSpec + +import java.util.concurrent.{CountDownLatch, TimeUnit} + +class WorkflowLifecycleManagerSpec extends AnyFlatSpec with BeforeAndAfterAll { + + // WorkflowLifecycleManager schedules through AmberRuntime's process-wide actor system. + // Preserve that shared reference because amber suites run concurrently in one JVM. + private lazy val testSystem: ActorSystem = + ActorSystem("WorkflowLifecycleManagerSpec-test", AmberRuntime.pekkoConfig) + + private var previousActorSystem: AnyRef = _ + + private def getAmberRuntimeField(name: String): AnyRef = { + val field = AmberRuntime.getClass.getDeclaredField(name) + field.setAccessible(true) + field.get(AmberRuntime) + } + + private def setAmberRuntimeField(name: String, value: AnyRef): Unit = { + val field = AmberRuntime.getClass.getDeclaredField(name) + field.setAccessible(true) + field.set(AmberRuntime, value) + } + + override protected def beforeAll(): Unit = { + super.beforeAll() + previousActorSystem = getAmberRuntimeField("_actorSystem") + setAmberRuntimeField("_actorSystem", testSystem) + } + + override protected def afterAll(): Unit = { + setAmberRuntimeField("_actorSystem", previousActorSystem) + TestKit.shutdownActorSystem(testSystem) + super.afterAll() + } + + private def managerWithCallback( + cleanUpTimeout: Int = 1 + ): (WorkflowLifecycleManager, CountDownLatch) = { + val cleaned = new CountDownLatch(1) + val manager = new WorkflowLifecycleManager( + id = "workflow-lifecycle-manager-spec", + cleanUpTimeout = cleanUpTimeout, + cleanUpCallback = () => cleaned.countDown() + ) + (manager, cleaned) + } + + private def assertCleanUpWithin(cleaned: CountDownLatch, seconds: Long): Unit = { + assert( + cleaned.await(seconds, TimeUnit.SECONDS), + "cleanup callback was not invoked before the deadline" + ) + } + + "WorkflowLifecycleManager" should "wait for the last user before scheduling cleanup" in { + val (manager, cleaned) = managerWithCallback() + + manager.increaseUserCount() + manager.increaseUserCount() + manager.decreaseUserCount(Some(COMPLETED)) + + assert( + !cleaned.await(1500, TimeUnit.MILLISECONDS), + "cleanup ran while a user was still present" + ) + + manager.decreaseUserCount(None) + assertCleanUpWithin(cleaned, seconds = 5) + } + + it should "cancel cleanup while the workflow is running and resume after completion" in { + val (manager, cleaned) = managerWithCallback() + val stateStore = new ExecutionStateStore + + manager.registerCleanUpOnStateChange(stateStore) + stateStore.metadataStore.updateState(_.withState(RUNNING)) + + assert(!cleaned.await(1500, TimeUnit.MILLISECONDS), "a running workflow must not be cleaned up") + + stateStore.metadataStore.updateState(_.withState(COMPLETED)) + assertCleanUpWithin(cleaned, seconds = 5) + } + + it should "refresh the deadline when a later terminal state arrives" in { + val (manager, cleaned) = managerWithCallback(cleanUpTimeout = 4) + val stateStore = new ExecutionStateStore + + manager.registerCleanUpOnStateChange(stateStore) + Thread.sleep(2000) + stateStore.metadataStore.updateState(_.withState(COMPLETED)) + + assert( + !cleaned.await(3000, TimeUnit.MILLISECONDS), + "a refreshed deadline must cancel the earlier one" + ) + assertCleanUpWithin(cleaned, seconds = 5) + } +} From 784191bd6a51d3b4b3d6734043eb81737018b433 Mon Sep 17 00:00:00 2001 From: Xinyuan Lin Date: Fri, 31 Jul 2026 16:25:29 -0700 Subject: [PATCH 2/2] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Xinyuan Lin --- .../org/apache/texera/web/WorkflowLifecycleManagerSpec.scala | 3 +++ 1 file changed, 3 insertions(+) diff --git a/amber/src/test/scala/org/apache/texera/web/WorkflowLifecycleManagerSpec.scala b/amber/src/test/scala/org/apache/texera/web/WorkflowLifecycleManagerSpec.scala index 1b7c219d9a2..6260de4ad6f 100644 --- a/amber/src/test/scala/org/apache/texera/web/WorkflowLifecycleManagerSpec.scala +++ b/amber/src/test/scala/org/apache/texera/web/WorkflowLifecycleManagerSpec.scala @@ -40,6 +40,7 @@ class WorkflowLifecycleManagerSpec extends AnyFlatSpec with BeforeAndAfterAll { ActorSystem("WorkflowLifecycleManagerSpec-test", AmberRuntime.pekkoConfig) private var previousActorSystem: AnyRef = _ + private var previousSerde: AnyRef = _ private def getAmberRuntimeField(name: String): AnyRef = { val field = AmberRuntime.getClass.getDeclaredField(name) @@ -56,10 +57,12 @@ class WorkflowLifecycleManagerSpec extends AnyFlatSpec with BeforeAndAfterAll { override protected def beforeAll(): Unit = { super.beforeAll() previousActorSystem = getAmberRuntimeField("_actorSystem") + previousSerde = getAmberRuntimeField("_serde") setAmberRuntimeField("_actorSystem", testSystem) } override protected def afterAll(): Unit = { + setAmberRuntimeField("_serde", previousSerde) setAmberRuntimeField("_actorSystem", previousActorSystem) TestKit.shutdownActorSystem(testSystem) super.afterAll()