6 files changed, 192 insertions, 16 deletions

diff --git a/etc/eventmq.conf-dist b/etc/eventmq.conf-dist
index 43ec6f2..1ff09c4 100644
--- a/etc/eventmq.conf-dist
+++ b/etc/eventmq.conf-dist
@@ -19,9 +19,22 @@ scheduler_addr=tcp://eventmq:47291
 
 [jobmanager]
 worker_addr=tcp://127.0.0.1:47290
-queues=[[50,"google"], [40,"pushes"], [10,"default"]]
+
+# Defines the weight and name of queues this worker deals with.
+queues=[[20,"heavy-cpu"], [30,"low-cpu"], [10,"default"]]
+
+# Specifies the number of of sub-processes to spawn to process jobs concurrently
 concurrent_jobs=2
 
+# This function is executed when EventMQ creates a new worker subprocess
+# subprocess_setup_func = 'path.to.my_setup_function'
+
+# This function is executed before every job
+# job_entry_func = 'path.to.my_job_setup_function'
+
+# This function is executed after every job
+# job_exit_func = 'path.to.my_job_teardown_function'
+
 [publisher]
 publisher_incoming_addr=tcp://0.0.0.0:47298
 publisher_outgoing_addr=tcp://0.0.0.0:47299
\ No newline at end of file
diff --git a/eventmq/__init__.py b/eventmq/__init__.py
index aefc3a2..2e180ed 100644
--- a/eventmq/__init__.py
+++ b/eventmq/__init__.py
@@ -1,5 +1,5 @@
 __author__ = 'EventMQ Contributors'
-__version__ = '0.3.5.6'
+__version__ = '0.3.6'
 
 PROTOCOL_VERSION = 'eMQP/1.0'
 
diff --git a/eventmq/conf.py b/eventmq/conf.py
index 8609716..715bbb2 100644
--- a/eventmq/conf.py
+++ b/eventmq/conf.py
@@ -89,9 +89,18 @@ RQ_PASSWORD = ''
 MAX_JOB_COUNT = 1024
 
 # Path/Callable to run on start of a worker process
+# These options are deprecated for the more user-friendly
+# SUBPROCESS_SETUP_FUNC which can be a full path to a function.
 SETUP_PATH = ''
 SETUP_CALLABLE = ''
 
+# Function to run on the start of a new worker subprocess
+SUBPROCESS_SETUP_FUNC = ''
+
+# function to be run before the execution of every job
+JOB_ENTRY_FUNC = ''
+# function to be run after the execution of every job
+JOB_EXIT_FUNC = ''
 # Time to wait after receiving SIGTERM to kill the workers in the jobmanager
 # forecfully
 KILL_GRACE_PERIOD = 300
diff --git a/eventmq/tests/test_worker.py b/eventmq/tests/test_worker.py
index 5e72f39..b8e4f3e 100644
--- a/eventmq/tests/test_worker.py
+++ b/eventmq/tests/test_worker.py
@@ -15,11 +15,19 @@
 
 import logging
 import time
+import sys
 
+import mock
 from .. import worker
 
-ADDR = 'inproc://pour_the_rice_in_the_thing'
+if sys.version[0] == '2':
+    import Queue
+else:
+    import queue as Queue
+
 
