Repository: DataDog/dd-trace-py

Issue:
ddtrace-run breaks Celery support, have to patch manually
The issue appears when trying to run Celery workers with `ddtrace-run`.

`dd-trace-py` library attempts to patch Celery but fails with the message - _module not installed_. Upon further examination we've discovered that it fails at `ddtrace/contrib/celery/__init.py__` line 51 when trying to import the Celery module.

The reason for that is that the import fails with Exception saying that `sys.argv` doesn't exists. Somewhere in the Celery module there is a reference towards `sys.argv`. As it seems, when importing Celery from `ddtrace-run`, this is not set (or `sys` is perhaps patched or muted) so the import fails and the patching fails silently saying that Celery is not installed.

The workaround is to manually call `patch(celery=True)` once your application regains control from `ddtrace-run`. In that case, Celery is normally imported as `sys.argv` exists.


Discussion:
Hello @apolloFER, I'm trying to reproduce this problem and I'm wondering if you can provide me more details. Currently I'm running with success the example you published in another issue/PR (https://gist.github.com/apolloFER/fc97831a69830e36ff83309e9f931d2f). The following command starts with success:
```
$ ddtrace-run celery -A worker worker --loglevel=info
```

Do you have an example that reproduces this specific issue? I would prefer to tackle this one before diving again in the other problem. Thank you very much!
Hi @palazzem 

The other example works because of this line
https://gist.github.com/apolloFER/fc97831a69830e36ff83309e9f931d2f#file-worker-py-L6
This line manually repatches Celery after startup. If you remove it you'll be able to reproduce this issue (ddtrace-run trying to patch celery but failing with an exception -- ImportError -- which is silenced).
Thank you @apolloFER ! Will continue from here fixing this issue before the other!
Hello @apolloFER! I'm still checking this example and this is exactly what is happening:
* If I use the example, I can reproduce the issue having also the following issue (#449):
```
TypeError: missing 1 required positional argument
```
* If I run a modified example that uses `ddtrace-run` without any `ddtrace` import (especially the `@patch_task` that is causing the problem above) I can send correctly traces to our Agent. Here a gist to replicate my current environment: https://gist.github.com/palazzem/615faaa7d26b69399ba681a01160dd8c

I'm still triaging the problem, but what I got is:
* while using the `@patch_task` and `patch(celery=True)` together, we have the `TypeError` issue
* if we remove all `ddtrace` imports and simply use the `ddtrace-run` it works as expected
* of course it's not working in your case as stated in this thread, so I'm wondering if I can slightly modify my gist to have exactly a case where `ddtrace-run` raises an `ImportError` avoiding patching Celery
Hi @palazzem 

I tried the modified example where I removed the `ddtrace` import (including all `@patch_task` calls).  Here's my current worker:
https://gist.github.com/apolloFER/a0e7ac6f5e0419c878d8bb120f3dfc69
Here's the startup log:
https://gist.github.com/apolloFER/7883745d5088595629f73f351a5f92a0
As you can see, celery is not patched (the reason is that ImportError which is caught so you won't get it in logs). Nothing crashes (the TypeError crash) but that's because Celery isn't patched and neither are the tasks. This is also seen in Datadog dashboard -- there are traces reported, but those are all Redis traces and no Celery traces.


With #469 we have fixed the issue about the `argv` error. Because the issue was only about using `ddtrace-run`, we can consider it addressed. I'll follow-up in other GitHub issues. Thank you very much!

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
@@ -1,15 +1,21 @@
-# Third party
 import celery
 
-# Project
+from wrapt import wrap_function_wrapper as _w
+
 from .app import patch_app, unpatch_app
+from .registry import _wrap_register
+from ...utils.wrappers import unwrap as _u
 
 
 def patch():
-    """ patch will add all available tracing to the celery library """
+    """Instrument Celery base application and the `TaskRegistry` so
+    that any new registered task is automatically instrumented
+    """
     setattr(celery, 'Celery', patch_app(celery.Celery))
+    _w('celery.app.registry', 'TaskRegistry.register', _wrap_register)
 
 
 def unpatch():
-    """ unpatch will remove tracing from the celery library """
+    """Removes instrumentation from Celery"""
     setattr(celery, 'Celery', unpatch_app(celery.Celery))
+    _u(celery.app.registry.TaskRegistry, 'register')
diff --git a/ddtrace/contrib/celery/registry.py b/ddtrace/contrib/celery/registry.py
new file mode 100644
--- /dev/null
+++ b/ddtrace/contrib/celery/registry.py
@@ -0,0 +1,15 @@
+from .task import patch_task
+
+
+def _wrap_register(func, instance, args, kwargs):
+    """Wraps the `TaskRegistry.register` function so that everytime
+    a `Task` is registered it is properly instrumented. This wrapper
+    is required because in old-style tasks (Celery 1.0+) we cannot
+    instrument the base class, otherwise a `Strategy` `KeyError`
+    exception is raised.
+    """
+    # the original signature requires one positional argument so the
+    # first and only parameter is the `Task` that must be instrumented
+    task = args[0]
+    patch_task(task)
+    func(*args, **kwargs)


Test patch: