perf(celery): reduce tracing overheads (#2579)

* perf(celery): reduce tracing overheads

This change tries to reduce the overheads introduced by the Celery
instrumentation. Tags are set straight from the context extraction
rather than via an intermediate dictionary. Tag names are pre-computed
to avoid the cost of string concats at runtime.

* Update ddtrace/contrib/celery/utils.py

Co-authored-by: Julien Danjou <[email protected]>

* guard del operation

Co-authored-by: Julien Danjou <[email protected]>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 3 people

ddtrace/contrib/celery/signals.py

@@ -14,7 +14,7 @@
 from .utils import detach_span
 from .utils import retrieve_span
 from .utils import retrieve_task_id
-from .utils import tags_from_context
+from .utils import set_tags_from_context
 
 
 log = get_logger(__name__)
@@ -70,8 +70,8 @@ def trace_postrun(*args, **kwargs):
     else:
         # request context tags
         span.set_tag(c.TASK_TAG_KEY, c.TASK_RUN)
-        span.set_tags(tags_from_context(kwargs))
-        span.set_tags(tags_from_context(task.request))
+        set_tags_from_context(span, kwargs)
+        set_tags_from_context(span, task.request.__dict__)
         span.finish()
         detach_span(task, task_id)
 
@@ -107,7 +107,7 @@ def trace_before_publish(*args, **kwargs):
     span.set_tag(SPAN_MEASURED_KEY)
     span.set_tag(c.TASK_TAG_KEY, c.TASK_APPLY_ASYNC)
     span.set_tag("celery.id", task_id)
-    span.set_tags(tags_from_context(kwargs))
+    set_tags_from_context(span, kwargs)
 
     # Note: adding tags from `traceback` or `state` calls will make an
     # API call to the backend for the properties so we should rely

ddtrace/contrib/celery/utils.py

@@ -1,57 +1,59 @@
+from typing import Any
+from typing import Dict
 from weakref import WeakValueDictionary
 
+from ddtrace.span import Span
+
 from .constants import CTX_KEY
 
 
-def tags_from_context(context):
+TAG_KEYS = frozenset(
+    [
+        ("compression", "celery.compression"),
+        ("correlation_id", "celery.correlation_id"),
+        ("countdown", "celery.countdown"),
+        ("delivery_info", "celery.delivery_info"),
+        ("eta", "celery.eta"),
+        ("exchange", "celery.exchange"),
+        ("expires", "celery.expires"),
+        ("hostname", "celery.hostname"),
+        ("id", "celery.id"),
+        ("priority", "celery.priority"),
+        ("queue", "celery.queue"),
+        ("reply_to", "celery.reply_to"),
+        ("retries", "celery.retries"),
+        ("routing_key", "celery.routing_key"),
+        ("serializer", "celery.serializer"),
+        ("timelimit", "celery.timelimit"),
+        # Celery 4.0 uses `origin` instead of `hostname`; this change preserves
+        # the same name for the tag despite Celery version
+        ("origin", "celery.hostname"),
+        ("state", "celery.state"),
+    ]
+)
+
+
+def set_tags_from_context(span, context):
+    # type: (Span, Dict[str, Any]) -> None
     """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:
+
+    for key, tag_name 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
+        # Skip `timelimit` if it is not set (its default/unset value is a
         # tuple or a list of `None` values
-        if key == "timelimit" and value in [(None, None), [None, None]]:
+        if key == "timelimit" and all(_ is None for _ in value):
             continue
 
-        # Skip `retries` if it's value is `0`
+        # Skip `retries` if its 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
+        span.set_tag(tag_name, value)
 
 
 def attach_span(task, task_id, span, is_publish=False):
@@ -64,7 +66,7 @@ def attach_span(task, task_id, span, is_publish=False):
          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
+         or when a task is manually retried (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
@@ -90,7 +92,10 @@ def detach_span(task, task_id, is_publish=False):
         return
 
     # DEV: See note in `attach_span` for key info
-    weak_dict.pop((task_id, is_publish), None)
+    try:
+        del weak_dict[(task_id, is_publish)]
+    except KeyError:
+        pass
 
 
 def retrieve_span(task, task_id, is_publish=False):

releasenotes/notes/perf-celery-reduced-overheads-6afe03acf1af703f.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    Performance of the Celery integration has been improved.

tests/contrib/celery/test_utils.py

@@ -4,7 +4,8 @@
 from ddtrace.contrib.celery.utils import detach_span
 from ddtrace.contrib.celery.utils import retrieve_span
 from ddtrace.contrib.celery.utils import retrieve_task_id
-from ddtrace.contrib.celery.utils import tags_from_context
+from ddtrace.contrib.celery.utils import set_tags_from_context
+from ddtrace.span import Span
 
 from .base import CeleryBaseTestCase
 
@@ -29,35 +30,42 @@ def test_tags_from_context(self):
             "custom_meta": "custom_value",
         }
 
-        metas = tags_from_context(context)
+        span = Span(None, "test")
+        set_tags_from_context(span, context)
+        metas = span.meta
+        metrics = span.metrics
+        sentinel = object()
         assert metas["celery.correlation_id"] == "44b7f305"
         assert metas["celery.delivery_info"] == '{"eager": "True"}'
         assert metas["celery.eta"] == "soon"
         assert metas["celery.expires"] == "later"
         assert metas["celery.hostname"] == "localhost"
         assert metas["celery.id"] == "44b7f305"
         assert metas["celery.reply_to"] == "44b7f305"
-        assert metas["celery.retries"] == 4
-        assert metas["celery.timelimit"] == ("now", "later")
-        assert metas.get("custom_meta", None) is None
+        assert metrics["celery.retries"] == 4
+        assert metas["celery.timelimit"] == "('now', 'later')"
+        assert metas.get("custom_meta", sentinel) is sentinel
+        assert metrics.get("custom_metric", sentinel) is sentinel
 
     def test_tags_from_context_empty_keys(self):
         # it should not extract empty keys
+        span = Span(None, "test")
         context = {
             "correlation_id": None,
             "exchange": "",
             "timelimit": (None, None),
             "retries": 0,
         }
+        tags = span.meta
 
-        tags = tags_from_context(context)
+        set_tags_from_context(span, context)
         assert {} == tags
         # edge case: `timelimit` can also be a list of None values
         context = {
             "timelimit": [None, None],
         }
 
-        tags = tags_from_context(context)
+        set_tags_from_context(span, context)
         assert {} == tags
 
     def test_span_propagation(self):