Repository: DataDog/dd-trace-py

Issue:
celery 4.1 worker tasks aren't being traced with ddtrace-py 0.11.1
```
ddtrace==0.11.1
celery==4.1.0
```

I'm manually patching celery due to #423 

I'm getting traces for 'defaultdb' (sqlite), 'redis', and 'postgres' when I run the celery worker

However, I'm not receiving any traces whatsoever for the task runs, and no 'celery' service appears in datadog.

Here's a gist with my worker startup in debug logging mode:
https://gist.github.com/eedwards-sk/f924548c043859901db6918aec95dada

I'm sending traces to a local running agent and I know that's working because I'm seeing the flush records in the logs, and I get traces for the above mentioned services when I'm running it.

(We have another application that is using celery 3.x and their traces show up for tasks, which we instrumented in more or less the same way)

Here's a gist showing how we're instantiating celery and the tracer:
https://gist.github.com/eedwards-sk/fe81f62406e7a245b7c26a9ada19c658


Discussion:
Thanks @eedwards-sk ! we're currently working on this problem and will be part of the next release (0.12). Will keep you updated when it's out!

Thanks for reporting the problem
Hello again @eedwards-sk ! I'm investigating #423 and I'm wondering if you can provide me more details, so that `ddtrace-run` works as expected in your configuration. If you don't see any error in your logs and everything seems simply "not instrumented", it's possible that Celery core functionalities (i.e. `celery.Celery` or `celery.task`) are imported before calling the `patch(celery=True)`, resulting in a missing instrumentation. They may be imported by the `celery` script (I still need to investigate that part, so consider it a first guess).

For now I would like to tackle one problem at a time, so it would be great if you can share in the other issue (#423) more details. Hopefully, we may fix both problems with a simple patch.

Thank you very much!
Hi @palazzem 

The gists I attached when I opened this issue show how I'm running the celery worker, debug output confirming when it's instrumented, and the source code for the entry point which shows instantiation of the tracer and the celery object.

If that wasn't enough detail already, can you please be more specific about what 'more details' you need?

Thanks!
Ah OK, I thought they were related more to "traces not reported in Celery 4.1" rather than `ddtrace-run` not working. I'll check again and see if it's enough to reproduce the problem! Thanks!
Hey @palazzem 

Today we noticed a celery task had been instrumented, but it was a built-in celery task (`celery.backend_cleanup`)

Thus, it appears that at least some aspect of the tracing is working, however our custom tasks are not being instrumented.

Possibly an issue with the datadog tracer and monkeypatching tasks when they are added using `app.autodiscover_tasks()` ?
@eedwards-sk I think that it's possible. Actually we found why `ddtrace-run` doesn't work and it will be part of a bugfix release. After writing a PR for the other problem, I would like to triage again this issue because I think it should be really related to the `autodiscover_tasks()`. Also:
* what version of Django are you using? not really relevant but I would like to reproduce the problem with a real app too
* are you using Celery 1.0 style tasks? just asking because we have a known problem for that
I work with @eedwards-sk 

```
>>> django.VERSION
(2, 0, 0, 'final', 0)
```
(but we can upgrade if necessary)

I do not believe we are using Celery 1.0 style tasks. The project is only a few months old.

```
>>> celery.VERSION
version_info_t(major=4, minor=1, micro=0, releaselevel='', serial='')
```

Here is an example of a celery call
```
from celery import shared_task

@shared_task
def foo_bar(baz):
    return 1
```
Hello @jallen1788 ! thanks for these details! Definitely not, so I need to check if we're experiencing issues with the `@shared_task` decorator. We're actively working on it right now, and will reach you once we have something stable that you can test (probably next week). Sounds good?
@palazzem thanks a bunch, sounds good!
@eedwards-sk just an update about the issue. We fixed some issues related to `ddtrace-run` + Celery and how we patch Celery tasks: https://github.com/DataDog/dd-trace-py/pull/469 and https://github.com/DataDog/dd-trace-py/pull/465

They're already on `master` but still not released. We're checking now if there are issues with `@shared_task` decorator so that we can consider this issue closed. Thanks and sorry for waiting more than expected.
@palazzem 

Thanks for the update! I'll keep an eye out for 0.12.1 which I assume will have this improvement.

Resolving patch:
diff --git a/ddtrace/contrib/celery/patch.py b/ddtrace/contrib/celery/patch.py
--- a/ddtrace/contrib/celery/patch.py
+++ b/ddtrace/contrib/celery/patch.py
@@ -3,19 +3,24 @@
 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():
     """Instrument Celery base application and the `TaskRegistry` so
-    that any new registered task is automatically instrumented
+    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.
     """
     setattr(celery, 'Celery', patch_app(celery.Celery))
     _w('celery.app.registry', 'TaskRegistry.register', _wrap_register)
+    _w('celery', 'shared_task', _wrap_shared_task)
 
 
 def unpatch():
     """Removes instrumentation from Celery"""
     setattr(celery, 'Celery', unpatch_app(celery.Celery))
     _u(celery.app.registry.TaskRegistry, 'register')
+    _u(celery, 'shared_task')
diff --git a/ddtrace/contrib/celery/task.py b/ddtrace/contrib/celery/task.py
--- a/ddtrace/contrib/celery/task.py
+++ b/ddtrace/contrib/celery/task.py
@@ -79,6 +79,14 @@ def unpatch_task(task):
     return task
 
 
+def _wrap_shared_task(decorator, instance, args, kwargs):
+    """Wrapper for Django-Celery shared tasks. `shared_task` is a decorator
+    that returns a `Task` from the given function.
+    """
+    task = decorator(*args, **kwargs)
+    return patch_task(task)
+
+
 def _task_init(func, task, args, kwargs):
     func(*args, **kwargs)
 


Test patch:
diff --git a/tests/contrib/celery/test_task.py b/tests/contrib/celery/test_task.py
--- a/tests/contrib/celery/test_task.py
+++ b/tests/contrib/celery/test_task.py
@@ -488,3 +488,19 @@ class CelerySubClass(CelerySuperClass):
         t.run()
         spans = self.tracer.writer.pop()
         self.assertEqual(len(spans), 4)
+
+    def test_celery_shared_task(self):
+        @celery.shared_task
+        def add(x ,y):
+            return x + y
+
+        res = add.run(2, 2)
+        self.assertEqual(res, 4)
+        spans = self.tracer.writer.pop()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+        self.assertEqual(span.service, 'celery-worker')
+        self.assertEqual(span.resource, 'tests.contrib.celery.test_task.add')
+        self.assertEqual(span.name, 'celery.run')
+        self.assertIsNone(span.parent_id)
+        self.assertEqual(span.error, 0)