Aurelien Larcher
2021-01-08 b5a8497ec9558d0061024ac4b49f445bbe7d1094
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
From 82dbfd5a04863d8b6363527e6a34a90c9aa5691b Mon Sep 17 00:00:00 2001
From: "Miss Islington (bot)"
 <31488909+miss-islington@users.noreply.github.com>
Date: Fri, 18 Dec 2020 11:19:10 -0800
Subject: [PATCH] bpo-41891: ensure asyncio.wait_for waits for task completion
 (GH-22461) (#23840)
 
(cherry picked from commit 17ef4319a34f5a2f95e7823dfb5f5b8cff11882d)
 
Co-authored-by: Richard Kojedzinszky <rkojedzinszky@users.noreply.github.com>
 
Co-authored-by: Richard Kojedzinszky <rkojedzinszky@users.noreply.github.com>
---
 Lib/asyncio/tasks.py                          |  5 +-
 Lib/test/test_asyncio/test_asyncio_waitfor.py | 61 +++++++++++++++++++
 .../2020-09-30-13-35-29.bpo-41891.pNAeYI.rst  |  1 +
 3 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 Lib/test/test_asyncio/test_asyncio_waitfor.py
 create mode 100644 Misc/NEWS.d/next/Library/2020-09-30-13-35-29.bpo-41891.pNAeYI.rst
 
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index f486b67229411..d6262ae75e656 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -471,7 +471,10 @@ async def wait_for(fut, timeout, *, loop=None):
                 return fut.result()
             else:
                 fut.remove_done_callback(cb)
-                fut.cancel()
+                # We must ensure that the task is not running
+                # after wait_for() returns.
+                # See https://bugs.python.org/issue32751
+                await _cancel_and_wait(fut, loop=loop)
                 raise
 
         if fut.done():
diff --git a/Lib/test/test_asyncio/test_asyncio_waitfor.py b/Lib/test/test_asyncio/test_asyncio_waitfor.py
new file mode 100644
index 0000000000000..2ca64abbeb527
--- /dev/null
+++ b/Lib/test/test_asyncio/test_asyncio_waitfor.py
@@ -0,0 +1,61 @@
+import asyncio
+import unittest
+import time
+
+def tearDownModule():
+    asyncio.set_event_loop_policy(None)
+
+
+class SlowTask:
+    """ Task will run for this defined time, ignoring cancel requests """
+    TASK_TIMEOUT = 0.2
+
+    def __init__(self):
+        self.exited = False
+
+    async def run(self):
+        exitat = time.monotonic() + self.TASK_TIMEOUT
+
+        while True:
+            tosleep = exitat - time.monotonic()
+            if tosleep <= 0:
+                break
+
+            try:
+                await asyncio.sleep(tosleep)
+            except asyncio.CancelledError:
+                pass
+
+        self.exited = True
+
+class AsyncioWaitForTest(unittest.TestCase):
+
+    async def atest_asyncio_wait_for_cancelled(self):
+        t  = SlowTask()
+
+        waitfortask = asyncio.create_task(asyncio.wait_for(t.run(), t.TASK_TIMEOUT * 2))
+        await asyncio.sleep(0)
+        waitfortask.cancel()
+        await asyncio.wait({waitfortask})
+
+        self.assertTrue(t.exited)
+
+    def test_asyncio_wait_for_cancelled(self):
+        asyncio.run(self.atest_asyncio_wait_for_cancelled())
+
+    async def atest_asyncio_wait_for_timeout(self):
+        t  = SlowTask()
+
+        try:
+            await asyncio.wait_for(t.run(), t.TASK_TIMEOUT / 2)
+        except asyncio.TimeoutError:
+            pass
+
+        self.assertTrue(t.exited)
+
+    def test_asyncio_wait_for_timeout(self):
+        asyncio.run(self.atest_asyncio_wait_for_timeout())
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/Misc/NEWS.d/next/Library/2020-09-30-13-35-29.bpo-41891.pNAeYI.rst b/Misc/NEWS.d/next/Library/2020-09-30-13-35-29.bpo-41891.pNAeYI.rst
new file mode 100644
index 0000000000000..75c2512780315
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2020-09-30-13-35-29.bpo-41891.pNAeYI.rst
@@ -0,0 +1 @@
+Ensure asyncio.wait_for waits for task completion