Worker optimizations now set up in celery.apps.worker not WorkerController, and tests reset optimiations

celery/apps/worker.py

@@ -27,6 +27,7 @@ from celery.app import app_or_default
 from celery.app.abstract import configurated, from_config
 from celery.exceptions import ImproperlyConfigured, SystemTerminate
 from celery.loaders.app import AppLoader
+from celery.task import trace
 from celery.utils import cry, isatty, worker_direct
 from celery.utils.imports import qualname
 from celery.utils.log import get_logger, mlevel, set_in_sighandler
@@ -150,6 +151,10 @@ class Worker(configurated):
     def run(self):
         self.init_queues()
         self.app.loader.init_worker()
+
+        # apply task execution optimizations
+        trace.setup_worker_optimizations(self.app)
+
         # this signal can be used to e.g. change queues after
         # the -Q option has been applied.
         signals.celeryd_after_setup.send(sender=self.hostname, instance=self,

celery/task/trace.py

@@ -49,10 +49,12 @@ EXCEPTION_STATES = states.EXCEPTION_STATES
 
 #: set by :func:`setup_worker_optimizations`
 _tasks = None
+_patched = {}
 
 
 def setup_worker_optimizations(app):
     global _tasks
+    global trace_task_ret
 
     # make sure custom Task.__call__ methods that calls super
     # will not mess up the request/task stack.
@@ -72,6 +74,21 @@ def setup_worker_optimizations(app):
     # set fast shortcut to task registry
     _tasks = app._tasks
 
+    trace_task_ret = _fast_trace_task
+
+
+def reset_worker_optimizations():
+    global trace_task_ret
+    trace_task_ret = trace_task
+    try:
+        delattr(BaseTask, '_stackprotected')
+    except AttributeError:
+        pass
+    try:
+        BaseTask.__call__ = _patched.pop('BaseTask.__call__')
+    except KeyError:
+        pass
+
 
 def _install_stack_protection():
     # Patches BaseTask.__call__ in the worker to handle the edge case
@@ -88,7 +105,7 @@ def _install_stack_protection():
     # will blow if a custom task class defines __call__ and also
     # calls super().
     if not getattr(BaseTask, '_stackprotected', False):
-        orig = BaseTask.__call__
+        _patched['BaseTask.__call__'] = orig = BaseTask.__call__
 
         def __protected_call__(self, *args, **kwargs):
             req, stack = self.request, self.request_stack
@@ -346,7 +363,15 @@ def trace_task(task, uuid, args, kwargs, request={}, **opts):
         return report_internal_error(task, exc)
 
 
-def trace_task_ret(task, uuid, args, kwargs, request={}):
+def _trace_task_ret(name, uuid, args, kwargs, request={}, **opts):
+    return trace_task(current_app.tasks[name],
+                      uuid, args, kwargs, request, **opts)
+trace_task_ret = _trace_task_ret
+
+
+def _fast_trace_task(task, uuid, args, kwargs, request={}):
+    # setup_worker_optimizations will point trace_task_ret to here,
+    # so this is the function used in the worker.
     return _tasks[task].__trace__(uuid, args, kwargs, request)[0]
 
 

celery/tests/bin/test_celeryd.py

@@ -20,6 +20,7 @@ from celery import current_app
 from celery.apps import worker as cd
 from celery.bin.celeryd import WorkerCommand, main as celeryd_main
 from celery.exceptions import ImproperlyConfigured, SystemTerminate
+from celery.task import trace
 from celery.utils.log import ensure_process_aware_logger
 from celery.worker import state
 
@@ -33,6 +34,14 @@ from celery.tests.utils import (
 ensure_process_aware_logger()
 
 
+
+class WorkerAppCase(AppCase):
+
+    def tearDown(self):
+        super(WorkerAppCase, self).tearDown()
+        trace.reset_worker_optimizations()
+
+
 def disable_stdouts(fun):
 
     @wraps(fun)
@@ -69,7 +78,7 @@ class Worker(cd.Worker):
         self.redirect_stdouts = False
 
 
-class test_Worker(AppCase):
+class test_Worker(WorkerAppCase):
 
     Worker = Worker
 
@@ -385,7 +394,7 @@ class test_Worker(AppCase):
         self.assertTrue(worker_ready_sent[0])
 
 
-class test_funs(AppCase):
+class test_funs(WorkerAppCase):
 
     def test_active_thread_count(self):
         self.assertTrue(cd.active_thread_count())
@@ -433,7 +442,7 @@ class test_funs(AppCase):
             sys.argv = s
 
 
-class test_signal_handlers(AppCase):
+class test_signal_handlers(WorkerAppCase):
 
     class _Worker(object):
         stopped = False

celery/worker/__init__.py

@@ -29,7 +29,6 @@ from celery import platforms
 from celery.app import app_or_default
 from celery.app.abstract import configurated, from_config
 from celery.exceptions import SystemTerminate, TaskRevokedError
-from celery.task.trace import setup_worker_optimizations
 from celery.utils.functional import noop
 from celery.utils.imports import qualname, reload_from_cwd
 from celery.utils.log import get_logger
@@ -308,7 +307,6 @@ class WorkController(configurated):
     def __init__(self, loglevel=None, hostname=None, ready_callback=noop,
             queues=None, app=None, pidfile=None, **kwargs):
         self.app = app_or_default(app or self.app)
-        setup_worker_optimizations(self.app)
 
         self._shutdown_complete = Event()
         self.setup_defaults(kwargs, namespace='celeryd')