core: Ensure add_handler() callbacks really receive _DEAD on shutdown

mitogen/core.py

@@ -1003,6 +1003,9 @@ class Router(object):
         for context in self._context_by_id.itervalues():
             context.on_shutdown(self.broker)
 
+        for _, func in self._handle_map.itervalues():
+            func(_DEAD)
+
     def add_route(self, target_id, via_id):
         _v and LOG.debug('%r.add_route(%r, %r)', self, target_id, via_id)
         try:

mitogen/master.py

@@ -660,6 +660,9 @@ class IdAllocator(object):
             self.lock.release()
 
     def on_allocate_id(self, msg):
+        if msg == mitogen.core._DEAD:
+            return
+
         id_ = self.allocate()
         requestee = self.router.context_by_id(msg.src_id)
         allocated = self.router.context_by_id(id_, msg.src_id)

tests/router_test.py

+import Queue
+import subprocess
+import time
+
+import unittest2
+
+import testlib
+import mitogen.master
+import mitogen.utils
+
+mitogen.utils.log_to_file()
+
+class AddHandlerTest(unittest2.TestCase):
+    klass = mitogen.master.Router
+
+    def test_invoked_at_shutdown(self):
+        router = self.klass()
+        queue = Queue.Queue()
+        handle = router.add_handler(queue.put)
+        router.broker.shutdown()
+        self.assertEquals(queue.get(timeout=5), mitogen.core._DEAD)
+
+
+if __name__ == '__main__':
+    unittest2.main()
+