Merge pull request #35 from sideshowdave7/feature/replies_and_latency0.2.3

Feature: job latencies

5 files changed, 107 insertions, 18 deletions

diff --git a/bin/send_request b/bin/send_request
new file mode 100755
index 0000000..aa84639
--- /dev/null
+++ b/bin/send_request
@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+"""
+Usage: ./send_msg <ipaddresss> <command> <frame w/ values space separated>
+"""
+import sys
+
+from eventmq.sender import Sender
+from eventmq.client.messages import send_request
+
+
+if __name__ == "__main__":
+    s = Sender()
+    s.connect(sys.argv[1])
+
+    msg = ['run', {
+           'path': 'eventmq.scheduler',
+           'callable': 'test_job',
+           'class_args': (),
+           'class_kwargs': {},
+           'args': (),
+           'kwargs': {}
+           }]
+
+    send_request(s, msg, guarantee=True, reply_requested=True)
diff --git a/eventmq/jobmanager.py b/eventmq/jobmanager.py
index 52b3c7e..6e4cd5e 100644
--- a/eventmq/jobmanager.py
+++ b/eventmq/jobmanager.py
@@ -17,7 +17,8 @@
 ================================
 Ensures things about jobs and spawns the actual tasks
 """
-from json import loads as serializer
+from json import loads as deserializer
+from json import dumps as serializer
 import logging
 import signal
 
@@ -67,14 +68,13 @@ class JobManager(HeartbeatMixin, EMQPService):
 
         #: keep track of workers
         concurrent_jobs = kwargs.pop('concurrent_jobs', None)
-        if concurrent_jobs is None:
-            concurrent_jobs = conf.CONCURRENT_JOBS
-        self.workers = Pool(processes=concurrent_jobs)
+        if concurrent_jobs is not None:
+            conf.CONCURRENT_JOBS = concurrent_jobs
 
         #: List of queues that this job manager is listening on
         self.queues = kwargs.pop('queues', None)
         if self.queues is None:
-            self.quques = conf.QUEUES
+            self.queues = conf.QUEUES
 
         if not kwargs.pop('skip_signal', False):
             # handle any sighups by reloading config
@@ -90,12 +90,23 @@ class JobManager(HeartbeatMixin, EMQPService):
 
         self._setup()
 
+    @property
+    def workers(self):
+        if not hasattr(self, '_workers'):
+            self._workers = Pool(processes=conf.CONCURRENT_JOBS)
+        elif self._workers.processes != conf.CONCURRENT_JOBS:
+            self._workers.close()
+            self._workers = Pool(processes=conf.CONCURRENT_JOBS)
+
+        return self._workers
+
     def _start_event_loop(self):
         """
         Starts the actual event loop. Usually called by :meth:`start`
         """
         # Acknowledgment has come
         # Send a READY for each available worker
+
         for i in range(0, conf.CONCURRENT_JOBS):
             self.send_ready()
 
@@ -145,16 +156,26 @@ class JobManager(HeartbeatMixin, EMQPService):
         # s_ indicates the string path vs the actual module and class
         # queue_name = msg[0]
 
-        # run callable
-        payload = serializer(msg[2])
-        # subcmd = payload[0]
+        # Parse REQUEST message values
+        headers = msg[1]
+        payload = deserializer(msg[2])
         params = payload[1]
 
+        if 'reply-requested' in headers:
+           callback = self.worker_done_with_reply
+        else:
+           callback = self.worker_done
+
+        # kick off the job asynchronously with an appropiate callback
         self.workers.apply_async(func=worker.run,
-                                 args=(params,),
-                                 callback=self.worker_done)
+                                 args=(params, msgid),
+                                 callback=callback)
+
+    def worker_done_with_reply(self, msgid):
+        self.send_reply(msgid)
+        self.send_ready()
 
-    def worker_done(self, result):
+    def worker_done(self, msgid):
         self.send_ready()
 
     def send_ready(self):
@@ -164,6 +185,23 @@ class JobManager(HeartbeatMixin, EMQPService):
         """
         sendmsg(self.outgoing, 'READY')
 
