Repository: elastic/apm-agent-python

Issue:
Django Celery trace.id integration
**Is your feature request related to a problem? Please describe.**
As of now, it is impossible to keep the same trace id between a Django view and a Celery task launched from the same view.

**Describe the solution you'd like**
Provide a way to easily pass a trace parent string to the Celery task. Preferably via it's headers field (introduced in Celery 3.1).

### What would it looks likes
**User side code (Django view)**
```python
def get(self, request):
    transaction = execution_context.get_transaction()
    trace_parent = transaction.trace_parent
    trace_parent_string = trace_parent.to_string()
    my_celery_task.apply_async(headers={"elasticapm": {"trace_parent_string": trace_parent_string} })
```

**Library side code (`elasticapm.contrib.celery.__init__.py`), rewrite of `begin_transaction()`, naïve implementation**
```python
def begin_transaction(*args, **kwargs):
    trace_parent = None
    try:
        trace_parent_string = kwargs["task"].request.headers["elasticapm"]["trace_parent_string"]
        trace_parent = TraceParent.from_string(trace_parent_string)
    except:
        pass
    client.begin_transaction("celery", trace_parent=trace_parent)
```

- **Why using Celery headers field ?** It seems the most unobstrusive way of doing it.
- **Why using a nested field (["elasticapm"]["trace_parent_string"]) ?** Seems "future proof", usefull future fields for elasticapm could be added under the "elasticapm" key. Users of the API shouldn't see their code break as they are aware that using this library, the headers Celery field has a reserved key "elasticapm" used for this integration.

**Additional context**
**For anyone wanting to try it, BEWARE !!** There is a Celery [bug](https://github.com/celery/celery/issues/4875) concerning it's headers field.
You might have to do this:
```python
my_celery_task.apply_async(headers={"headers": {"elasticapm": {"trace_parent_string": trace_parent_string} } })
```

Edits: fixed code error/typos


Discussion:
Hi @TPoret 

thanks for opening the detailed issue! This is something we're really interested in tackling, but haven't found the time to do it yet.

Ideally, injecting and extracting the tracing data should be done without the user having to write any custom code. Your code for extracting the header already does this, which is great. For injecting the trace data into the header, I think the [`before_task_publish`](https://docs.celeryproject.org/en/stable/userguide/signals.html#before-task-publish) signal is a great fit, as it explicitly allows for changing the headers before the task is published. This could be done in the same fashion as our current signal handlers: https://github.com/elastic/apm-agent-python/blob/f55f30dc9e0a66c237914bd100beee6f84a1e4ab/elasticapm/contrib/celery/__init__.py#L60-L85

Would you like to give it a try and open a PR? Your code already looks pretty good!
Thank for your suggestion ([before_task_publish()](https://docs.celeryproject.org/en/stable/userguide/signals.html#before-task-publish)), I had no idea how to do that without user intervention.

Although I have some reserve about doing everything automatically. Shouldn't we offer the user a way to choose when to propagate the trace ID to the Celery task ? And/or a way to disable this behavior for all Celery tasks ?

How I see it:
- Something like a settings in `ELASTIC_APM` dict to disable this behavior for all tasks (enabled by default).
- explicitely disable per task if the user provide a specific value to the header eg. `my_celery_task.apply_async(headers={"elasticapm": {"correlation": False} }`

I will gladly open a PR.


Resolving patch:
diff --git a/elasticapm/contrib/celery/__init__.py b/elasticapm/contrib/celery/__init__.py
--- a/elasticapm/contrib/celery/__init__.py
+++ b/elasticapm/contrib/celery/__init__.py
@@ -27,13 +27,15 @@
 #  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 #  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-
+from contextlib import suppress
 
 from celery import signals, states
 
 import elasticapm
 from elasticapm.conf import constants
+from elasticapm.traces import execution_context
 from elasticapm.utils import get_name_from_func
+from elasticapm.utils.disttracing import TraceParent
 
 
 class CeleryFilter(object):
@@ -57,9 +59,41 @@ def process_failure_signal(sender, task_id, exception, args, kwargs, traceback,
     _register_worker_signals(client)
 
 
+def set_celery_headers(headers=None, **kwargs):
+    """
+    Add elasticapm specific information to celery headers
+    """
+    headers = {} if headers is None else headers
+
+    transaction = execution_context.get_transaction()
+    if transaction is not None:
+        trace_parent = transaction.trace_parent
+        trace_parent_string = trace_parent.to_string()
+
+        headers.update({"elasticapm": {"trace_parent_string": trace_parent_string}})
+
+
+def get_trace_parent(celery_task):
+    """
+    Return a trace parent contained in the request headers of a Celery Task object or None
+    """
+    trace_parent = None
+    with suppress(AttributeError, KeyError, TypeError):
+        if celery_task.request.headers is not None:
+            trace_parent_string = celery_task.request.headers["elasticapm"]["trace_parent_string"]
+            trace_parent = TraceParent.from_string(trace_parent_string)
+        else:
+            trace_parent_string = celery_task.request.elasticapm["trace_parent_string"]
+            trace_parent = TraceParent.from_string(trace_parent_string)
+    return trace_parent
+
+
 def register_instrumentation(client):
     def begin_transaction(*args, **kwargs):
-        client.begin_transaction("celery")
+        task = kwargs["task"]
+
+        trace_parent = get_trace_parent(task)
+        client.begin_transaction("celery", trace_parent=trace_parent)
 
     def end_transaction(task_id, task, *args, **kwargs):
         name = get_name_from_func(task)
@@ -76,10 +110,12 @@ def end_transaction(task_id, task, *args, **kwargs):
     dispatch_uid = "elasticapm-tracing-%s"
 
     # unregister any existing clients
+    signals.before_task_publish.disconnect(set_celery_headers, dispatch_uid=dispatch_uid % "before-publish")
     signals.task_prerun.disconnect(begin_transaction, dispatch_uid=dispatch_uid % "prerun")
     signals.task_postrun.disconnect(end_transaction, dispatch_uid=dispatch_uid % "postrun")
 
     # register for this client
+    signals.before_task_publish.connect(set_celery_headers, dispatch_uid=dispatch_uid % "before-publish")
     signals.task_prerun.connect(begin_transaction, dispatch_uid=dispatch_uid % "prerun", weak=False)
     signals.task_postrun.connect(end_transaction, weak=False, dispatch_uid=dispatch_uid % "postrun")
     _register_worker_signals(client)


Test patch: