Merge 7f7c1820e5 into fea1b7b0a3

CHANGES/1032.feature.rst

@@ -0,0 +1 @@
+Cleanup will clear unused locks to prevent memory overflow.

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

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]

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