diff --git a/aiogram/fsm/storage/base.py b/aiogram/fsm/storage/base.py index 70c62f34..eea53560 100644 --- a/aiogram/fsm/storage/base.py +++ b/aiogram/fsm/storage/base.py @@ -94,6 +94,7 @@ class BaseEventIsolation(ABC): @abstractmethod @asynccontextmanager async def lock(self, bot: Bot, key: StorageKey) -> AsyncGenerator[None, None]: + self._locks: Dict[StorageKey, Any] """ Isolate events with lock. Will be used as context manager diff --git a/aiogram/fsm/storage/memory.py b/aiogram/fsm/storage/memory.py index 16c13157..c3dfbb23 100644 --- a/aiogram/fsm/storage/memory.py +++ b/aiogram/fsm/storage/memory.py @@ -47,6 +47,9 @@ class MemoryStorage(BaseStorage): class DisabledEventIsolation(BaseEventIsolation): + def __init__(self) -> None: + self._locks: DefaultDict[Hashable, Lock] + @asynccontextmanager async def lock(self, bot: Bot, key: StorageKey) -> AsyncGenerator[None, None]: yield @@ -66,5 +69,14 @@ class SimpleEventIsolation(BaseEventIsolation): async with lock: yield + self._cleanup(key) + async def close(self) -> None: self._locks.clear() + + def _cleanup(self, key: Hashable): + if self._locks[key]._waiters is None: + del self._locks[key] + + elif len(self._locks[key]._waiters) == 0: + del self._locks[key] diff --git a/tests/test_fsm/storage/test_isolation.py b/tests/test_fsm/storage/test_isolation.py index 5f20a3c9..bd3e680a 100644 --- a/tests/test_fsm/storage/test_isolation.py +++ b/tests/test_fsm/storage/test_isolation.py @@ -1,29 +1,66 @@ +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 pytestmark = pytest.mark.asyncio @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):