Add support for custom task classes

Author: lrafeei

newrelic/config.py

@@ -4007,18 +4007,25 @@ def _process_module_builtin_defaults():
         "instrument_rest_framework_decorators",
     )
 
-    _process_module_definition("celery.task.base", "newrelic.hooks.application_celery", "instrument_celery_app_task")
-    _process_module_definition("celery.app.task", "newrelic.hooks.application_celery", "instrument_celery_app_task")
+    _process_module_definition("celery.app.base", "newrelic.hooks.application_celery", "instrument_celery_app_base")
+    _process_module_definition("celery.local", "newrelic.hooks.application_celery", "instrument_celery_local")
+    # _process_module_definition("celery.task.base", "newrelic.hooks.application_celery", "instrument_celery_app_task")
+    # _process_module_definition("celery.app.task", "newrelic.hooks.application_celery", "instrument_celery_app_task")
     _process_module_definition("celery.app.trace", "newrelic.hooks.application_celery", "instrument_celery_app_trace")
     _process_module_definition("celery.worker", "newrelic.hooks.application_celery", "instrument_celery_worker")
+    # _process_module_definition(
+    #     "celery.concurrency.processes", "newrelic.hooks.application_celery", "instrument_celery_worker"
+    # )
     _process_module_definition(
-        "celery.concurrency.processes", "newrelic.hooks.application_celery", "instrument_celery_worker"
+        "celery.concurrency.prefork", "newrelic.hooks.application_celery", "instrument_celery_worker"
     )
+    # _process_module_definition(
+    #     "celery.concurrency.asynpool", "newrelic.hooks.application_celery", "instrument_celery_worker"
+    # )
     _process_module_definition(
-        "celery.concurrency.prefork", "newrelic.hooks.application_celery", "instrument_celery_worker"
+        "billiard.process", "newrelic.hooks.application_celery", "instrument_billiard_process"
     )
 
-    _process_module_definition("celery.app.base", "newrelic.hooks.application_celery", "instrument_celery_app_base")
     _process_module_definition("billiard.pool", "newrelic.hooks.application_celery", "instrument_billiard_pool")
 
     _process_module_definition("flup.server.cgi", "newrelic.hooks.adapter_flup", "instrument_flup_server_cgi")

newrelic/hooks/application_celery.py

@@ -27,14 +27,15 @@
 from newrelic.api.function_trace import FunctionTrace
 from newrelic.api.message_trace import MessageTrace
 from newrelic.api.pre_function import wrap_pre_function
+from newrelic.api.post_function import wrap_post_function
 from newrelic.api.transaction import current_transaction
 from newrelic.common.object_wrapper import FunctionWrapper, _NRBoundFunctionWrapper, wrap_function_wrapper
+from newrelic.common.signature import bind_args
 from newrelic.core.agent import shutdown_agent
 
 UNKNOWN_TASK_NAME = "<Unknown Task>"
 MAPPING_TASK_NAMES = {"celery.starmap", "celery.map"}
 
-
 def task_info(instance, *args, **kwargs):
     # Grab the current task, which can be located in either place
     if instance:
@@ -159,28 +160,6 @@ def run(self, *args, **kwargs):
     return wrapped_task
 
 
-def instrument_celery_app_task(module):
-    # Triggered for both 'celery.app.task' and 'celery.task.base'.
-
-    if hasattr(module, "BaseTask"):
-        # Need to add a wrapper for background task entry point.
-
-        # In Celery 2.2 the 'BaseTask' class actually resided in the
-        # module 'celery.task.base'. In Celery 2.3 the 'BaseTask' class
-        # moved to 'celery.app.task' but an alias to it was retained in
-        # the module 'celery.task.base'. We need to detect both module
-        # imports, but we check the module name associated with
-        # 'BaseTask' to ensure that we do not instrument the class via
-        # the alias in Celery 2.3 and later.
-
-        # In Celery 2.5+, although 'BaseTask' still exists execution of
-        # the task doesn't pass through it. For Celery 2.5+ need to wrap
-        # the tracer instead.
-
-        if module.BaseTask.__module__ == module.__name__:
-            module.BaseTask.__call__ = CeleryTaskWrapper(module.BaseTask.__call__)
-
-
 def wrap_Celery_send_task(wrapped, instance, args, kwargs):
     transaction = current_transaction()
     if not transaction:
@@ -222,22 +201,115 @@ def wrap_worker_optimizations(wrapped, instance, args, kwargs):
     return result
 
 
