From 91b84e63fa587930b9ad61b1e6cc1b49016aca3d Mon Sep 17 00:00:00 2001 From: andrew000 <11490628+andrew000@users.noreply.github.com> Date: Wed, 19 Oct 2022 12:17:16 +0300 Subject: [PATCH 1/4] Add cleanup to SimpleEventIsolation cleaunp will clear unused locks to prevent memory overflow --- aiogram/fsm/storage/base.py | 1 + aiogram/fsm/storage/memory.py | 12 +++++ tests/test_fsm/storage/test_isolation.py | 61 +++++++++++++++++++----- 3 files changed, 62 insertions(+), 12 deletions(-) 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): From fb2b685f1299c6f19ed29b4abf4a856c0171f916 Mon Sep 17 00:00:00 2001 From: andrew000 <11490628+andrew000@users.noreply.github.com> Date: Thu, 20 Oct 2022 00:19:23 +0300 Subject: [PATCH 2/4] Add changelog --- CHANGES/1032.feature.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 CHANGES/1032.feature.rst diff --git a/CHANGES/1032.feature.rst b/CHANGES/1032.feature.rst new file mode 100644 index 00000000..ed00922a --- /dev/null +++ b/CHANGES/1032.feature.rst @@ -0,0 +1 @@ +Add your info here \ No newline at end of file From 7125c2b39f323ed13f7954fe192d2eeb3d300a55 Mon Sep 17 00:00:00 2001 From: andrew000 <11490628+andrew000@users.noreply.github.com> Date: Thu, 20 Oct 2022 00:20:28 +0300 Subject: [PATCH 3/4] Refactor changelog --- CHANGES/1032.feature.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGES/1032.feature.rst b/CHANGES/1032.feature.rst index ed00922a..f4a0d5aa 100644 --- a/CHANGES/1032.feature.rst +++ b/CHANGES/1032.feature.rst @@ -1 +1 @@ -Add your info here \ No newline at end of file +Cleanup will clear unused locks to prevent memory overflow. From 7f7c1820e544ad33cae79152bbeb99efd041e7f2 Mon Sep 17 00:00:00 2001 From: andrew000 <11490628+andrew000@users.noreply.github.com> Date: Tue, 25 Oct 2022 23:42:10 +0300 Subject: [PATCH 4/4] Refactor code to pass the Lint --- aiogram/fsm/storage/base.py | 5 +++-- aiogram/fsm/storage/memory.py | 9 +++------ tests/test_fsm/storage/test_isolation.py | 3 ++- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/aiogram/fsm/storage/base.py b/aiogram/fsm/storage/base.py index eea53560..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,10 +91,11 @@ class BaseStorage(ABC): class BaseEventIsolation(ABC): + _locks: Dict[Hashable, Any] + @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 c3dfbb23..45213976 100644 --- a/aiogram/fsm/storage/memory.py +++ b/aiogram/fsm/storage/memory.py @@ -47,9 +47,6 @@ 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 @@ -74,9 +71,9 @@ class SimpleEventIsolation(BaseEventIsolation): async def close(self) -> None: self._locks.clear() - def _cleanup(self, key: Hashable): - if self._locks[key]._waiters is None: + 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: + 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 bd3e680a..417d7441 100644 --- a/tests/test_fsm/storage/test_isolation.py +++ b/tests/test_fsm/storage/test_isolation.py @@ -48,7 +48,8 @@ class TestLockIsolations: for _ in range(100): tasks.append( - asyncio.create_task(self._some_task(isolation, bot, self.random_storage_key(bot)))) + 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()])