Merge pull request #31 from sideshowdave7/0.4

Port master changes to 0.4

9 files changed, 194 insertions, 95 deletions

diff --git a/eventmq/__init__.py b/eventmq/__init__.py
index 43c0365..492f1b9 100644
--- a/eventmq/__init__.py
+++ b/eventmq/__init__.py
@@ -1,5 +1,5 @@
 __author__ = 'EventMQ Contributors'
-__version__ = '0.3.4'
+__version__ = '0.3.4.4'
 
 PROTOCOL_VERSION = 'eMQP/1.0'
 
diff --git a/eventmq/jobmanager.py b/eventmq/jobmanager.py
index d86c738..c490e68 100644
--- a/eventmq/jobmanager.py
+++ b/eventmq/jobmanager.py
@@ -84,12 +84,13 @@ class JobManager(HeartbeatMixin, EMQPService):
         self.name = conf.NAME or generate_device_name()
         logger.info('Initializing JobManager {}...'.format(self.name))
 
-        if skip_signal:
+        if not skip_signal:
             # handle any sighups by reloading config
             signal.signal(signal.SIGHUP, self.sighup_handler)
             signal.signal(signal.SIGTERM, self.sigterm_handler)
             signal.signal(signal.SIGINT, self.sigterm_handler)
             signal.signal(signal.SIGQUIT, self.sigterm_handler)
+            signal.signal(signal.SIGUSR1, self.handle_pdb)
 
         #: JobManager starts out by INFORMing the router of it's existence,
         #: then telling the router that it is READY. The reply will be the unit
@@ -99,11 +100,29 @@ class JobManager(HeartbeatMixin, EMQPService):
 
         self.poller = Poller()
 
+        #: Stats and monitoring information
+
+        #: Jobs in flight tracks all jobs currently executing.
+        #: Key: msgid, Value: The message with all the details of the job
+        self.jobs_in_flight = {}
+
+        #: Running total number of REQUEST messages received on the broker
+        self.total_requests = 0
+        #: Running total number of READY messages sent to the broker
+        self.total_ready_sent = 0
+        #: Keep track of what pids are servicing our requests
+        #: Key: pid, Value: # of jobs completed on the process with that pid
+        self.pid_distribution = {}
+
         #: Setup worker queues
         self.request_queue = mp_queue()
         self.finished_queue = mp_queue()
         self._setup()
 
+    def handle_pdb(self, sig, frame):
+        import pdb
+        pdb.Pdb().set_trace(frame)
+
     @property
     def workers(self):
         if not hasattr(self, '_workers'):
@@ -123,9 +142,6 @@ class JobManager(HeartbeatMixin, EMQPService):
         # Acknowledgment has come
         # Send a READY for each available worker
 
-        for i in range(0, len(self.workers)):
-            self.send_ready()
-
         self.status = STATUS.running
 
         try:
@@ -163,7 +179,7 @@ class JobManager(HeartbeatMixin, EMQPService):
                             sys.exit(0)
                     else:
                         try:
-                            events = self.poller.poll(10)
+                            events = self.poller.poll(1000)
                         except zmq.ZMQError:
                             logger.debug('Disconnecting due to ZMQError while'
                                          ' polling')
@@ -214,13 +230,21 @@ class JobManager(HeartbeatMixin, EMQPService):
 
         logger.debug(resp)
         pid = resp['pid']
+        msgid = resp['msgid']
+        callback = resp['callback']
+        death = resp['death']
 
-        callback = getattr(self, resp['callback'])
-        callback(resp['return'], resp['msgid'])
+        callback = getattr(self, callback)
+        callback(resp['return'], msgid, death, pid)
 
