Merge branch 'palazzem/celery-integration' into 'master'

Author: Emanuele Palazzetti

ddtrace/bootstrap/sitecustomize.py

@@ -70,9 +70,6 @@ def update_patched_modules():
     if opts:
         tracer.configure(**opts)
 
-    if not hasattr(sys, 'argv'):
-        sys.argv = ['']
-
     if patch:
         update_patched_modules()
         from ddtrace import patch_all; patch_all(**EXTRA_PATCHED_MODULES) # noqa

ddtrace/contrib/celery/__init__.py

@@ -1,6 +1,9 @@
 """
 The Celery integration will trace all tasks that are executed in the
-background. To trace your Celery application, call the patch method::
+background. Functions and class based tasks are traced only if the Celery API
+is used, so calling the function directly or via the ``run()`` method will not
+generate traces. On the other hand, calling ``apply()`` and ``apply_async()``
+will produce tracing data. To trace your Celery application, call the patch method::
 
     import celery
     from ddtrace import patch
@@ -12,36 +15,33 @@
     def my_task():
         pass
 
-
     class MyTask(app.Task):
         def run(self):
             pass
 
 
-If you don't need to patch all Celery tasks, you can patch individual
-applications or tasks using a fine grain patching method::
+To change Celery service name, you can update the attached ``Pin``
+instance::
 
-    import celery
-    from ddtrace.contrib.celery import patch_app, patch_task
+    from ddtrace import Pin
 
-    # patch only this application
     app = celery.Celery()
-    app = patch_app(app)
 
-    # or if you didn't patch the whole application, just patch
-    # a single function or class based Task
     @app.task
-    def fn_task():
+    def compute_stats():
         pass
 
+    # globally
+    Pin.override(app, service='background-jobs')
+
+    # by task
+    Pin.override(compute_stats, service='data-processing')
 
-    class BaseClassTask(celery.Task):
-        def run(self):
-            pass
 
+By default, reported service names are:
+    * ``celery-producer`` when tasks are enqueued for processing
+    * ``celery-worker`` when tasks are processed by a Celery process
 
-    BaseClassTask = patch_task(BaseClassTask)
-    fn_task = patch_task(fn_task)
 """
 from ...utils.importlib import require_modules
 

ddtrace/contrib/celery/app.py

@@ -1,74 +1,53 @@
-# Standard library
-import types
+from celery import signals
 
-# Third party
-import wrapt
-
-# Project
 from ddtrace import Pin
+from ddtrace.pin import _DD_PIN_NAME
 from ddtrace.ext import AppTypes
-from .task import patch_task, unpatch_task
-from .util import APP, WORKER_SERVICE, require_pin
+
+from .constants import APP, WORKER_SERVICE
+from .signals import (
+    trace_prerun,
+    trace_postrun,
+    trace_before_publish,
+    trace_after_publish,
+    trace_failure,
+)
 
 
 def patch_app(app, pin=None):
-    """ patch_app will add tracing to a celery app """
+    """Attach the Pin class to the application and connect
+    our handlers to Celery signals.
+    """
+    if getattr(app, '__datadog_patch', False):
+        return
+    setattr(app, '__datadog_patch', True)
+
+    # attach the PIN object
     pin = pin or Pin(service=WORKER_SERVICE, app=APP, app_type=AppTypes.worker)
-    patch_methods = [
-        ('task', _app_task),
-    ]
-    for method_name, wrapper in patch_methods:
-        # Get the original method
-        method = getattr(app, method_name, None)
-        if method is None:
-            continue
-
-        # Do not patch if method is already patched
-        if isinstance(method, wrapt.ObjectProxy):
-            continue
-
-        # Patch method
-        setattr(app, method_name, wrapt.FunctionWrapper(method, wrapper))
-
-    # patch the Task class if available
-    setattr(app, 'Task', patch_task(app.Task))
-
-    # Attach our pin to the app
     pin.onto(app)
+    # connect to the Signal framework
+    signals.task_prerun.connect(trace_prerun)
+    signals.task_postrun.connect(trace_postrun)
+    signals.before_task_publish.connect(trace_before_publish)
+    signals.after_task_publish.connect(trace_after_publish)
+    signals.task_failure.connect(trace_failure)
     return app
 
 
 def unpatch_app(app):
