diff --git a/CHANGES/1785.feature.rst b/CHANGES/1785.feature.rst new file mode 100644 index 00000000..3ec41245 --- /dev/null +++ b/CHANGES/1785.feature.rst @@ -0,0 +1 @@ +Add interface for tracing performance and integrate hooks, which allows native integration with OTel. diff --git a/aiogram/dispatcher/dispatcher.py b/aiogram/dispatcher/dispatcher.py index 5aff6a01..f0cfa81d 100644 --- a/aiogram/dispatcher/dispatcher.py +++ b/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: """ diff --git a/aiogram/dispatcher/event/telegram.py b/aiogram/dispatcher/event/telegram.py index 36b3843e..ca5695b8 100644 --- a/aiogram/dispatcher/event/telegram.py +++ b/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__( diff --git a/aiogram/dispatcher/middlewares/manager.py b/aiogram/dispatcher/middlewares/manager.py index e5cf73e7..f3078655 100644 --- a/aiogram/dispatcher/middlewares/manager.py +++ b/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 diff --git a/aiogram/tracer.py b/aiogram/tracer.py new file mode 100644 index 00000000..97488fe5 --- /dev/null +++ b/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)