+    def send_reply(self, res):
+         """
+         Sends an REPLY response
+
+         Args:
+             socket (socket): The socket to use for this ack
+             recipient (str): The recipient id for the ack
+             msgid: The unique id that we are acknowledging
+         """
+         msgid = res[0]
+
+         reply = res[1]
+
+         reply = serializer(reply)
+
+         sendmsg(self.outgoing, 'REPLY', [reply, msgid])
+
     def on_heartbeat(self, msgid, message):
         """
         a placeholder for a noop command. The actual 'logic' for HEARTBEAT is
diff --git a/eventmq/router.py b/eventmq/router.py
index fbd8d80..7dd2471 100644
--- a/eventmq/router.py
+++ b/eventmq/router.py
@@ -105,6 +105,12 @@ class Router(HeartbeatMixin):
         #:     }
         self.schedulers = {}
 
+
+        #: Latency tracking dictionary
+        #: Key: msgid of message each REQUEST received and forwarded to a worker
+        #: Value: (timestamp, queue_name)
+        self.job_latencies = {}
+
         #: Set to True when the router should die.
         self.received_disconnect = False
 
@@ -270,6 +276,25 @@ class Router(HeartbeatMixin):
             self.add_scheduler(sender)
             self.send_ack(self.incoming, sender, msgid)
 
+    def on_reply(self, sender, msgid, msg):
+        """
+        Handles an REPLY message. Replies are sent by the worker for latanecy
+        measurements
+        """
+
+        orig_msgid = msg[1]
+        logger.info('Received REPLY from {} (msgid: {}, ACK msgid: {})'.format(
+            sender, msgid, orig_msgid))
+
+        if orig_msgid in self.job_latencies:
+            logger.info("Completed {queue} job with msgid: {msgid} in "
+                        "{time:.2f}ms".\
+                format(queue=self.job_latencies[orig_msgid][1],
+                       msgid=orig_msgid,
+                       time=(monotonic()-self.job_latencies[orig_msgid][0])*1000.0))
+            del self.job_latencies[orig_msgid]
+
+
     def on_disconnect(self, msgid, msg):
         # Loops event loops should check for this and break out
         self.received_disconnect = True
@@ -364,6 +389,7 @@ class Router(HeartbeatMixin):
         try:
             # Rebuild the message to be sent to the worker. fwdmsg will
             # properly address the message.
+            self.job_latencies[msgid] = (monotonic(), queue_name)
             fwdmsg(self.outgoing, worker_addr, ['', constants.PROTOCOL_VERSION,
                                                 'REQUEST', msgid, ] + msg)
             self.workers[worker_addr]['available_slots'] -= 1
diff --git a/eventmq/tests/test_jobmanager.py b/eventmq/tests/test_jobmanager.py
index 6d3c663..483dde1 100644
--- a/eventmq/tests/test_jobmanager.py
+++ b/eventmq/tests/test_jobmanager.py
@@ -80,7 +80,7 @@ class TestCase(unittest.TestCase):
 
         jm.on_request(_msgid, _msg)
         apply_async_mock.assert_called_with(
-            args=({'a': 1},),
+            args=({'a': 1}, _msgid),
             callback=jm.worker_done,
             func=run_mock)
 
diff --git a/eventmq/worker.py b/eventmq/worker.py
index 78ea97a..284f478 100644
--- a/eventmq/worker.py
+++ b/eventmq/worker.py
@@ -26,7 +26,7 @@ from . import log
 logger = log.setup_logger(__name__)
 
 
-def run(payload):
+def run(payload, msgid):
     """
     process a run message and execute a job
 
@@ -47,7 +47,7 @@ def run(payload):
         reload(package)
     except Exception as e:
         logger.exception('Error importing module: {}'.format(str(e)))
-        return 'DONE'
+        return (msgid, str(e))
 
     if s_cls:
         cls = getattr(package, s_cls)
@@ -70,7 +70,7 @@ def run(payload):
         callable_ = getattr(obj, s_callable)
     except AttributeError as e:
         logger.exception('Error getting callable: {}'.format(str(e)))
-        return 'DONE'
+        return (msgid, str(e))
 
     if "args" in payload:
         args = payload["args"]
@@ -83,10 +83,11 @@ def run(payload):
         kwargs = {}
 
     try:
-        callable_(*args, **kwargs)
+        r = callable_(*args, **kwargs)
+        return (msgid, r)
     except Exception as e:
         logger.exception(e)
-        return 'DONE'
+        return (msgid, str(e))
 
     # Signal that we're done with this job
-    return 'DONE'
+    return (msgid, '')