mirror of
https://github.com/open-telemetry/opentelemetry-python-contrib.git
synced 2025-08-01 09:13:23 +08:00
107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
from weakref import WeakValueDictionary
|
|
|
|
from .constants import CTX_KEY
|
|
|
|
|
|
def tags_from_context(context):
|
|
"""Helper to extract meta values from a Celery Context"""
|
|
tag_keys = (
|
|
'compression', 'correlation_id', 'countdown', 'delivery_info', 'eta',
|
|
'exchange', 'expires', 'hostname', 'id', 'priority', 'queue', 'reply_to',
|
|
'retries', 'routing_key', 'serializer', 'timelimit', 'origin', 'state',
|
|
)
|
|
|
|
tags = {}
|
|
for key in tag_keys:
|
|
value = context.get(key)
|
|
|
|
# Skip this key if it is not set
|
|
if value is None or value == '':
|
|
continue
|
|
|
|
# Skip `timelimit` if it is not set (it's default/unset value is a
|
|
# tuple or a list of `None` values
|
|
if key == 'timelimit' and value in [(None, None), [None, None]]:
|
|
continue
|
|
|
|
# Skip `retries` if it's value is `0`
|
|
if key == 'retries' and value == 0:
|
|
continue
|
|
|
|
# Celery 4.0 uses `origin` instead of `hostname`; this change preserves
|
|
# the same name for the tag despite Celery version
|
|
if key == 'origin':
|
|
key = 'hostname'
|
|
|
|
# prefix the tag as 'celery'
|
|
tag_name = 'celery.{}'.format(key)
|
|
tags[tag_name] = value
|
|
return tags
|
|
|
|
|
|
def attach_span(task, task_id, span, is_publish=False):
|
|
"""Helper to propagate a `Span` for the given `Task` instance. This
|
|
function uses a `WeakValueDictionary` that stores a Datadog Span using
|
|
the `(task_id, is_publish)` as a key. This is useful when information must be
|
|
propagated from one Celery signal to another.
|
|
|
|
DEV: We use (task_id, is_publish) for the key to ensure that publishing a
|
|
task from within another task does not cause any conflicts.
|
|
|
|
This mostly happens when either a task fails and a retry policy is in place,
|
|
or when a task is manually retries (e.g. `task.retry()`), we end up trying
|
|
to publish a task with the same id as the task currently running.
|
|
|
|
Previously publishing the new task would overwrite the existing `celery.run` span
|
|
in the `weak_dict` causing that span to be forgotten and never finished.
|
|
|
|
NOTE: We cannot test for this well yet, because we do not run a celery worker,
|
|
and cannot run `task.apply_async()`
|
|
"""
|
|
weak_dict = getattr(task, CTX_KEY, None)
|
|
if weak_dict is None:
|
|
weak_dict = WeakValueDictionary()
|
|
setattr(task, CTX_KEY, weak_dict)
|
|
|
|
weak_dict[(task_id, is_publish)] = span
|
|
|
|
|
|
def detach_span(task, task_id, is_publish=False):
|
|
"""Helper to remove a `Span` in a Celery task when it's propagated.
|
|
This function handles tasks where the `Span` is not attached.
|
|
"""
|
|
weak_dict = getattr(task, CTX_KEY, None)
|
|
if weak_dict is None:
|
|
return
|
|
|
|
# DEV: See note in `attach_span` for key info
|
|
weak_dict.pop((task_id, is_publish), None)
|
|
|
|
|
|
def retrieve_span(task, task_id, is_publish=False):
|
|
"""Helper to retrieve an active `Span` stored in a `Task`
|
|
instance
|
|
"""
|
|
weak_dict = getattr(task, CTX_KEY, None)
|
|
if weak_dict is None:
|
|
return
|
|
else:
|
|
# DEV: See note in `attach_span` for key info
|
|
return weak_dict.get((task_id, is_publish))
|
|
|
|
|
|
def retrieve_task_id(context):
|
|
"""Helper to retrieve the `Task` identifier from the message `body`.
|
|
This helper supports Protocol Version 1 and 2. The Protocol is well
|
|
detailed in the official documentation:
|
|
http://docs.celeryproject.org/en/latest/internals/protocol.html
|
|
"""
|
|
headers = context.get('headers')
|
|
body = context.get('body')
|
|
if headers:
|
|
# Protocol Version 2 (default from Celery 4.0)
|
|
return headers.get('id')
|
|
else:
|
|
# Protocol Version 1
|
|
return body.get('id')
|