Merge pull request #530 from palazzem/celery-signals

[celery] use Celery signals instead of patch

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

@@ -0,0 +1,64 @@
+"""
+The Celery integration will trace all tasks that are executed in the
+background. To trace your Celery application, call the patch method::
+
+    import celery
+    from ddtrace import patch
+
+    patch(celery=True)
+    app = celery.Celery()
+
+    @app.task
+    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::
+
+    import celery
+    from ddtrace.contrib.celery import patch_app, patch_task
+
+    # 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():
+        pass
+
+
+    class BaseClassTask(celery.Task):
+        def run(self):
+            pass
+
+
+    BaseClassTask = patch_task(BaseClassTask)
+    fn_task = patch_task(fn_task)
+"""
+from ...utils.importlib import require_modules
+
+
+required_modules = ['celery']
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .app import patch_app, unpatch_app
+        from .patch import patch, unpatch
+        from .task import patch_task, unpatch_task
+
+        __all__ = [
+            'patch',
+            'patch_app',
+            'patch_task',
+            'unpatch',
+            'unpatch_app',
+            'unpatch_task',
+        ]

ddtrace/contrib/celery/app.py

@@ -0,0 +1,53 @@
+from celery import signals
+
+from ddtrace import Pin
+from ddtrace.pin import _DD_PIN_NAME
+from ddtrace.ext import AppTypes
+
+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):
+    """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)
+    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):
+    """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

@@ -0,0 +1,17 @@
+import celery
+
+from .app import patch_app, unpatch_app
+
+
+def patch():
+    """Instrument Celery base application and the `TaskRegistry` so
+    that any new registered task is automatically instrumented. In the
+    case of Django-Celery integration, also the `@shared_task` decorator
+    must be instrumented because Django doesn't use the Celery registry.
+    """
+    patch_app(celery.Celery)
+
+
+def unpatch():
+    """Disconnect all signals and remove Tracing capabilities"""
+    unpatch_app(celery.Celery)

ddtrace/contrib/celery/signals.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)

ddtrace/contrib/celery/task.py

@@ -0,0 +1,31 @@
+from .app import patch_app
+
+from ...utils.deprecation import deprecation
+
+
+def patch_task(task, pin=None):
+    """Deprecated API. The new API uses signals that can be activated via
+    patch(celery=True) or through `ddtrace-run` script. Using this API
+    enables instrumentation on all tasks.
+    """
+    deprecation(
+        name='ddtrace.contrib.celery.patch_task',
+        message='Use `patch(celery=True)` or `ddtrace-run` script instead',
+        version='1.0.0',
+    )
+
+    # Enable instrumentation everywhere
+    patch_app(task.app)
+    return task
+
+def unpatch_task(task):
+    """Deprecated API. The new API uses signals that can be deactivated
+    via unpatch() API. This API is now a no-op implementation so it doesn't
+    affect instrumented tasks.
+    """
+    deprecation(
+        name='ddtrace.contrib.celery.patch_task',
+        message='Use `unpatch()` instead',
+        version='1.0.0',
+    )
+    return task