Ensure user-allocated cancel scope just works!
Turns out the nursery doesn't have to care about allocating a per task
`CancelScope` since the user can just do that in the
`@task_scope_manager` if desired B) So just mask all the nursery cs
allocating with the intention of removal.
Also add a test for per-task-cancellation by starting the crash task as
a `trio.sleep_forever()` but then cancel it via the user allocated cs
and ensure the crash propagates as expected 💥
			
			
			
		
							parent
							
								
									e468f62c26
								
							
						
					
					
						commit
						e7b3254b7b
					
				|  | @ -126,18 +126,23 @@ class ScopePerTaskNursery(Struct): | ||||||
|         # task = new_tasks.pop() |         # task = new_tasks.pop() | ||||||
| 
 | 
 | ||||||
|         n: Nursery = self._n |         n: Nursery = self._n | ||||||
|         cs = CancelScope() | 
 | ||||||
|  |         sm = self.scope_manager | ||||||
|  |         # we do default behavior of a scope-per-nursery | ||||||
|  |         # if the user did not provide a task manager. | ||||||
|  |         if sm is None: | ||||||
|  |             return n.start_soon(async_fn, *args, name=None) | ||||||
|  | 
 | ||||||
|  |         # per_task_cs = CancelScope() | ||||||
|         new_task: Task | None = None |         new_task: Task | None = None | ||||||
|         to_return: tuple[Any] | None = None |         to_return: tuple[Any] | None = None | ||||||
| 
 | 
 | ||||||
