10 files changed, 124 insertions, 46 deletions

diff --git a/bin/emq-jobmanager b/bin/emq-jobmanager
index 4071ada..69c21ed 100755
--- a/bin/emq-jobmanager
+++ b/bin/emq-jobmanager
@@ -34,7 +34,10 @@ if __name__ == "__main__":
     parser.add_argument('--jobs', '-J', type=int, nargs='?',
                         help='the max number of concurrent jobs to manage at '
                         'a time')
-
+    parser.add_argument('--name', '-n', type=str, default=None,
+                        help="A unique ame to give this node. If one "
+                             "isn't provided a random uuid will be "
+                             "generated")
     args = parser.parse_args()
 
     # Overwrite the default config location with the one passed to the app
@@ -48,6 +51,8 @@ if __name__ == "__main__":
 
     broker_addr = args.broker_addr
     concurrent_jobs = args.jobs
+    name = args.name
 
-    j = JobManager(queues=queues, concurrent_jobs=concurrent_jobs)
+    j = JobManager(queues=queues, concurrent_jobs=concurrent_jobs,
+                   name=name)
     j.jobmanager_main(broker_addr=broker_addr)
diff --git a/bin/send_msg b/bin/send_msg
index 47de45f..c3e88fa 100755
--- a/bin/send_msg
+++ b/bin/send_msg
@@ -22,11 +22,11 @@ if __name__ == "__main__":
             'callable': 'work_job',
             'class_args': ('blurp',),
             'class_kwargs': {'kwarg1': True},
-            'args': (50, ),
+            'args': (10, ),
             'kwargs': {}
             }]
 
-        msgid = send_request(s, msg, guarantee=True, reply_requested=True, timeout=10)
+        msgid = send_request(s, msg, guarantee=True, reply_requested=True, timeout=10, queue='other')
         msgid = send_request(s, msg, guarantee=True, reply_requested=True)
         # print 'Sent message, use msgid={} to track responses'.format(msgid)
         # events = dict(poller.poll(500))
diff --git a/eventmq/__init__.py b/eventmq/__init__.py
index 30572f3..60bb61c 100644
--- a/eventmq/__init__.py
+++ b/eventmq/__init__.py
@@ -1,5 +1,5 @@
 __author__ = 'EventMQ Contributors'
-__version__ = '0.3.4.3'
+__version__ = '0.3.4.6'
 
 PROTOCOL_VERSION = 'eMQP/1.0'
 
diff --git a/eventmq/conf.py b/eventmq/conf.py
index 5710933..4749e6b 100644
--- a/eventmq/conf.py
+++ b/eventmq/conf.py
@@ -95,5 +95,6 @@ SETUP_CALLABLE = ''
 # Time to wait after receiving SIGTERM to kill the workers in the jobmanager
 # forecfully
 KILL_GRACE_PERIOD = 300
+GLOBAL_TIMEOUT = 300
 
 # }}}
diff --git a/eventmq/jobmanager.py b/eventmq/jobmanager.py
index d4a0ddc..4840fac 100644
--- a/eventmq/jobmanager.py
+++ b/eventmq/jobmanager.py
@@ -75,9 +75,11 @@ class JobManager(HeartbeatMixin, EMQPService):
         """
         super(JobManager, self).__init__(*args, **kwargs)
 
+        setup_logger("eventmq")
+
         #: Define the name of this JobManager instance. Useful to know when
         #: referring to the logs.
-        self.name = kwargs.pop('name', generate_device_name())
+        self.name = kwargs.pop('name', None) or generate_device_name()
         logger.info('Initializing JobManager {}...'.format(self.name))
 
         #: keep track of workers
@@ -94,6 +96,7 @@ class JobManager(HeartbeatMixin, EMQPService):
             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
@@ -103,11 +106,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'):
@@ -127,9 +148,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:
@@ -167,7 +185,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')
@@ -224,13 +242,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):
         """
@@ -262,6 +288,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]
@@ -279,6 +308,8 @@ 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):
@@ -287,31 +318,36 @@ class JobManager(HeartbeatMixin, EMQPService):
         """
         return
 
-    def worker_death(self, reply, msgid):
+    def worker_death(self, reply, msgid, death, pid):
         """
