Test commit

Author: rbagd

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/__init__.py

@@ -1,3 +1,4 @@
+# type: ignore
 # Copyright The OpenTelemetry Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -60,7 +61,7 @@ def add(x, y):
 """
 
 import logging
-from timeit import default_timer
+import time
 from typing import Collection, Iterable
 
 from billiard import VERSION
@@ -76,6 +77,7 @@ def add(x, y):
 from opentelemetry.metrics import get_meter
 from opentelemetry.propagate import extract, inject
 from opentelemetry.propagators.textmap import Getter
+from opentelemetry.semconv._incubating.metrics import messaging_metrics
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace.status import Status, StatusCode
 
@@ -96,6 +98,12 @@ def add(x, y):
 _TASK_REVOKED_TERMINATED_SIGNAL_KEY = "celery.terminated.signal"
 _TASK_NAME_KEY = "celery.task_name"
 
+# Metric names
+_TASK_COUNT_ACTIVE = "messaging.client.active_tasks"
+_TASK_COUNT_PREFETCHED = "messaging.client.prefetched_tasks"
+_TASK_PROCESSING_TIME = messaging_metrics.MESSAGING_PROCESS_DURATION
+_TASK_PREFETCH_TIME = "messaging.prefetch.duration"
+
 
 class CeleryGetter(Getter):
     def get(self, carrier, key):
@@ -113,9 +121,37 @@ def keys(self, carrier):
 celery_getter = CeleryGetter()
 
 
+class TaskDurationTracker:
+    def __init__(self, metrics):
+        self.metrics = metrics
+        self.tracker = {}
+
+    def record_start(self, key, step):
+        self.tracker.setdefault(key, {})[step] = time.perf_counter()
+
+    def record_finish(self, key, metric_name, attributes):
+
+        try:
+            time_elapsed = self._time_elapsed(key, metric_name)
+            self.metrics[metric_name].record(
+                max(0, time_elapsed), attributes=attributes
+            )
+        except KeyError:
+            logger.warning("Failed to record %s for task %s", metric_name, key)
+
+    def _time_elapsed(self, key, step):
+        end_time = time.perf_counter()
+        try:
+            start_time = self.tracker.get(key, {}).pop(step)
+            time_elapsed = end_time - start_time
+            return time_elapsed
+        finally:
+            # Cleanup operation
+            if key in self.tracker and not self.tracker.get(key):
+                self.tracker.pop(key)
+
+
 class CeleryInstrumentor(BaseInstrumentor):
-    metrics = None
-    task_id_to_start_time = {}
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
@@ -139,60 +175,89 @@ def _instrument(self, **kwargs):
             schema_url="https://opentelemetry.io/schemas/1.11.0",
         )
 
-        self.create_celery_metrics(meter)
+        self.metrics = _create_celery_worker_metrics(meter)
+        self.time_tracker = TaskDurationTracker(self.metrics)
 
+        signals.task_received.connect(self._trace_received, weak=False)
         signals.task_prerun.connect(self._trace_prerun, weak=False)
         signals.task_postrun.connect(self._trace_postrun, weak=False)
-        signals.before_task_publish.connect(
-            self._trace_before_publish, weak=False
-        )
-        signals.after_task_publish.connect(
-            self._trace_after_publish, weak=False
-        )
+        signals.before_task_publish.connect(self._trace_before_publish, weak=False)
+        signals.after_task_publish.connect(self._trace_after_publish, weak=False)
         signals.task_failure.connect(self._trace_failure, weak=False)
         signals.task_retry.connect(self._trace_retry, weak=False)
 
     def _uninstrument(self, **kwargs):
+        signals.task_received.disconnect(self._trace_received)
         signals.task_prerun.disconnect(self._trace_prerun)
         signals.task_postrun.disconnect(self._trace_postrun)
         signals.before_task_publish.disconnect(self._trace_before_publish)
         signals.after_task_publish.disconnect(self._trace_after_publish)
         signals.task_failure.disconnect(self._trace_failure)
         signals.task_retry.disconnect(self._trace_retry)
 
+    def _trace_received(self, *args, **kwargs):
+        """
+        On prerun signal, task is prefetched and prefetch timer starts
+        """
+
+        request = utils.retrieve_task_from_request(kwargs)
+
+        metrics_attributes = utils.get_metrics_attributes_from_request(request)
+        self.metrics[_TASK_COUNT_PREFETCHED].add(1, attributes=metrics_attributes)
+        self.time_tracker.record_start(request.task_id, _TASK_PREFETCH_TIME)
+
     def _trace_prerun(self, *args, **kwargs):
+        """
+        On prerun signal, task is no longer prefetched, and execution timer
+        starts along with the task span
+        """
+
         task = utils.retrieve_task(kwargs)
         task_id = utils.retrieve_task_id(kwargs)
 
         if task is None or task_id is None:
             return
 
-        self.update_task_duration_time(task_id)
+        metrics_attributes = utils.get_metrics_attributes_from_task(task)
+        self.metrics[_TASK_COUNT_PREFETCHED].add(-1, attributes=metrics_attributes)
+        self.time_tracker.record_finish(
+            task_id, _TASK_PREFETCH_TIME, metrics_attributes
+        )
+        self.time_tracker.record_start(task_id, _TASK_PROCESSING_TIME)
+
         request = task.request
         tracectx = extract(request, getter=celery_getter) or None
         token = context_api.attach(tracectx) if tracectx is not None else None
 
-        logger.debug("prerun signal start task_id=%s", task_id)
-
         operation_name = f"{_TASK_RUN}/{task.name}"
         span = self._tracer.start_span(
             operation_name, context=tracectx, kind=trace.SpanKind.CONSUMER
         )
 
         activation = trace.use_span(span, end_on_exit=True)
-        activation.__enter__()  # pylint: disable=E1101
+        activation.__enter__()
         utils.attach_context(task, task_id, span, activation, token)
 
+        self.metrics[_TASK_COUNT_ACTIVE].add(1, attributes=metrics_attributes)
+
     def _trace_postrun(self, *args, **kwargs):
+        """
+        On postrun signal, task is no longer being executed
+        """
+
         task = utils.retrieve_task(kwargs)
         task_id = utils.retrieve_task_id(kwargs)
 
         if task is None or task_id is None:
             return
 
-        logger.debug("postrun signal task_id=%s", task_id)
+        metrics_attributes = utils.get_metrics_attributes_from_task(task)
+        self.metrics[_TASK_COUNT_ACTIVE].add(-1, attributes=metrics_attributes)
+        self.time_tracker.record_finish(
+            task_id, _TASK_PROCESSING_TIME, metrics_attributes
+        )
 
-        # retrieve and finish the Span
+        # Retrieve and finish the Span
         ctx = utils.retrieve_context(task, task_id)
 
         if ctx is None:
@@ -201,7 +266,7 @@ def _trace_postrun(self, *args, **kwargs):
 
         span, activation, token = ctx
 
-        # request context tags
+        # Request context tags
         if span.is_recording():
             span.set_attribute(_TASK_TAG_KEY, _TASK_RUN)
             utils.set_attributes_from_context(span, kwargs)
@@ -210,10 +275,8 @@ def _trace_postrun(self, *args, **kwargs):
 
         activation.__exit__(None, None, None)
         utils.detach_context(task, task_id)
-        self.update_task_duration_time(task_id)
-        labels = {"task": task.name, "worker": task.request.hostname}
-        self._record_histograms(task_id, labels)
-        # if the process sending the task is not instrumented
+
+        # If the process sending the task is not instrumented,
         # there's no incoming context and no token to detach
         if token is not None:
             context_api.detach(token)
@@ -226,18 +289,16 @@ def _trace_before_publish(self, *args, **kwargs):
             return
 
         if task is None:
-            # task is an anonymous task send using send_task or using canvas workflow
+            # task is an anonymous task sent using send_task or using canvas workflow
             # Signatures() to send to a task not in the current processes dependency
             # tree
             task_name = kwargs.get("sender", "unknown")
         else:
             task_name = task.name
         operation_name = f"{_TASK_APPLY_ASYNC}/{task_name}"
-        span = self._tracer.start_span(
-            operation_name, kind=trace.SpanKind.PRODUCER
-        )
+        span = self._tracer.start_span(operation_name, kind=trace.SpanKind.PRODUCER)
 
-        # apply some attributes here because most of the data is not available
+        # Apply some attributes here because most of the data is not available
         if span.is_recording():
             span.set_attribute(_TASK_TAG_KEY, _TASK_APPLY_ASYNC)
             span.set_attribute(SpanAttributes.MESSAGING_MESSAGE_ID, task_id)
@@ -247,9 +308,7 @@ def _trace_before_publish(self, *args, **kwargs):
         activation = trace.use_span(span, end_on_exit=True)
         activation.__enter__()  # pylint: disable=E1101
 
-        utils.attach_context(
-            task, task_id, span, activation, None, is_publish=True
-        )
+        utils.attach_context(task, task_id, span, activation, None, is_publish=True)
 
         headers = kwargs.get("headers")
         if headers:
@@ -263,7 +322,7 @@ def _trace_after_publish(*args, **kwargs):
         if task is None or task_id is None:
             return
 
-        # retrieve and finish the Span
+        # Retrieve and finish the Span
         ctx = utils.retrieve_context(task, task_id, is_publish=True)
 
         if ctx is None:
@@ -272,7 +331,7 @@ def _trace_after_publish(*args, **kwargs):
 
         _, activation, _ = ctx
 
-        activation.__exit__(None, None, None)  # pylint: disable=E1101
+        activation.__exit__(None, None, None)
         utils.detach_context(task, task_id, is_publish=True)
 
     @staticmethod
@@ -345,29 +404,30 @@ def _trace_retry(*args, **kwargs):
         # something that isn't an `Exception`
         span.set_attribute(_TASK_RETRY_REASON_KEY, str(reason))
 
-    def update_task_duration_time(self, task_id):
-        cur_time = default_timer()
-        task_duration_time_until_now = (
-            cur_time - self.task_id_to_start_time[task_id]
-            if task_id in self.task_id_to_start_time
-            else cur_time
-        )
-        self.task_id_to_start_time[task_id] = task_duration_time_until_now
-
-    def _record_histograms(self, task_id, metric_attributes):
-        if task_id is None:
-            return
-
-        self.metrics["flower.task.runtime.seconds"].record(
-            self.task_id_to_start_time.get(task_id),
-            attributes=metric_attributes,
-        )
 
-    def create_celery_metrics(self, meter) -> None:
-        self.metrics = {
-            "flower.task.runtime.seconds": meter.create_histogram(
-                name="flower.task.runtime.seconds",
-                unit="seconds",
-                description="The time it took to run the task.",
-            )
-        }
+def _create_celery_worker_metrics(meter) -> None:
+
+    metrics = {
+        _TASK_PROCESSING_TIME: meter.create_histogram(
+            name=_TASK_PROCESSING_TIME,
+            unit="s",
+            description="The time it took to run the task.",
+        ),
+        _TASK_COUNT_ACTIVE: meter.create_up_down_counter(
+            name=_TASK_COUNT_ACTIVE,
+            unit="{message}",
+            description="Number of tasks currently being executed by the worker",
+        ),
+        _TASK_COUNT_PREFETCHED: meter.create_up_down_counter(
+            name=_TASK_COUNT_PREFETCHED,
+            unit="{message}",
+            description="Number of tasks prefetched by the worker",
+        ),
+        _TASK_PREFETCH_TIME: meter.create_histogram(
+            name=_TASK_PREFETCH_TIME,
+            unit="s",
+            description="The time the task spent in prefetch mode",
+        ),
+    }
+
+    return metrics

instrumentation/opentelemetry-instrumentation-celery/src/opentelemetry/instrumentation/celery/utils.py

@@ -1,3 +1,5 @@
+# type: ignore
+
 # Copyright The OpenTelemetry Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,6 +22,10 @@
 from celery import registry  # pylint: disable=no-name-in-module
 from celery.app.task import Task
 
+from opentelemetry.semconv._incubating.attributes.messaging_attributes import (
+    MESSAGING_CLIENT_ID,
+    MESSAGING_OPERATION_NAME,
+)
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import Span
 
@@ -91,9 +97,7 @@ def set_attributes_from_context(span, context):
             routing_key = value.get("routing_key")
 
             if routing_key is not None:
-                span.set_attribute(
-                    SpanAttributes.MESSAGING_DESTINATION, routing_key
-                )
+                span.set_attribute(SpanAttributes.MESSAGING_DESTINATION, routing_key)
 
             value = str(value)
 
@@ -217,6 +221,15 @@ def retrieve_task_id(kwargs):
     return task_id
 
 
+def retrieve_task_from_request(kwargs):
+    # retry signal does not include task_id as argument so use request argument
+    request = kwargs.get("request")
+    if request is None:
+        logger.debug("Unable to retrieve the request from signal arguments")
+
+    return request
+
+
 def retrieve_task_id_from_request(kwargs):
     # retry signal does not include task_id as argument so use request argument
     request = kwargs.get("request")
@@ -250,3 +263,17 @@ def retrieve_reason(kwargs):
     if not reason:
         logger.debug("Unable to retrieve the retry reason")
     return reason
+
+
+def get_metrics_attributes_from_request(request):
+    return {
+        MESSAGING_OPERATION_NAME: request.task.name,
+        MESSAGING_CLIENT_ID: request.hostname,
+    }
+
+
+def get_metrics_attributes_from_task(task):
+    return {
+        MESSAGING_OPERATION_NAME: task.name,
+        MESSAGING_CLIENT_ID: task.request.hostname,
+    }

instrumentation/opentelemetry-instrumentation-celery/test-requirements-0.txt

@@ -2,7 +2,7 @@ amqp==5.2.0
 asgiref==3.8.1
 backports.zoneinfo==0.2.1
 billiard==4.2.0
-celery==5.3.6
+celery[redis]==5.3.6
 click==8.1.7
 click-didyoumean==0.3.0
 click-plugins==1.1.1
@@ -24,5 +24,6 @@ vine==5.1.0
 wcwidth==0.2.13
 wrapt==1.16.0
 zipp==3.19.2
+sqlalchemy
 -e opentelemetry-instrumentation
 -e instrumentation/opentelemetry-instrumentation-celery