From 7a60b20dc6a5151761d1b3fc5b096f0e3fe9b375 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 21 Sep 2017 13:41:12 +0530
Subject: [PATCH] core: Generalize/duplicate the call/send_await code using
 Receiver.

---
 mitogen/core.py   | 40 +++++++++++++++++++---------------------
 mitogen/master.py | 44 ++------------------------------------------
 2 files changed, 21 insertions(+), 63 deletions(-)

diff --git a/mitogen/core.py b/mitogen/core.py
index cdff71d2..1c1d2e27 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -235,10 +235,11 @@ class Sender(object):
 
 
 class Receiver(object):
-    def __init__(self, router, handle=None):
+    def __init__(self, router, handle=None, persist=True, respondent=None):
         self.router = router
         self.handle = handle  # Avoid __repr__ crash in add_handler()
-        self.handle = router.add_handler(self._on_receive, handle)
+        self.handle = router.add_handler(self._on_receive, handle,
+                                         persist, respondent)
         self._queue = Queue.Queue()
 
     def __repr__(self):
@@ -266,7 +267,7 @@ class Receiver(object):
                 continue
 
         if msg is None:
-            return
+            raise TimeoutError('deadline exceeded.')
 
         IOLOG.debug('%r.on_receive() got %r', self, msg)
         if msg == _DEAD:
@@ -282,6 +283,9 @@ class Receiver(object):
 
         return msg, data
 
+    def get_data(self, timeout=None):
+        return self.get(timeout)[1]
+
     def __iter__(self):
         """Yield objects from this channel until it is closed."""
         while True:
@@ -393,7 +397,7 @@ class Importer(object):
             self._cache[fullname] = ret = (
                 self._context.send_await(
                     Message(data=fullname, handle=GET_MODULE)
-                ).unpickle()
+                )
             )
 
         if ret is None:
@@ -687,29 +691,23 @@ class Context(object):
             msg.src_id = mitogen.context_id
         self.router.route(msg)
 
-    def send_await(self, msg, deadline=None):
-        """Send `msg` and wait for a response with an optional timeout."""
+    def send_async(self, msg, persist=False):
         if self.router.broker._thread == threading.currentThread():  # TODO
             raise SystemError('Cannot making blocking call on broker thread')
 
-        queue = Queue.Queue()
-        msg.reply_to = self.router.add_handler(queue.put,
-            persist=False,
-            respondent=self)
-        LOG.debug('%r.send_await(%r)', self, msg)
+        receiver = Receiver(self.router, persist=persist, respondent=self)
+        msg.reply_to = receiver.handle
 
+        LOG.debug('%r.send_async(%r)', self, msg)
         self.send(msg)
-        try:
-            msg = queue.get(True, deadline)
-        except Queue.Empty:
-            # self.broker.defer(self.stream.on_disconnect, self.broker)
-            raise TimeoutError('deadline exceeded.')
-
-        if msg == _DEAD:
-            raise StreamError('lost connection during call.')
+        return receiver
 
-        IOLOG.debug('%r._send_await() -> %r', self, msg)
-        return msg
+    def send_await(self, msg, deadline=None):
+        """Send `msg` and wait for a response with an optional timeout."""
+        receiver = self.send_async(msg)
+        response = receiver.get_data(deadline)
+        IOLOG.debug('%r._send_await() -> %r', self, response)
+        return response
 
     def __repr__(self):
         return 'Context(%s, %r)' % (self.context_id, self.name)
diff --git a/mitogen/master.py b/mitogen/master.py
index 7c664997..d2fe66c8 100644
--- a/mitogen/master.py
+++ b/mitogen/master.py
@@ -649,16 +649,6 @@ class Context(mitogen.core.Context):
         """
         mitogen.core.fire(self, 'disconnect')
 
-    def _discard_result(self, msg):
-        data = msg.unpickle()
-        if isinstance(data, Exception):
-            try:
-                raise data
-            except Exception:
-                LOG.exception('_discard_result')
-        else:
-            LOG.debug('_discard_result: %r', data)
-
     def call_async(self, fn, *args, **kwargs):
         LOG.debug('%r.call_async(%r, *%r, **%r)',
                   self, fn, args, kwargs)
@@ -669,46 +659,16 @@ class Context(mitogen.core.Context):
         else:
             klass = None
 
-        self.send(
+        return self.send_async(
             mitogen.core.Message.pickled(
                 (fn.__module__, klass, fn.__name__, args, kwargs),
                 handle=mitogen.core.CALL_FUNCTION,
-                reply_to=self.router.add_handler(self._discard_result),
             )
         )
 
-    def call_with_deadline(self, deadline, fn, *args, **kwargs):
-        """Invoke `fn([context,] *args, **kwargs)` in the external context.
-
-        If `deadline` is not ``None``, expire the call after `deadline`
-        seconds. If `deadline` is ``None``, the invocation may block
-        indefinitely."""
-        LOG.debug('%r.call_with_deadline(%r, %r, *%r, **%r)',
-                  self, deadline, fn, args, kwargs)
-
-        if isinstance(fn, types.MethodType) and \
-           isinstance(fn.im_self, (type, types.ClassType)):
-            klass = fn.im_self.__name__
-        else:
-            klass = None
-
-        response = self.send_await(
-            mitogen.core.Message.pickled(
-                (fn.__module__, klass, fn.__name__, args, kwargs),
-                handle=mitogen.core.CALL_FUNCTION
-            ),
-            deadline
-        )
-
-        decoded = response.unpickle()
-        if isinstance(decoded, mitogen.core.CallError):
-            raise decoded
-        return decoded
-
     def call(self, fn, *args, **kwargs):
         """Invoke `fn(*args, **kwargs)` in the external context."""
-        return self.call_with_deadline(None, fn, *args, **kwargs)
-
+        return self.call_async(fn, *args, **kwargs).get_data()
 
 
 def _local_method():