Merge 0dfec1f24a into e4d3692ac2

CHANGES/1785.feature.rst

@@ -0,0 +1 @@
+Add interface for tracing performance and integrate hooks, which allows native integration with OTel.

aiogram/dispatcher/dispatcher.py

@@ -17,7 +17,8 @@ from aiogram.fsm.storage.base import BaseEventIsolation, BaseStorage
 from aiogram.fsm.storage.memory import DisabledEventIsolation, MemoryStorage
 from aiogram.fsm.strategy import FSMStrategy
 from aiogram.methods import GetUpdates, TelegramMethod
-from aiogram.types import Update, User
+from aiogram.tracer import AbstractTracer, tracer
+from aiogram.types import TelegramObject, Update, User
 from aiogram.types.base import UNSET, UNSET_TYPE
 from aiogram.types.update import UpdateTypeLookupError
 from aiogram.utils.backoff import Backoff, BackoffConfig
@@ -48,6 +49,7 @@ class Dispatcher(Router):
         events_isolation: BaseEventIsolation | None = None,
         disable_fsm: bool = False,
         name: str | None = None,
+        tracer: AbstractTracer | None = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -73,6 +75,7 @@ class Dispatcher(Router):
             router=self,
             event_name="update",
         )
+        self.tracer = tracer
         self.update.register(self._listen_update)
 
         # Error handlers should work is out of all other functions
@@ -159,7 +162,6 @@ class Dispatcher(Router):
             # The preferred way is that pass already mounted Bot instance to this update
             # before call feed_update method
             update = Update.model_validate(update.model_dump(), context={"bot": bot})
