[celery] add Config object to change Worker/Producer service names (#540)

* [celery] add Config object to change Worker/Producer service names

* [celery] update documentation after the Config API change

Author: Emanuele Palazzetti

ddtrace/contrib/celery/__init__.py

@@ -2,7 +2,7 @@
 The Celery integration will trace all tasks that are executed in the
 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()``
+generate traces. On the other hand, calling ``apply()``, ``apply_async()`` and ``delay()``
 will produce tracing data. To trace your Celery application, call the patch method::
 
     import celery
@@ -20,8 +20,7 @@ def run(self):
             pass
 
 
-To change Celery service name, you can update the attached ``Pin``
-instance::
+To change Celery service name, you can use the ``Config`` API as follows::
 
     from ddtrace import Pin
 
@@ -31,12 +30,9 @@ def run(self):
     def compute_stats():
         pass
 
-    # globally
-    Pin.override(app, service='background-jobs')
-
-    # by task
-    Pin.override(compute_stats, service='data-processing')
-
+    # change service names for producers and workers
+    config.celery['producer_service_name'] = 'task-queue'
+    config.celery['worker_service_name'] = 'worker-notify'
 
 By default, reported service names are:
     * ``celery-producer`` when tasks are enqueued for processing

ddtrace/contrib/celery/app.py

@@ -1,10 +1,10 @@
 from celery import signals
 
-from ddtrace import Pin
+from ddtrace import Pin, config
 from ddtrace.pin import _DD_PIN_NAME
 from ddtrace.ext import AppTypes
 
-from .constants import APP, WORKER_SERVICE
+from .constants import APP
 from .signals import (
     trace_prerun,
     trace_postrun,
@@ -23,7 +23,12 @@ def patch_app(app, pin=None):
     setattr(app, '__datadog_patch', True)
 
     # attach the PIN object
-    pin = pin or Pin(service=WORKER_SERVICE, app=APP, app_type=AppTypes.worker)
+    pin = pin or Pin(
+        service=config.celery['worker_service_name'],
+        app=APP,
+        app_type=AppTypes.worker,
+        _config=config.celery,
+    )
     pin.onto(app)
     # connect to the Signal framework
     signals.task_prerun.connect(trace_prerun)

ddtrace/contrib/celery/constants.py

@@ -15,5 +15,7 @@
 
 # Service info
 APP = 'celery'
+# `getenv()` call must be kept for backward compatibility; we may remove it
+# later when we do a full migration to the `Config` class
 PRODUCER_SERVICE = getenv('DATADOG_SERVICE_NAME') or 'celery-producer'
 WORKER_SERVICE = getenv('DATADOG_SERVICE_NAME') or 'celery-worker'

ddtrace/contrib/celery/patch.py

@@ -1,6 +1,17 @@
 import celery
 
+from ddtrace import config
+
 from .app import patch_app, unpatch_app
+from .constants import PRODUCER_SERVICE, WORKER_SERVICE
+from ...utils.formats import get_env
+
+
+# Celery default settings
+config._add('celery', {
+    'producer_service_name': get_env('celery', 'producer_service_name', PRODUCER_SERVICE),
+    'worker_service_name': get_env('celery', 'worker_service_name', WORKER_SERVICE),
+})
 
 
 def patch():

ddtrace/contrib/celery/signals.py

@@ -1,6 +1,6 @@
 import logging
 
-from ddtrace import Pin
+from ddtrace import Pin, config
 
 from celery import registry
 
@@ -32,7 +32,8 @@ def trace_prerun(*args, **kwargs):
         return
 
     # propagate the `Span` in the current task Context
-    span = pin.tracer.trace(c.WORKER_ROOT_SPAN, service=c.WORKER_SERVICE, resource=task.name)
+    service = config.celery['worker_service_name']
+    span = pin.tracer.trace(c.WORKER_ROOT_SPAN, service=service, resource=task.name)
     attach_span(task, task_id, span)
 
 
@@ -79,7 +80,8 @@ def trace_before_publish(*args, **kwargs):
 
     # 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)
+    service = config.celery['producer_service_name']
+    span = pin.tracer.trace(c.PRODUCER_ROOT_SPAN, service=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))

tests/contrib/celery/base.py

@@ -2,7 +2,7 @@
 
 from celery import Celery
 
-from ddtrace import Pin
+from ddtrace import Pin, config
 from ddtrace.compat import PY2
 from ddtrace.contrib.celery import patch, unpatch
 
@@ -21,19 +21,23 @@ class CeleryBaseTestCase(unittest.TestCase):
     """
 
     def setUp(self):
+        # keep track of original config
+        self._config = dict(config.celery)
         # instrument Celery and create an app with Broker and Result backends
         patch()
         self.tracer = get_dummy_tracer()
         self.pin = Pin(service='celery-unittest', tracer=self.tracer)
         self.app = Celery('celery.test_app', broker=BROKER_URL, backend=BACKEND_URL)
         # override pins to use our Dummy Tracer
         Pin.override(self.app, tracer=self.tracer)
-        Pin.override(self.app.task, tracer=self.tracer)
-        Pin.override(self.app.Task, tracer=self.tracer)
 
     def tearDown(self):
+        # remove instrumentation from Celery
         unpatch()
         self.app = None
+        # restore the global configuration
+        config.celery.update(self._config)
+        self._config = None
 
     def assert_items_equal(self, a, b):
         if PY2:

tests/contrib/celery/test_integration.py

@@ -2,6 +2,7 @@
 
 from nose.tools import eq_, ok_
 
+from ddtrace import config
 from ddtrace.contrib.celery import patch, unpatch
 
 from .base import CeleryBaseTestCase
@@ -272,3 +273,40 @@ def add(x ,y):
         eq_(span.get_tag('celery.id'), res.task_id)
         eq_(span.get_tag('celery.action'), 'run')
         eq_(span.get_tag('celery.state'), 'SUCCESS')
+
+    def test_worker_service_name(self):
+        # Ensure worker service name can be changed via
+        # configuration object
+        config.celery['worker_service_name'] = 'worker-notify'
+
+        @self.app.task
+        def fn_task():
+            return 42
+
+        t = fn_task.apply()
+        ok_(t.successful())
+        eq_(42, t.result)
+
+        traces = self.tracer.writer.pop_traces()
+        eq_(1, len(traces))
+        eq_(1, len(traces[0]))
+        span = traces[0][0]
+        eq_(span.service, 'worker-notify')
+
+    def test_producer_service_name(self):
+        # Ensure producer service name can be changed via
+        # configuration object
+        config.celery['producer_service_name'] = 'task-queue'
+
+        @self.app.task
+        def fn_task():
+            return 42
+
+        t = fn_task.delay()
+        eq_('PENDING', t.status)
+
+        traces = self.tracer.writer.pop_traces()
+        eq_(1, len(traces))
+        eq_(1, len(traces[0]))
+        span = traces[0][0]
+        eq_(span.service, 'task-queue')