From fe614aa9661650a8cd8985d3677ba4412b666eef Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 29 Mar 2018 19:32:31 +0545
Subject: [PATCH] core: cleanup handlers on broker crash; closes #112.

---
 mitogen/core.py      |  8 +++++++-
 tests/router_test.py | 30 ++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/mitogen/core.py b/mitogen/core.py
index 29e04a2a..b0fb8603 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -1201,6 +1201,7 @@ class Router(object):
 
     def __init__(self, broker):
         self.broker = broker
+        listen(broker, 'crash', self._cleanup_handlers)
         listen(broker, 'shutdown', self.on_broker_shutdown)
 
         # Here seems as good a place as any.
@@ -1230,7 +1231,11 @@ class Router(object):
         for context in self._context_by_id.itervalues():
             context.on_shutdown(self.broker)
 
-        for _, func in self._handle_map.itervalues():
+        self._cleanup_handlers()
+
+    def _cleanup_handlers(self):
+        while self._handle_map:
+            _, (_, func) = self._handle_map.popitem()
             func(_DEAD)
 
     def register(self, context, stream):
@@ -1415,6 +1420,7 @@ class Broker(object):
                 side.stream.on_disconnect(self)
         except Exception:
             LOG.exception('_broker_main() crashed')
+            fire(self, 'crash')
 
         fire(self, 'exit')
 
diff --git a/tests/router_test.py b/tests/router_test.py
index c6b4e2df..053037fb 100644
--- a/tests/router_test.py
+++ b/tests/router_test.py
@@ -21,6 +21,36 @@ def send_n_sized_reply(sender, n):
     return 123
 
 
+class CrashTest(testlib.BrokerMixin, unittest2.TestCase):
+    # This is testing both Broker's ability to crash nicely, and Router's
+    # ability to respond to the crash event.
+    klass = mitogen.master.Router
+
+    def _naughty(self):
+        raise ValueError('eek')
+
+    def test_shutdown(self):
+        router = self.klass(self.broker)
+
+        sem = mitogen.core.Latch()
+        router.add_handler(sem.put)
+
+        log = testlib.LogCapturer('mitogen')
+        log.start()
+
+        # Force a crash and ensure it wakes up.
+        self.broker._loop_once = self._naughty
+        self.broker.defer(lambda: None)
+
+        # sem should have received _DEAD.
+        self.assertEquals(mitogen.core._DEAD, sem.get())
+
+        # Ensure it was logged.
+        expect = '_broker_main() crashed'
+        self.assertTrue(expect in log.stop())
+
+
+
 class AddHandlerTest(unittest2.TestCase):
     klass = mitogen.master.Router