Celery contrib: read custom parent span id from the task headers (#1500)

* Celery contrib: read parent span id from celery headers if provided

* Simplify setting custom parent span if by copying from the current trace parent

* Should attempt to get trace parent from celery task headers before trying the inner headers

* CHANGELOG and lint

Co-authored-by: Colton Myers <[email protected]>

Author: namoshizun

CHANGELOG.asciidoc

@@ -40,6 +40,7 @@ endif::[]
 * Change default for `sanitize_field_names` to sanitize `*auth*` instead of `authorization` {pull}1494[#1494]
 * Add `span_stack_trace_min_duration` to replace deprecated `span_frames_min_duration` {pull}1498[#1498]
 * Enable exact_match span compression by default {pull}1504[#1504]
+* Allow parent celery tasks to specify the downstream `parent_span_id` in celery headers {pull}1500[#1500]
 
 [float]
 ===== Bug fixes

elasticapm/contrib/celery/__init__.py

@@ -68,24 +68,30 @@ def set_celery_headers(headers=None, **kwargs):
     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}})
+        # Customize parent span id (if provided)
+        apm_headers = headers.get("elasticapm", dict())
+        if "parent_span_id" in apm_headers:
+            trace_parent = trace_parent.copy_from()
+            trace_parent.span_id = apm_headers["parent_span_id"]
+
+        apm_headers["trace_parent_string"] = trace_parent.to_string()
+        headers.update(elasticapm=apm_headers)
 
 
 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
+    read_from_inner_headers = lambda: celery_task.request.headers["elasticapm"]["trace_parent_string"]
+    read_from_request = lambda: celery_task.request.elasticapm["trace_parent_string"]
+
+    for read_fun in (read_from_request, read_from_inner_headers):
+        with suppress(AttributeError, KeyError, TypeError):
+            trace_parent_string = read_fun()
+            return TraceParent.from_string(trace_parent_string)
+
+    return None
 
 
 def register_instrumentation(client):