diff --git a/helpers/defer.py b/helpers/defer.py index 052c264afe..5db56e727d 100644 --- a/helpers/defer.py +++ b/helpers/defer.py @@ -8,11 +8,22 @@ THREAD_BACKGROUND = "Background" +# How long to wait for a loop's pending tasks to finish cancelling before it is +# torn down. Reached only when a task ignores cancellation; the bound is what +# keeps that from wedging the caller. +DRAIN_TIMEOUT = 10.0 + class EventLoopThread: _instances: dict[str, "EventLoopThread"] = {} _lock = threading.Lock() + loop: Optional[asyncio.AbstractEventLoop] + thread: Optional[threading.Thread] + # How many live ``DeferredTask``s share this thread. Instances are keyed by + # name, so the count is what decides whether a task may stop the loop. + _users: int + def __init__(self, thread_name: str = THREAD_BACKGROUND) -> None: """Initialize the event loop thread.""" self.thread_name = thread_name @@ -22,23 +33,70 @@ def __new__(cls, thread_name: str = THREAD_BACKGROUND): with cls._lock: if thread_name not in cls._instances: instance = super(EventLoopThread, cls).__new__(cls) + # Set here rather than in ``__init__``: instances are shared by + # name, so ``__init__`` runs again for every task that joins an + # existing thread and would reset the count. + instance._users = 0 cls._instances[thread_name] = instance return cls._instances[thread_name] + def acquire(self) -> None: + """Registers a task as a user of this shared thread.""" + with self.__class__._lock: + self._users += 1 + + def release(self, terminate: bool) -> bool: + """Drops a user, reporting whether the caller may tear the loop down. + + Only the last user may: the instance is shared by name, so stopping the + loop while another task still runs on it would cancel that task's work. + + Unregistering happens under the same lock that makes the decision, so a + task created after this point gets a fresh thread rather than one that + is on its way out. + """ + with self.__class__._lock: + if self._users > 0: + self._users -= 1 + if not (terminate and self._users == 0): + return False + if self.__class__._instances.get(self.thread_name) is self: + del self.__class__._instances[self.thread_name] + return True + def _start(self): - if not hasattr(self, "loop") or not self.loop: - self.loop = asyncio.new_event_loop() - if not hasattr(self, "thread") or not self.thread: + loop = getattr(self, "loop", None) + thread = getattr(self, "thread", None) + # A closed loop counts as absent, not just a null one: ``terminate()`` + # called from the loop's own thread cannot null these attributes out -- + # it is running inside the callback the loop still has to return from -- + # so a torn-down instance keeps a stale loop and thread attached. + # + # Both are rebuilt together. Replacing only the loop would leave it + # unattended, since the surviving thread runs the loop it was handed. + if loop is None or loop.is_closed() or thread is None or not thread.is_alive(): + self.loop = loop = asyncio.new_event_loop() self.thread = threading.Thread( - target=self._run_event_loop, daemon=True, name=self.thread_name + target=self._run_event_loop, + args=(loop,), + daemon=True, + name=self.thread_name, ) self.thread.start() - def _run_event_loop(self): - if not self.loop: - raise RuntimeError("Event loop is not initialized") - asyncio.set_event_loop(self.loop) - self.loop.run_forever() + def _run_event_loop(self, loop: asyncio.AbstractEventLoop): + # The loop is passed in rather than read from ``self.loop``: a restarted + # instance replaces that attribute, and this thread must keep running + # the loop it was created for. + asyncio.set_event_loop(loop) + try: + loop.run_forever() + finally: + # Closed here rather than by ``terminate()``, which may itself be + # running *on* this thread, where ``close()`` would raise because + # the loop is still running. + if not loop.is_closed(): + loop.close() def terminate(self): loop = getattr(self, "loop", None) @@ -47,25 +105,32 @@ def terminate(self): if not loop: return - if loop.is_running(): - if thread and thread is threading.current_thread(): - loop.stop() - else: + # ``terminate()`` can reach here from inside its own loop, via a task's + # done-callback killing a child that shares this thread. Joining would + # then wait on the current thread and deadlock. + on_own_thread = thread is not None and thread is threading.current_thread() + + # Scheduled rather than conditional on ``is_running()``: a thread that + # has started but not yet entered ``run_forever`` reports False, and + # joining it would block for the full timeout. + if on_own_thread: + loop.call_soon(loop.stop) + else: + try: loop.call_soon_threadsafe(loop.stop) - if thread: - thread.join() - elif thread and thread.is_alive() and thread is not threading.current_thread(): - thread.join() - - if not loop.is_closed(): - loop.close() + except RuntimeError: + # Already closed. + pass + if thread is not None and thread.is_alive(): + thread.join(timeout=DRAIN_TIMEOUT) with self.__class__._lock: - if self.thread_name in self.__class__._instances: + if self.__class__._instances.get(self.thread_name) is self: del self.__class__._instances[self.thread_name] - self.loop = None - self.thread = None + if not on_own_thread: + self.loop = None + self.thread = None def run_coroutine(self, coro): self._start() @@ -86,6 +151,8 @@ def __init__( thread_name: str = THREAD_BACKGROUND, ): self.event_loop_thread = EventLoopThread(thread_name) + self.event_loop_thread.acquire() + self._released = False self._future: Optional[Future] = None self.children: list[ChildTask] = [] self.func: Optional[Callable[..., Coroutine[Any, Any, Any]]] = None @@ -108,6 +175,14 @@ def _start_task(self): if self.func is None: raise RuntimeError("Task callable is no longer available") + if self._released: + # ``restart()`` goes through ``kill()``, which drops this task's + # claim on the thread. Re-register before using it again, or a + # sibling's later ``kill(terminate_thread=True)`` would see a count + # that no longer includes this task and stop the loop under it. + self._released = False + self.event_loop_thread.acquire() + self._future = self.event_loop_thread.run_coroutine( self._run(self.func, self.args, self.kwargs) ) @@ -161,23 +236,64 @@ def _get_result(): return await loop.run_in_executor(None, _get_result) def kill(self, terminate_thread: bool = False) -> None: - """Kill the task and optionally terminate its thread.""" + """Kill the task and optionally terminate its thread. + + ``terminate_thread`` is honoured only once this is the *last* task using + the thread. ``EventLoopThread`` instances are shared by name, so + stopping the loop while a sibling task still runs on it would cancel + that sibling's work -- silently, since the cancellation surfaces only on + the abandoned task's own future. + """ self.kill_children() if self._future and not self._future.done(): self._future.cancel() self._clear_call() - if terminate_thread and self.event_loop_thread.loop: - if self.event_loop_thread.loop.is_running(): + # ``release`` is idempotent per task: killing twice must not drop the + # count twice and let the thread go while a sibling still needs it. + may_terminate = False + if not self._released: + self._released = True + may_terminate = self.event_loop_thread.release(terminate_thread) + + if not may_terminate: + return + + event_loop_thread = self.event_loop_thread + loop = event_loop_thread.loop + if loop is None: + return + + if event_loop_thread.thread is threading.current_thread(): + # Waiting on the drain from inside the loop would deadlock: the + # coroutine can only advance on this thread, which the wait would be + # occupying. The drain and the teardown are chained into a task + # instead, so the loop is not stopped out from under the drain. + # ``release()`` has already unregistered the instance, so a task + # created before this finishes gets a fresh thread rather than this + # one on its way out. + async def drain_then_terminate() -> None: try: - cleanup_future = asyncio.run_coroutine_threadsafe( - self._drain_event_loop_tasks(), self.event_loop_thread.loop - ) - cleanup_future.result() - except Exception: - pass - - self.event_loop_thread.terminate() + await self._drain_event_loop_tasks() + finally: + event_loop_thread.terminate() + + loop.create_task(drain_then_terminate()) + return + + if loop.is_running(): + try: + cleanup_future = asyncio.run_coroutine_threadsafe( + self._drain_event_loop_tasks(), loop + ) + # Bounded: an unbounded wait hangs the caller outright if a task + # swallows cancellation, and this runs on the request thread + # that serves chat deletion and task deletion. + cleanup_future.result(timeout=DRAIN_TIMEOUT) + except Exception: + pass + + event_loop_thread.terminate() def kill_children(self) -> None: for child in self.children: diff --git a/helpers/defer.py.dox.md b/helpers/defer.py.dox.md index 4e234c89fa..87290e7fbb 100644 --- a/helpers/defer.py.dox.md +++ b/helpers/defer.py.dox.md @@ -12,6 +12,8 @@ - `defer.py.dox.md` owns durable notes about responsibilities, contracts, side effects, and verification for that implementation. - Classes: - `EventLoopThread` (no explicit base class) + - `acquire(self) -> None` + - `release(self, terminate: bool) -> bool` - `terminate(self)` - `run_coroutine(self, coro)` - `ChildTask` (no explicit base class) @@ -24,13 +26,17 @@ - `kill_children(self) -> None` - `is_alive(self) -> bool` - `restart(self, terminate_thread: bool=...) -> None` -- Notable constants/configuration names: `T`, `THREAD_BACKGROUND`. +- Notable constants/configuration names: `T`, `THREAD_BACKGROUND`, `DRAIN_TIMEOUT`. ## Runtime Contracts - Helper modules own reusable framework APIs and must preserve public callers unless all callers, tests, and docs are updated together. - `DeferredTask` retains its callable and arguments only while an invocation is active; completion and `kill()` clear those references after the running coroutine has taken its own snapshot. - Task results remain available after completion. `restart()` can restart an active invocation, but a completed invocation has no retained call recipe and must be started again explicitly. +- `EventLoopThread` instances are shared per thread name, so a name is a shared resource rather than a private one. Every `DeferredTask` registers with `acquire()` on construction and drops its claim in `kill()`; `kill(terminate_thread=True)` tears the loop down only once the count reaches zero. Ignoring the count cancels every sibling task on the same name, and the cancellation surfaces only on the abandoned task's own future, so it is otherwise silent. `TaskScheduler` shares one name across all scheduled tasks and `THREAD_BACKGROUND` is shared process-wide. +- `kill()` and `terminate()` may run on the loop's own thread, since a task's done callback runs there and kills its children. Neither may join that thread or wait on a coroutine that only that thread can advance; the in-loop path chains the drain and the stop onto the loop instead and leaves the loop to close itself as `run_forever` returns. +- Because the in-loop path cannot clear the instance's `loop`/`thread` attributes, `_start` treats a closed loop or a dead thread as absent and rebuilds both together, so reusing a torn-down thread name yields a working loop. +- Waits on teardown are bounded by `DRAIN_TIMEOUT`: `kill(terminate_thread=True)` is reached from request handlers (chat deletion, scheduler task deletion), so a task that ignores cancellation must not hang them. - Update this file whenever public functions, classes, persistence behavior, path/security assumptions, side effects, or cross-module contracts change. - Observed side-effect areas: scheduler state. - Imported dependency areas include: `asyncio`, `concurrent.futures`, `dataclasses`, `threading`, `typing`. @@ -49,7 +55,9 @@ ## Verification - Run targeted tests for changed helper behavior; run security regressions for auth, filesystem, WebSocket, tunnel, upload, or secret-handling helpers. +- Shared-thread and teardown behavior is covered by `tests/test_defer_lifecycle.py`. Run those tests with a per-test timeout: a regression in the in-loop teardown path wedges the loop thread permanently rather than failing, which stalls the run instead of reporting. - Related tests observed by source search: + - `tests/test_defer_lifecycle.py` - `tests/test_office_document_store.py` ## Child DOX Index diff --git a/tests/test_defer_lifecycle.py b/tests/test_defer_lifecycle.py index 879029aaf4..6b2f571065 100644 --- a/tests/test_defer_lifecycle.py +++ b/tests/test_defer_lifecycle.py @@ -1,5 +1,6 @@ import asyncio import threading +import time import uuid import weakref @@ -113,3 +114,283 @@ async def run(value): assert task.args == ("argument",) finally: task.kill(terminate_thread=True) + + +def wait_until(predicate, timeout: float = 5.0) -> bool: + """Polls ``predicate`` until it holds, so teardown need not be raced. + + Teardown finishes on the loop's own thread in one case, so there is no + future to wait on -- but polling still fails within the timeout rather than + passing on a sleep that happened to be long enough. + """ + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.02) + return predicate() + + +def test_killing_one_task_does_not_cancel_siblings_on_the_shared_thread(): + """``EventLoopThread`` instances are shared by name, so one task's + ``kill(terminate_thread=True)`` used to stop a loop that its siblings were + still running on, cancelling their work. + + ``TaskScheduler`` passes a single fixed thread name for every scheduled + task (`helpers/task_scheduler.py`), so deleting one running task took down + every other running task in the process. The cancellation surfaces only on + the abandoned task's own future, so nothing reported it. + """ + name = f"defer-shared-{uuid.uuid4()}" + victim, survivor = DeferredTask(name), DeferredTask(name) + assert victim.event_loop_thread is survivor.event_loop_thread + + victim_started = threading.Event() + survivor_started = threading.Event() + progress: list[int] = [] + + async def blocks_forever(): + victim_started.set() + await asyncio.Future() + + async def keeps_working(): + survivor_started.set() + for step in range(10): + await asyncio.sleep(0.05) + progress.append(step) + return "finished" + + try: + victim.start_task(blocks_forever) + survivor.start_task(keeps_working) + assert victim_started.wait(5) and survivor_started.wait(5) + + victim.kill(terminate_thread=True) + + # The survivor's own result, not a progress count: a cancelled task + # would raise here, while asserting only that the counter advanced + # could pass on work done before the kill landed. + assert survivor.result_sync(timeout=5) == "finished" + finally: + survivor.kill(terminate_thread=True) + victim.kill(terminate_thread=True) + + +def test_last_task_to_be_killed_still_tears_the_thread_down(): + """Deferring teardown to the last user must not mean never tearing down. + + Otherwise the fix for the shared-thread cancellation would trade it for a + thread leak, which matters most for the per-request thread names + (``BrowserRuntime-``) that are created and discarded constantly. + """ + name = f"defer-teardown-{uuid.uuid4()}" + first, second = DeferredTask(name), DeferredTask(name) + event_loop_thread = first.event_loop_thread + loop, thread = event_loop_thread.loop, event_loop_thread.thread + assert loop is not None and thread is not None + + async def idle(): + await asyncio.Future() + + first.start_task(idle) + second.start_task(idle) + + first.kill(terminate_thread=True) + assert not loop.is_closed(), "the loop went down while a sibling still held it" + + second.kill(terminate_thread=True) + assert wait_until(lambda: loop.is_closed()), "the loop outlived its last user" + assert wait_until(lambda: not thread.is_alive()), "the thread outlived its loop" + + +def test_a_new_task_after_teardown_gets_a_working_thread(): + """Reusing a thread name after it was torn down must not hand back a dead + loop: instances are cached by name, and a terminated one stays cached until + it is replaced. + + This one passes on the unfixed code too. It is here as a guard on the fix + rather than on the bug: teardown now leaves a closed loop attached in the + in-loop case, so ``_start`` has to treat a closed loop as absent. The + reuse-after-teardown path is what would break if it did not. + """ + name = f"defer-reuse-{uuid.uuid4()}" + first = DeferredTask(name) + + async def done(): + return "first" + + first.start_task(done) + assert first.result_sync(timeout=5) == "first" + first.kill(terminate_thread=True) + + second = DeferredTask(name) + + async def again(): + return "second" + + try: + second.start_task(again) + assert second.result_sync(timeout=5) == "second" + finally: + second.kill(terminate_thread=True) + + +def test_killing_from_the_loops_own_thread_does_not_deadlock(): + """A task's done-callback runs *on* the loop thread, and it kills children. + + A child that shares the parent's thread therefore reached + ``kill(terminate_thread=True)`` from inside the loop, where the old code + waited on ``run_coroutine_threadsafe(...).result()`` with no timeout -- a + coroutine that only that thread can advance. The thread wedged permanently, + taking every task on it with it. + """ + name = f"defer-own-thread-{uuid.uuid4()}" + child = DeferredTask(name) + parent = DeferredTask(name) + parent.add_child_task(child, terminate_thread=True) + loop = parent.event_loop_thread.loop + assert loop is not None + + child_started = threading.Event() + + async def child_body(): + child_started.set() + await asyncio.Future() + + async def parent_body(): + return "parent done" + + try: + child.start_task(child_body) + assert child_started.wait(5) + parent.start_task(parent_body) + assert parent.result_sync(timeout=5) == "parent done" + + # Probed by scheduling onto the loop rather than by elapsed time: the + # parent's result arrives before the done-callback runs, so a wedge + # would otherwise go unnoticed here. + assert wait_until(lambda: child._future is not None and child._future.done()) + probe = asyncio.run_coroutine_threadsafe(asyncio.sleep(0), loop) + probe.result(timeout=5) + finally: + parent.kill(terminate_thread=True) + + +def test_teardown_requested_from_inside_the_loop_still_completes(): + """The in-loop teardown path cannot join its own thread or close a loop that + is still running, so it chains the drain and the stop onto the loop. It must + still finish, and it must leave the instance reusable rather than holding a + closed loop. + + ``kill()`` is invoked on the loop directly here. The refcount checked above + is what keeps a shared thread from reaching this branch, so within the + current tree there is no indirect route to it -- but ``__del__`` calls + ``kill()`` from wherever the last reference happens to be dropped, so the + branch has to hold on its own. + """ + name = f"defer-in-loop-{uuid.uuid4()}" + task = DeferredTask(name) + event_loop_thread = task.event_loop_thread + loop, thread = event_loop_thread.loop, event_loop_thread.thread + assert loop is not None and thread is not None + + started = threading.Event() + + async def idle(): + started.set() + await asyncio.Future() + + task.start_task(idle) + assert started.wait(5) + + loop.call_soon_threadsafe(lambda: task.kill(terminate_thread=True)) + + assert wait_until(lambda: loop.is_closed()), "the loop never closed" + assert wait_until(lambda: not thread.is_alive()), "the thread never stopped" + + # The in-loop path cannot null out ``loop``/``thread``, so a later task on + # the same name would otherwise be handed the closed loop. + successor = DeferredTask(name) + + async def done(): + return "reusable" + + try: + successor.start_task(done) + assert successor.result_sync(timeout=5) == "reusable" + finally: + successor.kill(terminate_thread=True) + + +def test_a_restarted_task_keeps_its_claim_on_the_thread(): + """``restart()`` goes through ``kill()``, which drops the task's claim. + + Without re-registering, a sibling's later ``kill(terminate_thread=True)`` + would see a count that no longer includes the restarted task and stop the + loop underneath it. + """ + name = f"defer-restart-{uuid.uuid4()}" + restarted = DeferredTask(name) + sibling = DeferredTask(name) + starts = [threading.Event(), threading.Event()] + run_count = 0 + progress: list[int] = [] + + async def run(value): + nonlocal run_count + current = run_count + run_count += 1 + starts[current].set() + if current == 0: + await asyncio.Future() + for step in range(10): + await asyncio.sleep(0.05) + progress.append(step) + return f"restarted:{value}" + + async def idle(): + await asyncio.Future() + + try: + sibling.start_task(idle) + restarted.start_task(run, "argument") + assert starts[0].wait(5) + restarted.restart() + assert starts[1].wait(5) + + sibling.kill(terminate_thread=True) + + assert restarted.result_sync(timeout=5) == "restarted:argument" + finally: + restarted.kill(terminate_thread=True) + sibling.kill(terminate_thread=True) + + +def test_killing_the_same_task_twice_does_not_release_the_thread_twice(): + """``kill()`` is called from ``__del__`` as well as explicitly, and + ``close_runtime_sync`` kills in a ``finally`` after a task that may already + have been killed. A double release would drop a sibling's claim.""" + name = f"defer-double-kill-{uuid.uuid4()}" + killed_twice = DeferredTask(name) + sibling = DeferredTask(name) + progress: list[int] = [] + + async def idle(): + await asyncio.Future() + + async def keeps_working(): + for step in range(10): + await asyncio.sleep(0.05) + progress.append(step) + return "finished" + + try: + killed_twice.start_task(idle) + sibling.start_task(keeps_working) + + killed_twice.kill(terminate_thread=True) + killed_twice.kill(terminate_thread=True) + + assert sibling.result_sync(timeout=5) == "finished" + finally: + sibling.kill(terminate_thread=True)