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 55f72c35d81..fd16c654f10 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 @@ -221,6 +221,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 b1d647d035c..ce88a46e6ec 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 @@ -42,10 +42,14 @@ 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 import org.scalatest.matchers.should.Matchers +import org.scalatest.time.{Millis, Span} import java.time.Instant +import java.util.concurrent.ExecutorService import scala.collection.mutable.ListBuffer import scala.reflect.ClassTag @@ -441,4 +445,26 @@ class ExecutionConsoleServiceSpec 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 + } + } + } }