-    """ unpatch_app will remove tracing from a celery app """
-    patched_methods = [
-        'task',
-    ]
-    for method_name in patched_methods:
-        # Get the wrapped method
-        wrapper = getattr(app, method_name, None)
-        if wrapper is None:
-            continue
-
-        # Only unpatch if the wrapper is an `ObjectProxy`
-        if not isinstance(wrapper, wrapt.ObjectProxy):
-            continue
-
-        # Restore original method
-        setattr(app, method_name, wrapper.__wrapped__)
-
-    # restore the original Task class
-    setattr(app, 'Task', unpatch_task(app.Task))
-    return app
-
-
-@require_pin
-def _app_task(pin, func, app, args, kwargs):
-    task = func(*args, **kwargs)
-
-    # `app.task` is a decorator which may return a function wrapper
-    if isinstance(task, types.FunctionType):
-        def wrapper(func, instance, args, kwargs):
-            return patch_task(func(*args, **kwargs), pin=pin)
-        return wrapt.FunctionWrapper(task, wrapper)
-
-    return patch_task(task, pin=pin)
+    """Remove the Pin instance from the application and disconnect
+    our handlers from Celery signal framework.
+    """
+    if not getattr(app, '__datadog_patch', False):
+        return
+    setattr(app, '__datadog_patch', False)
+
+    pin = Pin.get_from(app)
+    if pin is not None:
+        delattr(app, _DD_PIN_NAME)
+
+    signals.task_prerun.disconnect(trace_prerun)
+    signals.task_postrun.disconnect(trace_postrun)
+    signals.before_task_publish.disconnect(trace_before_publish)
+    signals.after_task_publish.disconnect(trace_after_publish)
+    signals.task_failure.disconnect(trace_failure)

ddtrace/contrib/celery/constants.py

@@ -0,0 +1,19 @@
+from os import getenv
+
+# Celery Context key
+CTX_KEY = '__dd_task_span'
+
+# Span names
+PRODUCER_ROOT_SPAN = 'celery.apply'
+WORKER_ROOT_SPAN = 'celery.run'
+
+# Task operations
+TASK_TAG_KEY = 'celery.action'
+TASK_APPLY = 'apply'
+TASK_APPLY_ASYNC = 'apply_async'
+TASK_RUN = 'run'
+
+# Service info
+APP = 'celery'
+PRODUCER_SERVICE = getenv('DATADOG_SERVICE_NAME') or 'celery-producer'
+WORKER_SERVICE = getenv('DATADOG_SERVICE_NAME') or 'celery-worker'

ddtrace/contrib/celery/patch.py

@@ -1,11 +1,6 @@
 import celery
 
-from wrapt import wrap_function_wrapper as _w
-
 from .app import patch_app, unpatch_app
-from .task import _wrap_shared_task
-from .registry import _wrap_register
-from ...utils.wrappers import unwrap as _u
 
 
 def patch():
