req-context, handle early close request

bombela · bombela · commit 9b9cfa0afe93 · 2012-05-17T20:06:03.000-07:00
diff --git a/zerorpc/core.py b/zerorpc/core.py
@@ -43,7 +43,7 @@
 class ServerBase(object):
 
     def __init__(self, channel, methods=None, name=None, context=None,
-            pool_size=None, heartbeat=5):
+            pool_size=None, heartbeat=5, allow_remote_stop=False):
         self._multiplexer = ChannelMultiplexer(channel)
 
         if methods is None:
@@ -58,6 +58,9 @@ def __init__(self, channel, methods=None, name=None, context=None,
         self._inject_builtins()
         self._heartbeat_freq = heartbeat
 
+        if allow_remote_stop:
+            self._methods['_zerorpc_stop'] = self.stop
+
         for (k, functor) in self._methods.items():
             if not isinstance(functor, DecoratorBase):
                 self._methods[k] = rep(functor)
diff --git a/zerorpc/patterns.py b/zerorpc/patterns.py
@@ -77,7 +77,8 @@ def process_call(self, context, bufchan, event, functor):
         context_gen =context.middleware_call_procedure(functor, *event.args)
         methods = context_gen.next()
         wchannel = WrappedEvents(bufchan)
-        server = core.ServerBase(wchannel, methods, heartbeat=None)
+        server = core.ServerBase(wchannel, methods, heartbeat=None,
+                allow_remote_stop=True)
         bufchan.emit('CTX', (None,))
         try:
             server.run()
@@ -105,13 +106,32 @@ def process_answer(self, context, bufchan, event, method,
         wchannel = WrappedEvents(bufchan)
 
         class ContextClient(core.ClientBase):
+            def __init__(self, channel):
+                self._closed = False
+                super(ContextClient, self).__init__(channel, heartbeat=None)
+
             def close(self):
-                self._channel.emit('CTX_CLOSE', (None,))
+                if self._closed:
+                    return
+                self('_zerorpc_stop')
                 super(ContextClient, self).close()
+                wchannel.close()
                 bufchan.close()
                 bufchan.channel.close()
                 bufchan.channel.channel.close()
+                self._closed = True
+
+            def __call__(self, method, *args, **kargs):
+                if self._closed:
+                    raise ValueError('I/O operation on closed context')
+                return super(ContextClient, self).__call__(method, *args, **kargs)
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *args):
+                self.close()
 
-        return ContextClient(wchannel, heartbeat=None)
+        return ContextClient(wchannel)
 
 patterns_list = [ReqContext(), ReqStream(), ReqRep()]