+ADDR = 'inproc://pour_the_rice_in_the_thing'
+SETUP_SUCCESS_RETVAL = 'job setup success'
 
 def test_run_with_timeout():
     payload = {
@@ -33,11 +41,124 @@ def test_run_with_timeout():
     assert msgid
 
 
+@mock.patch('eventmq.worker.callable_from_name')
+def test_run_job_setup_hook(callable_from_name_mock):
+    from eventmq import conf
+
+    setup_func_str = 'eventmq.tests.test_worker.job_setup_hook'
+    callable_from_name_mock.return_value = mock.Mock()
+
+    payload = {
+        'path': 'eventmq.tests.test_worker',
+        'callable': 'job',
+        'args': [2]
+    }
+
+    q, res_q = Queue.Queue(), Queue.Queue()
+
+    q.put(payload)
+    q.put('DONE')
+
+    try:
+        conf.JOB_ENTRY_FUNC = setup_func_str
+        worker._run(q, res_q, logging.getLogger())
+    finally:
+        conf.JOB_ENTRY_FUNC = ''
+
+    callable_from_name_mock.assert_called_with(setup_func_str)
+    assert callable_from_name_mock.return_value.call_count == 1
+
+
+@mock.patch('eventmq.worker.callable_from_name')
+def test_run_job_teardown_hook(callable_from_name_mock):
+    from eventmq import conf
+
+    teardown_func_str = 'eventmq.tests.test_worker.job_teardown_hook'
+    callable_from_name_mock.return_value = mock.Mock()
+
+    payload = {
+        'path': 'eventmq.tests.test_worker',
+        'callable': 'job',
+        'args': [2]
+    }
+
+    q, res_q = Queue.Queue(), Queue.Queue()
+
+    q.put(payload)
+    q.put('DONE')
+
+    try:
+        conf.JOB_EXIT_FUNC = teardown_func_str
+        worker._run(q, res_q, logging.getLogger())
+    finally:
+        conf.JOB_EXIT_FUNC = ''
+
+    callable_from_name_mock.assert_called_with(teardown_func_str)
+    assert callable_from_name_mock.return_value.call_count == 1
+
+
+@mock.patch('eventmq.worker.callable_from_name')
+def test_run_subprocess_setup_func(callable_from_name_mock):
+    from eventmq import conf
+
+    setup_func_str = 'eventmq.tests.test_worker.process_setup_hook'
+    callable_from_name_mock.return_value = mock.Mock()
+
+    payload = {
+        'path': 'eventmq.tests.test_worker',
+        'callable': 'job',
+        'args': [2]
+    }
+
+    q, res_q = Queue.Queue(), Queue.Queue()
+
+    q.put(payload)
+    q.put('DONE')
+
+    try:
+        conf.SUBPROCESS_SETUP_FUNC = setup_func_str
+        worker._run(q, res_q, logging.getLogger())
+    finally:
+        conf.SUBPROCESS_SETUP_FUNC = ''
+
+    callable_from_name_mock.assert_called_with(setup_func_str)
+    assert callable_from_name_mock.return_value.call_count == 1
+
+
+@mock.patch('eventmq.worker.run_setup')
+def test_run_run_setup_func(run_setup_mock):
+    from eventmq import conf
+
+    setup_func_path = 'eventmq.tests.test_worker'
+    setup_func_callable = 'process_setup_hook'
+
+    payload = {
+        'path': 'eventmq.tests.test_worker',
+        'callable': 'job',
+        'args': [2]
+    }
+
+    q, res_q = Queue.Queue(), Queue.Queue()
+
+    q.put(payload)
+    q.put('DONE')
+
+    try:
+        conf.SETUP_PATH = setup_func_path
+        conf.SETUP_CALLABLE = setup_func_callable
+        worker._run(q, res_q, logging.getLogger())
+    finally:
+        conf.SETUP_PATH = ''
+        conf.SETUP_CALLABLE = ''
+
+    run_setup_mock.assert_called_with(setup_func_path, setup_func_callable)
+
+
 def test_run_setup():
-    setup_callable = 'pre_hook'
+    setup_callable = 'process_setup_hook'
     setup_path = 'eventmq.tests.test_worker'
 
-    worker.run_setup(setup_path, setup_callable)
+    assert worker.run_setup(setup_path, setup_callable)
 
 
 def job(sleep_time=0):
@@ -46,9 +167,16 @@ def job(sleep_time=0):
     return True
 
 
-def pre_hook():
-    return 1
+def process_setup_hook():
+    print 'process setup hook executed'
+    return True
 
 
-def post_hook():
-    return 1
+def job_setup_hook():
+    print 'job setup hook executed'
+    return SETUP_SUCCESS_RETVAL
+
+
+def job_teardown_hook():
+    print 'job teardown hook executed'
+    return True
diff --git a/eventmq/worker.py b/eventmq/worker.py
index 1621578..aa03284 100644
--- a/eventmq/worker.py
+++ b/eventmq/worker.py
@@ -18,16 +18,13 @@
 Defines different short-lived workers that execute jobs
 """
 from importlib import import_module
-
 import logging
-
 from multiprocessing import Process
-
 import os
 import sys
-
 from threading import Thread
 
+from .utils.functions import callable_from_name
 from . import conf
 
 if sys.version[0] == '2':
@@ -180,7 +177,19 @@ def _run(queue, result_queue, logger):
         "class_kwargs": {"value": 2}
     }
     """
-    if any(conf.SETUP_CALLABLE) and any(conf.SETUP_PATH):
+    if any(conf.SUBPROCESS_SETUP_FUNC):
+        try:
+            logger.debug("Running setup ({}) for worker id {}".format(
+                conf.SUBPROCESS_SETUP_FUNC, os.getpid()))
+            setup_func = callable_from_name(conf.SUBPROCESS_SETUP_FUNC)
+            setup_func()
+        except Exception as e:
+            logger.warning('Unable to do setup task ({}): {}'
+                           .format(conf.SUBPROCESS_SETUP_FUNC, str(e)))
+
+    elif any(conf.SETUP_CALLABLE) and any(conf.SETUP_PATH):
+        logger.warning("SETUP_CALLABLE and SETUP_PATH deprecated in favor for "
+                       "SUBPROCESS_SETUP_FUNC")
         try:
             logger.debug("Running setup ({}.{}) for worker id {}"
                          .format(
@@ -193,6 +202,16 @@ def _run(queue, result_queue, logger):
                            .format(conf.SETUP_PATH,
                                    conf.SETUP_CALLABLE, str(e)))
 
+    if conf.JOB_ENTRY_FUNC:
+        job_entry_func = callable_from_name(conf.JOB_ENTRY_FUNC)
+    else:
+        job_entry_func = None
+
+    if conf.JOB_EXIT_FUNC:
+        job_exit_func = callable_from_name(conf.JOB_EXIT_FUNC)
+    else:
+        job_exit_func = None
+
     while True:
         # Blocking get so we don't spin cycles reading over and over
         try:
@@ -204,7 +223,14 @@ def _run(queue, result_queue, logger):
         if payload == 'DONE':
             break
 
+        if job_entry_func:
+            job_entry_func()
+
         return_val = _run_job(payload, logger)
+
+        if job_exit_func:
+            job_exit_func()
+
         # Signal that we're done with this job and put its return value on the
         # result queue
         result_queue.put(return_val)
@@ -278,4 +304,4 @@ def run_setup(setup_path, setup_callable):
 
         setup_callable_ = getattr(setup_package, setup_callable)
 
-        setup_callable_()
+        return setup_callable_()
diff --git a/setup.py b/setup.py
index 64afbd9..9bacfbf 100644
--- a/setup.py
+++ b/setup.py
@@ -7,7 +7,7 @@ from setuptools import find_packages, setup
 
 setup(
     name='eventmq',
-    version='0.3.5.6',
+    version='0.3.6',
     description='EventMQ job execution and messaging system based on ZeroMQ',
     packages=find_packages(),
     install_requires=['pyzmq==15.4.0',