-        if resp['return'] == 'DEATH':
-            if pid in self._workers.keys():
-                del self._workers[pid]
+        if msgid in self.jobs_in_flight:
+            del self.jobs_in_flight[msgid]
+
+        if not death:
+            if pid not in self.pid_distribution:
+                self.pid_distribution[pid] = 1
+            else:
+                self.pid_distribution[pid] += 1
 
     def on_request(self, msgid, msg):
         """
@@ -258,6 +282,9 @@ class JobManager(HeartbeatMixin, EMQPService):
         # queue_name = msg[0]
 
         # Parse REQUEST message values
+
+        self.total_requests += 1
+
         headers = msg[1]
         payload = deserializer(msg[2])
         params = payload[1]
@@ -275,27 +302,54 @@ class JobManager(HeartbeatMixin, EMQPService):
         payload['msgid'] = msgid
         payload['callback'] = callback
 
+        self.jobs_in_flight[msgid] = (monotonic(), payload)
+
         self.request_queue.put(payload)
 
     def premature_death(self, reply, msgid):
+        """
+        Worker died before running any jobs
+        """
         return
 
-    def worker_death(self, reply, msgid):
-        return
+    def worker_death(self, reply, msgid, death, pid):
+        """
+        Worker died of natural causes, ensure its death and
+        remove from tracking, will be replaced on next heartbeat
+        """
+        if pid in self._workers.keys():
+            del self._workers[pid]
 
-    def worker_done_with_reply(self, reply, msgid):
-        reply = serializer(reply)
-        self.send_reply(reply, msgid)
+    def worker_ready(self, reply, msgid, death, pid):
         self.send_ready()
 
-    def worker_done(self, msgid):
-        self.send_ready()
+    def worker_done_with_reply(self, reply, msgid, death, pid):
+        """
+        Worker finished a job and requested the return value
+        """
+        try:
+            reply = serializer(reply)
+        except TypeError as e:
+            reply = serializer({"value": str(e)})
+
+        self.send_reply(reply, msgid)
+
+        if self.status != STATUS.stopping and not death:
+            self.send_ready()
+
+    def worker_done(self, reply, msgid, death, pid):
+        """
+        Worker finished a job, notify broker of an additional slot opening
+        """
+        if self.status != STATUS.stopping and not death:
+            self.send_ready()
 
     def send_ready(self):
         """
         send the READY command upstream to indicate that JobManager is ready
         for another REQUEST message.
         """
+        self.total_ready_sent += 1
         sendmsg(self.frontend, 'READY')
 
     def send_reply(self, reply, msgid):
@@ -323,6 +377,9 @@ class JobManager(HeartbeatMixin, EMQPService):
         necessary
         """
         # Kill workers that aren't alive
+        logger.debug("Jobs in flight: {}".format(len(self.jobs_in_flight)))
+        logger.debug("Total requests: {}".format(self.total_requests))
+        logger.debug("Total ready sent: {}".format(self.total_ready_sent))
         try:
             [self.kill_worker(w.pid, signal.SIGKILL) for w in self.workers
              if not w.is_alive]
@@ -345,7 +402,7 @@ class JobManager(HeartbeatMixin, EMQPService):
         try:
             os.kill(pid, signal)
         except Exception as e:
