Skip to content

Commit ed736ed

Browse files
committed
extmod/uasyncio: Backport py3.11 asyncio's taskgroups.
This patch adds a mostly-compatible backport of CPython 3.11's `TaskGroup` class to uasyncio. Also, there is a new `run_server` method in uasyncio/stream.py which supports task groups and doesn't run in the background (no need). TODO: Write a couple of tests. Closes #8508.
1 parent f12754a commit ed736ed

File tree

4 files changed

+296
-2
lines changed

4 files changed

+296
-2
lines changed

extmod/uasyncio/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@
1414
"Lock": "lock",
1515
"open_connection": "stream",
1616
"start_server": "stream",
17+
"run_server": "stream",
1718
"StreamReader": "stream",
1819
"StreamWriter": "stream",
20+
"TaskGroup": "taskgroup",
1921
}
2022

2123
# Lazy loader, effectively does:

extmod/uasyncio/core.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,14 @@
1515
# Exceptions
1616

1717

18+
class BaseExceptionGroup(BaseException):
19+
pass
20+
21+
22+
class ExceptionGroup(Exception): # TODO cannot also inherit from BaseExceptionGroup
23+
pass
24+
25+
1826
class CancelledError(BaseException):
1927
pass
2028

extmod/uasyncio/stream.py

Lines changed: 42 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,8 @@ async def __aenter__(self):
1717
return self
1818

1919
async def __aexit__(self, exc_type, exc, tb):
20-
await self.close()
20+
self.s.close()
21+
pass
2122

2223
def close(self):
2324
pass
@@ -127,7 +128,7 @@ async def _serve(self, s, cb):
127128

128129

129130
# Helper function to start a TCP stream server, running as a new task
130-
# TODO could use an accept-callback on socket read activity instead of creating a task
131+
# DOES NOT USE TASKGROUPS. Use run_server instead
131132
async def start_server(cb, host, port, backlog=5):
132133
import usocket as socket
133134

@@ -145,6 +146,45 @@ async def start_server(cb, host, port, backlog=5):
145146
return srv
146147

147148

149+
# Helper task to run a TCP stream server
150+
# Callbacks may run in a different taskgroup
151+
async def run_server(cb, host, port, backlog=5, taskgroup=None):
152+
import usocket as socket
153+
154+
# Create and bind server socket.
155+
host = socket.getaddrinfo(host, port)[0] # TODO this is blocking!
156+
s = socket.socket()
157+
s.setblocking(False)
158+
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
159+
s.bind(host[-1])
160+
s.listen(backlog)
161+
162+
if taskgroup is None:
163+
async with TaskGroup() as tg:
164+
await _run_server(tg, s, cb)
165+
else:
166+
await _run_server(taskgroup, s, cb)
167+
168+
169+
async def _run_server(tg, s, cb):
170+
while True:
171+
try:
172+
yield core._io_queue.queue_read(s)
173+
except core.CancelledError:
174+
# Shutdown server
175+
s.close()
176+
return
177+
try:
178+
s2, addr = s.accept()
179+
except Exception:
180+
# Ignore a failed accept
181+
continue
182+
183+
s2.setblocking(False)
184+
s2s = Stream(s2, {"peername": addr})
185+
tg.create_task(cb(s2s, s2s))
186+
187+
148188
################################################################################
149189
# Legacy uasyncio compatibility
150190

extmod/uasyncio/taskgroup.py

