[BUG] Worker is cancelled only after child widgets are removed from the DOM #6361
Description
Issue description
Consider the following textual app, consisting of two screens:
import asyncio
import logging
import time
from textual import work
from textual.app import App, ComposeResult
from textual.events import Key
from textual.screen import Screen
from textual.widgets import Static, Header
logging.basicConfig(
filename="demo.log",
level=logging.DEBUG,
format="%(asctime)s %(levelname)s %(message)s",
)
log = logging.getLogger(__name__)
class ScreenA(Screen):
def compose(self) -> ComposeResult:
yield Header()
yield Static("Screen A\n\nPress 'i' to push Screen B")
def on_key(self, event: Key):
if event.key == "i":
self.app.push_screen(ScreenB())
class ScreenB(Screen):
def compose(self) -> ComposeResult:
yield Header()
yield Static("Screen B\n\nPress 'u' to pop back to Screen A")
yield Static("0.000s", id="stopwatch")
def on_mount(self):
log.info("ScreenB mounted")
self.run_stopwatch()
def on_key(self, event: Key):
if event.key == "u":
log.info("ScreenB: pop_screen() called")
self.app.pop_screen()
@work(exclusive=True)
async def run_stopwatch(self):
start = time.monotonic()
while True:
await asyncio.sleep(0.01)
elapsed = time.monotonic() - start
widget = self.screen.query_one("#stopwatch", Static)
widget.update(f"{elapsed:.3f}s")
class DemoApp(App):
SCREENS = {"screen_a": ScreenA}
def on_mount(self):
self.push_screen(ScreenA())
if __name__ == "__main__":
app = DemoApp()
app.run()Running the app and toggling between the two screens by mashing the "u" and "i" keys will trigger a crash:
╭─────────────────────────────────────────────────────────────────────────────────────────────────── Traceback (most recent call last) ────────────────────────────────────────────────────────────────────────────────────────────────────╮
│ /usr/local/dev/textual-demo/.venv/lib64/python3.9/site-packages/textual/worker.py:370 in _run │
│ │
│ 367 │ │ │ self.state = WorkerState.RUNNING ╭───────────────────────────────────────────────────── locals ─────────────────────────────────────────────────────╮ │
│ 368 │ │ │ app.log.worker(self) │ app = DemoApp(title='DemoApp', classes={'-dark-mode'}, pseudo_classes={'dark', 'focus'}) │ │
│ 369 │ │ │ try: │ error = NoMatches("No nodes match '#stopwatch' on ScreenB()") │ │
│ ❱ 370 │ │ │ │ self._result = await self.run() │ self = <Worker ERROR name='run_stopwatch' description='run_stopwatch()'> │ │
│ 371 │ │ │ except asyncio.CancelledError as error: │ worker_failed = WorkerFailed('Worker raised exception: NoMatches("No nodes match \'#stopwatch\' on ScreenB()")') │ │
│ 372 │ │ │ │ self.state = WorkerState.CANCELLED ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ │
│ 373 │ │ │ │ self._error = error │
│ │
│ /usr/local/dev/textual-demo/.venv/lib64/python3.9/site-packages/textual/worker.py:354 in run │
│ │
│ 351 │ │ Returns: ╭───────────────────────────────── locals ─────────────────────────────────╮ │
│ 352 │ │ │ Return value of the work. │ self = <Worker ERROR name='run_stopwatch' description='run_stopwatch()'> │ │
│ 353 │ │ """ ╰──────────────────────────────────────────────────────────────────────────╯ │
│ ❱ 354 │ │ return await ( │
│ 355 │ │ │ self._run_threaded() if self._thread_worker else self._run_async() │
│ 356 │ │ ) │
│ 357 │
│ │
│ /usr/local/dev/textual-demo/.venv/lib64/python3.9/site-packages/textual/worker.py:339 in _run_async │
│ │
│ 336 │ │ │ or hasattr(self._work, "func") ╭───────────────────────────────── locals ─────────────────────────────────╮ │
│ 337 │ │ │ and inspect.iscoroutinefunction(self._work.func) │ self = <Worker ERROR name='run_stopwatch' description='run_stopwatch()'> │ │
│ 338 │ │ ): ╰──────────────────────────────────────────────────────────────────────────╯ │
│ ❱ 339 │ │ │ return await self._work() │
│ 340 │ │ elif inspect.isawaitable(self._work): │
│ 341 │ │ │ return await self._work │
│ 342 │ │ elif callable(self._work): │
│ │
│ /usr/local/dev/textual-demo/app.py:59 in run_stopwatch │
│ │
│ 56 │ │ │ await asyncio.sleep(0.01) ╭───────────── locals ─────────────╮ │
│ 57 │ │ │ elapsed = time.monotonic() - start │ elapsed = 0.048568773083388805 │ │
│ 58 │ │ │ │ self = ScreenB() │ │
│ ❱ 59 │ │ │ widget = self.screen.query_one("#stopwatch", Static) │ start = 1652181.945696306 │ │
│ 60 │ │ │ widget.update(f"{elapsed:.3f}s") │ widget = Static(id='stopwatch') │ │
│ 61 ╰──────────────────────────────────╯ │
│ 62 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
NoMatches: No nodes match '#stopwatch' on ScreenB()
Analysis
It seems like the child widgets of ScreenB are removed from the DOM before the worker is fully cancelled when popping the screen. The worker tries to access the Static widget, but it has already been removed, which leads to a crash.
Looking closer, the _message_loop_exit() method runs after a node is pruned from the tree. As you can see, this prunes the children as well, and it waits for them to fully exit before cancelling the workers on the current node by emitting the Unmount event (also see _on_unmount()).
Discussion
Personally, I expected the worker to be closed before the pruning of the DOM elements.
Right now, the code in the workers must be written in a way that accounts for the DOM being in a "pruning" state, using query_one_optional and handling the cases in which it returns None.
I guess that one solution would be to cancel the workers before posting the Prune message to the children.