diff --git a/mitogen/core.py b/mitogen/core.py index c2c5fcb1..60198973 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -868,11 +868,6 @@ class Router(object): return handle - def allocate_id(self): - master = Context(self, 0) - reply = master.send_await(Message(dst_id=0, handle=ALLOCATE_ID)) - return reply.unpickle() - def on_shutdown(self, broker): """Called during :py:meth:`Broker.shutdown`, informs callbacks registered with :py:meth:`add_handle_cb` the connection is dead.""" diff --git a/mitogen/master.py b/mitogen/master.py index 4d17eed1..b21a8478 100644 --- a/mitogen/master.py +++ b/mitogen/master.py @@ -739,6 +739,7 @@ METHOD_NAMES = { def upgrade_router(econtext): if not isinstance(econtext.router, Router): # TODO econtext.router.__class__ = Router # TODO + econtext.router.id_allocator = ChildIdAllocator(econtext.router) LOG.debug('_proxy_connect(): constructing ModuleForwarder') ModuleForwarder(econtext.router, econtext.parent, econtext.importer) @@ -785,6 +786,18 @@ class IdAllocator(object): ) +class ChildIdAllocator(object): + def __init__(self, router): + self.router = router + + def allocate(self): + master = Context(self.router, 0) + reply = master.send_await( + mitogen.core.Message(dst_id=0, handle=mitogen.core.ALLOCATE_ID) + ) + return reply.unpickle() + + class Router(mitogen.core.Router): debug = False diff --git a/tests/data/id_allocation.py b/tests/data/id_allocation.py index dcfd5230..ce83f3fc 100644 --- a/tests/data/id_allocation.py +++ b/tests/data/id_allocation.py @@ -1,8 +1,10 @@ import mitogen.core +import mitogen.master -@mitogen.core.takes_router -def allocate_an_id(router): - return router.allocate_id() +@mitogen.core.takes_econtext +def allocate_an_id(econtext): + mitogen.master.upgrade_router(econtext) + return econtext.router.allocate_id()