From 17f47cd398a495e9caeef228e0afa3ec3549bda3 Mon Sep 17 00:00:00 2001
From: Martin Vu <127065329+VuMartin@users.noreply.github.com>
Date: Fri, 28 Aug 2026 03:54:24 +0000
Subject: [PATCH] fix(WorkflowExecutionService): shutdown console writer thread
on unsubscribe (#7914)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
### What changes were proposed in this PR?
Fixes a console writer thread leak in `ExecutionConsoleService`.
The console writer executor was not being shut down when an execution
service was unsubscribed. This could leave `texera-console-writer`
threads alive after workflow execution finished.
Before:
After:
This PR:
- Shuts down the console writer executor during `unsubscribeAll()`.
- Waits for termination and falls back to `shutdownNow()` if necessary.
- Closes active console message writers and clears the writer map.
- Adds a test verifying the console writer executor is shut down and
terminated.
### Any related issues, documentation, discussions?
Fixes #7455
### How was this PR tested?
Ran:
```bash
sbt "project WorkflowExecutionService" "testOnly *ExecutionConsoleServiceSpec -- -z unsubscribeAll"
```
Manual testing:
Ran workflows multiple times and checked the console writer threads
with:
```bash
jcmd 57436 Thread.print | grep "texera-console-writer"
```
```bash
jcmd 67548 Thread.print | grep "texera-console-writer"
```
Verified that the console writer threads are terminated after workflow
execution completes and unsubscribeAll() is called.
### Was this PR authored or co-authored using generative AI tooling?
Generated-by: ChatGPT (5.5 mini)
(backported from commit bdc6d2a90eb1415014a4f705b7bb7cca31bb9688)
---
.../web/service/ExecutionConsoleService.scala | 25 ++
.../service/ExecutionConsoleServiceSpec.scala | 215 ++++++++++++++++++
2 files changed, 240 insertions(+)
diff --git a/amber/src/main/scala/org/apache/texera/web/service/ExecutionConsoleService.scala b/amber/src/main/scala/org/apache/texera/web/service/ExecutionConsoleService.scala
index 1678494e937..64c664bcb4f 100644
--- a/amber/src/main/scala/org/apache/texera/web/service/ExecutionConsoleService.scala
+++ b/amber/src/main/scala/org/apache/texera/web/service/ExecutionConsoleService.scala
@@ -216,6 +216,31 @@ class ExecutionConsoleService(
}
)
+ override def unsubscribeAll(): Unit = {
+ consoleMessageOpIdToWriterMap.values.foreach { writer =>
+ try {
+ writer.close()
+ } catch {
+ case e: Exception =>
+ logger.error("Failed to close console message writer during unsubscribeAll", e)
+ }
+ }
+ consoleMessageOpIdToWriterMap.clear()
+
+ super.unsubscribeAll()
+
+ consoleWriterThread.shutdown()
+ try {
+ if (!consoleWriterThread.awaitTermination(5, java.util.concurrent.TimeUnit.SECONDS)) {
+ consoleWriterThread.shutdownNow()
+ }
+ } catch {
+ case _: InterruptedException =>
+ consoleWriterThread.shutdownNow()
+ Thread.currentThread().interrupt()
+ }
+ }
+
/**
* Processes a console message for display, performing truncation if needed.
* This method uses the shared implementation in ConsoleMessageProcessor.
diff --git a/amber/src/test/scala/org/apache/texera/web/service/ExecutionConsoleServiceSpec.scala b/amber/src/test/scala/org/apache/texera/web/service/ExecutionConsoleServiceSpec.scala
index d4753984cf1..c1ccd35cb60 100644
--- a/amber/src/test/scala/org/apache/texera/web/service/ExecutionConsoleServiceSpec.scala
+++ b/amber/src/test/scala/org/apache/texera/web/service/ExecutionConsoleServiceSpec.scala
@@ -25,10 +25,30 @@ import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
ConsoleMessageType
}
import org.apache.texera.amber.engine.common.executionruntimestate.ExecutionConsoleStore
+<<<<<<< HEAD
import org.scalatest.flatspec.AnyFlatSpec
+=======
+import org.apache.texera.web.WebsocketInput
+import org.apache.texera.web.model.websocket.event.TexeraWebSocketEvent
+import org.apache.texera.web.model.websocket.event.python.ConsoleUpdateEvent
+import org.apache.texera.web.model.websocket.request.python.DebugCommandRequest
+import org.apache.texera.web.storage.ExecutionStateStore
+import org.scalamock.scalatest.MockFactory
+import org.scalatest.BeforeAndAfterAll
+import org.scalatest.concurrent.Eventually.eventually
+import org.scalatest.concurrent.PatienceConfiguration.{Interval, Timeout}
+import org.scalatest.flatspec.AnyFlatSpecLike
+>>>>>>> bdc6d2a90 (fix(WorkflowExecutionService): shutdown console writer thread on unsubscribe (#7914))
import org.scalatest.matchers.should.Matchers
+import org.scalatest.time.{Millis, Span}
import java.time.Instant
+<<<<<<< HEAD
+=======
+import java.util.concurrent.ExecutorService
+import scala.collection.mutable.ListBuffer
+import scala.reflect.ClassTag
+>>>>>>> bdc6d2a90 (fix(WorkflowExecutionService): shutdown console writer thread on unsubscribe (#7914))
class ExecutionConsoleServiceSpec extends AnyFlatSpec with Matchers {
@@ -231,4 +251,199 @@ class ExecutionConsoleServiceSpec extends AnyFlatSpec with Matchers {
val expectedTruncatedTitle = "a" * (messageDisplayLength - 3) + "..."
opInfo.consoleMessages.head.title shouldBe expectedTruncatedTitle
}
+<<<<<<< HEAD
+=======
+ // ---------------------------------------------------------------- instance
+
+ /** Empty-plan client that captures the ConsoleMessage callback the service registers. */
+ private final class TestAmberClient(
+ override val coordinatorInterface: CoordinatorServiceFs2Grpc[TwitterFuture, Unit]
+ ) extends AmberClient(
+ system,
+ new WorkflowContext(),
+ PhysicalPlan(Set.empty, Set.empty),
+ CoordinatorConfig(None, None, None, None),
+ _ => ()
+ ) {
+ var consoleCallback: ConsoleMessage => Unit = _
+
+ override def registerCallback[T](callback: T => Unit)(implicit ct: ClassTag[T]): Disposable = {
+ if (ct.runtimeClass == classOf[ConsoleMessage]) {
+ consoleCallback = callback.asInstanceOf[ConsoleMessage => Unit]
+ }
+ Disposable.empty()
+ }
+
+ def dispose(): Unit = super.shutdown()
+ }
+
+ private final class Fixture(
+ val client: TestAmberClient,
+ val coordinator: CoordinatorServiceFs2Grpc[TwitterFuture, Unit],
+ val stateStore: ExecutionStateStore,
+ val wsInput: WebsocketInput,
+ val service: ExecutionConsoleService
+ ) {
+ def close(): Unit = {
+ service.unsubscribeAll()
+ client.dispose()
+ }
+ }
+
+ private def fixture(): Fixture = {
+ val coordinator = mock[CoordinatorServiceFs2Grpc[TwitterFuture, Unit]]
+ val client = new TestAmberClient(coordinator)
+ val stateStore = new ExecutionStateStore
+ val wsInput = new WebsocketInput(ListBuffer.empty[Throwable] += _)
+ val service = new ExecutionConsoleService(client, stateStore, wsInput, new WorkflowContext())
+ new Fixture(client, coordinator, stateStore, wsInput, service)
+ }
+
+ private def message(
+ workerId: String = "Worker:WF1-udf1-main-0",
+ title: String = "hello",
+ msgType: ConsoleMessageType = ConsoleMessageType.PRINT
+ ): ConsoleMessage =
+ new ConsoleMessage(workerId, Timestamp(Instant.now), msgType, "src", title, "content")
+
+ private def withFixture(body: Fixture => Unit): Unit = {
+ val f = fixture()
+ try body(f)
+ finally f.close()
+ }
+
+ "processConsoleMessage" should "leave a debugger message untouched however long it is" in {
+ // The debugger's output is the payload the user asked to see; truncating it would cut off the
+ // frame or variable they are inspecting.
+ withFixture { f =>
+ val long = "a" * (f.service.consoleMessageDisplayLength + 50)
+
+ val processed = f.service.processConsoleMessage(
+ message(title = long, msgType = ConsoleMessageType.DEBUGGER)
+ )
+
+ processed.title shouldBe long
+ }
+ }
+
+ it should "still truncate an ordinary message to the configured length" in {
+ withFixture { f =>
+ val long = "a" * (f.service.consoleMessageDisplayLength + 50)
+
+ val processed = f.service.processConsoleMessage(message(title = long))
+
+ processed.title.length shouldBe f.service.consoleMessageDisplayLength
+ processed.title should endWith("...")
+ }
+ }
+
+ "the console diff handler" should "report only the messages added since the last state" in {
+ // The frontend appends what it is sent. Emitting the whole buffer instead of the delta would
+ // duplicate every earlier line on each update.
+ withFixture { f =>
+ // One batch is published per state update. Subscribing up front and reading the batch for the
+ // SECOND message is what shows the delta: the first message must not appear in it again.
+ val batches = ListBuffer.empty[Iterable[TexeraWebSocketEvent]]
+ val sub = f.stateStore.consoleStore.getWebsocketEventObservable
+ .subscribe((batch: Iterable[TexeraWebSocketEvent]) => batches += batch)
+
+ try {
+ f.client.consoleCallback(message(title = "first"))
+ f.client.consoleCallback(message(title = "second"))
+ } finally sub.dispose()
+
+ val titles =
+ batches.last.collect { case e: ConsoleUpdateEvent => e.messages.map(_.title) }.flatten
+ titles shouldBe Seq("second")
+ }
+ }
+
+ "the console message callback" should "file a message under the logical operator id" in {
+ // The worker id carries the physical layer and worker index; the frontend console is keyed by
+ // the logical operator, so anything else silently strands the output.
+ withFixture { f =>
+ f.client.consoleCallback(message(workerId = "Worker:WF1-udf1-main-0"))
+
+ f.stateStore.consoleStore.getState.operatorConsole.keys should contain("udf1")
+ }
+ }
+
+ it should "store the truncated form, not the original" in {
+ withFixture { f =>
+ val long = "b" * (f.service.consoleMessageDisplayLength + 50)
+
+ f.client.consoleCallback(message(title = long))
+
+ val stored = f.stateStore.consoleStore.getState.operatorConsole("udf1").consoleMessages
+ stored.map(_.title.length) shouldBe Seq(f.service.consoleMessageDisplayLength)
+ }
+ }
+
+ "a debug command" should "be attributed to the user that issued it" in {
+ withFixture { f =>
+ (f.coordinator.debugCommand _)
+ .expects(AmberDebugCommandRequest("Worker:WF1-udf1-main-0", "break 12"), ())
+ .returning(TwitterFuture.value(EmptyReturn()))
+
+ f.wsInput.onNext(
+ DebugCommandRequest("udf1", "Worker:WF1-udf1-main-0", "break 12"),
+ Some(7)
+ )
+
+ val stored = f.stateStore.consoleStore.getState.operatorConsole("udf1").consoleMessages
+ stored.map(_.source) shouldBe Seq("USER-7")
+ stored.map(_.title) shouldBe Seq("break 12")
+ }
+ }
+
+ it should "fall back to UNKNOWN when there is no session user" in {
+ withFixture { f =>
+ (f.coordinator.debugCommand _)
+ .expects(*, *)
+ .returning(TwitterFuture.value(EmptyReturn()))
+
+ f.wsInput.onNext(DebugCommandRequest("udf1", "Worker:WF1-udf1-main-0", "cont"), None)
+
+ f.stateStore.consoleStore.getState
+ .operatorConsole("udf1")
+ .consoleMessages
+ .map(_.source) shouldBe Seq("USER-UNKNOWN")
+ }
+ }
+
+ it should "file the command under the operator, not the worker" in {
+ // req carries both; keying by workerId would scatter the command across per-worker consoles.
+ withFixture { f =>
+ (f.coordinator.debugCommand _).expects(*, *).returning(TwitterFuture.value(EmptyReturn()))
+
+ f.wsInput.onNext(DebugCommandRequest("udf1", "Worker:WF1-udf1-main-0", "cont"), Some(1))
+
+ val keys = f.stateStore.consoleStore.getState.operatorConsole.keys
+ keys should contain("udf1")
+ keys should not contain "Worker:WF1-udf1-main-0"
+ }
+ }
+
+ "unsubscribeAll" should "shutdown consoleWriterThread" in {
+ withFixture { f =>
+ f.client.consoleCallback(message(title = "test"))
+
+ val threadField = classOf[ExecutionConsoleService].getDeclaredField("consoleWriterThread")
+ threadField.setAccessible(true)
+ val executor = threadField.get(f.service).asInstanceOf[ExecutorService]
+
+ // Verify it is initially active
+ executor.isShutdown shouldBe false
+
+ // Trigger the teardown
+ f.service.unsubscribeAll()
+
+ // Use Eventually to wait for async termination without blocking arbitrarily
+ eventually(Timeout(Span(2000, Millis)), Interval(Span(50, Millis))) {
+ executor.isShutdown shouldBe true
+ executor.isTerminated shouldBe true
+ }
+ }
+ }
+>>>>>>> bdc6d2a90 (fix(WorkflowExecutionService): shutdown console writer thread on unsubscribe (#7914))
}