bpo-32751: Wait for task cancel in asyncio.wait_for() when timeout <= 0 (#21895)

elprans · web-flow · commit c517fc712105 · 2020-08-26T09:42:22.000-07:00
When I was fixing bpo-32751 back in GH-7216 I missed the case when *timeout* is zero or negative. This takes care of that. Props to @aaliddell for noticing the inconsistency.
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
@@ -445,8 +445,13 @@ async def wait_for(fut, timeout, *, loop=None):
 if fut.done():
 return fut.result()
 
- fut.cancel()
- raise exceptions.TimeoutError()
+ await _cancel_and_wait(fut, loop=loop)
+ try:
+ fut.result()
+ except exceptions.CancelledError as exc:
+ raise exceptions.TimeoutError() from exc
+ else:
+ raise exceptions.TimeoutError()
 
 waiter = loop.create_future()
 timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
@@ -1131,6 +1131,9 @@ async def inner():
 nonlocal task_done
 try:
 await asyncio.sleep(0.2)
+ except asyncio.CancelledError:
+ await asyncio.sleep(_EPSILON)
+ raise
 finally:
 task_done = True
 
@@ -1145,6 +1148,34 @@ async def inner():
 chained = cm.exception.__context__
 self.assertEqual(type(chained), asyncio.CancelledError)
 
+ def test_wait_for_waits_for_task_cancellation_w_timeout_0(self):
+ loop = asyncio.new_event_loop()
+ self.addCleanup(loop.close)
+
+ task_done = False
+
+ async def foo():
+ async def inner():
+ nonlocal task_done
+ try:
+ await asyncio.sleep(10)
+ except asyncio.CancelledError:
+ await asyncio.sleep(_EPSILON)
+ raise
+ finally:
+ task_done = True
+
+ inner_task = self.new_task(loop, inner())
+ await asyncio.sleep(_EPSILON)
+ await asyncio.wait_for(inner_task, timeout=0)
+
+ with self.assertRaises(asyncio.TimeoutError) as cm:
+ loop.run_until_complete(foo())
+
+ self.assertTrue(task_done)
+ chained = cm.exception.__context__
+ self.assertEqual(type(chained), asyncio.CancelledError)
+
 def test_wait_for_reraises_exception_during_cancellation(self):
 loop = asyncio.new_event_loop()
 self.addCleanup(loop.close)
diff --git a/Misc/NEWS.d/next/Library/2020-08-15-15-50-12.bpo-32751.85je5X.rst b/Misc/NEWS.d/next/Library/2020-08-15-15-50-12.bpo-32751.85je5X.rst
@@ -0,0 +1,3 @@
+When cancelling the task due to a timeout, :meth:`asyncio.wait_for` will now
+wait until the cancellation is complete also in the case when *timeout* is
+<= 0, like it does with positive timeouts.

-Original file line number
+Diff line change
 iffut.done():
 returnfut.result()
 -fut.cancel()
 -raiseexceptions.TimeoutError()
 +await_cancel_and_wait(fut, loop=loop)
 +try:
 +fut.result()
 +exceptexceptions.CancelledErrorasexc:
 +raiseexceptions.TimeoutError() fromexc
 +else:
 +raiseexceptions.TimeoutError()
 waiter=loop.create_future()
 timeout_handle=loop.call_later(timeout, _release_waiter, waiter)
-Original file line number
+Diff line change
 nonlocaltask_done
 try:
 awaitasyncio.sleep(0.2)
 +exceptasyncio.CancelledError:
 +awaitasyncio.sleep(_EPSILON)
 +raise
 finally:
 task_done=True
 chained=cm.exception.__context__
 self.assertEqual(type(chained), asyncio.CancelledError)
 +deftest_wait_for_waits_for_task_cancellation_w_timeout_0(self):
 +loop=asyncio.new_event_loop()
 +self.addCleanup(loop.close)
++
 +task_done=False
++
 +asyncdeffoo():
 +asyncdefinner():
 +nonlocaltask_done
 +try:
 +awaitasyncio.sleep(10)
 +exceptasyncio.CancelledError:
 +awaitasyncio.sleep(_EPSILON)
 +raise
 +finally:
 +task_done=True
++
 +inner_task=self.new_task(loop, inner())
 +awaitasyncio.sleep(_EPSILON)
 +awaitasyncio.wait_for(inner_task, timeout=0)
++
 +withself.assertRaises(asyncio.TimeoutError) ascm:
 +loop.run_until_complete(foo())
++
 +self.assertTrue(task_done)
 +chained=cm.exception.__context__
 +self.assertEqual(type(chained), asyncio.CancelledError)
++
 deftest_wait_for_reraises_exception_during_cancellation(self):
 loop=asyncio.new_event_loop()
 self.addCleanup(loop.close)
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+When cancelling the task due to a timeout, :meth:`asyncio.wait_for` will now
	`2`	`+wait until the cancellation is complete also in the case when timeout is`
	`3`	`+<= 0, like it does with positive timeouts.`