fix(WorkflowExecutionService, v1.2): shutdown console writer thread on unsubscribe

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.

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))
 }