Tear down blocked task chains iteratively

crates/monty/src/bytecode/vm/scheduler.rs

@@ -9,6 +9,7 @@
 use std::{collections::VecDeque, mem};
 
 use ahash::AHashMap;
+use smallvec::{SmallVec, smallvec};
 
 use crate::{
     asyncio::{Awaiter, CallId, ExternalFutureState, TaskId},
@@ -384,8 +385,24 @@ impl Scheduler {
     /// result, and tears down any inner gather it was blocked on. After this
     /// call the task no longer exists in `Scheduler::tasks`; its owning
     /// references to its coroutine and (outer) gather are released by the
-    /// `Task::drop_with` call at the end.
+    /// `Task::drop_with` in [`Scheduler::cancel_one`].
+    ///
+    /// Drains a worklist rather than recursing into inner gathers: a chain of
+    /// blocked tasks costs no native stack to *build*, so recursive teardown
+    /// turned that stored depth back into frames and overflowed.
     pub fn cancel_task(&mut self, task_id: TaskId, heap: &mut HeapReader<'_>) {
+        let mut pending: SmallVec<[TaskId; 4]> = smallvec![task_id];
+        while let Some(task_id) = pending.pop() {
+            self.cancel_one(task_id, heap, &mut pending);
+        }
+    }
+
+    /// Cancels one task, queueing the tasks spawned under any gather it was
+    /// blocked on for [`Scheduler::cancel_task`] to drain.
+    ///
+    /// Dropping this task ahead of the children it queued is sound: each owns
+    /// an inc_ref on that same gather (see [`Scheduler::spawn`]).
+    fn cancel_one(&mut self, task_id: TaskId, heap: &mut HeapReader<'_>, pending: &mut SmallVec<[TaskId; 4]>) {
         // No-op if the task has already been removed (idempotent — finalization
         // sites may iterate task ids that include already-cancelled siblings).
         let Some(task) = self.tasks.remove(&task_id) else {
@@ -407,37 +424,51 @@ impl Scheduler {
         if !task.is_finished() {
             self.ready_queue.retain(|&id| id != task_id);
 
-            // If blocked on an awaitable, dispatch by kind via `heap.read`.
-            // For a gather: recursively cancel its task children — external
-            // children manage themselves via the owning `Awaiter::GatherSlot`
-            // (the gather stays alive until each external resolves and
-            // releases its inc_ref), but spawned tasks have no such anchor
-            // and would otherwise linger in `self.tasks` holding inc_refs.
-            // For an external future: no extra teardown.
-            if let TaskState::Blocked(blocked_id) = task.state
-                && let HeapReadOutput::GatherFuture(gather) = heap.read(blocked_id)
-            {
-                let inner_task_ids: Vec<TaskId> = gather
-                    .get(heap)
-                    .as_awaited()
-                    .map(|awaited| {
-                        awaited
-                            .pending_children
-                            .keys()
-                            .filter_map(|id| self.coroutine_to_task.get(id).copied())
-                            .collect()
-                    })
-                    .unwrap_or_default();
-                drop(gather);
-                for inner_task_id in inner_task_ids {
-                    self.cancel_task(inner_task_id, heap);
-                }
+            // Blocked on a gather: queue the tasks spawned under it. An
+            // external future needs no extra teardown.
+            if let TaskState::Blocked(blocked_id) = task.state {
+                self.queue_gather_tasks(blocked_id, heap, pending);
             }
         }
 
         task.drop_with(heap);
     }
 
+    /// Queues every task spawned under the gather `root`, walking nested
+    /// gathers iteratively.
+    ///
+    /// A gather item can itself be a gather (`gather(gather(coro()))`), whose
+    /// tasks are just as orphaned as direct coroutine children if left in
+    /// `self.tasks` — they would keep running and then deliver a result to the
+    /// task cancelled here. External children *are* left alone: the owning
+    /// `Awaiter::GatherSlot` anchors them.
+    ///
+    /// Must run while the cancelled task still holds its `Blocked` inc_ref on
+    /// `root`, since the walk takes no references of its own: each nested
+    /// gather is kept alive by its parent's `items`, and the parent in turn by
+    /// the `Awaiter::GatherSlot` inc_ref that nested child holds.
+    fn queue_gather_tasks(&self, root: HeapId, heap: &HeapReader<'_>, pending: &mut SmallVec<[TaskId; 4]>) {
+        // Gathers nest as a tree — a gather may only be awaited once, so the
+        // walk cannot revisit a node and terminates.
+        let mut gathers: SmallVec<[HeapId; 4]> = smallvec![root];
+        while let Some(gather_id) = gathers.pop() {
+            // Coroutine and external children land here too; only gathers have
+            // children of their own to walk.
+            let HeapReadOutput::GatherFuture(gather) = heap.read(gather_id) else {
+                continue;
+            };
+            if let Some(awaited) = gather.get(heap).as_awaited() {
+                for child_id in awaited.pending_children.keys() {
+                    match self.coroutine_to_task.get(child_id) {
+                        Some(&task_id) => pending.push(task_id),
+                        None => gathers.push(*child_id),
+                    }
+                }
+            }
+            drop(gather);
+        }
+    }
+
     /// Records a host-side failure for `call_id` and returns the awaiter the
     /// caller should walk to propagate the error.
     ///

crates/monty/test_cases/async__gather_all.py

@@ -201,3 +201,27 @@ async def boom():
 nested_2 = asyncio.gather(nested_1, nested_1)
 
 assert await nested_2 == [[1, 1], [1, 1]]  # pyright: ignore
+
+
+# === Sibling failure while blocked on a gather of gathers ===
+# The failing sibling cancels the task blocked on `gather(gather(...))`, whose
+# inner gather owns a task of its own; the scheduler must keep running after.
+async def leaf():
+    return 1
+
+
+async def blocked_on_nested():
+    return await asyncio.gather(asyncio.gather(leaf()))
+
+
+async def detonate_sibling():
+    raise ValueError('boom')
+
+
+try:
+    await asyncio.gather(blocked_on_nested(), detonate_sibling())  # pyright: ignore
+    assert False, 'expected the failing sibling to propagate'
+except ValueError as e:
+    assert str(e) == 'boom'
+
+assert await asyncio.gather(leaf(), leaf()) == [1, 1]  # pyright: ignore

crates/monty/test_cases/refcount__gather_nested_gather_cancel.py

@@ -0,0 +1,23 @@
+# Test that a nested GatherFuture *item* is cleaned up when the task awaiting it
+# is cancelled. Unlike refcount__gather_nested_cancel, the inner gather is a direct
+# item of the gather the cancelled task is blocked on, not reached via a coroutine.
+import asyncio
+
+
+async def leaf():
+    return 1
+
+
+async def task_with_gather_item():
+    return await asyncio.gather(asyncio.gather(leaf()))
+
+
+async def task_fail():
+    raise ValueError('outer task failed')
+
+
+try:
+    result = await asyncio.gather(task_with_gather_item(), task_fail())  # pyright: ignore
+except ValueError:
+    pass
+# ref-counts={'asyncio': 1}

crates/monty/tests/asyncio.rs

@@ -3,6 +3,8 @@
 //! These tests verify the behavior of the async execution model, specifically around
 //! resolving external futures incrementally via `ResolveFutures::resume()`.
 
+use std::thread;
+
 use monty::{MontyRun, ResolveFutures, RunProgress};
 use monty_types::{
     CompileOptions, ExcType, ExtFunctionResult, MontyException, MontyObject, NameLookupResult, PrintWriter,
@@ -969,6 +971,70 @@ await main()
     assert_eq!(result, MontyObject::Int(333));
 }
 
+// === Test: Deep blocked task chains are torn down without recursing ===
+
+/// A chain of blocked tasks costs no native stack to *build*, so teardown must
+/// not turn that stored depth back into frames.
+///
+/// Runs on a 2 MiB thread — a worker's budget, and where the abort was seen;
+/// libtest's 8 MiB would need a far deeper, slower chain to prove the same.
+#[test]
+fn deep_blocked_task_chain_teardown_does_not_overflow_the_stack() {
+    thread::Builder::new()
+        .stack_size(2 * 1024 * 1024)
+        .spawn(fail_sibling_of_deep_task_chain)
+        .expect("spawning the bounded-stack thread")
+        .join()
+        .expect("tearing down a deep blocked chain must not overflow the stack");
+}
+
+/// Wraps `leaf()` in 20,000 nested gathers and awaits that chain alongside a
+/// `sibling()`, so both park on external calls: the chain's `parked` (never
+/// resolved) and the sibling's `doomed`. Resolving `doomed` with an error
+/// fails the outer gather, which cancels all 20,000 blocked tasks in one walk,
+/// and asserts that error surfaces as the run's `ValueError`.
+///
+/// Failing the *sibling* is what makes it a single deep walk — failing the
+/// chain's own future would instead unwind it level by level.
+fn fail_sibling_of_deep_task_chain() {
+    let code = r"
+import asyncio
+
+async def leaf():
+    return await parked(1)
+
+async def wrap(g):
+    return await g
+
+async def sibling():
+    return await doomed(2)
+
+g = leaf()
+for _ in range(20000):
+    g = asyncio.gather(wrap(g))
+await asyncio.gather(g, sibling())
+";
+    let runner = MontyRun::new(code.to_owned(), "test.py", vec![], CompileOptions::default()).unwrap();
+    let progress = runner
+        .start(vec![], ResourceTracker::default(), PrintWriter::Stdout)
+        .unwrap();
+
+    let (state, calls) = drive_collecting_calls(progress);
+    let doomed_id = calls
+        .iter()
+        .find_map(|(id, name)| (name == "doomed").then_some(*id))
+        .expect("the sibling should have parked on an external call");
+    assert_eq!(calls.len(), 2, "the chain's leaf and the sibling should both park");
+
+    // Failing the sibling tears down the enclosing gather, cancelling the
+    // chain top-down; the exception itself only walks up to the main task.
+    let error = MontyException::new(ExcType::ValueError, Some("sibling failed".to_string()));
+    let result = state.resume(vec![(doomed_id, ExtFunctionResult::Error(error))], PrintWriter::Stdout);
+
+    let exc = result.expect_err("the failed sibling should surface as an exception");
+    assert_eq!(exc.exc_type(), ExcType::ValueError);
+}
+
 /// Propagating a failure through deeply nested gather waiters must not recurse
 /// on the native Rust stack.
 #[test]

limitations/asyncio.md

@@ -53,3 +53,21 @@ Concurrency is cooperative and host-driven. `gather` suspends Monty whenever
 every branch is blocked on an external call, hands the pending calls to the
 host, and resumes when the host returns results. There is no preemption, no
 threads, and no in-sandbox scheduler.
+
+### A failing `gather` cancels its siblings
+
+When one child of a `gather` raises, every sibling still running is cancelled where it is blocked and never resumes.
+That includes the tasks of any gather a sibling was itself awaiting.
+CPython leaves those siblings running as tasks on the loop, so:
+
+```python
+async def worker():
+    for _ in range(3):
+        await asyncio.gather(step())
+    done.append('finished')
+```
+
+appends `'finished'` under CPython after a sibling of `worker()` raises, but not under Monty.
+
+External calls already passed to the host are not cancelled.
+The host still resolves them and the results are discarded.