Lines changed: 244 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,244 @@
1+
# Adapted with permission from the EdgeDB project.
2+
# Adapted for MicroPython
3+
4+
5+
__all__ = ["TaskGroup"]
6+
7+
from . import core
8+
from . import event
9+
from . import task
10+
11+
DEBUG = False
12+
13+
14+
class TaskGroup:
15+
def __init__(self):
16+
self._entered = False
17+
self._exiting = False
18+
self._aborting = False
19+
self._loop = None
20+
self._parent_task = None
21+
self._parent_cancel_requested = False
22+
self._tasks = set()
23+
self._errors = []
24+
self._base_error = None
25+
self._on_completed = None
26+
27+
def __repr__(self):
28+
info = [""]
29+
if self._tasks:
30+
info.append(f"tasks={len(self._tasks)}")
31+
if self._errors:
32+
info.append(f"errors={len(self._errors)}")
33+
if self._aborting:
34+
info.append("cancelling")
35+
elif self._entered:
36+
info.append("entered")
37+
38+
info_str = " ".join(info)
39+
return f"<TaskGroup{info_str}>"
40+
41+
async def __aenter__(self):
42+
if self._entered:
43+
raise RuntimeError(f"TaskGroup {repr(self)} has been already entered")
44+
self._entered = True
45+
46+
if self._loop is None:
47+
self._loop = core.get_event_loop()
48+
49+
self._parent_task = core.current_task()
50+
if self._parent_task is None:
51+
raise RuntimeError(f"TaskGroup {repr(self)} cannot determine the parent task")
52+
53+
return self
54+
55+
async def __aexit__(self, et, exc, tb):
56+
self._exiting = True
57+
propagate_cancellation_error = None
58+
59+
if exc is not None and self._is_base_error(exc) and self._base_error is None:
60+
self._base_error = exc
61+
62+
if et is not None:
63+
if et is core.CancelledError:
64+
# micropython doesn't have an uncancel counter
65+
# if self._parent_cancel_requested and not self._parent_task.uncancel():
66+
# Do nothing, i.e. swallow the error.
67+
# pass
68+
# else:
69+
propagate_cancellation_error = exc
70+
71+
if not self._aborting:
72+
# Our parent task is being cancelled:
73+
#
74+
# async with TaskGroup() as g:
75+
# g.create_task(...)
76+
# await ... # <- CancelledError
77+
#
78+
# or there's an exception in "async with":
79+
#
80+
# async with TaskGroup() as g:
81+
# g.create_task(...)
82+
# 1 / 0
83+
#
84+
self._abort()
85+
86+
# We use while-loop here because "self._on_completed"
87+
# can be cancelled multiple times if our parent task
88+
# is being cancelled repeatedly (or even once, when
89+
# our own cancellation is already in progress)
90+
while self._tasks:
91+
if self._on_completed is None:
92+
self._on_completed = event.Event()
93+
94+
try:
95+
await self._on_completed.wait()
96+
except core.CancelledError as ex:
97+
if not self._aborting:
98+
# Our parent task is being cancelled:
99+
#
100+
# async def wrapper():
101+
# async with TaskGroup() as g:
102+
# g.create_task(foo)
103+
#
104+
# "wrapper" is being cancelled while "foo" is
105+
# still running.
106+
propagate_cancellation_error = ex
107+
self._abort()
108+
109+
self._on_completed = None
110+
111+
assert not self._tasks
112+
113+
if self._base_error is not None:
114+
raise self._base_error
115+
116+
if propagate_cancellation_error is not None:
117+
# The wrapping task was cancelled; since we're done with
118+
# closing all child tasks, just propagate the cancellation
119+
# request now.
120+
raise propagate_cancellation_error
121+
122+
if et is not None and et is not core.CancelledError:
123+
self._errors.append(exc)
124+
125+
if self._errors:
126+
# Exceptions are heavy objects that can have object
127+
# cycles (bad for GC); let's not keep a reference to
128+
# a bunch of them.
129+
errors = self._errors
130+
self._errors = None
131+
132+
if len(errors) == 1:
133+
me = errors[0]
134+
else:
135+
EGroup = core.ExceptionGroup
136+
for err in errors:
137+
if not isinstance(err, Exception):
138+
EGroup = core.baseExceptionGroup
139+
break
140+
me = EGroup("unhandled errors in a TaskGroup", errors)
141+
raise me
142+
143+
def create_task(self, coro):
144+
if not self._entered:
145+
raise RuntimeError(f"TaskGroup {repr(self)} has not been entered")
146+
if self._exiting and not self._tasks:
147+
raise RuntimeError(f"TaskGroup {repr(self)} is finished")
148+
149+
k = [None]
150+
t = self._loop.create_task(self._run_task(k, coro))
151+
k[0] = t # sigh
152+
self._tasks.add(t)
153+
return t
154+
155+
def cancel(self):
156+
# Extension (not in CPython): kill off a whole taskgroup
157+
# TODO this waits for the parent to die before killing the child
158+
# tasks. Shouldn't that be the other way round?
159+
try:
160+
self._parent_task.cancel()
161+
except RuntimeError:
162+
raise core.CancelledError()
163+
164+
# Since Python 3.8 Tasks propagate all exceptions correctly,
165+
# except for KeyboardInterrupt and SystemExit which are
166+
# still considered special.
167+
168+
def _is_base_error(self, exc: BaseException) -> bool:
169+
assert isinstance(exc, BaseException)
170+
return isinstance(exc, (SystemExit, KeyboardInterrupt))
171+
172+
def _abort(self):
173+
self._aborting = True
174+
175+
for t in self._tasks:
176+
if not t.done():
177+
t.cancel()
178+
179+
async def _run_task(self, k, coro):
180+
task = k[0]
181+
assert task is not None
182+
183+
try:
184+
await coro
185+
except core.CancelledError as e:
186+
exc = e
187+
except BaseException as e:
188+
if DEBUG:
189+
import sys
190+
191+
sys.print_exception(e)
192+
193+
exc = e
194+
else:
195+
exc = None
196+
finally:
197+
self._tasks.discard(task)
198+
if self._on_completed is not None and not self._tasks:
199+
self._on_completed.set()
200+
201+
if type(exc) is core.CancelledError:
202+
return
203+
204+
if exc is None:
205+
return
206+
207+
self._errors.append(exc)
208+
if self._is_base_error(exc) and self._base_error is None:
209+
self._base_error = exc
210+
211+
if self._parent_task.done():
212+
# Not sure if this case is possible, but we want to handle
213+
# it anyways.
214+
self._loop.call_exception_handler(
215+
{
216+
"message": f"Task {repr(task)} has errored out but its parent task {self._parent_task} is already completed",
217+
"exception": exc,
218+
"task": task,
219+
}
220+
)
221+
return
222+
223+
if not self._aborting and not self._parent_cancel_requested:
224+
# If parent task *is not* being cancelled, it means that we want
225+
# to manually cancel it to abort whatever is being run right now
226+
# in the TaskGroup. But we want to mark parent task as
227+
# "not cancelled" later in __aexit__. Example situation that
228+
# we need to handle:
229+
#
230+
# async def foo():
231+
# try:
232+
# async with TaskGroup() as g:
233+
# g.create_task(crash_soon())
234+
# await something # <- this needs to be canceled
235+
# # by the TaskGroup, e.g.
236+
# # foo() needs to be cancelled
237+
# except Exception:
238+
# # Ignore any exceptions raised in the TaskGroup
239+
# pass
240+
# await something_else # this line has to be called
241+
# # after TaskGroup is finished.
242+
self._abort()
243+
self._parent_cancel_requested = True
244+
self._parent_task.cancel()

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy