feat: Add OpenTelemetry instrumentation

docs/README.md

@@ -14,13 +14,13 @@ actions:
 head:
  - - meta
  - name: "google-site-verification"
- content: "hQCR5w2tmeuOvYIYXsOYU3u4kLNwT86lnqltANYlRQ0"
+ content: "hQCR5w2tmeuOvYIYXsOYU3u4kLNwT86lnqltANYlRQ0" # pragma: allowlist secret
  - - meta
  - name: "msvalidate.01"
- content: "97DC185FE0A2F5B123861F0790FDFB26"
+ content: "97DC185FE0A2F5B123861F0790FDFB26" # pragma: allowlist secret
  - - meta
  - name: "yandex-verification"
- content: "9b105f7c58cbc920"
+ content: "9b105f7c58cbc920" # pragma: allowlist secret
 highlights:
  - features:
  - title: Production ready

pyproject.toml

@@ -36,6 +36,10 @@ pycron = "^3.0.0"
 taskiq_dependencies = ">=1.3.1,<2"
 anyio = ">=3"
 packaging = ">=19"
+# For opentelemetry instrumentation
+opentelemetry-api = { version = "^1.38.0", optional = true }
+opentelemetry-instrumentation = { version = "^0.59b0", optional = true}
+opentelemetry-semantic-conventions = { version = "^0.59b0", optional = true}
 # For prometheus metrics
 prometheus_client = { version = "^0", optional = true }
 # For ZMQBroker
@@ -69,10 +73,12 @@ pytest-mock = "^3.11.1"
 tzlocal = "^5.0.1"
 types-tzlocal = "^5.0.1.1"
 types-pytz = "^2023.3.1.1"
+opentelemetry-test-utils = "^0.59b0"
 
 [tool.poetry.extras]
 zmq = ["pyzmq"]
 uv = ["uvloop"]
+opentelemetry = ["opentelemetry-api", "opentelemetry-instrumentation", "opentelemetry-semantic-conventions"]
 metrics = ["prometheus_client"]
 reload = ["watchdog", "gitignore-parser"]
 orjson = ["orjson"]
@@ -86,6 +92,9 @@ taskiq = "taskiq.__main__:main"
 worker = "taskiq.cli.worker.cmd:WorkerCMD"
 scheduler = "taskiq.cli.scheduler.cmd:SchedulerCMD"
 
+[tool.poetry.plugins.opentelemetry_instrumentor]
+taskiq = "taskiq.instrumentation:TaskiqInstrumentor"
+
 [tool.mypy]
 strict = true
 ignore_missing_imports = true
@@ -174,6 +183,7 @@ line-length = 88
  "SLF001", # Private member accessed
  "S311", # Standard pseudo-random generators are not suitable for security/cryptographic purposes
  "D101", # Missing docstring in public class
+ "D102", # Missing docstring in public method
 ]
 
 [tool.ruff.lint.pydocstyle]

taskiq/instrumentation.py

@@ -0,0 +1,138 @@
+"""
+Instrument `taskiq`_ to trace Taskiq applications.
+
+.. _taskiq: https://pypi.org/project/taskiq/
+
+Usage
+-----
+
+* Run instrumented task
+
+.. code:: python
+
+ import asyncio
+
+ from taskiq import InMemoryBroker, TaskiqEvents, TaskiqState
+ from taskiq.instrumentation import TaskiqInstrumentor
+
+ broker = InMemoryBroker()
+
+ @broker.on_event(TaskiqEvents.WORKER_STARTUP)
+ async def startup(state: TaskiqState) -> None:
+ TaskiqInstrumentor().instrument()
+
+ @broker.task
+ async def add(x, y):
+ return x + y
+
+ async def main():
+ await broker.startup()
+ await my_task.kiq(1, 2)
+ await broker.shutdown()
+
+ if __name__ == "__main__":
+ asyncio.run(main())
+
+API
+---
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Collection, Optional
+from weakref import WeakSet as _WeakSet
+
+try:
+ import opentelemetry # noqa: F401
+except ImportError as exc:
+ raise ImportError(
+ "Cannot instrument. Please install 'taskiq[opentelemetry]'.",
+ ) from exc
+
+
+from opentelemetry.instrumentation.instrumentor import ( # type: ignore[attr-defined]
+ BaseInstrumentor,
+)
+from opentelemetry.instrumentation.utils import unwrap
+from opentelemetry.metrics import MeterProvider
+from opentelemetry.trace import TracerProvider
+from wrapt import wrap_function_wrapper
+
+from taskiq import AsyncBroker
+from taskiq.middlewares.opentelemetry_middleware import OpenTelemetryMiddleware
+
+if TYPE_CHECKING:
+ pass
+
+logger = logging.getLogger("taskiq.opentelemetry")
+
+
+class TaskiqInstrumentor(BaseInstrumentor):
+ """OpenTelemetry instrumentor for Taskiq."""
+
+ _instrumented_brokers: _WeakSet[AsyncBroker] = _WeakSet()
+
+ def __init__(self) -> None:
+ super().__init__()
+ self._middleware = None
+
+ def instrument_broker(
+ self,
+ broker: AsyncBroker,
+ tracer_provider: Optional[TracerProvider] = None,
+ meter_provider: Optional[MeterProvider] = None,
+ ) -> None:
+ """Instrument broker."""
+ if not hasattr(broker, "_is_instrumented_by_opentelemetry"):
+ broker._is_instrumented_by_opentelemetry = False # type: ignore[attr-defined] # noqa: SLF001
+
+ if not getattr(broker, "is_instrumented_by_opentelemetry", False):
+ broker.middlewares.insert(
+ 0,
+ OpenTelemetryMiddleware(
+ tracer_provider=tracer_provider,
+ meter_provider=meter_provider,
+ ),
+ )
+ broker._is_instrumented_by_opentelemetry = True # type: ignore[attr-defined] # noqa: SLF001
+ if broker not in self._instrumented_brokers:
+ self._instrumented_brokers.add(broker)
+ else:
+ logger.warning(
+ "Attempting to instrument taskiq broker while already instrumented",
+ )
+
+ def uninstrument_broker(self, broker: AsyncBroker) -> None:
+ """Uninstrument broker."""
+ broker.middlewares = [
+ middleware
+ for middleware in broker.middlewares
+ if not isinstance(middleware, OpenTelemetryMiddleware)
+ ]
+ broker._is_instrumented_by_opentelemetry = False # type: ignore[attr-defined] # noqa: SLF001
+ self._instrumented_brokers.discard(broker)
+
+ def instrumentation_dependencies(self) -> Collection[str]:
+ """This function tells which library this instrumentor instruments."""
+ return ("taskiq >= 0.0.1",)
+
+ def _instrument(self, **kwargs: Any) -> None:
+ def broker_init(
+ init: Callable[[Any], Any],
+ broker: AsyncBroker,
+ args: Any,
+ kwargs: Any,
+ ) -> None:
+ result = init(*args, **kwargs)
+ self.instrument_broker(broker)
+ return result
+
+ wrap_function_wrapper("taskiq", "AsyncBroker.__init__", broker_init)
+
+ def _uninstrument(self, **kwargs: Any) -> None:
+ instances_to_uninstrument = list(self._instrumented_brokers)
+ for broker in instances_to_uninstrument:
+ self.uninstrument_broker(broker)
+ self._instrumented_brokers.clear()
+ unwrap(AsyncBroker, "__init__")