Pass the request event to the client_after_request hook

lopter · lopter · commit a96fdc93da36 · 2013-01-01T18:56:17.000+01:00
This is useful to get some context on the current task when the reply
event is None (either because we are timing out or because we are
talking to a PUSH/PULL or PUB/SUB).

It also replaces the useless method argument from the patterns.
diff --git a/tests/test_middleware_client.py b/tests/test_middleware_client.py
@@ -95,9 +95,11 @@ def client_before_request(self, event):
 class ClientAfterRequestMiddleware(object):
     def __init__(self):
         self.called = False
-    def client_after_request(self, event, exception):
+    def client_after_request(self, req_event, rep_event, exception):
         self.called = True
-        self.retcode = event.name
+        assert req_event is not None
+        assert req_event.name == "echo" or req_event.name == "echoes"
+        self.retcode = rep_event.name
         assert exception is None
 
 def test_hook_client_after_request():
@@ -158,9 +160,11 @@ def test_hook_client_after_request_timeout():
     class ClientAfterRequestMiddleware(object):
         def __init__(self):
             self.called = False
-        def client_after_request(self, event, exception):
+        def client_after_request(self, req_event, rep_event, exception):
             self.called = True
-            assert event is None
+            assert req_event is not None
+            assert req_event.name == "timeout"
+            assert rep_event is None
 
     zero_ctx = zerorpc.Context()
     test_middleware = ClientAfterRequestMiddleware()
@@ -187,7 +191,9 @@ def client_after_request(self, event, exception):
 class ClientAfterFailedRequestMiddleware(object):
     def __init__(self):
         self.called = False
-    def client_after_request(self, rep_event, exception):
+    def client_after_request(self, req_event, rep_event, exception):
+        assert req_event is not None
+        assert req_event.name == "crash" or req_event.name == "echoes_crash"
         self.called = True
         assert isinstance(exception, zerorpc.RemoteError)
         assert exception.name == 'RuntimeError'
@@ -338,7 +344,9 @@ def client_handle_remote_error(self, event):
             etype = eval(name)
             e = etype(tb)
             return e
-        def client_after_request(self, rep_event, exception):
+        def client_after_request(self, req_event, rep_event, exception):
+            assert req_event is not None
+            assert req_event.name == "crash"
             self.called = True
             assert isinstance(exception, RuntimeError)
 
diff --git a/zerorpc/context.py b/zerorpc/context.py
@@ -151,7 +151,7 @@ def hook_client_before_request(self, event):
         for functor in self._hooks['client_before_request']:
             functor(event)
 
-    def hook_client_after_request(self, reply_event, exception=None):
+    def hook_client_after_request(self, request_event, reply_event, exception=None):
         """Called when an answer or a timeout has been received from the server.
 
         This hook is called right before the answer is returned to the client.
@@ -171,4 +171,4 @@ def hook_client_after_request(self, reply_event, exception=None):
 
         """
         for functor in self._hooks['client_after_request']:
-            functor(reply_event, exception)
+            functor(request_event, reply_event, exception)
diff --git a/zerorpc/core.py b/zerorpc/core.py
@@ -208,19 +208,17 @@ def _select_pattern(self, event):
         msg = 'Unable to find a pattern for: {0}'.format(event)
         raise RuntimeError(msg)
 
-    def _process_response(self, method, bufchan, timeout):
+    def _process_response(self, request_event, bufchan, timeout):
         try:
-            event = bufchan.recv(timeout)
-            pattern = self._select_pattern(event)
-            return pattern.process_answer(self._context, bufchan, event, method,
-                    self._handle_remote_error)
+            reply_event = bufchan.recv(timeout)
+            pattern = self._select_pattern(reply_event)
+            return pattern.process_answer(self._context, bufchan, request_event,
+                    reply_event, self._handle_remote_error)
         except TimeoutExpired:
             bufchan.close()
-            # FIXME: Add the method name as a separate argument? (so you
-            # could get it back using the args attribute in the middleware).
             ex = TimeoutExpired(timeout,
-                    'calling remote method {0}'.format(method))
-            self._context.hook_client_after_request(None, ex)
+                    'calling remote method {0}'.format(request_event.name))
+            self._context.hook_client_after_request(request_event, None, ex)
             raise ex
         except:
             bufchan.close()
@@ -234,13 +232,13 @@ def __call__(self, method, *args, **kargs):
         bufchan = BufferedChannel(hbchan, inqueue_size=kargs.get('slots', 100))
 
         xheader = self._context.hook_get_task_context()
-        event = bufchan.create_event(method, args, xheader)
-        self._context.hook_client_before_request(event)
-        bufchan.emit_event(event)
+        request_event = bufchan.create_event(method, args, xheader)
+        self._context.hook_client_before_request(request_event)
+        bufchan.emit_event(request_event)
 
         try:
             if kargs.get('async', False) is False:
-                return self._process_response(method, bufchan, timeout)
+                return self._process_response(request_event, bufchan, timeout)
 
             async_result = gevent.event.AsyncResult()
             gevent.spawn(self._process_response, method, bufchan,
diff --git a/zerorpc/patterns.py b/zerorpc/patterns.py
@@ -36,13 +36,13 @@ def process_call(self, context, bufchan, req_event, functor):
     def accept_answer(self, event):
         return True
 
-    def process_answer(self, context, bufchan, rep_event, method,
+    def process_answer(self, context, bufchan, req_event, rep_event,
             handle_remote_error):
         if rep_event.name == 'ERR':
             exception = handle_remote_error(rep_event)
-            context.hook_client_after_request(rep_event, exception)
+            context.hook_client_after_request(req_event, rep_event, exception)
             raise exception
-        context.hook_client_after_request(rep_event)
+        context.hook_client_after_request(req_event, rep_event)
         bufchan.close()
         result = rep_event.args[0]
         return result
@@ -66,27 +66,27 @@ def process_call(self, context, bufchan, req_event, functor):
     def accept_answer(self, event):
         return event.name in ('STREAM', 'STREAM_DONE')
 
-    def process_answer(self, context, bufchan, rep_event, method,
+    def process_answer(self, context, bufchan, req_event, rep_event,
             handle_remote_error):
 
         def is_stream_done(rep_event):
             return rep_event.name == 'STREAM_DONE'
         bufchan.on_close_if = is_stream_done
 
-        def iterator(rep_event):
+        def iterator(req_event, rep_event):
             while rep_event.name == 'STREAM':
                 # Like in process_call, we made the choice to call the
                 # after_exec hook only when the stream is done.
                 yield rep_event.args
                 rep_event = bufchan.recv()
             if rep_event.name == 'ERR':
                 exception = handle_remote_error(rep_event)
-                context.hook_client_after_request(rep_event, exception)
+                context.hook_client_after_request(req_event, rep_event, exception)
                 raise exception
-            context.hook_client_after_request(rep_event)
+            context.hook_client_after_request(req_event, rep_event)
             bufchan.close()
 
-        return iterator(rep_event)
+        return iterator(req_event, rep_event)
 
 
 patterns_list = [ReqStream(), ReqRep()]
