From 7316c08237c2407b7c66327f385301ee66d38995 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 29 Apr 2018 06:02:19 +0100
Subject: [PATCH] core: fix _tls_init() race.

The GIL could be lost between the check for an empty list and popping a
socket off the list. Previously _tls_init (per its name) used per-thread
storage, hence the bug.
---
 mitogen/core.py   | 12 +++++++-----
 mitogen/parent.py |  2 +-
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/mitogen/core.py b/mitogen/core.py
index 8b81fc45..1bb4a1ee 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -1008,12 +1008,14 @@ class Latch(object):
         return len(self._queue) == 0
 
     def _tls_init(self):
-        if self._sockets:
+        # pop() must be atomic, which is true for GIL-equipped interpreters.
+        try:
             return self._sockets.pop()
-        rsock, wsock = socket.socketpair()
-        set_cloexec(rsock.fileno())
-        set_cloexec(wsock.fileno())
-        return rsock, wsock
+        except IndexError:
+            rsock, wsock = socket.socketpair()
+            set_cloexec(rsock.fileno())
+            set_cloexec(wsock.fileno())
+            return rsock, wsock
 
     def get(self, timeout=None, block=True):
         _vv and IOLOG.debug('%r.get(timeout=%r, block=%r)',
diff --git a/mitogen/parent.py b/mitogen/parent.py
index eae9ee6c..497ccd0a 100644
--- a/mitogen/parent.py
+++ b/mitogen/parent.py
@@ -798,7 +798,6 @@ class Context(mitogen.core.Context):
         LOG.debug('%r.shutdown() sending SHUTDOWN', self)
         latch = mitogen.core.Latch()
         mitogen.core.listen(self, 'disconnect', lambda: latch.put(None))
-
         self.send(
             mitogen.core.Message(
                 handle=mitogen.core.SHUTDOWN,
@@ -807,6 +806,7 @@ class Context(mitogen.core.Context):
 
         if wait:
             latch.get()
+        return latch
 
 
 class RouteMonitor(object):