@@ -14,13 +9,9 @@ def patch():
     case of Django-Celery integration, also the `@shared_task` decorator
     must be instrumented because Django doesn't use the Celery registry.
     """
-    setattr(celery, 'Celery', patch_app(celery.Celery))
-    _w('celery.app.registry', 'TaskRegistry.register', _wrap_register)
-    _w('celery', 'shared_task', _wrap_shared_task)
+    patch_app(celery.Celery)
 
 
 def unpatch():
-    """Removes instrumentation from Celery"""
-    setattr(celery, 'Celery', unpatch_app(celery.Celery))
-    _u(celery.app.registry.TaskRegistry, 'register')
-    _u(celery, 'shared_task')
+    """Disconnect all signals and remove Tracing capabilities"""
+    unpatch_app(celery.Celery)

ddtrace/contrib/celery/registry.py

@@ -0,0 +1,130 @@
+import logging
+
+from ddtrace import Pin
+
+from celery import registry
+
+from . import constants as c
+from .utils import (
+    tags_from_context,
+    retrieve_task_id,
+    attach_span,
+    detach_span,
+    retrieve_span,
+)
+
+
+log = logging.getLogger(__name__)
+
+
+def trace_prerun(*args, **kwargs):
+    # safe-guard to avoid crashes in case the signals API
+    # changes in Celery
+    task = kwargs.get('sender')
+    task_id = kwargs.get('task_id')
+    if task is None or task_id is None:
+        log.debug('unable to extract the Task and the task_id. This version of Celery may not be supported.')
+        return
+
+    # retrieve the task Pin or fallback to the global one
+    pin = Pin.get_from(task) or Pin.get_from(task.app)
+    if pin is None:
+        return
+
+    # propagate the `Span` in the current task Context
+    span = pin.tracer.trace(c.WORKER_ROOT_SPAN, service=c.WORKER_SERVICE, resource=task.name)
+    attach_span(task, task_id, span)
+
+
+def trace_postrun(*args, **kwargs):
+    # safe-guard to avoid crashes in case the signals API
+    # changes in Celery
+    task = kwargs.get('sender')
+    task_id = kwargs.get('task_id')
+    if task is None or task_id is None:
+        log.debug('unable to extract the Task and the task_id. This version of Celery may not be supported.')
+        return
+
+    # retrieve and finish the Span
+    span = retrieve_span(task, task_id)
+    if span is None:
+        return
+    else:
+        # request context tags
+        span.set_tag(c.TASK_TAG_KEY, c.TASK_RUN)
+        span.set_tags(tags_from_context(kwargs))
+        span.set_tags(tags_from_context(task.request))
+        span.finish()
+        detach_span(task, task_id)
+
+
+def trace_before_publish(*args, **kwargs):
+    # `before_task_publish` signal doesn't propagate the task instance so
+    # we need to retrieve it from the Celery Registry to access the `Pin`. The
+    # `Task` instance **does not** include any information about the current
+    # execution, so it **must not** be used to retrieve `request` data.
+    task_name = kwargs.get('sender')
+    task = registry.tasks.get(task_name)
+    task_id = retrieve_task_id(kwargs)
+    # safe-guard to avoid crashes in case the signals API
+    # changes in Celery
+    if task is None or task_id is None:
+        log.debug('unable to extract the Task and the task_id. This version of Celery may not be supported.')
+        return
+
+    # propagate the `Span` in the current task Context
+    pin = Pin.get_from(task) or Pin.get_from(task.app)
+    if pin is None:
+        return
+
+    # apply some tags here because most of the data is not available
+    # in the task_after_publish signal
+    span = pin.tracer.trace(c.PRODUCER_ROOT_SPAN, service=c.PRODUCER_SERVICE, resource=task_name)
+    span.set_tag(c.TASK_TAG_KEY, c.TASK_APPLY_ASYNC)
+    span.set_tag('celery.id', task_id)
+    span.set_tags(tags_from_context(kwargs))
+    # Note: adding tags from `traceback` or `state` calls will make an
+    # API call to the backend for the properties so we should rely
+    # only on the given `Context`
+    attach_span(task, task_id, span)
+
+
+def trace_after_publish(*args, **kwargs):
+    task_name = kwargs.get('sender')
+    task = registry.tasks.get(task_name)
+    task_id = retrieve_task_id(kwargs)
+    # safe-guard to avoid crashes in case the signals API
+    # changes in Celery
+    if task is None or task_id is None:
+        log.debug('unable to extract the Task and the task_id. This version of Celery may not be supported.')
+        return
+
+    # retrieve and finish the Span
+    span = retrieve_span(task, task_id)
+    if span is None:
+        return
+    else:
+        span.finish()
+        detach_span(task, task_id)
+
+
+def trace_failure(*args, **kwargs):
+    # safe-guard to avoid crashes in case the signals API
+    # changes in Celery
+    task = kwargs.get('sender')
+    task_id = kwargs.get('task_id')
+    if task is None or task_id is None:
+        log.debug('unable to extract the Task and the task_id. This version of Celery may not be supported.')
+        return
+
+    # retrieve and finish the Span
+    span = retrieve_span(task, task_id)
+    if span is None:
+        return
+    else:
+        # add Exception tags; post signals are still called
+        # so we don't need to attach other tags here
+        ex = kwargs.get('einfo')
+        if ex is None:
+            return
+        span.set_exc_info(ex.type, ex.exception, ex.tb)