diff --git a/mitogen/core.py b/mitogen/core.py index b2b7e0df..fa274324 100644 --- a/mitogen/core.py +++ b/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: diff --git a/mitogen/master.py b/mitogen/master.py index a0a85a8e..273c6920 100644 --- a/mitogen/master.py +++ b/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) diff --git a/tests/router_test.py b/tests/router_test.py new file mode 100644 index 00000000..3f460c3f --- /dev/null +++ b/tests/router_test.py @@ -0,0 +1,26 @@ +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() +