Add cleanup to SimpleEventIsolation

cleaunp will clear unused locks to prevent memory overflow

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

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]

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):