|
| 1 | +""" |
| 2 | +Instrument `taskiq`_ to trace Taskiq applications. |
| 3 | +
|
| 4 | +.. _taskiq: https://pypi.org/project/taskiq/ |
| 5 | +
|
| 6 | +Usage |
| 7 | +----- |
| 8 | +
|
| 9 | +* Run instrumented task |
| 10 | +
|
| 11 | +.. code:: python |
| 12 | +
|
| 13 | + import asyncio |
| 14 | +
|
| 15 | + from taskiq import InMemoryBroker, TaskiqEvents, TaskiqState |
| 16 | + from taskiq.instrumentation import TaskiqInstrumentor |
| 17 | +
|
| 18 | + broker = InMemoryBroker() |
| 19 | +
|
| 20 | + @broker.on_event(TaskiqEvents.WORKER_STARTUP) |
| 21 | + async def startup(state: TaskiqState) -> None: |
| 22 | + TaskiqInstrumentor().instrument() |
| 23 | +
|
| 24 | + @broker.task |
| 25 | + async def add(x, y): |
| 26 | + return x + y |
| 27 | +
|
| 28 | + async def main(): |
| 29 | + await broker.startup() |
| 30 | + await my_task.kiq(1, 2) |
| 31 | + await broker.shutdown() |
| 32 | +
|
| 33 | + if __name__ == "__main__": |
| 34 | + asyncio.run(main()) |
| 35 | +
|
| 36 | +API |
| 37 | +--- |
| 38 | +""" |
| 39 | + |
| 40 | +from __future__ import annotations |
| 41 | + |
| 42 | +import logging |
| 43 | +from typing import TYPE_CHECKING, Any, Callable, Collection, Optional |
| 44 | +from weakref import WeakSet as _WeakSet |
| 45 | + |
| 46 | +try: |
| 47 | + import opentelemetry # noqa: F401 |
| 48 | +except ImportError as exc: |
| 49 | + raise ImportError( |
| 50 | + "Cannot instrument. Please install 'taskiq[opentelemetry]'.", |
| 51 | + ) from exc |
| 52 | + |
| 53 | + |
| 54 | +from opentelemetry.instrumentation.instrumentor import ( # type: ignore[attr-defined] |
| 55 | + BaseInstrumentor, |
| 56 | +) |
| 57 | +from opentelemetry.instrumentation.utils import unwrap |
| 58 | +from opentelemetry.metrics import MeterProvider |
| 59 | +from opentelemetry.trace import TracerProvider |
| 60 | +from wrapt import wrap_function_wrapper |
| 61 | + |
| 62 | +from taskiq import AsyncBroker |
| 63 | +from taskiq.middlewares.opentelemetry_middleware import OpenTelemetryMiddleware |
| 64 | + |
| 65 | +if TYPE_CHECKING: |
| 66 | + pass |
| 67 | + |
| 68 | +logger = logging.getLogger("taskiq.opentelemetry") |
| 69 | + |
| 70 | + |
| 71 | +class TaskiqInstrumentor(BaseInstrumentor): |
| 72 | + """OpenTelemetry instrumentor for Taskiq.""" |
| 73 | + |
| 74 | + _instrumented_brokers: _WeakSet[AsyncBroker] = _WeakSet() |
| 75 | + |
| 76 | + def __init__(self) -> None: |
| 77 | + super().__init__() |
| 78 | + self._middleware = None |
| 79 | + |
| 80 | + def instrument_broker( |
| 81 | + self, |
| 82 | + broker: AsyncBroker, |
| 83 | + tracer_provider: Optional[TracerProvider] = None, |
| 84 | + meter_provider: Optional[MeterProvider] = None, |
| 85 | + ) -> None: |
| 86 | + """Instrument broker.""" |
| 87 | + if not hasattr(broker, "_is_instrumented_by_opentelemetry"): |
| 88 | + broker._is_instrumented_by_opentelemetry = False # type: ignore[attr-defined] # noqa: SLF001 |
| 89 | + |
| 90 | + if not getattr(broker, "is_instrumented_by_opentelemetry", False): |
| 91 | + broker.middlewares.insert( |
| 92 | + 0, |
| 93 | + OpenTelemetryMiddleware( |
| 94 | + tracer_provider=tracer_provider, |
| 95 | + meter_provider=meter_provider, |
| 96 | + ), |
| 97 | + ) |
| 98 | + broker._is_instrumented_by_opentelemetry = True # type: ignore[attr-defined] # noqa: SLF001 |
| 99 | + if broker not in self._instrumented_brokers: |
| 100 | + self._instrumented_brokers.add(broker) |
| 101 | + else: |
| 102 | + logger.warning( |
| 103 | + "Attempting to instrument taskiq broker while already instrumented", |
| 104 | + ) |
| 105 | + |
| 106 | + def uninstrument_broker(self, broker: AsyncBroker) -> None: |
| 107 | + """Uninstrument broker.""" |
| 108 | + broker.middlewares = [ |
| 109 | + middleware |
| 110 | + for middleware in broker.middlewares |
| 111 | + if not isinstance(middleware, OpenTelemetryMiddleware) |
| 112 | + ] |
| 113 | + broker._is_instrumented_by_opentelemetry = False # type: ignore[attr-defined] # noqa: SLF001 |
| 114 | + self._instrumented_brokers.discard(broker) |
| 115 | + |
| 116 | + def instrumentation_dependencies(self) -> Collection[str]: |
| 117 | + """This function tells which library this instrumentor instruments.""" |
| 118 | + return ("taskiq >= 0.0.1",) |
| 119 | + |
| 120 | + def _instrument(self, **kwargs: Any) -> None: |
| 121 | + def broker_init( |
| 122 | + init: Callable[[Any], Any], |
| 123 | + broker: AsyncBroker, |
| 124 | + args: Any, |
| 125 | + kwargs: Any, |
| 126 | + ) -> None: |
| 127 | + result = init(*args, **kwargs) |
| 128 | + self.instrument_broker(broker) |
| 129 | + return result |
| 130 | + |
| 131 | + wrap_function_wrapper("taskiq", "AsyncBroker.__init__", broker_init) |
| 132 | + |
| 133 | + def _uninstrument(self, **kwargs: Any) -> None: |
| 134 | + instances_to_uninstrument = list(self._instrumented_brokers) |
| 135 | + for broker in instances_to_uninstrument: |
| 136 | + self.uninstrument_broker(broker) |
| 137 | + self._instrumented_brokers.clear() |
| 138 | + unwrap(AsyncBroker, "__init__") |
0 commit comments