Skip to content

Commit

Permalink
Prevent update handling task pointers from being garbage collected, b…
Browse files Browse the repository at this point in the history
…ackport of #1328 (#1331)

* Preserve update handling task pointers, backport of #1328

* Changelog

* Typing improvements
  • Loading branch information
Kylmakalle committed Oct 8, 2023
1 parent f681afb commit cad4258
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 4 deletions.
1 change: 1 addition & 0 deletions CHANGES/1331.misc
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Prevent update handling task pointers from being garbage collected, backport from 2.x
7 changes: 5 additions & 2 deletions aiogram/dispatcher/dispatcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import warnings
from asyncio import CancelledError, Event, Future, Lock
from contextlib import suppress
from typing import Any, AsyncGenerator, Dict, List, Optional, Union
from typing import Any, AsyncGenerator, Dict, List, Optional, Set, Union

from .. import loggers
from ..client.bot import Bot
Expand Down Expand Up @@ -95,6 +95,7 @@ def __init__(
self._running_lock = Lock()
self._stop_signal: Optional[Event] = None
self._stopped_signal: Optional[Event] = None
self._handle_update_tasks: Set[asyncio.Task[Any]] = set()

def __getitem__(self, item: str) -> Any:
return self.workflow_data[item]
Expand Down Expand Up @@ -349,7 +350,9 @@ async def _polling(
):
handle_update = self._process_update(bot=bot, update=update, **kwargs)
if handle_as_tasks:
asyncio.create_task(handle_update)
handle_update_task = asyncio.create_task(handle_update)
self._handle_update_tasks.add(handle_update_task)
handle_update_task.add_done_callback(self._handle_update_tasks.discard)
else:
await handle_update
finally:
Expand Down
7 changes: 5 additions & 2 deletions aiogram/webhook/aiohttp_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import secrets
from abc import ABC, abstractmethod
from asyncio import Transport
from typing import Any, Awaitable, Callable, Dict, Optional, Tuple, cast
from typing import Any, Awaitable, Callable, Dict, Optional, Set, Tuple, cast

from aiohttp import MultipartWriter, web
from aiohttp.abc import Application
Expand Down Expand Up @@ -98,6 +98,7 @@ def __init__(
self.dispatcher = dispatcher
self.handle_in_background = handle_in_background
self.data = data
self._background_feed_update_tasks: Set[asyncio.Task[Any]] = set()

def register(self, app: Application, /, path: str, **kwargs: Any) -> None:
"""
Expand Down Expand Up @@ -139,11 +140,13 @@ async def _background_feed_update(self, bot: Bot, update: Dict[str, Any]) -> Non
await self.dispatcher.silent_call_request(bot=bot, result=result)

async def _handle_request_background(self, bot: Bot, request: web.Request) -> web.Response:
asyncio.create_task(
feed_update_task = asyncio.create_task(
self._background_feed_update(
bot=bot, update=await request.json(loads=bot.session.json_loads)
)
)
self._background_feed_update_tasks.add(feed_update_task)
feed_update_task.add_done_callback(self._background_feed_update_tasks.discard)
return web.json_response({}, dumps=bot.session.json_dumps)

def _build_response_writer(
Expand Down

0 comments on commit cad4258

Please sign in to comment.