From 1d7ece30766c7076bdbac84286870d958769a8ab Mon Sep 17 00:00:00 2001 From: jakkdl Date: Tue, 8 Oct 2024 12:37:15 +0200 Subject: [PATCH] make broken_by attribute a list, clean up tests --- src/trio/_core/_parking_lot.py | 20 ++-- src/trio/_core/_run.py | 3 +- src/trio/_core/_tests/test_parking_lot.py | 123 +++++++++++++--------- 3 files changed, 82 insertions(+), 64 deletions(-) diff --git a/src/trio/_core/_parking_lot.py b/src/trio/_core/_parking_lot.py index 3f2edc5ba..af6ff610e 100644 --- a/src/trio/_core/_parking_lot.py +++ b/src/trio/_core/_parking_lot.py @@ -73,7 +73,6 @@ import inspect import math -import warnings from collections import OrderedDict from typing import TYPE_CHECKING @@ -152,7 +151,7 @@ class ParkingLot: # {task: None}, we just want a deque where we can quickly delete random # items _parked: OrderedDict[Task, None] = attrs.field(factory=OrderedDict, init=False) - broken_by: Task | None = None + broken_by: list[Task] = attrs.field(factory=list, init=False) def __len__(self) -> int: """Returns the number of parked tasks.""" @@ -176,7 +175,7 @@ async def park(self) -> None: breaks before we get to unpark. """ - if self.broken_by is not None: + if self.broken_by: raise _core.BrokenResourceError( f"Attempted to park in parking lot broken by {self.broken_by}", ) @@ -289,16 +288,13 @@ def break_lot(self, task: Task | None = None) -> None: """ if task is None: task = _core.current_task() - if self.broken_by is not None: - if self.broken_by != task: - warnings.warn( - RuntimeWarning( - f"{task} attempted to break parking lot {self} already broken by {self.broken_by}", - ), - stacklevel=2, - ) + + # if lot is already broken, just mark this as another breaker and return + if self.broken_by: + self.broken_by.append(task) return - self.broken_by = task + + self.broken_by.append(task) for parked_task in self._parked: _core.reschedule( diff --git a/src/trio/_core/_run.py b/src/trio/_core/_run.py index 871016963..defbaa6a8 100644 --- a/src/trio/_core/_run.py +++ b/src/trio/_core/_run.py @@ -1900,8 +1900,7 @@ def task_exited(self, task: Task, outcome: Outcome[Any]) -> None: # break parking lots associated with the task exiting if task in GLOBAL_PARKING_LOT_BREAKER: for lot in GLOBAL_PARKING_LOT_BREAKER[task]: - if lot.broken_by is None: - lot.break_lot(task) + lot.break_lot(task) del GLOBAL_PARKING_LOT_BREAKER[task] if ( diff --git a/src/trio/_core/_tests/test_parking_lot.py b/src/trio/_core/_tests/test_parking_lot.py index b8cb97712..cad3e7b43 100644 --- a/src/trio/_core/_tests/test_parking_lot.py +++ b/src/trio/_core/_tests/test_parking_lot.py @@ -1,5 +1,6 @@ from __future__ import annotations +import re from typing import TypeVar import pytest @@ -233,6 +234,53 @@ async def dummy_task( async def test_parking_lot_breaker_basic() -> None: + """Test basic functionality for breaking lots.""" + lot = ParkingLot() + task = current_task() + + # defaults to current task + lot.break_lot() + assert lot.broken_by == [task] + + # breaking the lot again with the same task appends another copy in `broken_by` + lot.break_lot() + assert lot.broken_by == [task, task] + + # trying to park in broken lot errors + broken_by_str = re.escape(str([task, task])) + with pytest.raises( + _core.BrokenResourceError, + match=f"^Attempted to park in parking lot broken by {broken_by_str}$", + ): + await lot.park() + + +async def test_parking_lot_break_parking_tasks() -> None: + """Checks that tasks currently waiting to park raise an error when the breaker exits.""" + + async def bad_parker(lot: ParkingLot, scope: _core.CancelScope) -> None: + add_parking_lot_breaker(current_task(), lot) + with scope: + await trio.sleep_forever() + + lot = ParkingLot() + cs = _core.CancelScope() + + # check that parked task errors + with RaisesGroup( + Matcher(_core.BrokenResourceError, match="^Parking lot broken by"), + ): + async with _core.open_nursery() as nursery: + nursery.start_soon(bad_parker, lot, cs) + await wait_all_tasks_blocked() + + nursery.start_soon(lot.park) + await wait_all_tasks_blocked() + + cs.cancel() + + +async def test_parking_lot_breaker_registration() -> None: lot = ParkingLot() task = current_task() @@ -254,58 +302,60 @@ async def test_parking_lot_breaker_basic() -> None: ): remove_parking_lot_breaker(task, lot) - # defaults to current task - lot.break_lot() - assert lot.broken_by == task - - # breaking the lot again with the same task is a no-op + # registering a task as breaker on an already broken lot is fine lot.break_lot() - - # registering a task as a breaker on an already broken lot is a no-op. child_task = None async with trio.open_nursery() as nursery: child_task = await nursery.start(dummy_task) add_parking_lot_breaker(child_task, lot) nursery.cancel_scope.cancel() + assert lot.broken_by == [task, child_task] # manually breaking a lot with an already exited task is fine lot = ParkingLot() lot.break_lot(child_task) + assert lot.broken_by == [child_task] -async def test_parking_lot_breaker_warnings() -> None: +async def test_parking_lot_breaker_rebreak() -> None: lot = ParkingLot() task = current_task() lot.break_lot() - warn_str = "attempted to break parking .* already broken by .*" - # breaking an already broken lot with a different task gives a warning + # breaking an already broken lot with a different task is allowed # The nursery is only to create a task we can pass to lot.break_lot async with trio.open_nursery() as nursery: child_task = await nursery.start(dummy_task) - with pytest.warns( - RuntimeWarning, - match=warn_str, - ): - lot.break_lot(child_task) + lot.break_lot(child_task) nursery.cancel_scope.cancel() - # and doesn't change broken_by - assert lot.broken_by == task + # and appends the task + assert lot.broken_by == [task, child_task] + +async def test_parking_lot_multiple_breakers_exit() -> None: # register multiple tasks as lot breakers, then have them all exit # No warning is given on task exit, even if the lot is already broken. lot = ParkingLot() - child_task = None async with trio.open_nursery() as nursery: - child_task = await nursery.start(dummy_task) + child_task1 = await nursery.start(dummy_task) child_task2 = await nursery.start(dummy_task) child_task3 = await nursery.start(dummy_task) - add_parking_lot_breaker(child_task, lot) + add_parking_lot_breaker(child_task1, lot) add_parking_lot_breaker(child_task2, lot) add_parking_lot_breaker(child_task3, lot) nursery.cancel_scope.cancel() + # I think the order is guaranteed currently, but doesn't hurt to be safe. + assert set(lot.broken_by) == {child_task1, child_task2, child_task3} + + +async def test_parking_lot_breaker_register_exited_task() -> None: + lot = ParkingLot() + child_task = None + async with trio.open_nursery() as nursery: + child_task = await nursery.start(dummy_task) + nursery.cancel_scope.cancel() # trying to register an exited task as lot breaker errors with pytest.raises( trio.BrokenResourceError, @@ -314,34 +364,7 @@ async def test_parking_lot_breaker_warnings() -> None: add_parking_lot_breaker(child_task, lot) -async def test_parking_lot_breaker_bad_parker() -> None: - async def bad_parker(lot: ParkingLot, scope: _core.CancelScope) -> None: - add_parking_lot_breaker(current_task(), lot) - with scope: - await trio.sleep_forever() - - lot = ParkingLot() - cs = _core.CancelScope() - - # check that parked task errors - with RaisesGroup( - Matcher(_core.BrokenResourceError, match="^Parking lot broken by"), - ): - async with _core.open_nursery() as nursery: - nursery.start_soon(bad_parker, lot, cs) - await wait_all_tasks_blocked() - - nursery.start_soon(lot.park) - await wait_all_tasks_blocked() - - cs.cancel() - - # check that trying to park in broken lot errors - with pytest.raises(_core.BrokenResourceError): - await lot.park() - - -async def test_parking_lot_weird() -> None: +async def test_parking_lot_break_itself() -> None: """Break a parking lot, where the breakee is parked. Doing this is weird, but should probably be supported. """ @@ -359,5 +382,5 @@ async def return_me_and_park( Matcher(_core.BrokenResourceError, match="^Parking lot broken by"), ): async with _core.open_nursery() as nursery: - task = await nursery.start(return_me_and_park, lot) - lot.break_lot(task) + child_task = await nursery.start(return_me_and_park, lot) + lot.break_lot(child_task)