asyncio: Fix race in FastChildWatcher (by its original author, Anthony Baire).

1 files changed, 16 insertions, 20 deletions

diff --git a/Lib/asyncio/unix_events.py b/Lib/asyncio/unix_events.py
index 7a6546d..2418642 100644
--- a/Lib/asyncio/unix_events.py
+++ b/Lib/asyncio/unix_events.py
@@ -641,22 +641,16 @@ class FastChildWatcher(BaseChildWatcher):
 
     def add_child_handler(self, pid, callback, *args):
         assert self._forks, "Must use the context manager"
+        with self._lock:
+            try:
+                returncode = self._zombies.pop(pid)
+            except KeyError:
+                # The child is running.
+                self._callbacks[pid] = callback, args
+                return
 
-        self._callbacks[pid] = callback, args
-
-        try:
-            # Ensure that the child is not already terminated.
-            # (raise KeyError if still alive)
-            returncode = self._zombies.pop(pid)
-
-            # Child is dead, therefore we can fire the callback immediately.
-            # First we remove it from the dict.
-            # (raise KeyError if .remove_child_handler() was called in-between)
-            del self._callbacks[pid]
-        except KeyError:
-            pass
-        else:
-            callback(pid, returncode, *args)
+        # The child is dead already. We can fire the callback.
+        callback(pid, returncode, *args)
 
     def remove_child_handler(self, pid):
         try:
@@ -681,16 +675,18 @@ class FastChildWatcher(BaseChildWatcher):
 
                 returncode = self._compute_returncode(status)
 
-            try:
-                callback, args = self._callbacks.pop(pid)
-            except KeyError:
-                # unknown child
-                with self._lock:
+            with self._lock:
+                try:
+                    callback, args = self._callbacks.pop(pid)
+                except KeyError:
+                    # unknown child
                     if self._forks:
                         # It may not be registered yet.
                         self._zombies[pid] = returncode
                         continue
+                    callback = None
 
+            if callback is None:
                 logger.warning(
                     "Caught subprocess termination from unknown pid: "
                     "%d -> %d", pid, returncode)