8000 bpo-46771: Implement asyncio context managers for handling timeouts (… · python/cpython@f537b2a · GitHub
[go: up one dir, main page]
Skip to content

Commit f537b2a

Browse files
asvetlovgvanrossum
asvetlov
and
gvanrossum
authored
bpo-46771: Implement asyncio context managers for handling timeouts (GH-31394)
Example: async with asyncio.timeout(5): await some_task() Will interrupt the await and raise TimeoutError if some_task() takes longer than 5 seconds. Co-authored-by: Guido van Rossum <guido@python.org>
1 parent 32bf359 commit f537b2a

File tree

4 files changed

+384
-0
lines changed

4 files changed

+384
-0
lines changed

‎Lib/asyncio/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from .subprocess import *
1919
from .tasks import *
2020
from .taskgroups import *
21+
from .timeouts import *
2122
from .threads import *
2223
from .transports import *
2324

@@ -34,6 +35,7 @@
3435
subprocess.__all__ +
353 EDBE 6
tasks.__all__ +
3637
threads.__all__ +
38+
timeouts.__all__ +
3739
transports.__all__)
3840

3941
if sys.platform == 'win32': # pragma: no cover

‎Lib/asyncio/timeouts.py

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
import enum
2+
3+
from types import TracebackType
4+
from typing import final, Optional, Type
5+
6+
from . import events
7+
from . import exceptions
8+
from . import tasks
9+
10+
11+
__all__ = (
12+
"Timeout",
13+
"timeout",
14+
"timeout_at",
15+
)
16+
17+
18+
class _State(enum.Enum):
19+
CREATED = "created"
20+
ENTERED = "active"
21+
EXPIRING = "expiring"
22+
EXPIRED = "expired"
23+
EXITED = "finished"
24+
25+
26+
@final
27+
class Timeout:
28+
29+
def __init__(self, when: Optional[float]) -> None:
30+
self._state = _State.CREATED
31+
32+
self._timeout_handler: Optional[events.TimerHandle] = None
33+
self._task: Optional[tasks.Task] = None
34+
self._when = when
35+
36+
def when(self) -> Optional[float]:
37+
return self._when
38+
39+
def reschedule(self, when: Optional[float]) -> None:
40+
assert self._state is not _State.CREATED
41+
if self._state is not _State.ENTERED:
42+
raise RuntimeError(
43+
f"Cannot change state of {self._state.value} Timeout",
44+
)
45+
46+
self._when = when
47+
48+
if self._timeout_handler is not None:
49+
self._timeout_handler.cancel()
50+
51+
if when is None:
52+
self._timeout_handler = None
53+
else:
54+
loop = events.get_running_loop()
55+
self._timeout_handler = loop.call_at(
56+
when,
57+
self._on_timeout,
58+
)
59+
60+
def expired(self) -> bool:
61+
"""Is timeout expired during execution?"""
62+
return self._state in (_State.EXPIRING, _State.EXPIRED)
63+
64+
def __repr__(self) -> str:
65+
info = ['']
66+
if self._state is _State.ENTERED:
67+
when = round(self._when, 3) if self._when is not None else None
68+
info.append(f"when={when}")
69+
info_str = ' '.join(info)
70+
return f"<Timeout [{self._state.value}]{info_str}>"
71+
72+
async def __aenter__(self) -> "Timeout":
73+
self._state = _State.ENTERED
74+
self._task = tasks.current_task()
75+
if self._task is None:
76+
raise RuntimeError("Timeout should be used inside a task")
77+
self.reschedule(self._when)
78+
return self
79+
80+
async def __aexit__(
81+
self,
82+
exc_type: Optional[Type[BaseException]],
83+
exc_val: Optional[BaseException],
84+
exc_tb: Optional[TracebackType],
85+
) -> Optional[bool]:
86+
assert self._state in (_State.ENTERED, _State.EXPIRING)
87+
88+
if self._timeout_handler is not None:
89+
self._timeout_handler.cancel()
90+
self._timeout_handler = None
91+
92+
if self._state is _State.EXPIRING:
93+
self._state = _State.EXPIRED
94+
95+
if self._task.uncancel() == 0 and exc_type is exceptions.CancelledError:
96+
# Since there are no outstanding cancel requests, we're
97+
# handling this.
98+
raise TimeoutError
99+
elif self._state is _State.ENTERED:
100+
self._state = _State.EXITED
101+
102+
return None
103+
104+
def _on_timeout(self) -> None:
105+
assert self._state is _State.ENTERED
106+
self._task.cancel()
107+
self._state = _State.EXPIRING
108+
# drop the reference early
109+
self._timeout_handler = None
110+
111+
112+
def timeout(delay: Optional[float]) -> Timeout:
113+
"""Timeout async context manager.
114+
115+
Useful in cases when you want to apply timeout logic around block
116+
of code or in cases when asyncio.wait_for is not suitable. For example:
117+
118+
>>> async with asyncio.timeout(10): # 10 seconds timeout
119+
... await long_running_task()
120+
121+
122+
delay - value in seconds or None to disable timeout logic
123+
124+
long_running_task() is interrupted by raising asyncio.CancelledError,
125+
the top-most affected timeout() context manager converts CancelledError
126+
into TimeoutError.
127+
"""
128+
loop = events.get_running_loop()
129+
return Timeout(loop.time() + delay if delay is not None else None)
130+
131+
132+
def timeout_at(when: Optional[float]) -> Timeout:
133+
"""Schedule the timeout at absolute time.
134+
135+
Like timeout() but argument gives absolute time in the same clock system
136+
as loop.time().
137+
138+
Please note: it is not POSIX time but a time with
139+
undefined starting base, e.g. the time of the system power on.
140+
141+
>>> async with asyncio.timeout_at(loop.time() + 10):
142+
... await long_running_task()
143+
144+
145+
when - a deadline when timeout occurs or None to disable timeout logic
146+
147+
long_running_task() is interrupted by raising asyncio.CancelledError,
148+
the top-most affected timeout() context manager converts CancelledError
149+
into TimeoutError.
150+
"""
151+
return Timeout(when)

