prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN 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
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `elasticapm/contrib/celery/__init__.py`
  Content:
  ```
  1 #  BSD 3-Clause License
  2 #
  3 #  Copyright (c) 2012, the Sentry Team, see AUTHORS for more details
  4 #  Copyright (c) 2019, Elasticsearch BV
  5 #  All rights reserved.
  6 #
  7 #  Redistribution and use in source and binary forms, with or without
  8 #  modification, are permitted provided that the following conditions are met:
  9 #
  10 #  * Redistributions of source code must retain the above copyright notice, this
  11 #    list of conditions and the following disclaimer.
  12 #
  13 #  * Redistributions in binary form must reproduce the above copyright notice,
  14 #    this list of conditions and the following disclaimer in the documentation
  15 #    and/or other materials provided with the distribution.
  16 #
  17 #  * Neither the name of the copyright holder nor the names of its
  18 #    contributors may be used to endorse or promote products derived from
  19 #    this software without specific prior written permission.
  20 #
  21 #  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  22 #  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  23 #  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
  24 #  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
  25 #  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  26 #  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
  27 #  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  28 #  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
  29 #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
  30 
  31 
  32 from celery import signals, states
  33 
  34 import elasticapm
  35 from elasticapm.conf import constants
  36 from elasticapm.utils import get_name_from_func
  37 
  38 
  39 class CeleryFilter(object):
  40     def filter(self, record):
  41         if record.funcName in ("_log_error",):
  42             return 0
  43         else:
  44             return 1
  45 
  46 
  47 def register_exception_tracking(client):
  48     dispatch_uid = "elasticapm-exc-tracking"
  49 
  50     def process_failure_signal(sender, task_id, exception, args, kwargs, traceback, einfo, **kw):
  51         client.capture_exception(
  52             extra={"task_id": task_id, "task": sender, "args": args, "kwargs": kwargs}, handled=False
  53         )
  54 
  55     signals.task_failure.disconnect(process_failure_signal, dispatch_uid=dispatch_uid)
  56     signals.task_failure.connect(process_failure_signal, weak=False, dispatch_uid=dispatch_uid)
  57     _register_worker_signals(client)
  58 
  59 
  60 def register_instrumentation(client):
  61     def begin_transaction(*args, **kwargs):
  62         client.begin_transaction("celery")
  63 
  64     def end_transaction(task_id, task, *args, **kwargs):
  65         name = get_name_from_func(task)
  66         state = kwargs.get("state", "None")
  67         if state == states.SUCCESS:
  68             outcome = constants.OUTCOME.SUCCESS
  69         elif state in states.EXCEPTION_STATES:
  70             outcome = constants.OUTCOME.FAILURE
  71         else:
  72             outcome = constants.OUTCOME.UNKNOWN
  73         elasticapm.set_transaction_outcome(outcome, override=False)
  74         client.end_transaction(name, state)
  75 
  76     dispatch_uid = "elasticapm-tracing-%s"
  77 
  78     # unregister any existing clients
  79     signals.task_prerun.disconnect(begin_transaction, dispatch_uid=dispatch_uid % "prerun")
  80     signals.task_postrun.disconnect(end_transaction, dispatch_uid=dispatch_uid % "postrun")
  81 
  82     # register for this client
  83     signals.task_prerun.connect(begin_transaction, dispatch_uid=dispatch_uid % "prerun", weak=False)
  84     signals.task_postrun.connect(end_transaction, weak=False, dispatch_uid=dispatch_uid % "postrun")
  85     _register_worker_signals(client)
  86 
  87 
  88 def _register_worker_signals(client):
  89     def worker_shutdown(*args, **kwargs):
  90         client.close()
  91 
  92     def connect_worker_process_init(*args, **kwargs):
  93         signals.worker_process_shutdown.connect(worker_shutdown, dispatch_uid="elasticapm-shutdown-worker", weak=False)
  94 
  95     signals.worker_init.connect(
  96         connect_worker_process_init, dispatch_uid="elasticapm-connect-start-threads", weak=False
  97     )
  98
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: 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 @@
       _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 @@
       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)
  
verification_info: {"golden_diff": "diff --git a/elasticapm/contrib/celery/__init__.py b/elasticapm/contrib/celery/__init__.py\n--- a/elasticapm/contrib/celery/__init__.py\n+++ b/elasticapm/contrib/celery/__init__.py\n@@ -27,13 +27,15 @@\n #  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\n #  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\n #  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n-\n+from contextlib import suppress\n \n from celery import signals, states\n \n import elasticapm\n from elasticapm.conf import constants\n+from elasticapm.traces import execution_context\n from elasticapm.utils import get_name_from_func\n+from elasticapm.utils.disttracing import TraceParent\n \n \n class CeleryFilter(object):\n@@ -57,9 +59,41 @@\n     _register_worker_signals(client)\n \n \n+def set_celery_headers(headers=None, **kwargs):\n+    \"\"\"\n+    Add elasticapm specific information to celery headers\n+    \"\"\"\n+    headers = {} if headers is None else headers\n+\n+    transaction = execution_context.get_transaction()\n+    if transaction is not None:\n+        trace_parent = transaction.trace_parent\n+        trace_parent_string = trace_parent.to_string()\n+\n+        headers.update({\"elasticapm\": {\"trace_parent_string\": trace_parent_string}})\n+\n+\n+def get_trace_parent(celery_task):\n+    \"\"\"\n+    Return a trace parent contained in the request headers of a Celery Task object or None\n+    \"\"\"\n+    trace_parent = None\n+    with suppress(AttributeError, KeyError, TypeError):\n+        if celery_task.request.headers is not None:\n+            trace_parent_string = celery_task.request.headers[\"elasticapm\"][\"trace_parent_string\"]\n+            trace_parent = TraceParent.from_string(trace_parent_string)\n+        else:\n+            trace_parent_string = celery_task.request.elasticapm[\"trace_parent_string\"]\n+            trace_parent = TraceParent.from_string(trace_parent_string)\n+    return trace_parent\n+\n+\n def register_instrumentation(client):\n     def begin_transaction(*args, **kwargs):\n-        client.begin_transaction(\"celery\")\n+        task = kwargs[\"task\"]\n+\n+        trace_parent = get_trace_parent(task)\n+        client.begin_transaction(\"celery\", trace_parent=trace_parent)\n \n     def end_transaction(task_id, task, *args, **kwargs):\n         name = get_name_from_func(task)\n@@ -76,10 +110,12 @@\n     dispatch_uid = \"elasticapm-tracing-%s\"\n \n     # unregister any existing clients\n+    signals.before_task_publish.disconnect(set_celery_headers, dispatch_uid=dispatch_uid % \"before-publish\")\n     signals.task_prerun.disconnect(begin_transaction, dispatch_uid=dispatch_uid % \"prerun\")\n     signals.task_postrun.disconnect(end_transaction, dispatch_uid=dispatch_uid % \"postrun\")\n \n     # register for this client\n+    signals.before_task_publish.connect(set_celery_headers, dispatch_uid=dispatch_uid % \"before-publish\")\n     signals.task_prerun.connect(begin_transaction, dispatch_uid=dispatch_uid % \"prerun\", weak=False)\n     signals.task_postrun.connect(end_transaction, weak=False, dispatch_uid=dispatch_uid % \"postrun\")\n     _register_worker_signals(client)\n", "issue": "Django Celery trace.id integration\n**Is your feature request related to a problem? Please describe.**\r\nAs of now, it is impossible to keep the same trace id between a Django view and a Celery task launched from the same view.\r\n\r\n**Describe the solution you'd like**\r\nProvide a way to easily pass a trace parent string to the Celery task. Preferably via it's headers field (introduced in Celery 3.1).\r\n\r\n### What would it looks likes\r\n**User side code (Django view)**\r\n```python\r\ndef get(self, request):\r\n    transaction = execution_context.get_transaction()\r\n    trace_parent = transaction.trace_parent\r\n    trace_parent_string = trace_parent.to_string()\r\n    my_celery_task.apply_async(headers={\"elasticapm\": {\"trace_parent_string\": trace_parent_string} })\r\n```\r\n\r\n**Library side code (`elasticapm.contrib.celery.__init__.py`), rewrite of `begin_transaction()`, na\u00efve implementation**\r\n```python\r\ndef begin_transaction(*args, **kwargs):\r\n    trace_parent = None\r\n    try:\r\n        trace_parent_string = kwargs[\"task\"].request.headers[\"elasticapm\"][\"trace_parent_string\"]\r\n        trace_parent = TraceParent.from_string(trace_parent_string)\r\n    except:\r\n        pass\r\n    client.begin_transaction(\"celery\", trace_parent=trace_parent)\r\n```\r\n\r\n- **Why using Celery headers field ?** It seems the most unobstrusive way of doing it.\r\n- **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.\r\n\r\n**Additional context**\r\n**For anyone wanting to try it, BEWARE !!** There is a Celery [bug](https://github.com/celery/celery/issues/4875) concerning it's headers field.\r\nYou might have to do this:\r\n```python\r\nmy_celery_task.apply_async(headers={\"headers\": {\"elasticapm\": {\"trace_parent_string\": trace_parent_string} } })\r\n```\r\n\r\nEdits: fixed code error/typos\n", "before_files": [{"content": "#  BSD 3-Clause License\n#\n#  Copyright (c) 2012, the Sentry Team, see AUTHORS for more details\n#  Copyright (c) 2019, Elasticsearch BV\n#  All rights reserved.\n#\n#  Redistribution and use in source and binary forms, with or without\n#  modification, are permitted provided that the following conditions are met:\n#\n#  * Redistributions of source code must retain the above copyright notice, this\n#    list of conditions and the following disclaimer.\n#\n#  * Redistributions in binary form must reproduce the above copyright notice,\n#    this list of conditions and the following disclaimer in the documentation\n#    and/or other materials provided with the distribution.\n#\n#  * Neither the name of the copyright holder nor the names of its\n#    contributors may be used to endorse or promote products derived from\n#    this software without specific prior written permission.\n#\n#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\n#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\n#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\n#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\n#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\n#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\n\n\nfrom celery import signals, states\n\nimport elasticapm\nfrom elasticapm.conf import constants\nfrom elasticapm.utils import get_name_from_func\n\n\nclass CeleryFilter(object):\n    def filter(self, record):\n        if record.funcName in (\"_log_error\",):\n            return 0\n        else:\n            return 1\n\n\ndef register_exception_tracking(client):\n    dispatch_uid = \"elasticapm-exc-tracking\"\n\n    def process_failure_signal(sender, task_id, exception, args, kwargs, traceback, einfo, **kw):\n        client.capture_exception(\n            extra={\"task_id\": task_id, \"task\": sender, \"args\": args, \"kwargs\": kwargs}, handled=False\n        )\n\n    signals.task_failure.disconnect(process_failure_signal, dispatch_uid=dispatch_uid)\n    signals.task_failure.connect(process_failure_signal, weak=False, dispatch_uid=dispatch_uid)\n    _register_worker_signals(client)\n\n\ndef register_instrumentation(client):\n    def begin_transaction(*args, **kwargs):\n        client.begin_transaction(\"celery\")\n\n    def end_transaction(task_id, task, *args, **kwargs):\n        name = get_name_from_func(task)\n        state = kwargs.get(\"state\", \"None\")\n        if state == states.SUCCESS:\n            outcome = constants.OUTCOME.SUCCESS\n        elif state in states.EXCEPTION_STATES:\n            outcome = constants.OUTCOME.FAILURE\n        else:\n            outcome = constants.OUTCOME.UNKNOWN\n        elasticapm.set_transaction_outcome(outcome, override=False)\n        client.end_transaction(name, state)\n\n    dispatch_uid = \"elasticapm-tracing-%s\"\n\n    # unregister any existing clients\n    signals.task_prerun.disconnect(begin_transaction, dispatch_uid=dispatch_uid % \"prerun\")\n    signals.task_postrun.disconnect(end_transaction, dispatch_uid=dispatch_uid % \"postrun\")\n\n    # register for this client\n    signals.task_prerun.connect(begin_transaction, dispatch_uid=dispatch_uid % \"prerun\", weak=False)\n    signals.task_postrun.connect(end_transaction, weak=False, dispatch_uid=dispatch_uid % \"postrun\")\n    _register_worker_signals(client)\n\n\ndef _register_worker_signals(client):\n    def worker_shutdown(*args, **kwargs):\n        client.close()\n\n    def connect_worker_process_init(*args, **kwargs):\n        signals.worker_process_shutdown.connect(worker_shutdown, dispatch_uid=\"elasticapm-shutdown-worker\", weak=False)\n\n    signals.worker_init.connect(\n        connect_worker_process_init, dispatch_uid=\"elasticapm-connect-start-threads\", weak=False\n    )\n", "path": "elasticapm/contrib/celery/__init__.py"}], "after_files": [{"content": "#  BSD 3-Clause License\n#\n#  Copyright (c) 2012, the Sentry Team, see AUTHORS for more details\n#  Copyright (c) 2019, Elasticsearch BV\n#  All rights reserved.\n#\n#  Redistribution and use in source and binary forms, with or without\n#  modification, are permitted provided that the following conditions are met:\n#\n#  * Redistributions of source code must retain the above copyright notice, this\n#    list of conditions and the following disclaimer.\n#\n#  * Redistributions in binary form must reproduce the above copyright notice,\n#    this list of conditions and the following disclaimer in the documentation\n#    and/or other materials provided with the distribution.\n#\n#  * Neither the name of the copyright holder nor the names of its\n#    contributors may be used to endorse or promote products derived from\n#    this software without specific prior written permission.\n#\n#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n#  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\n#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\n#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\n#  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\n#  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\n#  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\nfrom contextlib import suppress\n\nfrom celery import signals, states\n\nimport elasticapm\nfrom elasticapm.conf import constants\nfrom elasticapm.traces import execution_context\nfrom elasticapm.utils import get_name_from_func\nfrom elasticapm.utils.disttracing import TraceParent\n\n\nclass CeleryFilter(object):\n    def filter(self, record):\n        if record.funcName in (\"_log_error\",):\n            return 0\n        else:\n            return 1\n\n\ndef register_exception_tracking(client):\n    dispatch_uid = \"elasticapm-exc-tracking\"\n\n    def process_failure_signal(sender, task_id, exception, args, kwargs, traceback, einfo, **kw):\n        client.capture_exception(\n            extra={\"task_id\": task_id, \"task\": sender, \"args\": args, \"kwargs\": kwargs}, handled=False\n        )\n\n    signals.task_failure.disconnect(process_failure_signal, dispatch_uid=dispatch_uid)\n    signals.task_failure.connect(process_failure_signal, weak=False, dispatch_uid=dispatch_uid)\n    _register_worker_signals(client)\n\n\ndef set_celery_headers(headers=None, **kwargs):\n    \"\"\"\n    Add elasticapm specific information to celery headers\n    \"\"\"\n    headers = {} if headers is None else headers\n\n    transaction = execution_context.get_transaction()\n    if transaction is not None:\n        trace_parent = transaction.trace_parent\n        trace_parent_string = trace_parent.to_string()\n\n        headers.update({\"elasticapm\": {\"trace_parent_string\": trace_parent_string}})\n\n\ndef get_trace_parent(celery_task):\n    \"\"\"\n    Return a trace parent contained in the request headers of a Celery Task object or None\n    \"\"\"\n    trace_parent = None\n    with suppress(AttributeError, KeyError, TypeError):\n        if celery_task.request.headers is not None:\n            trace_parent_string = celery_task.request.headers[\"elasticapm\"][\"trace_parent_string\"]\n            trace_parent = TraceParent.from_string(trace_parent_string)\n        else:\n            trace_parent_string = celery_task.request.elasticapm[\"trace_parent_string\"]\n            trace_parent = TraceParent.from_string(trace_parent_string)\n    return trace_parent\n\n\ndef register_instrumentation(client):\n    def begin_transaction(*args, **kwargs):\n        task = kwargs[\"task\"]\n\n        trace_parent = get_trace_parent(task)\n        client.begin_transaction(\"celery\", trace_parent=trace_parent)\n\n    def end_transaction(task_id, task, *args, **kwargs):\n        name = get_name_from_func(task)\n        state = kwargs.get(\"state\", \"None\")\n        if state == states.SUCCESS:\n            outcome = constants.OUTCOME.SUCCESS\n        elif state in states.EXCEPTION_STATES:\n            outcome = constants.OUTCOME.FAILURE\n        else:\n            outcome = constants.OUTCOME.UNKNOWN\n        elasticapm.set_transaction_outcome(outcome, override=False)\n        client.end_transaction(name, state)\n\n    dispatch_uid = \"elasticapm-tracing-%s\"\n\n    # unregister any existing clients\n    signals.before_task_publish.disconnect(set_celery_headers, dispatch_uid=dispatch_uid % \"before-publish\")\n    signals.task_prerun.disconnect(begin_transaction, dispatch_uid=dispatch_uid % \"prerun\")\n    signals.task_postrun.disconnect(end_transaction, dispatch_uid=dispatch_uid % \"postrun\")\n\n    # register for this client\n    signals.before_task_publish.connect(set_celery_headers, dispatch_uid=dispatch_uid % \"before-publish\")\n    signals.task_prerun.connect(begin_transaction, dispatch_uid=dispatch_uid % \"prerun\", weak=False)\n    signals.task_postrun.connect(end_transaction, weak=False, dispatch_uid=dispatch_uid % \"postrun\")\n    _register_worker_signals(client)\n\n\ndef _register_worker_signals(client):\n    def worker_shutdown(*args, **kwargs):\n        client.close()\n\n    def connect_worker_process_init(*args, **kwargs):\n        signals.worker_process_shutdown.connect(worker_shutdown, dispatch_uid=\"elasticapm-shutdown-worker\", weak=False)\n\n    signals.worker_init.connect(\n        connect_worker_process_init, dispatch_uid=\"elasticapm-connect-start-threads\", weak=False\n    )\n", "path": "elasticapm/contrib/celery/__init__.py"}]}