update 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,41 +175,60 @@ 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
@@ -183,14 +238,24 @@ def _trace_prerun(self, *args, **kwargs):
         activation.__enter__()  # pylint: disable=E1101
         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
         ctx = utils.retrieve_context(task, task_id)
@@ -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)
@@ -345,29 +408,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_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",
+        ),
+        _TASK_PROCESSING_TIME: meter.create_histogram(
+            name=_TASK_PROCESSING_TIME,
+            unit="s",
+            description="The time it took to run the task.",
+        ),
+    }
+
+    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
 
@@ -217,6 +223,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 +265,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/tests/celery_test_tasks.py

@@ -1,3 +1,5 @@
+# type: ignore
+
 # Copyright The OpenTelemetry Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,14 +14,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import time
+
 from celery import Celery
 
 from opentelemetry import baggage
 
 
 class Config:
-    result_backend = "rpc"
-    broker_backend = "memory"
+    result_backend = "rpc://"
+    without_gossip = True
+    without_heartbeat = True
+    without_mingle = True
 
 
 app = Celery(broker="memory:///")
@@ -31,8 +37,14 @@ class CustomError(Exception):
 
 
 @app.task
-def task_add(num_a, num_b):
-    return num_a + num_b
+def task_add(x=1, y=2):
+    return x + y
+
+
+@app.task
+def task_sleep(sleep_time):
+    time.sleep(sleep_time)
+    return 1
 
 
 @app.task

instrumentation/opentelemetry-instrumentation-celery/tests/test_duplicate.py

@@ -26,5 +26,3 @@ def test_duplicate_instrumentaion(self):
         CeleryInstrumentor().uninstrument()
         self.assertIsNotNone(first.metrics)
         self.assertIsNotNone(second.metrics)
-        self.assertEqual(first.task_id_to_start_time, {})
-        self.assertEqual(second.task_id_to_start_time, {})