From 6377f2d69c95691eaa3965a6660670984deb41c9 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sun, 10 Jun 2018 03:06:59 +0000
Subject: [PATCH] issue #257: split pool shutdown and join.

---
 ansible_mitogen/process.py | 19 +++++++++++++++----
 mitogen/service.py         |  6 +++++-
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/ansible_mitogen/process.py b/ansible_mitogen/process.py
index d61de0a2..ff637299 100644
--- a/ansible_mitogen/process.py
+++ b/ansible_mitogen/process.py
@@ -30,6 +30,7 @@ from __future__ import absolute_import
 import errno
 import logging
 import os
+import signal
 import socket
 import sys
 
@@ -140,6 +141,7 @@ class MuxProcess(object):
         self.router.responder.whitelist_prefix('ansible')
         self.router.responder.whitelist_prefix('ansible_mitogen')
         mitogen.core.listen(self.router.broker, 'shutdown', self.on_broker_shutdown)
+        mitogen.core.listen(self.router.broker, 'exit', self.on_broker_exit)
         self.listener = mitogen.unix.Listener(
             router=self.router,
             path=self.unix_listener_path,
@@ -168,14 +170,23 @@ class MuxProcess(object):
 
     def on_broker_shutdown(self):
         """
-        Respond to the Router shutdown (indirectly triggered through exit of
-        the main thread) by unlinking the listening socket. Ideally this would
-        happen explicitly, but Ansible provides no hook to allow it.
+        Respond to broker shutdown by beginning service pool shutdown. Do not
+        join on the pool yet, since that would block the broker thread which
+        then cannot clean up pending handlers, which is required for the
+        threads to exit gracefully.
         """
-        self.pool.stop()
+        self.pool.stop(join=False)
         try:
             os.unlink(self.listener.path)
         except OSError, e:
             # Prevent a shutdown race with the parent process.
             if e.args[0] != errno.ENOENT:
                 raise
+
+    def on_broker_exit(self):
+        """
+        Respond to the broker thread about to exit by sending SIGTERM to
+        ourself. In future this should gracefully join the pool, but TERM is
+        fine for now.
+        """
+        os.kill(os.getpid(), signal.SIGTERM)
diff --git a/mitogen/service.py b/mitogen/service.py
index 62180e33..dd8ac5cb 100644
--- a/mitogen/service.py
+++ b/mitogen/service.py
@@ -456,9 +456,13 @@ class Pool(object):
 
     closed = False
 
-    def stop(self):
+    def stop(self, join=True):
         self.closed = True
         self._select.close()
+        if join:
+            self.join()
+
+    def join(self):
         for th in self._threads:
             th.join()
         for invoker in self._invoker_by_name.itervalues():