Closes #21886, #21447: Fix a race condition in asyncio when setting the result

vstinner · vstinner · commit a9acbe82e782 · 2014-07-05T15:29:41.000+02:00
of a Future with call_soon(). Add an helper, a private method, to set the
result only if the future was not cancelled.
diff --git a/Lib/asyncio/coroutines.py b/Lib/asyncio/coroutines.py
@@ -64,6 +64,12 @@ def __init__(self, gen, func):
         self.gen = gen
         self.func = func
         self._source_traceback = traceback.extract_stack(sys._getframe(1))
+        # __name__, __qualname__, __doc__ attributes are set by the coroutine()
+        # decorator
+
+    def __repr__(self):
+        return ('<%s %s>'
+                % (self.__class__.__name__, _format_coroutine(self)))
 
     def __iter__(self):
         return self
diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py
@@ -316,6 +316,12 @@ def remove_done_callback(self, fn):
 
     # So-called internal methods (note: no set_running_or_notify_cancel()).
 
+    def _set_result_unless_cancelled(self, result):
+        """Helper setting the result only if the future was not cancelled."""
+        if self.cancelled():
+            return
+        self.set_result(result)
+
     def set_result(self, result):
         """Mark the future done and set its result.
 
diff --git a/Lib/asyncio/proactor_events.py b/Lib/asyncio/proactor_events.py
@@ -38,7 +38,7 @@ def __init__(self, loop, sock, protocol, waiter=None,
             self._server.attach(self)
         self._loop.call_soon(self._protocol.connection_made, self)
         if waiter is not None:
-            self._loop.call_soon(waiter.set_result, None)
+            self._loop.call_soon(waiter._set_result_unless_cancelled, None)
 
     def _set_extra(self, sock):
         self._extra['pipe'] = sock
diff --git a/Lib/asyncio/queues.py b/Lib/asyncio/queues.py
@@ -173,7 +173,7 @@ def get(self):
             # run, we need to defer the put for a tick to ensure that
             # getters and putters alternate perfectly. See
             # ChannelTest.test_wait.
-            self._loop.call_soon(putter.set_result, None)
+            self._loop.call_soon(putter._set_result_unless_cancelled, None)
 
             return self._get()
 
diff --git a/Lib/asyncio/selector_events.py b/Lib/asyncio/selector_events.py
@@ -481,7 +481,7 @@ def __init__(self, loop, sock, protocol, waiter=None,
         self._loop.add_reader(self._sock_fd, self._read_ready)
         self._loop.call_soon(self._protocol.connection_made, self)
         if waiter is not None:
-            self._loop.call_soon(waiter.set_result, None)
+            self._loop.call_soon(waiter._set_result_unless_cancelled, None)
 
     def pause_reading(self):
         if self._closing:
@@ -690,7 +690,8 @@ def _on_handshake(self):
         self._loop.add_reader(self._sock_fd, self._read_ready)
         self._loop.call_soon(self._protocol.connection_made, self)
         if self._waiter is not None:
-            self._loop.call_soon(self._waiter.set_result, None)
+            self._loop.call_soon(self._waiter._set_result_unless_cancelled,
+                                 None)
 
     def pause_reading(self):
         # XXX This is a bit icky, given the comment at the top of
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
@@ -487,7 +487,8 @@ def _wait_for_one():
 def sleep(delay, result=None, *, loop=None):
     """Coroutine that completes after a given time (in seconds)."""
     future = futures.Future(loop=loop)
-    h = future._loop.call_later(delay, future.set_result, result)
+    h = future._loop.call_later(delay,
+                                future._set_result_unless_cancelled, result)
     try:
         return (yield from future)
     finally:
diff --git a/Lib/asyncio/unix_events.py b/Lib/asyncio/unix_events.py
@@ -269,7 +269,7 @@ def __init__(self, loop, pipe, protocol, waiter=None, extra=None):
         self._loop.add_reader(self._fileno, self._read_ready)
         self._loop.call_soon(self._protocol.connection_made, self)
         if waiter is not None:
-            self._loop.call_soon(waiter.set_result, None)
+            self._loop.call_soon(waiter._set_result_unless_cancelled, None)
 
     def _read_ready(self):
         try:
@@ -353,7 +353,7 @@ def __init__(self, loop, pipe, protocol, waiter=None, extra=None):
 
         self._loop.call_soon(self._protocol.connection_made, self)
         if waiter is not None:
-            self._loop.call_soon(waiter.set_result, None)
+            self._loop.call_soon(waiter._set_result_unless_cancelled, None)
 
     def get_write_buffer_size(self):
         return sum(len(data) for data in self._buffer)
diff --git a/Lib/test/test_asyncio/test_futures.py b/Lib/test/test_asyncio/test_futures.py
@@ -343,6 +343,12 @@ def memroy_error():
         message = m_log.error.call_args[0][0]
         self.assertRegex(message, re.compile(regex, re.DOTALL))
 
+    def test_set_result_unless_cancelled(self):
+        fut = asyncio.Future(loop=self.loop)
+        fut.cancel()
+        fut._set_result_unless_cancelled(2)
+        self.assertTrue(fut.cancelled())
+
 
 class FutureDoneCallbackTests(test_utils.TestCase):
 
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
@@ -211,6 +211,10 @@ def notmuch():
         coro = ('%s() at %s:%s'
                 % (coro_qualname, code.co_filename, code.co_firstlineno))
 
+        # test repr(CoroWrapper)
+        if coroutines._DEBUG:
+            self.assertEqual(repr(gen), '<CoroWrapper %s>' % coro)
+
         # test pending Task
         t = asyncio.Task(gen, loop=self.loop)
         t.add_done_callback(Dummy())