+def wrap_task_trace_wrapper(wrapped, instance, args, kwargs):
+    bound_args = bind_args(wrapped, args, kwargs)
+    task_name = bound_args.get("task", None)
+    localized_args = bound_args.get("_loc", None) or wrapped.__globals__['_localized']
+    try:
+        tasks, accept, hostname = localized_args if localized_args else (None, None, None)
+        task_function = tasks.get(task_name, None)
+        tasks[task_name] = task_function
+        bound_args["_loc"] = (tasks, accept, hostname)
+        transaction = current_transaction(active_only=False)
+
+        # Grab task name and source
+        _name = task_name
+        _source = task_function
+
+        # A Celery Task can be called either outside of a transaction, or
+        # within the context of an existing transaction. There are 3
+        # possibilities we need to handle:
+        #
+        #   1. In an inactive transaction
+        #
+        #      If the end_of_transaction() or ignore_transaction() API calls
+        #      have been invoked, this task may be called in the context
+        #      of an inactive transaction. In this case, don't wrap the task
+        #      in any way. Just run the original function.
+        #
+        #   2. In an active transaction
+        #
+        #      Run the original function inside a FunctionTrace.
+        #
+        #   3. Outside of a transaction
+        #
+        #      This is the typical case for a celery Task. Since it's not
+        #      running inside of an existing transaction, we want to create
+        #      a new background transaction for it.
+
+        if transaction and (transaction.ignore_transaction or transaction.stopped):
+            return wrapped(*args, **kwargs)
+
+        elif transaction:
+            with FunctionTrace(_name, source=_source):
+                return wrapped(*args, **kwargs)
+
+        else:
+            with BackgroundTask(application_instance(), _name, "Celery", source=_source) as transaction:
+                # Attempt to grab distributed tracing headers
+                try:
+                    # Headers on earlier versions of Celery may end up as attributes
+                    # on the request context instead of as custom headers. Handler this
+                    # by defaulting to using vars() if headers is not available
+                    
+                    request = task_function.request
+                    headers = getattr(request, "headers", None) or vars(request)
+
+                    settings = transaction.settings
+                    if headers is not None and settings is not None:
+                        if settings.distributed_tracing.enabled:
+                            transaction.accept_distributed_trace_headers(headers, transport_type="AMQP")
+                        elif transaction.settings.cross_application_tracer.enabled:
+                            transaction._process_incoming_cat_headers(
+                                headers.get(MessageTrace.cat_id_key, None),
+                                headers.get(MessageTrace.cat_transaction_key, None),
+                            )
+                except Exception:
+                    pass
+
+                return wrapped(*args, **kwargs)
+            
+    except Exception as e:
+        return wrapped(*args, **kwargs)
+
+
+def wrap_request_wrapper(wrapped, instance, args, kwargs):
+    try:
+        result = wrapped(*args, **kwargs)
+        ready, request = result
+        job, i, task_tracer_function, _args, _kwargs = request
+        task_tracer_function = FunctionWrapper(task_tracer_function, wrap_task_trace_wrapper)
+        request = (job, i, task_tracer_function, _args, _kwargs)
+        result = (ready, request)
+    except Exception as e:
+        return wrapped(*args, **kwargs)
+
+    return result   
+
+
+def wrap_protected_receive_wrapper(wrapped, instance, args, kwargs):
+    receive_function = wrapped(*args, **kwargs)
+    wrapped_receive_function = FunctionWrapper(receive_function, wrap_request_wrapper)
+    return wrapped_receive_function
+
+
+def instrument_celery_local(module):
+    if hasattr(module, "Proxy"):
+        # This is used in the case where the function is 
+        # called directly on the Proxy object (rather than 
+        # using "delay" or "apply_async")
+        module.Proxy.__call__ = CeleryTaskWrapper(module.Proxy.__call__)
+
+
 def instrument_celery_app_base(module):
     if hasattr(module, "Celery") and hasattr(module.Celery, "send_task"):
         wrap_function_wrapper(module, "Celery.send_task", wrap_Celery_send_task)
 
 
 def instrument_celery_worker(module):
-    # Triggered for 'celery.worker' and 'celery.concurrency.processes'.
-
     if hasattr(module, "process_initializer"):
-        # We try and force registration of default application after
-        # fork of worker process rather than lazily on first request.
-
-        # Originally the 'process_initializer' function was located in
-        # 'celery.worker'. In Celery 2.5 the function 'process_initializer'
-        # was moved to the module 'celery.concurrency.processes'.
-
+        # We try and force activation of the agent before 
+        # the worker process starts.
         _process_initializer = module.process_initializer
 
         @functools.wraps(module.process_initializer)
@@ -246,6 +318,26 @@ def process_initializer(*args, **kwargs):
             return _process_initializer(*args, **kwargs)
 
         module.process_initializer = process_initializer
+        
+    
+    if hasattr(module, "process_destructor"):
+        # We try and force shutdown of the agent before
+        # the worker process exits.
+        _process_destructor = module.process_destructor
+
+        @functools.wraps(module.process_destructor)
+        def process_destructor(*args, **kwargs):
+            shutdown_agent()
+            return _process_destructor(*args, **kwargs)
+
+        module.process_destructor = process_destructor
+        
+
+def instrument_billiard_process(module):
+    def force_application_activation(*args, **kwargs):
+        application_instance().activate()
+
+        wrap_pre_function(module, "Process.run", force_application_activation)
 
 
 def instrument_celery_loaders_base(module):
@@ -259,8 +351,15 @@ def instrument_billiard_pool(module):
     def force_agent_shutdown(*args, **kwargs):
         shutdown_agent()
 
-    if hasattr(module, "Worker"):
+    if hasattr(module, "Worker") and hasattr(module.Worker, "_make_protected_receive"):
+        # `_make_protected_receive` returns the `receive` function
+        # `receive` returns a tuple of (`ready`, `request`)
+        # where `request` is a tuple of (job, i, task_tracer_function, _args, _kwargs)
+        # and `task_tracer_function` is a function that executes the task
+        wrap_function_wrapper(module, "Worker._make_protected_receive", wrap_protected_receive_wrapper)
+    if hasattr(module, "Worker") and hasattr(module.Worker, "_do_exit"):
         wrap_pre_function(module, "Worker._do_exit", force_agent_shutdown)
+    
 
 
 def instrument_celery_app_trace(module):
@@ -270,3 +369,4 @@ def instrument_celery_app_trace(module):
 
     if hasattr(module, "reset_worker_optimizations"):
         wrap_function_wrapper(module, "reset_worker_optimizations", wrap_worker_optimizations)
+