-            logger.excpetion(
+            logger.exception(
                 "Encountered exception trying to send {} to worker {}: {}"
                 .format(signal, pid, str(e)))
 
@@ -366,8 +423,9 @@ class JobManager(HeartbeatMixin, EMQPService):
         self.received_disconnect = True
 
     def sigterm_handler(self, signum, frame):
-        logger.info('Shutting down..')
-        sendmsg(self.frontend, KBYE)
+        if not self.received_disconnect:
+            logger.info('Shutting down..')
+            sendmsg(self.frontend, KBYE)
 
-        self.awaiting_startup_ack = False
-        self.received_disconnect = True
+            self.awaiting_startup_ack = False
+            self.received_disconnect = True
diff --git a/eventmq/sender.py b/eventmq/sender.py
index 063cd74..f0cc09c 100644
--- a/eventmq/sender.py
+++ b/eventmq/sender.py
@@ -18,6 +18,7 @@
 The sender is responsible for sending messages
 """
 import logging
+import sys
 import uuid
 
 import zmq
@@ -136,9 +137,11 @@ class Sender(ZMQSendMixin, ZMQReceiveMixin):
         if self.zsocket:
             self.zsocket.close()
 
-        self.name = kwargs.pop('name', str(uuid.uuid4()).encode('ascii'))
         self.zsocket = kwargs.pop('socket', self.zcontext.socket(zmq.DEALER))
-        self.zsocket.setsockopt(zmq.IDENTITY, self.name)
+        if sys.version[0] == '2':
+            self.zsocket.setsockopt(zmq.IDENTITY, self.name)
+        else:
+            self.zsocket.setsockopt_string(zmq.IDENTITY, str(self.name))
 
         self.status = constants.STATUS.ready
 
diff --git a/eventmq/settings.py b/eventmq/settings.py
index ee15742..18ec306 100644
--- a/eventmq/settings.py
+++ b/eventmq/settings.py
@@ -185,6 +185,13 @@ _CONFIG_DEFS = {
                     "weight from the name with a comma. For example "
                     "-Q 10,default 20,high"
         },
+        'MAX_JOB_COUNT': {
+            'default': 1024,
+            'long-arg': '--max-job-count',
+            'type': int,
+            'help': "Maximum number of jobs each worker process executes "
+                    "before resetting",
+        },
         'CONCURRENT_JOBS': {
             'default': 4,
             'long-arg': '--concurrent-jobs',
@@ -209,7 +216,15 @@ _CONFIG_DEFS = {
             'long-arg': '--kill-grace-period',
             'type': int,
             'help': 'Seconds to wait before forefully killing worker '
-                    'processes after receiving a SIGTERM'
+                    'processes after receiving a SIGTERM',
+        },
+        'GLOBAL_TIMEOUT': {
+            'default': 300,
+            'long-arg': '--global-timeout',
+            'short-arg': '-t',
+            'type': int,
+            'help': "Global default timeout for all jobs unless the request "
+                    "specifies otherwise",
         }
     },
     'scheduler': {
diff --git a/eventmq/tests/test_jobmanager.py b/eventmq/tests/test_jobmanager.py
index 105f76c..1ace379 100644
--- a/eventmq/tests/test_jobmanager.py
+++ b/eventmq/tests/test_jobmanager.py
@@ -18,7 +18,6 @@ import unittest
 import mock
 
 from .. import constants, jobmanager
-from ..settings import conf
 
 ADDR = 'inproc://pour_the_rice_in_the_thing'
 
@@ -53,8 +52,7 @@ class TestCase(unittest.TestCase):
     @mock.patch('eventmq.jobmanager.Sender.recv_multipart')
     @mock.patch('eventmq.jobmanager.Poller.poll')
     @mock.patch('eventmq.jobmanager.JobManager.maybe_send_heartbeat')
-    @mock.patch('eventmq.jobmanager.JobManager.send_ready')
-    def test__start_event_loop(self, send_ready_mock, maybe_send_hb_mock,
+    def test__start_event_loop(self, maybe_send_hb_mock,
                                poll_mock, sender_mock, process_msg_mock,
                                pool_close_mock):
         jm = jobmanager.JobManager(override_settings={
@@ -66,9 +64,6 @@ class TestCase(unittest.TestCase):
 
         jm._start_event_loop()
 
-        # send int(conf.CONCURRENT_JOBS) ready messages
-        self.assertEqual(conf.CONCURRENT_JOBS, send_ready_mock.call_count)
-
         process_msg_mock.assert_called_with(
             sender_mock.return_value)
 
diff --git a/eventmq/tests/test_sender.py b/eventmq/tests/test_sender.py
index b2570a3..f292696 100644
--- a/eventmq/tests/test_sender.py
+++ b/eventmq/tests/test_sender.py
@@ -35,7 +35,7 @@ class TestCase(unittest.TestCase):
 
         self.assert_(socket.poll() != 0)
         msg = socket.recv_multipart()
-        self.assertEqual(msg[0], self.sender.name)
+        self.assertEqual(msg[0].decode('ascii'), self.sender.name)
         self.assertEqual(msg[1], b'')
         self.assertEqual(msg[2], b'1')
         self.assertEqual(msg[3], b'Hello!')
diff --git a/eventmq/tests/test_worker.py b/eventmq/tests/test_worker.py
index 115aa52..66808fd 100644
--- a/eventmq/tests/test_worker.py
+++ b/eventmq/tests/test_worker.py
@@ -39,7 +39,7 @@ def test_run_with_timeout():
         'args': [2]
     }
 
-    msgid = worker._run(payload, logging.getLogger())
+    msgid = worker._run_job(payload, logging.getLogger())
 
     assert msgid
 
diff --git a/eventmq/worker.py b/eventmq/worker.py
index 04a35eb..14d4dfe 100644
--- a/eventmq/worker.py
+++ b/eventmq/worker.py
@@ -26,7 +26,7 @@ from multiprocessing import Process
 import os
 import sys
 
-from threading import Event, Thread
+from threading import Thread
 
 from .settings import conf
 
@@ -36,33 +36,6 @@ else:
     import queue as Queue
 
 
-class StoppableThread(Thread):
-    """Thread class with a stop() method. The thread itself has to check
-    regularly for the stopped() condition."""
-
-    def __init__(self, target, name=None, args=()):
-        super(StoppableThread, self).__init__(name=name, target=target,
-                                              args=args)
-        self._return = None
-        self._stop = Event()
-
-    def stop(self):
-        self._stop.set()
-
-    def stopped(self):
-        return self._stop.isSet()
-
-    def run(self):
-        """
-        Overrides default run to have a side effect of saving the result
-        in self._return that will be accessible when the job completes,
-        or remain None after a timeout
-        """
-        if self._Thread__target is not None:
-            self._return = self._Thread__target(*self._Thread__args,
-                                                **self._Thread__kwargs)
-
-
 class MultiprocessWorker(Process):
     """
     Defines a worker that spans the job in a multiprocessing task
@@ -78,10 +51,7 @@ class MultiprocessWorker(Process):
 
     @property
     def logger(self):
-        if not hasattr(self, '_logger'):
-            self._logger = logging.getLogger(__name__ + '.' + str(os.getpid()))
-
-        return self._logger
+        return logging.getLogger(__name__ + '.' + str(os.getpid()))
 
     def run(self):
         """
@@ -89,31 +59,39 @@ class MultiprocessWorker(Process):
 
         This is designed to run in a seperate process.
         """
-        if self.run_setup:
-            self.run_setup = False
-            if any(conf.SETUP_CALLABLE) and any(conf.SETUP_PATH):
-                try:
-                    self.logger.debug("Running setup ({}.{}) for worker id {}"
-                                      .format(
-                                          conf.SETUP_PATH,
-                                          conf.SETUP_CALLABLE,
-                                          os.getpid()))
-                    run_setup(conf.SETUP_PATH, conf.SETUP_CALLABLE)
-                except Exception as e:
-                    self.logger.warning('Unable to do setup task ({}.{}): {}'
-                                        .format(conf.SETUP_PATH,
-                                                conf.SETUP_CALLABLE, str(e)))
+        # Define the 2 queues for communicating with the worker thread
+        logger = self.logger
+
+        worker_queue = Queue.Queue(1)
+        worker_result_queue = Queue.Queue(1)
+        worker_thread = Thread(target=_run,
+                               args=(worker_queue,
+                                     worker_result_queue,
+                                     logger))
 
         import zmq
         zmq.Context.instance().term()
 
+        self.output_queue.put(
+            {'msgid': None,
+             'return': None,
+             'death': False,
+             'pid': os.getpid(),
+             'callback': 'worker_ready'}
+        )
+
+        callback = 'premature_death'
+
+        worker_thread.start()
+
         # Main execution loop, only break in cases that we can't recover from
         # or we reach job count limit
         while True:
             try:
-                payload = self.input_queue.get(block=False, timeout=1000)
+                payload = self.input_queue.get(timeout=1)
                 if payload == 'DONE':
                     break
+
             except Queue.Empty:
                 if os.getppid() != self.ppid:
                     break
@@ -121,53 +99,71 @@ class MultiprocessWorker(Process):
             except Exception as e:
                 break
             finally:
+                # If I'm an orphan, die
                 if os.getppid() != self.ppid:
                     break
 
             try:
                 return_val = 'None'
                 self.job_count += 1
-                timeout = payload.get("timeout", None)
+                timeout = payload.get("timeout") or conf.GLOBAL_TIMEOUT
                 msgid = payload.get('msgid', '')
                 callback = payload.get('callback', '')
+                logger.debug("Putting on thread queue msgid: {}".format(
+                    msgid))
+
+                worker_queue.put(payload['params'])
 
-                worker_thread = StoppableThread(target=_run,
-                                                args=(payload['params'],
-                                                      self.logger))
-                worker_thread.start()
-                worker_thread.join(timeout)
-                return_val = {"value": worker_thread._return}
+                try:
+                    return_val = worker_result_queue.get(timeout=timeout)
 
-                if worker_thread.isAlive():
-                    worker_thread.stop()
+                    logger.debug("Got from result queue msgid: {}".format(
+                        msgid))
+                except Queue.Empty:
                     return_val = 'TimeoutError'
 
+                return_val = {"value": return_val}
+
                 try:
                     self.output_queue.put_nowait(
                         {'msgid': msgid,
                          'return': return_val,
+                         'death': self.job_count >= conf.MAX_JOB_COUNT or
+                         return_val == 'TimeoutError',
                          'pid': os.getpid(),
                          'callback': callback}
                     )
                 except Exception:
                     break
 
+                if return_val["value"] == 'TimeoutError':
+                    break
+
             except Exception as e:
                 return_val = str(e)
 
             if self.job_count >= conf.MAX_JOB_COUNT:
+                logger.debug("Worker reached job limit, exiting")
                 break
 
+        worker_queue.put('DONE')
+        worker_thread.join(timeout=5)
+
         self.output_queue.put(
             {'msgid': None,
              'return': 'DEATH',
+             'death': True,
              'pid': os.getpid(),
              'callback': 'worker_death'}
             )
-        self.logger.debug("Worker death")
+
+        logger.debug("Worker death")
+
+        if worker_thread.is_alive():
+            logger.debug("Worker thread did not die gracefully")
 
 
-def _run(payload, logger):
+def _run(queue, result_queue, logger):
     """
     Takes care of actually executing the code given a message payload
 
@@ -181,6 +177,39 @@ def _run(payload, logger):
         "class_kwargs": {"value": 2}
     }
     """
+    if any(conf.SETUP_CALLABLE) and any(conf.SETUP_PATH):
+        try:
+            logger.debug("Running setup ({}.{}) for worker id {}"
+                         .format(
+                             conf.SETUP_PATH,
+                             conf.SETUP_CALLABLE,
+                             os.getpid()))
+            run_setup(conf.SETUP_PATH, conf.SETUP_CALLABLE)
+        except Exception as e:
+            logger.warning('Unable to do setup task ({}.{}): {}'
+                           .format(conf.SETUP_PATH,
+                                   conf.SETUP_CALLABLE, str(e)))
+
+    while True:
+        # Blocking get so we don't spin cycles reading over and over
+        try:
+            payload = queue.get()
+        except Exception as e:
+            logger.exception(e)
+            continue
+
+        if payload == 'DONE':
+            break
+
+        return_val = _run_job(payload, logger)
+        # Signal that we're done with this job and put its return value on the
+        # result queue
+        result_queue.put(return_val)
+
+    logger.debug("Worker thread death")
+
+
+def _run_job(payload, logger):
     try:
         if ":" in payload["path"]:
             _pkgsplit = payload["path"].split(':')
@@ -224,9 +253,8 @@ def _run(payload, logger):
         return_val = callable_(*args, **kwargs)
     except Exception as e:
         logger.exception(e)
-        return str(e)
+        return_val = str(e)
 
-    # Signal that we're done with this job
     return return_val
 
 
diff --git a/setup.py b/setup.py
index 694c6cc..72e8e35 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@ from setuptools import find_packages, setup
 
 setup(
     name='eventmq',
-    version='0.3.4',
+    version='0.3.4.4',
     description='EventMQ job execution and messaging system based on ZeroMQ',
     packages=find_packages(),
     install_requires=['pyzmq==16.0.2',