core: cleanup handlers on broker crash; closes #112.

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')
 

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