diff --git a/CHANGES/1032.feature.rst b/CHANGES/1032.feature.rst new file mode 100644 index 00000000..f4a0d5aa --- /dev/null +++ b/CHANGES/1032.feature.rst @@ -0,0 +1 @@ +Cleanup will clear unused locks to prevent memory overflow. diff --git a/aiogram/fsm/storage/base.py b/aiogram/fsm/storage/base.py index 70c62f34..08aadda6 100644 --- a/aiogram/fsm/storage/base.py +++ b/aiogram/fsm/storage/base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any, AsyncGenerator, Dict, Optional, Union +from typing import Any, AsyncGenerator, Dict, Hashable, Optional, Union from aiogram import Bot from aiogram.fsm.state import State @@ -91,6 +91,8 @@ class BaseStorage(ABC): class BaseEventIsolation(ABC): + _locks: Dict[Hashable, Any] + @abstractmethod @asynccontextmanager async def lock(self, bot: Bot, key: StorageKey) -> AsyncGenerator[None, None]: diff --git a/aiogram/fsm/storage/memory.py b/aiogram/fsm/storage/memory.py index 8be371b6..e20eb071 100644 --- a/aiogram/fsm/storage/memory.py +++ b/aiogram/fsm/storage/memory.py @@ -71,5 +71,14 @@ class SimpleEventIsolation(BaseEventIsolation): async with lock: yield + self._cleanup(key) + async def close(self) -> None: self._locks.clear() + + def _cleanup(self, key: Hashable) -> None: + if self._locks[key]._waiters is None: # type: ignore[attr-defined] + del self._locks[key] + + elif len(self._locks[key]._waiters) == 0: # type: ignore[attr-defined] + del self._locks[key] diff --git a/tests/test_fsm/storage/test_isolation.py b/tests/test_fsm/storage/test_isolation.py index 042cc305..d3c4021c 100644 --- a/tests/test_fsm/storage/test_isolation.py +++ b/tests/test_fsm/storage/test_isolation.py @@ -1,27 +1,65 @@ +import asyncio +from random import randint, uniform + import pytest -from aiogram.fsm.storage.base import BaseEventIsolation, StorageKey +from aiogram.fsm.storage.base import StorageKey +from aiogram.fsm.storage.memory import DisabledEventIsolation, SimpleEventIsolation +from aiogram.fsm.storage.redis import RedisEventIsolation from tests.mocked_bot import MockedBot @pytest.fixture(name="storage_key") -def create_storate_key(bot: MockedBot): +def create_storage_key(bot: MockedBot): return StorageKey(chat_id=-42, user_id=42, bot_id=bot.id) -@pytest.mark.parametrize( - "isolation", - [ - pytest.lazy_fixture("redis_isolation"), - pytest.lazy_fixture("lock_isolation"), - pytest.lazy_fixture("disabled_isolation"), - ], -) -class TestIsolations: +@pytest.mark.parametrize("isolation", [pytest.lazy_fixture("disabled_isolation")]) +class TestDisabledIsolation: async def test_lock( self, bot: MockedBot, - isolation: BaseEventIsolation, + isolation: DisabledEventIsolation, + storage_key: StorageKey, + ): + async with isolation.lock(bot=bot, key=storage_key): + assert True, "You are kidding me?" + + +@pytest.mark.parametrize("isolation", [pytest.lazy_fixture("lock_isolation")]) +class TestLockIsolations: + @staticmethod + async def _some_task(isolation: SimpleEventIsolation, bot: MockedBot, key: StorageKey): + async with isolation.lock(bot=bot, key=key): + await asyncio.sleep(uniform(0, 1)) + + @staticmethod + def random_storage_key(bot: MockedBot): + return StorageKey(chat_id=randint(-44, -40), user_id=randint(40, 44), bot_id=bot.id) + + async def test_lock( + self, + bot: MockedBot, + isolation: SimpleEventIsolation, + ): + tasks = [] + + for _ in range(100): + tasks.append( + asyncio.create_task(self._some_task(isolation, bot, self.random_storage_key(bot))) + ) + await asyncio.sleep(0.01) + + await asyncio.gather(*[task for task in tasks if not task.done()]) + assert len(isolation._locks) == 0 + + +@pytest.mark.parametrize("isolation", [pytest.lazy_fixture("redis_isolation")]) +class TestRedisIsolation: + async def test_lock( + self, + bot: MockedBot, + isolation: RedisEventIsolation, storage_key: StorageKey, ): async with isolation.lock(bot=bot, key=storage_key):