|
19 | 19 | get_cancelled_exc_class,
|
20 | 20 | sleep,
|
21 | 21 | wait_readable,
|
| 22 | + ClosedResourceError, |
22 | 23 | )
|
23 | 24 | from anyio.abc import TaskGroup, TaskStatus
|
24 |
| -from anyioutils import FIRST_COMPLETED, Future, create_task, wait |
| 25 | +from anyioutils import FIRST_COMPLETED, Future, create_task |
25 | 26 |
|
26 | 27 | import zmq
|
27 | 28 | from zmq import EVENTS, POLLIN, POLLOUT
|
@@ -890,31 +891,24 @@ async def _start(self, *, task_status: TaskStatus[None] = TASK_STATUS_IGNORED):
|
890 | 891 | task_status.started()
|
891 | 892 | self.started.set()
|
892 | 893 | self._thread = get_ident()
|
| 894 | + |
| 895 | + async def wait_or_cancel() -> None: |
| 896 | + assert self.stopped is not None |
| 897 | + await self.stopped.wait() |
| 898 | + tg.cancel_scope.cancel() |
| 899 | + |
893 | 900 | try:
|
894 | 901 | while True:
|
895 |
| - wait_stopped_task = create_task( |
896 |
| - self.stopped.wait(), |
897 |
| - self._task_group, |
898 |
| - exception_handler=ignore_exceptions, |
899 |
| - ) |
900 |
| - tasks = [ |
901 |
| - create_task( |
902 |
| - wait_readable(self._shadow_sock), # type: ignore[arg-type] |
903 |
| - self._task_group, |
904 |
| - exception_handler=ignore_exceptions, |
905 |
| - ), |
906 |
| - wait_stopped_task, |
907 |
| - ] |
908 |
| - done, pending = await wait( |
909 |
| - tasks, self._task_group, return_when=FIRST_COMPLETED |
910 |
| - ) |
911 |
| - for task in pending: |
912 |
| - task.cancel() |
913 |
| - if wait_stopped_task in done: |
| 902 | + async with create_task_group() as tg: |
| 903 | + tg.start_soon(wait_or_cancel) |
| 904 | + await wait_readable(self._shadow_sock) |
| 905 | + if self.stopped.is_set(): |
914 | 906 | break
|
915 | 907 | await self._handle_events()
|
916 |
| - except BaseException: |
917 |
| - pass |
| 908 | + except get_cancelled_exc_class(): |
| 909 | + raise |
| 910 | + except ClosedResourceError: |
| 911 | + self._task_group.cancel_scope.cancel() |
918 | 912 | finally:
|
919 | 913 | self._exited.set()
|
920 | 914 |
|
|
0 commit comments