-        Worker died of natural causes
+        Worker died of natural causes, ensure its death and
+        remove from tracking, will be replaced on next heartbeat
         """
-        return
+        if pid in self._workers.keys():
+            del self._workers[pid]
+
+    def worker_ready(self, reply, msgid, death, pid):
+        self.send_ready()
 
-    def worker_done_with_reply(self, reply, msgid):
+    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 = {"value": str(e)}
+            reply = serializer({"value": str(e)})
 
         self.send_reply(reply, msgid)
 
-        if self.status != STATUS.stopping:
+        if self.status != STATUS.stopping and not death:
             self.send_ready()
 
-    def worker_done(self, reply, msgid):
+    def worker_done(self, reply, msgid, death, pid):
         """
         Worker finished a job, notify broker of an additional slot opening
         """
-        if self.status != STATUS.stopping:
+        if self.status != STATUS.stopping and not death:
             self.send_ready()
 
     def send_ready(self):
@@ -319,6 +355,7 @@ class JobManager(HeartbeatMixin, EMQPService):
         send the READY command upstream to indicate that JobManager is ready
         for another REQUEST message.
         """
+        self.total_ready_sent += 1
         sendmsg(self.outgoing, 'READY')
 
     def send_reply(self, reply, msgid):
@@ -346,6 +383,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]
@@ -368,7 +408,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)))
 
@@ -404,7 +444,6 @@ class JobManager(HeartbeatMixin, EMQPService):
         Args:
             broker_addr (str): The address of the broker to connect to.
         """
-        setup_logger('')
         import_settings()
         import_settings(section='jobmanager')
 
diff --git a/eventmq/sender.py b/eventmq/sender.py
index 063cd74..12dec05 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,14 @@ 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)
+
+        self.name = kwargs.pop('name', str(uuid.uuid4()))
+
+        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/tests/test_jobmanager.py b/eventmq/tests/test_jobmanager.py
index 927e5f6..566c583 100644
--- a/eventmq/tests/test_jobmanager.py
+++ b/eventmq/tests/test_jobmanager.py
@@ -49,8 +49,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()
@@ -60,9 +59,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/worker.py b/eventmq/worker.py
index 4f42a2c..859b831 100644
--- a/eventmq/worker.py
+++ b/eventmq/worker.py
@@ -51,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):
         """
@@ -63,16 +60,26 @@ class MultiprocessWorker(Process):
         This is designed to run in a seperate process.
         """
         # Define the 2 queues for communicating with the worker thread
-        worker_queue = Queue.Queue()
-        worker_result_queue = Queue.Queue()
+        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,
-                                     self.logger))
+                                     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()
@@ -99,14 +106,19 @@ class MultiprocessWorker(Process):
             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'])
 
                 try:
                     return_val = worker_result_queue.get(timeout=timeout)
+
+                    logger.debug("Got from result queue msgid: {}".format(
+                        msgid))
                 except Queue.Empty:
                     return_val = 'TimeoutError'
 
@@ -116,6 +128,8 @@ class MultiprocessWorker(Process):
                     self.output_queue.put_nowait(
                         {'msgid': msgid,
                          'return': return_val,
+                         'death': self.job_count >= conf.MAX_JOB_COUNT or
+                         return_val["value"] == 'TimeoutError',
                          'pid': os.getpid(),
                          'callback': callback}
                     )
@@ -129,15 +143,24 @@ class MultiprocessWorker(Process):
                 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(queue, result_queue, logger):
@@ -169,14 +192,22 @@ def _run(queue, result_queue, logger):
 
     while True:
         # Blocking get so we don't spin cycles reading over and over
-        payload = queue.get()
-        return_val = _run_job(payload, logger)
+        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
-        queue.task_done()
         result_queue.put(return_val)
 
+    logger.debug("Worker thread death")
+
 
 def _run_job(payload, logger):
     try:
diff --git a/setup.py b/setup.py
index ee81dae..9ac6089 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@ from setuptools import find_packages, setup
 
 setup(
     name='eventmq',
-    version='0.3.4.3',
+    version='0.3.4.6',
     description='EventMQ job execution and messaging system based on ZeroMQ',
     packages=find_packages(),
     install_requires=['pyzmq==15.4.0',