From 0bc12ef585913daf050e5707624785b1f6a230a1 Mon Sep 17 00:00:00 2001 From: John Belmonte Date: Sun, 19 Aug 2018 00:18:36 +0900 Subject: [PATCH 01/10] implement open_nursery with explicit context manager --- trio/_core/_run.py | 86 ++++++++++++++++++---------------------------- 1 file changed, 34 insertions(+), 52 deletions(-) diff --git a/trio/_core/_run.py b/trio/_core/_run.py index 6c8651ceb8..24a5a22092 100644 --- a/trio/_core/_run.py +++ b/trio/_core/_run.py @@ -16,9 +16,7 @@ from sniffio import current_async_library_cvar import attr -from async_generator import ( - async_generator, yield_, asynccontextmanager, isasyncgen -) +from async_generator import isasyncgen from sortedcontainers import SortedDict from outcome import Error, Value @@ -295,57 +293,41 @@ def started(self, value=None): self._old_nursery._check_nursery_closed() -@asynccontextmanager -@async_generator -@enable_ki_protection -async def open_nursery(): - """Returns an async context manager which creates a new nursery. - - This context manager's ``__aenter__`` method executes synchronously. Its - ``__aexit__`` method blocks until all child tasks have exited. +class NurseryManager: + @enable_ki_protection + async def __aenter__(self): + assert currently_ki_protected() + self._scope_manager = open_cancel_scope() + scope = self._scope_manager.__enter__() + self._nursery = Nursery(current_task(), scope) + return self._nursery - """ - assert currently_ki_protected() - with open_cancel_scope() as scope: - nursery = Nursery(current_task(), scope) - nested_child_exc = None - try: - await yield_(nursery) - except BaseException as exc: - nested_child_exc = exc + @enable_ki_protection + async def __aexit__(self, etype, exc, tb): assert currently_ki_protected() - await nursery._nested_child_finished(nested_child_exc) - - -# I *think* this is equivalent to the above, and it gives *much* nicer -# exception tracebacks... but I'm a little nervous about it because it's much -# trickier code :-( -# -# class NurseryManager: -# @enable_ki_protection -# async def __aenter__(self): -# self._scope_manager = open_cancel_scope() -# scope = self._scope_manager.__enter__() -# self._parent_nursery = Nursery(current_task(), scope) -# return self._parent_nursery -# -# @enable_ki_protection -# async def __aexit__(self, etype, exc, tb): -# try: -# await self._parent_nursery._clean_up(exc) -# except BaseException as new_exc: -# if not self._scope_manager.__exit__( -# type(new_exc), new_exc, new_exc.__traceback__): -# if exc is new_exc: -# return False -# else: -# raise -# else: -# self._scope_manager.__exit__(None, None, None) -# return True -# -# def open_nursery(): -# return NurseryManager() + try: + await self._nursery._nested_child_finished(exc) + except BaseException as new_exc: + if not self._scope_manager.__exit__( + type(new_exc), new_exc, new_exc.__traceback__): + if new_exc is exc: + return False + else: + raise + else: + self._scope_manager.__exit__(None, None, None) + return True + + def __enter__(self): + raise RuntimeError( + "use 'async with open_nursery(...)', not 'with open_nursery(...)'") + + def __exit__(self): # pragma: no cover + assert False, """Never called, but should be defined""" + + +def open_nursery(): + return NurseryManager() class Nursery: From 2c12889233abc7672f271134a1ad223f00d9bf27 Mon Sep 17 00:00:00 2001 From: "Nathaniel J. Smith" Date: Sun, 19 Aug 2018 16:41:15 -0700 Subject: [PATCH 02/10] Well, that was obscure --- trio/_core/_run.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/trio/_core/_run.py b/trio/_core/_run.py index 24a5a22092..085dc2c75c 100644 --- a/trio/_core/_run.py +++ b/trio/_core/_run.py @@ -314,6 +314,8 @@ async def __aexit__(self, etype, exc, tb): return False else: raise + else: + return True else: self._scope_manager.__exit__(None, None, None) return True From 292516d98096bc28215967c4a62cf18d6b67114c Mon Sep 17 00:00:00 2001 From: "Nathaniel J. Smith" Date: Sun, 19 Aug 2018 16:55:16 -0700 Subject: [PATCH 03/10] Flip the branches of an 'if' to make the flow more clear --- trio/_core/_run.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/trio/_core/_run.py b/trio/_core/_run.py index 085dc2c75c..ac3e5ac29c 100644 --- a/trio/_core/_run.py +++ b/trio/_core/_run.py @@ -308,14 +308,14 @@ async def __aexit__(self, etype, exc, tb): try: await self._nursery._nested_child_finished(exc) except BaseException as new_exc: - if not self._scope_manager.__exit__( + if self._scope_manager.__exit__( type(new_exc), new_exc, new_exc.__traceback__): + return True + else: if new_exc is exc: return False else: raise - else: - return True else: self._scope_manager.__exit__(None, None, None) return True From f38a50b5b5e7cc6b86b63b9c6bc5c2297b726720 Mon Sep 17 00:00:00 2001 From: John Belmonte Date: Mon, 20 Aug 2018 10:56:07 +0900 Subject: [PATCH 04/10] formatting --- trio/_core/_run.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/trio/_core/_run.py b/trio/_core/_run.py index ac3e5ac29c..c90ee966b3 100644 --- a/trio/_core/_run.py +++ b/trio/_core/_run.py @@ -309,7 +309,8 @@ async def __aexit__(self, etype, exc, tb): await self._nursery._nested_child_finished(exc) except BaseException as new_exc: if self._scope_manager.__exit__( - type(new_exc), new_exc, new_exc.__traceback__): + type(new_exc), new_exc, new_exc.__traceback__ + ): return True else: if new_exc is exc: @@ -322,7 +323,8 @@ async def __aexit__(self, etype, exc, tb): def __enter__(self): raise RuntimeError( - "use 'async with open_nursery(...)', not 'with open_nursery(...)'") + "use 'async with open_nursery(...)', not 'with open_nursery(...)'" + ) def __exit__(self): # pragma: no cover assert False, """Never called, but should be defined""" From 3ce95efec80ad6326fcb019391a2df94bbdf3d6c Mon Sep 17 00:00:00 2001 From: John Belmonte Date: Mon, 20 Aug 2018 22:46:52 +0900 Subject: [PATCH 05/10] simplify trace of explicit exception in nursery scope, add test --- trio/_core/_run.py | 17 +++++++++-------- trio/_core/tests/test_run.py | 6 ++++++ 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/trio/_core/_run.py b/trio/_core/_run.py index c90ee966b3..5c8b5d3c08 100644 --- a/trio/_core/_run.py +++ b/trio/_core/_run.py @@ -308,15 +308,16 @@ async def __aexit__(self, etype, exc, tb): try: await self._nursery._nested_child_finished(exc) except BaseException as new_exc: - if self._scope_manager.__exit__( - type(new_exc), new_exc, new_exc.__traceback__ - ): - return True - else: - if new_exc is exc: + try: + if self._scope_manager.__exit__( + type(new_exc), new_exc, new_exc.__traceback__ + ): + return True + except BaseException as scope_manager_exc: + if scope_manager_exc == exc: return False - else: - raise + raise # scope_manager_exc + raise # new_exc else: self._scope_manager.__exit__(None, None, None) return True diff --git a/trio/_core/tests/test_run.py b/trio/_core/tests/test_run.py index 49ffb76baf..c6dc36ad98 100644 --- a/trio/_core/tests/test_run.py +++ b/trio/_core/tests/test_run.py @@ -1823,6 +1823,12 @@ async def start_sleep_then_crash(nursery): assert _core.current_time() - t0 == 7 +async def test_nursery_explicit_exception(): + with pytest.raises(KeyError): + async with _core.open_nursery(): + raise KeyError() + + def test_contextvar_support(): var = contextvars.ContextVar("test") var.set("before") From 9ead5611baa875cca39d3dfb5c4c733b37e5e188 Mon Sep 17 00:00:00 2001 From: John Belmonte Date: Tue, 21 Aug 2018 21:54:07 +0900 Subject: [PATCH 06/10] add test for StopIteration bug; document NurseryManager rationale --- trio/_core/_run.py | 7 +++++++ trio/_core/tests/test_run.py | 12 ++++++++++++ 2 files changed, 19 insertions(+) diff --git a/trio/_core/_run.py b/trio/_core/_run.py index 5c8b5d3c08..6cd48324f3 100644 --- a/trio/_core/_run.py +++ b/trio/_core/_run.py @@ -294,6 +294,13 @@ def started(self, value=None): class NurseryManager: + """Nursery context manager. + + Note we explicitly avoid @asynccontextmanager and @async_generator + since they add a lot of extraneous stack frames to exceptions, as + well as cause problematic behavior with handling of StopIteration. + + """ @enable_ki_protection async def __aenter__(self): assert currently_ki_protected() diff --git a/trio/_core/tests/test_run.py b/trio/_core/tests/test_run.py index c6dc36ad98..45eab55744 100644 --- a/trio/_core/tests/test_run.py +++ b/trio/_core/tests/test_run.py @@ -1829,6 +1829,18 @@ async def test_nursery_explicit_exception(): raise KeyError() +async def test_nursery_stop_iteration(): + async def fail(): + raise ValueError + + try: + async with _core.open_nursery() as nursery: + nursery.start_soon(fail) + raise StopIteration + except _core.MultiError as e: + assert tuple(map(type, e.exceptions)) == (StopIteration, ValueError) + + def test_contextvar_support(): var = contextvars.ContextVar("test") var.set("before") From eaf60649d69dc354af132b9d49b2c1b633f8187c Mon Sep 17 00:00:00 2001 From: John Belmonte Date: Tue, 21 Aug 2018 22:02:35 +0900 Subject: [PATCH 07/10] add test for StopAsyncIteration bug --- trio/_core/tests/test_run.py | 55 ++++++++++++++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/trio/_core/tests/test_run.py b/trio/_core/tests/test_run.py index 45eab55744..770641b29e 100644 --- a/trio/_core/tests/test_run.py +++ b/trio/_core/tests/test_run.py @@ -1841,6 +1841,61 @@ async def fail(): assert tuple(map(type, e.exceptions)) == (StopIteration, ValueError) +async def test_nursery_stop_async_iteration(): + class it(object): + def __init__(self, count): + self.count = count + self.val = 0 + + def __aiter__(self): + return self + + async def __anext__(self): + await sleep(0) + val = self.val + if val >= self.count: + raise StopAsyncIteration + self.val += 1 + return val + + class async_zip(object): + def __init__(self, *largs): + self.nexts = [obj.__anext__ for obj in largs] + + async def _accumulate(self, f, items, i): + items[i] = await f() + + def __aiter__(self): + return self + + async def __anext__(self): + nexts = self.nexts + items = [None, ] * len(nexts) + got_stop = False + + def handle(exc): + nonlocal got_stop + if isinstance(exc, StopAsyncIteration): + got_stop = True + return None + else: + return exc + + with _core.MultiError.catch(handle): + async with _core.open_nursery() as nursery: + for i, f in enumerate(nexts): + nursery.start_soon(self._accumulate, f, items, i) + + if got_stop: + raise StopAsyncIteration + return items + + result = [] + async for vals in async_zip(it(4), it(2)): + result.append(vals) + assert result == [[0, 0], [1, 1]] + + def test_contextvar_support(): var = contextvars.ContextVar("test") var.set("before") From 0fa15dfc8ebe0bd6f7ae5d20093262e31377c15e Mon Sep 17 00:00:00 2001 From: John Belmonte Date: Tue, 21 Aug 2018 22:46:58 +0900 Subject: [PATCH 08/10] formatting and docs; fix __aiter__ compatibility of test --- trio/_core/_run.py | 4 +++- trio/_core/tests/test_run.py | 7 ++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/trio/_core/_run.py b/trio/_core/_run.py index 6cd48324f3..8421f3be0a 100644 --- a/trio/_core/_run.py +++ b/trio/_core/_run.py @@ -298,9 +298,11 @@ class NurseryManager: Note we explicitly avoid @asynccontextmanager and @async_generator since they add a lot of extraneous stack frames to exceptions, as - well as cause problematic behavior with handling of StopIteration. + well as cause problematic behavior with handling of StopIteration + and StopAsyncIteration. """ + @enable_ki_protection async def __aenter__(self): assert currently_ki_protected() diff --git a/trio/_core/tests/test_run.py b/trio/_core/tests/test_run.py index 770641b29e..4e02c1dece 100644 --- a/trio/_core/tests/test_run.py +++ b/trio/_core/tests/test_run.py @@ -17,6 +17,7 @@ from .tutil import check_sequence_matches, gc_collect_harder from ... import _core from ..._timeouts import sleep +from ..._util import aiter_compat from ...testing import ( wait_all_tasks_blocked, Sequencer, @@ -1847,6 +1848,7 @@ def __init__(self, count): self.count = count self.val = 0 + @aiter_compat def __aiter__(self): return self @@ -1865,12 +1867,15 @@ def __init__(self, *largs): async def _accumulate(self, f, items, i): items[i] = await f() + @aiter_compat def __aiter__(self): return self async def __anext__(self): nexts = self.nexts - items = [None, ] * len(nexts) + items = [ + None, + ] * len(nexts) got_stop = False def handle(exc): From 9afcb58bf4aa8d2fecda7b8bb844442b8bf8b644 Mon Sep 17 00:00:00 2001 From: John Belmonte Date: Tue, 21 Aug 2018 22:49:20 +0900 Subject: [PATCH 09/10] newsfragment for nursery context manager --- newsfragments/612.bugfix.rst | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 newsfragments/612.bugfix.rst diff --git a/newsfragments/612.bugfix.rst b/newsfragments/612.bugfix.rst new file mode 100644 index 0000000000..d9a9bdfb1d --- /dev/null +++ b/newsfragments/612.bugfix.rst @@ -0,0 +1,4 @@ +The nursery context manager was rewritten to avoid use of +`@asynccontextmanager` and `@async_generator`. This reduces extraneous frames +in exception traces and addresses bugs regarding `StopIteration` and +`StopAsyncIteration` exceptions not propagating correctly. From 4be156d96a96a6d37a57cc6c3d7636ad2eea535b Mon Sep 17 00:00:00 2001 From: John Belmonte Date: Tue, 21 Aug 2018 23:07:41 +0900 Subject: [PATCH 10/10] code coverage --- trio/_core/tests/test_run.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/trio/_core/tests/test_run.py b/trio/_core/tests/test_run.py index 4e02c1dece..c264c39e75 100644 --- a/trio/_core/tests/test_run.py +++ b/trio/_core/tests/test_run.py @@ -1848,10 +1848,6 @@ def __init__(self, count): self.count = count self.val = 0 - @aiter_compat - def __aiter__(self): - return self - async def __anext__(self): await sleep(0) val = self.val @@ -1883,7 +1879,7 @@ def handle(exc): if isinstance(exc, StopAsyncIteration): got_stop = True return None - else: + else: # pragma: no cover return exc with _core.MultiError.catch(handle):