-
         try:
             response = await self.update.wrap_outer_middleware(
                 self.update.trigger,
@@ -279,9 +281,26 @@ class Dispatcher(Router):
             raise SkipHandler() from e
 
         kwargs.update(event_update=update)
-
+        if self.tracer is not None:
+            tracer.set(self.tracer)
         return await self.propagate_event(update_type=update_type, event=event, **kwargs)
 
+    async def _propagate_event(
+        self,
+        observer: TelegramEventObserver | None,
+        update_type: str,
+        event: TelegramObject,
+        **kwargs: Any,
+    ) -> Any:
+        if (
+            self.tracer is None
+            or (tracer_manager := self.tracer.get_trigger_span_manager(event)) is None
+        ):
+            return await super()._propagate_event(observer, update_type, event, **kwargs)
+
+        async with tracer_manager:
+            return await super()._propagate_event(observer, update_type, event, **kwargs)
+
     @classmethod
     async def silent_call_request(cls, bot: Bot, result: TelegramMethod[Any]) -> None:
         """

aiogram/dispatcher/event/telegram.py

@@ -13,6 +13,10 @@ from .handler import CallbackType, FilterObject, HandlerObject
 if TYPE_CHECKING:
     from aiogram.dispatcher.router import Router
     from aiogram.types import TelegramObject
+import contextlib
+import functools
+
+from aiogram.tracer import AbstractTracer, tracer
 
 
 class TelegramEventObserver:
@@ -113,20 +117,39 @@ class TelegramEventObserver:
         Propagate event to handlers and stops propagation on first match.
         Handler will be called when all its filters are pass.
         """
+        tracer_instance = tracer.get()
         for handler in self.handlers:
+            if tracer_instance is not None and (
+                filter_manager := tracer_instance.get_filter_span_manager(handler)
+            ):
+                async with filter_manager:
+                    result, data = await handler.check(event, **kwargs)
+            else:
+                result, data = await handler.check(event, **kwargs)
+            if not result:
+                continue
             kwargs["handler"] = handler
-            result, data = await handler.check(event, **kwargs)
-            if result:
-                kwargs.update(data)
-                try:
-                    wrapped_inner = self.outer_middleware.wrap_middlewares(
-                        self._resolve_middlewares(),
-                        handler.call,
-                    )
-                    return await wrapped_inner(event, kwargs)
-                except SkipHandler:
-                    continue
+            kwargs.update(data)
+            try:
+                handler_call = handler.call
+                if tracer_instance is not None and (
+                    handler_manager := tracer_instance.get_handler_span_manager(handler)
+                ):
 
+                    @functools.wraps(handler.call)
+                    async def handler_wrapper(
+                        event: TelegramObject, *args: Any, **kwargs: Any
+                    ) -> Any:
+                        async with handler_manager:  # noqa: B023
+                            return await handler.call(event, *args, **kwargs)  # noqa: B023
+
+                    handler_call = handler_wrapper
+                wrapped_inner = self.outer_middleware.wrap_middlewares(
+                    self._resolve_middlewares(), handler_call
+                )
+                return await wrapped_inner(event, kwargs)
+            except SkipHandler:
+                continue
         return UNHANDLED
 
     def __call__(

aiogram/dispatcher/middlewares/manager.py

@@ -8,6 +8,7 @@ from aiogram.dispatcher.event.bases import (
     NextMiddlewareType,
 )
 from aiogram.dispatcher.event.handler import CallbackType
+from aiogram.tracer import AbstractTracer, tracer
 from aiogram.types import TelegramObject
 
 
@@ -54,15 +55,35 @@ class MiddlewareManager(Sequence[MiddlewareType[TelegramObject]]):
         return len(self._middlewares)
 
     @staticmethod
-    def wrap_middlewares(
-        middlewares: Sequence[MiddlewareType[MiddlewareEventType]],
+    async def middleware_step(
+        tracer_instance: AbstractTracer,
+        middleware: MiddlewareType[TelegramObject],
         handler: CallbackType,
+        event: TelegramObject,
+        kwargs: dict[str, Any],
+    ) -> Any:
+        manager = tracer_instance.get_middleware_span_manager(middleware)
+        if manager is None:
+            return await middleware(handler, event, kwargs)
+        async with manager:
+            return await middleware(handler, event, kwargs)
+
+    @staticmethod
+    def wrap_middlewares(
+        middlewares: Sequence[MiddlewareType[MiddlewareEventType]], handler: CallbackType
     ) -> NextMiddlewareType[MiddlewareEventType]:
         @functools.wraps(handler)
         def handler_wrapper(event: TelegramObject, kwargs: dict[str, Any]) -> Any:
             return handler(event, **kwargs)
 
         middleware = handler_wrapper
-        for m in reversed(middlewares):
-            middleware = functools.partial(m, middleware)  # type: ignore[assignment]
+        tracer_instance = tracer.get()
+        if tracer_instance is None:
+            for m in reversed(middlewares):
+                middleware = functools.partial(m, middleware)  # type: ignore[assignment]
+        else:
+            for m in reversed(middlewares):
+                middleware = functools.partial(  # type: ignore[assignment]
+                    MiddlewareManager.middleware_step, tracer_instance, m, middleware
+                )
         return middleware

aiogram/tracer.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from contextlib import AbstractAsyncContextManager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING
+
+from aiogram.dispatcher.event.bases import MiddlewareEventType, MiddlewareType
+
+if TYPE_CHECKING:
+    from aiogram.dispatcher.event.handler import HandlerObject
+    from aiogram.types import TelegramObject
+
+
+class AbstractTracer(ABC):
+    @abstractmethod
+    def get_middleware_span_manager(
+        self, middleware: MiddlewareType[MiddlewareEventType]
+    ) -> AbstractAsyncContextManager[None] | None:
+        pass
+
+    @abstractmethod
+    def get_handler_span_manager(
+        self, handler: HandlerObject
+    ) -> AbstractAsyncContextManager[None] | None:
+        pass
+
+    @abstractmethod
+    def get_trigger_span_manager(
+        self, event: TelegramObject
+    ) -> AbstractAsyncContextManager[None] | None:
+        pass
+
+    @abstractmethod
+    def get_filter_span_manager(
+        self, handler: HandlerObject
+    ) -> AbstractAsyncContextManager[None] | None:
+        pass
+
+
+tracer: ContextVar[AbstractTracer | None] = ContextVar("tracer", default=None)