|         sm = self.scope_manager |  | ||||||
|         if sm is None: |  | ||||||
|             mngr = nullcontext([cs]) |  | ||||||
|         else: |  | ||||||
|         # NOTE: what do we enforce as a signature for the |         # NOTE: what do we enforce as a signature for the | ||||||
|         # `@task_scope_manager` here? |         # `@task_scope_manager` here? | ||||||
|             mngr = sm(nursery=n) |         mngr = sm( | ||||||
| 
 |             nursery=n, | ||||||
|  |             # scope=per_task_cs, | ||||||
|  |         ) | ||||||
|         async def _start_wrapped_in_scope( |         async def _start_wrapped_in_scope( | ||||||
|             task_status: TaskStatus[ |             task_status: TaskStatus[ | ||||||
|                 tuple[CancelScope, Task] |                 tuple[CancelScope, Task] | ||||||
|  | @ -148,10 +153,12 @@ class ScopePerTaskNursery(Struct): | ||||||
|             # TODO: this was working before?! |             # TODO: this was working before?! | ||||||
|             # nonlocal to_return |             # nonlocal to_return | ||||||
| 
 | 
 | ||||||
|             with cs: |  | ||||||
| 
 |  | ||||||
|             task = trio.lowlevel.current_task() |             task = trio.lowlevel.current_task() | ||||||
|                 self._scopes[cs] = task |             # self._scopes[per_task_cs] = task | ||||||
|  | 
 | ||||||
|  |             # NOTE: we actually don't need this since the user can | ||||||
|  |             # just to it themselves inside mngr! | ||||||
|  |             # with per_task_cs: | ||||||
| 
 | 
 | ||||||
|             # execute up to the first yield |             # execute up to the first yield | ||||||
|             try: |             try: | ||||||
|  | @ -178,8 +185,7 @@ class ScopePerTaskNursery(Struct): | ||||||
|             try: |             try: | ||||||
|                 mngr.send(outcome) |                 mngr.send(outcome) | ||||||
| 
 | 
 | ||||||
|                     # NOTE: this will instead send the underlying | 
 | ||||||
|                     # `.value`? Not sure if that's better or not? |  | ||||||
|                 # I would presume it's better to have a handle to |                 # I would presume it's better to have a handle to | ||||||
|                 # the `Outcome` entirely? This method sends *into* |                 # the `Outcome` entirely? This method sends *into* | ||||||
|                 # the mngr this `Outcome.value`; seems like kinda |                 # the mngr this `Outcome.value`; seems like kinda | ||||||
|  | @ -200,7 +206,6 @@ class ScopePerTaskNursery(Struct): | ||||||
|         return to_return |         return to_return | ||||||
| 
 | 
 | ||||||
| 
 | 
 | ||||||
| 
 |  | ||||||
| # TODO: you could wrap your output task handle in this? | # TODO: you could wrap your output task handle in this? | ||||||
| # class TaskHandle(Struct): | # class TaskHandle(Struct): | ||||||
| #     task: Task | #     task: Task | ||||||
|  | @ -214,6 +219,11 @@ class ScopePerTaskNursery(Struct): | ||||||
| def add_task_handle_and_crash_handling( | def add_task_handle_and_crash_handling( | ||||||
|     nursery: Nursery, |     nursery: Nursery, | ||||||
| 
 | 
 | ||||||
|  |     # TODO: is this the only way we can have a per-task scope | ||||||
|  |     # allocated or can we allow the user to somehow do it if | ||||||
|  |     # they want below? | ||||||
|  |     # scope: CancelScope, | ||||||
|  | 
 | ||||||
| ) -> Generator[None, list[Any]]: | ) -> Generator[None, list[Any]]: | ||||||
| 
 | 
 | ||||||
|     task_outcome = TaskOutcome() |     task_outcome = TaskOutcome() | ||||||
|  | @ -222,8 +232,12 @@ def add_task_handle_and_crash_handling( | ||||||
|     task: Task = trio.lowlevel.current_task() |     task: Task = trio.lowlevel.current_task() | ||||||
|     print(f'Spawning task: {task.name}') |     print(f'Spawning task: {task.name}') | ||||||
| 
 | 
 | ||||||
|  |     # yields back when task is terminated, cancelled, returns. | ||||||
|     try: |     try: | ||||||
|         # yields back when task is terminated, cancelled, returns? |         # XXX: wait, this isn't doing anything right since we'd have to | ||||||
|  |         # manually activate this scope using something like: | ||||||
|  |         # `task._activate_cancel_status(cs._cancel_status)` ?? | ||||||
|  |         # oh wait, but `.__enter__()` does all that already? | ||||||
|         with CancelScope() as cs: |         with CancelScope() as cs: | ||||||
| 
 | 
 | ||||||
|             # the yielded value(s) here are what are returned to the |             # the yielded value(s) here are what are returned to the | ||||||
|  | @ -260,6 +274,19 @@ async def sleep_then_return_val(val: str): | ||||||
|     return val |     return val | ||||||
| 
 | 
 | ||||||
| 
 | 
 | ||||||
|  | async def ensure_cancelled(): | ||||||
|  |     try: | ||||||
|  |         await trio.sleep_forever() | ||||||
|  | 
 | ||||||
|  |     except trio.Cancelled: | ||||||
|  |         task = trio.lowlevel.current_task() | ||||||
|  |         print(f'heyyo ONLY {task.name} was cancelled as expected B)') | ||||||
|  |         assert 0 | ||||||
|  | 
 | ||||||
|  |     except BaseException: | ||||||
|  |         raise RuntimeError("woa woa woa this ain't right!") | ||||||
|  | 
 | ||||||
|  | 
 | ||||||
| if __name__ == '__main__': | if __name__ == '__main__': | ||||||
| 
 | 
 | ||||||
|     async def main(): |     async def main(): | ||||||
|  | @ -267,17 +294,22 @@ if __name__ == '__main__': | ||||||
|             scope_manager=add_task_handle_and_crash_handling, |             scope_manager=add_task_handle_and_crash_handling, | ||||||
|         ) as sn: |         ) as sn: | ||||||
|             for _ in range(3): |             for _ in range(3): | ||||||
|                 outcome, cs = await sn.start_soon(trio.sleep_forever) |                 outcome, _ = await sn.start_soon(trio.sleep_forever) | ||||||
| 
 | 
 | ||||||
|             # extra task we want to engage in debugger post mortem. |             # extra task we want to engage in debugger post mortem. | ||||||
|             err_outcome, *_ = await sn.start_soon(sleep_then_err) |             err_outcome, cs = await sn.start_soon(ensure_cancelled) | ||||||
| 
 | 
 | ||||||
|             val: str = 'yoyoyo' |             val: str = 'yoyoyo' | ||||||
|             val_outcome, cs = await sn.start_soon(sleep_then_return_val, val) |             val_outcome, _ = await sn.start_soon( | ||||||
|  |                 sleep_then_return_val, | ||||||
|  |                 val, | ||||||
|  |             ) | ||||||
|             res = await val_outcome.wait_for_result() |             res = await val_outcome.wait_for_result() | ||||||
|             assert res == val |             assert res == val | ||||||
|             print(f'GOT EXPECTED TASK VALUE: {res}') |             print(f'{res} -> GOT EXPECTED TASK VALUE') | ||||||
| 
 | 
 | ||||||
|             print('WAITING FOR CRASH..') |             await trio.sleep(0.6) | ||||||
|  |             print('Cancelling and waiting for CRASH..') | ||||||
|  |             cs.cancel() | ||||||
| 
 | 
 | ||||||
|     trio.run(main) |     trio.run(main) | ||||||
|  |  | ||||||
		Loading…
	
		Reference in New Issue