[Python-checkins] cpython (3.4): asyncio, Tulip issue 201: Fix a race condition in wait_for()

[victor.stinner]

http://hg.python.org/cpython/rev/20244eda2946
changeset:   92257:20244eda2946
branch:      3.4
parent:      92253:a48ef19911ea
user:        Victor Stinner <victor.stinner at gmail.com>
date:        Thu Aug 28 11:19:25 2014 +0200
summary:
  asyncio, Tulip issue 201: Fix a race condition in wait_for()

Don't raise a TimeoutError if we reached the timeout and the future completed
in the same iteration of the event loop. A side effect of the bug is that
Queue.get() looses items.

files:
  Lib/asyncio/tasks.py                |  15 +++++++++------
  Lib/test/test_asyncio/test_tasks.py |  15 +++++++++++++++
  2 files changed, 24 insertions(+), 6 deletions(-)


diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -330,9 +330,9 @@
     return (yield from _wait(fs, timeout, return_when, loop))
 
 
-def _release_waiter(waiter, value=True, *args):
+def _release_waiter(waiter, *args):
     if not waiter.done():
-        waiter.set_result(value)
+        waiter.set_result(None)
 
 
 @coroutine
@@ -357,14 +357,17 @@
         return (yield from fut)
 
     waiter = futures.Future(loop=loop)
-    timeout_handle = loop.call_later(timeout, _release_waiter, waiter, False)
-    cb = functools.partial(_release_waiter, waiter, True)
+    timeout_handle = loop.call_later(timeout, _release_waiter, waiter)
+    cb = functools.partial(_release_waiter, waiter)
 
     fut = async(fut, loop=loop)
     fut.add_done_callback(cb)
 
     try:
-        if (yield from waiter):
+        # wait until the future completes or the timeout
+        yield from waiter
+
+        if fut.done():
             return fut.result()
         else:
             fut.remove_done_callback(cb)
@@ -397,7 +400,7 @@
             if timeout_handle is not None:
                 timeout_handle.cancel()
             if not waiter.done():
-                waiter.set_result(False)
+                waiter.set_result(None)
 
     for f in fs:
         f.add_done_callback(_on_completion)
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -552,6 +552,21 @@
         self.assertTrue(fut.done())
         self.assertTrue(fut.cancelled())
 
+    def test_wait_for_race_condition(self):
+
+        def gen():
+            yield 0.1
+            yield 0.1
+            yield 0.1
+
+        loop = self.new_test_loop(gen)
+
+        fut = asyncio.Future(loop=loop)
+        task = asyncio.wait_for(fut, timeout=0.2, loop=loop)
+        loop.call_later(0.1, fut.set_result, "ok")
+        res = loop.run_until_complete(task)
+        self.assertEqual(res, "ok")
+
     def test_wait(self):
 
         def gen():

-- 
Repository URL: http://hg.python.org/cpython