‎Lib/test/test_asyncio/test_timeouts.py

Lines changed: 229 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,229 @@
1+
"""Tests for asyncio/timeouts.py"""
2+
3+
import unittest
4+
import time
5+
6+
import asyncio
7+
from asyncio import tasks
8+
9+
10+
def tearDownModule():
11+
asyncio.set_event_loop_policy(None)
12+
13+
14+
class TimeoutTests(unittest.IsolatedAsyncioTestCase):
15+
16+
async def test_timeout_basic(self):
17+
with self.assertRaises(TimeoutError):
18+
async with asyncio.timeout(0.01) as cm:
19+
await asyncio.sleep(10)
20+
self.assertTrue(cm.expired())
21+
22+
async def test_timeout_at_basic(self):
23+
loop = asyncio.get_running_loop()
24+
25+
with self.assertRaises(TimeoutError):
26+
deadline = loop.time() + 0.01
27+
async with asyncio.timeout_at(deadline) as cm:
28+
await asyncio.sleep(10)
29+
self.assertTrue(cm.expired())
30+
self.assertEqual(deadline, cm.when())
31+
32+
async def test_nested_timeouts(self):
33+
loop = asyncio.get_running_loop()
34+
cancelled = False
35+
with self.assertRaises(TimeoutError):
36+
deadline = loop.time() + 0.01
37+
async with asyncio.timeout_at(deadline) as cm1:
38+
# Only the topmost context manager should raise TimeoutError
39+
try:
40+
async with asyncio.timeout_at(deadline) as cm2:
41+
await asyncio.sleep(10)
42+
except asyncio.CancelledError:
43+
cancelled = True
44+
raise
45+
self.assertTrue(cancelled)
46+
self.assertTrue(cm1.expired())
47+
self.assertTrue(cm2.expired())
48+
49+
async def test_waiter_cancelled(self):
50+
loop = asyncio.get_running_loop()
51+
cancelled = False
52+
with self.assertRaises(TimeoutError):
53+
async with asyncio.timeout(0.01):
54+
try:
55+
await asyncio.sleep(10)
56+
except asyncio.CancelledError:
57+
cancelled = True
58+
raise
59+
self.assertTrue(cancelled)
60+
61+
async def test_timeout_not_called(self):
62+
loop = asyncio.get_running_loop()
63+
t0 = loop.time()
64+
async with asyncio.timeout(10) as cm:
65+
await asyncio.sleep(0.01)
66+
t1 = loop.time()
67+
68+
self.assertFalse(cm.expired())
69+
# 2 sec for slow CI boxes
70+
self.assertLess(t1-t0, 2)
71+
self.assertGreater(cm.when(), t1)
72+
73+
async def test_timeout_disabled(self):
74+
loop = asyncio.get_running_loop()
75+
t0 = loop.time()
76+
async with asyncio.timeout(None) as cm:
77+
await asyncio.sleep(0.01)
78+
t1 = loop.time()
79+
80+
self.assertFalse(cm.expired())
81+
self.assertIsNone(cm.when())
82+
# 2 sec for slow CI boxes
83+
self.assertLess(t1-t0, 2)
84+
85+
async def test_timeout_at_disabled(self):
86+
loop = asyncio.get_running_loop()
87+
t0 = loop.time()
88+
async with asyncio.timeout_at(None) as cm:
89+
await asyncio.sleep(0.01)
90+
t1 = loop.time()
91+
92+
self.assertFalse(cm.expired())
93+
self.assertIsNone(cm.when())
94+
# 2 sec for slow CI boxes
95+
self.assertLess(t1-t0, 2)
96+
97+
async def test_timeout_zero(self):
98+
loop = asyncio.get_running_loop()
99+
t0 = loop.time()
100+
with self.assertRaises(TimeoutError):
101+
async with asyncio.timeout(0) as cm:
102+
await asyncio.sleep(10)
103+
t1 = loop.time()
104+
self.assertTrue(cm.expired())
105+
# 2 sec for slow CI boxes
106+
self.assertLess(t1-t0, 2)
107+
self.assertTrue(t0 <= cm.when() <= t1)
108+
109+
async def test_foreign_exception_passed(self):
110+
with self.assertRaises(KeyError):
111+
async with asyncio.timeout(0.01) as cm:
112+
raise KeyError
113+
self.assertFalse(cm.expired())
114+
115+
async def test_foreign_exception_on_timeout(self):
116+
async def crash():
117+
try:
118+
await asyncio.sleep(1)
119+
finally:
120+
1/0
121+
with self.assertRaises(ZeroDivisionError):
122+
async with asyncio.timeout(0.01):
123+
await crash()
124+
125+
async def test_foreign_cancel_doesnt_timeout_if_not_expired(self):
126+
with self.assertRaises(asyncio.CancelledError):
127+
async with asyncio.timeout(10) as cm:
128+
asyncio.current_task().cancel()
129+
await asyncio.sleep(10)
130+
self.assertFalse(cm.expired())
131+
132+
async def test_outer_task_is_not_cancelled(self):
133+
async def outer() -> None:
134+
with self.assertRaises(TimeoutError):
135+
async with asyncio.timeout(0.001):
136+
await asyncio.sleep(10)
137+
138+
task = asyncio.create_task(outer())
139+
await task
140+
self.assertFalse(task.cancelled())
141+
self.assertTrue(task.done())
142+
143+
async def test_nested_timeouts_concurrent(self):
144+
with self.assertRaises(TimeoutError):
145+
async with asyncio.timeout(0.002):
146+
with self.assertRaises(TimeoutError):
147+
async with asyncio.timeout(0.1):
148+
# Pretend we crunch some numbers.
149+
time.sleep(0.01)
150+
await asyncio.sleep(1)
151+
152+
async def test_nested_timeouts_loop_busy(self):
153+
# After the inner timeout is an expensive operation which should
154+
# be stopped by the outer timeout.
155+
loop = asyncio.get_running_loop()
156+
# Disable a message about long running task
157+
loop.slow_callback_duration = 10
158+
t0 = loop.time()
159+
with self.assertRaises(TimeoutError):
160+
async with asyncio.timeout(0.1): # (1)
161+
with self.assertRaises(TimeoutError):
162+
async with asyncio.timeout(0.01): # (2)
163+
# Pretend the loop is busy for a while.
164+
time.sleep(0.1)
165+
await asyncio.sleep(1)
166+
# TimeoutError was cought by (2)
167+
await asyncio.sleep(10) # This sleep should be interrupted by (1)
168+
t1 = loop.time()
169+
self.assertTrue(t0 <= t1 <= t0 + 1)
170+
171+
async def test_reschedule(self):
172+
loop = asyncio.get_running_loop()
173+
fut = loop.create_future()
174+
deadline1 = loop.time() + 10
175+
deadline2 = deadline1 + 20
176+
177+
async def f():
178+
async with asyncio.timeout_at(deadline1) as cm:
179+
fut.set_result(cm)
180+
await asyncio.sleep(50)
181+
182+
task = asyncio.create_task(f())
183+
cm = await fut
184+
185+
self.assertEqual(cm.when(), deadline1)
186+
cm.reschedule(deadline2)
187+
self.assertEqual(cm.when(), deadline2)
188+
cm.reschedule(None)
189+
self.assertIsNone(cm.when())
190+
191+
task.cancel()
192+
193+
with self.assertRaises(asyncio.CancelledError):
194+
await task
195+
self.assertFalse(cm.expired())
196+
197+
async def test_repr_active(self):
198+
async with asyncio.timeout(10) as cm:
199+
self.assertRegex(repr(cm), r"<Timeout \[active\] when=\d+\.\d*>")
200+
201+
async def test_repr_expired(self):
202+
with self.assertRaises(TimeoutError):
203+
async with asyncio.timeout(0.01) as cm:
204+
await asyncio.sleep(10)
205+
self.assertEqual(repr(cm), "<Timeout [expired]>")
206+
207+
async def test_repr_finished(self):
208+
async with asyncio.timeout(10) as cm:
209+
await asyncio.sleep(0)
210+
211+
self.assertEqual(repr(cm), "<Timeout [finished]>")
212+
213+
async def test_repr_disabled(self):
214+
async with asyncio.timeout(None) as cm:
215+
self.assertEqual(repr(cm), r"<Timeout [active] when=None>")
216+
217+
async def test_nested_timeout_in_finally(self):
218+
with self.assertRaises(TimeoutError):
219+
async with asyncio.timeout(0.01):
220+
try:
221+
await asyncio.sleep(1)
222+
finally:
223+
with self.assertRaises(TimeoutError):
224+
async with asyncio.timeout(0.01):
225+
await asyncio.sleep(10)
226+
227+
228+
if __name__ == '__main__':
229+
unittest.main()

0 commit comments

Comments
 (0)
0