diff --git a/README.md b/README.md index 840c5ca..b4583b5 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ dp = Dispatcher() # Local / single process setup_input(dp, storage=MemoryInputStorage()) -# Production (multi-worker): Redis-backed wait markers +# Production: Redis markers + TTL (futures still resolve on the owning worker) # redis = Redis.from_url("redis://localhost:6379/0") # setup_input(dp, storage=RedisInputStorage(redis, ttl=300)) diff --git a/pyproject.toml b/pyproject.toml index f0ea9b9..ac49447 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "aiogram-input" -version = "4.2.0" +version = "4.2.1" description = "Await user replies in aiogram bots with Dispatcher-scoped setup and pluggable storage." readme = "README.md" license = "MIT" diff --git a/src/aiogram_input/registry.py b/src/aiogram_input/registry.py index fc80682..76606fc 100644 --- a/src/aiogram_input/registry.py +++ b/src/aiogram_input/registry.py @@ -1,9 +1,11 @@ from __future__ import annotations import logging -from asyncio import Lock +from asyncio import InvalidStateError, Lock from typing import Dict, Optional +from aiogram.types import Message + from .types import PendingWait logger = logging.getLogger(__name__) @@ -38,6 +40,27 @@ async def pop_if(self, chat_id: int, wait_id: str, /) -> Optional[PendingWait]: return None return self._waits.pop(chat_id) + async def resolve_if( + self, chat_id: int, wait_id: str, message: Message, / + ) -> bool: + """ + Atomically complete ``wait_id`` with ``message`` if it is still current. + + Returns True only when this call consumed the wait. + """ + async with self._lock: + wait = self._waits.get(chat_id) + if wait is None or wait.wait_id != wait_id: + return False + future = wait.future + if future.done(): + return False + try: + future.set_result(message) + except InvalidStateError: + return False + return True + async def contains(self, chat_id: int, /) -> bool: async with self._lock: return chat_id in self._waits @@ -48,5 +71,8 @@ def cancel_wait(wait: PendingWait, *, chat_id: int) -> None: future = wait.future if future.done(): return - future.set_result(None) + try: + future.set_result(None) + except InvalidStateError: + return logger.debug("[REGISTRY] Aborted leftover wait chat=%s", chat_id) diff --git a/src/aiogram_input/session.py b/src/aiogram_input/session.py index 34bd4ec..1d00165 100644 --- a/src/aiogram_input/session.py +++ b/src/aiogram_input/session.py @@ -76,16 +76,14 @@ async def feed(self, message: Message) -> bool: message.message_id, ) - if not await self._storage.contains(chat_id): - logger.debug("[SESSION] No pending entry chat=%s", chat_id) - return False - wait = await self._registry.get(chat_id) if wait is None: - logger.debug("[SESSION] Marker without registry wait chat=%s", chat_id) - await self._storage.pop(chat_id) + # Do NOT delete storage markers here: with Redis another worker may + # own the wait. Orphan Redis keys should expire via TTL. + logger.debug("[SESSION] No local pending wait chat=%s", chat_id) return False + wait_id = wait.wait_id if not await self._check_filter(wait.filter, message): filter_name = wait.filter.__class__.__name__ if wait.filter else str(None) logger.debug( @@ -95,11 +93,15 @@ async def feed(self, message: Message) -> bool: ) return False - if wait.future.done(): - logger.debug("[SESSION] Future already done chat=%s", chat_id) + consumed = await self._registry.resolve_if(chat_id, wait_id, message) + if not consumed: + logger.debug( + "[SESSION] Resolve skipped stale/done wait chat=%s wait_id=%s", + chat_id, + wait_id, + ) return False - wait.future.set_result(message) logger.debug( "[SESSION] Future resolved chat=%s, message_id=%s", chat_id, @@ -122,7 +124,11 @@ async def _await_future( async def _check_filter(filter: FilterObjectType, message: Message) -> bool: if filter is None: return True - return await filter.call(message) + try: + return bool(await filter.call(message)) + except Exception: + logger.exception("[SESSION] Filter raised; treating as reject") + return False async def _register_pending( self, diff --git a/src/aiogram_input/storage/redis.py b/src/aiogram_input/storage/redis.py index 407e231..80f8719 100644 --- a/src/aiogram_input/storage/redis.py +++ b/src/aiogram_input/storage/redis.py @@ -1,10 +1,13 @@ from __future__ import annotations import json +import logging from typing import Any, Optional from ..types import WaitRecord +logger = logging.getLogger(__name__) + _POP_IF_SCRIPT = """ local raw = redis.call('GET', KEYS[1]) if not raw then @@ -23,8 +26,12 @@ class RedisInputStorage: """ Redis-backed InputStorage for wait markers. - Futures/filters stay in-process; only ``WaitRecord`` is persisted. - Requires the optional dependency: ``pip install aiogram-input[redis]``. + Futures/filters stay in-process on the worker that called ``wait()``. + Redis only stores markers so TTLs and cross-process visibility of "someone + is waiting" work; resolving still requires the owning process (or a future + pub/sub bridge). Use ``ttl`` so abandoned markers cannot grow forever. + + Requires: ``pip install aiogram-input[redis]``. """ def __init__( @@ -35,6 +42,8 @@ def __init__( key_prefix: str = "aiogram_input:wait:", ttl: Optional[int] = None, ) -> None: + if not isinstance(key_prefix, str) or not key_prefix: + raise ValueError("key_prefix must be a non-empty string") if ttl is not None and ttl <= 0: raise ValueError("ttl must be a positive integer or None") self._redis = redis @@ -52,18 +61,29 @@ def _dumps(record: WaitRecord) -> str: separators=(",", ":"), ) - @staticmethod - def _loads(raw: Any) -> WaitRecord: - if isinstance(raw, bytes): - raw = raw.decode() - data = json.loads(raw) - return WaitRecord(wait_id=data["wait_id"], created_at=float(data["created_at"])) + @classmethod + def _loads(cls, raw: Any) -> Optional[WaitRecord]: + try: + if isinstance(raw, bytes): + raw = raw.decode() + data = json.loads(raw) + wait_id = data["wait_id"] + created_at = float(data["created_at"]) + if not isinstance(wait_id, str) or not wait_id: + raise ValueError("invalid wait_id") + return WaitRecord(wait_id=wait_id, created_at=created_at) + except (TypeError, ValueError, KeyError, json.JSONDecodeError, UnicodeError): + logger.warning("[REDIS] Ignoring corrupt wait marker payload") + return None async def get(self, chat_id: int, /) -> Optional[WaitRecord]: raw = await self._redis.get(self._key(chat_id)) if raw is None: return None - return self._loads(raw) + record = self._loads(raw) + if record is None: + await self._redis.delete(self._key(chat_id)) + return record async def set(self, chat_id: int, record: WaitRecord, /) -> None: key = self._key(chat_id) diff --git a/tests/test_concurrency_hardening.py b/tests/test_concurrency_hardening.py new file mode 100644 index 0000000..ec8071b --- /dev/null +++ b/tests/test_concurrency_hardening.py @@ -0,0 +1,91 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from aiogram_input.registry import WaitRegistry +from aiogram_input.session import SessionManager +from aiogram_input.storage import MemoryInputStorage, RedisInputStorage +from aiogram_input.types import PendingWait, WaitRecord +from tests.helpers import make_message + + +class BoomFilter: + async def call(self, _message) -> bool: + raise RuntimeError("filter exploded") + + +@pytest.mark.asyncio +async def test_filter_exception_does_not_consume(session: SessionManager) -> None: + chat_id = 1 + task = asyncio.create_task( + session.start_waiting(chat_id, timeout=0.2, filter=BoomFilter()) # type: ignore[arg-type] + ) + await asyncio.sleep(0) + assert await session.feed(make_message(chat_id)) is False + assert await task is None + + +@pytest.mark.asyncio +async def test_resolve_if_rejects_stale_wait_id(registry: WaitRegistry) -> None: + loop = asyncio.get_running_loop() + future = loop.create_future() + await registry.set( + 1, PendingWait(wait_id="current", future=future, filter=None) + ) + assert await registry.resolve_if(1, "stale", make_message(1)) is False + assert not future.done() + assert await registry.resolve_if(1, "current", make_message(1)) is True + assert future.done() + assert future.result().message_id == 1 + + +@pytest.mark.asyncio +async def test_concurrent_feeds_only_one_consumes(session: SessionManager) -> None: + chat_id = 88 + task = asyncio.create_task(session.start_waiting(chat_id, timeout=2, filter=None)) + await asyncio.sleep(0) + results = await asyncio.gather( + session.feed(make_message(chat_id, message_id=1)), + session.feed(make_message(chat_id, message_id=2)), + ) + assert results.count(True) == 1 + assert results.count(False) == 1 + got = await task + assert got is not None + assert got.message_id in {1, 2} + + +@pytest.mark.asyncio +async def test_redis_corrupt_payload_is_dropped() -> None: + from fakeredis import FakeAsyncRedis + + redis = FakeAsyncRedis(decode_responses=True) + storage = RedisInputStorage(redis, key_prefix="bad:") + await redis.set("bad:1", "{not-json") + assert await storage.get(1) is None + assert await redis.exists("bad:1") == 0 + await redis.aclose() + + +@pytest.mark.asyncio +async def test_redis_rejects_empty_key_prefix() -> None: + from fakeredis import FakeAsyncRedis + + with pytest.raises(ValueError, match="key_prefix"): + RedisInputStorage(FakeAsyncRedis(), key_prefix="") + + +@pytest.mark.asyncio +async def test_feed_after_overwrite_goes_to_latest_wait( + session: SessionManager, +) -> None: + chat_id = 77 + first = asyncio.create_task(session.start_waiting(chat_id, timeout=2, filter=None)) + await asyncio.sleep(0) + second = asyncio.create_task(session.start_waiting(chat_id, timeout=2, filter=None)) + assert await first is None + msg = make_message(chat_id, message_id=9) + assert await session.feed(msg) is True + assert await second is msg diff --git a/tests/test_edge_cases.py b/tests/test_edge_cases.py index 1e6f066..bdf3604 100644 --- a/tests/test_edge_cases.py +++ b/tests/test_edge_cases.py @@ -36,12 +36,13 @@ async def handler(event, data): @pytest.mark.asyncio -async def test_orphan_storage_marker_is_cleaned_on_feed( +async def test_storage_marker_without_local_wait_is_not_deleted( session: SessionManager, storage: MemoryInputStorage ) -> None: + """Redis multi-worker safety: never delete another process's marker.""" await storage.set(55, WaitRecord(wait_id="orphan", created_at=1.0)) assert await session.feed(make_message(55)) is False - assert await storage.contains(55) is False + assert await storage.contains(55) is True @pytest.mark.asyncio @@ -55,6 +56,8 @@ async def test_feed_when_future_already_done(session: SessionManager) -> None: ) await session._storage.set(chat_id, WaitRecord(wait_id="x", created_at=1.0)) assert await session.feed(make_message(chat_id)) is False + # Marker may remain until owning cleanup/TTL; feed must not crash. + assert await session._storage.contains(chat_id) is True @pytest.mark.asyncio diff --git a/tests/test_redis_storage.py b/tests/test_redis_storage.py index 8bb5650..7558b95 100644 --- a/tests/test_redis_storage.py +++ b/tests/test_redis_storage.py @@ -78,3 +78,9 @@ async def test_redis_loads_bytes_payload() -> None: record = WaitRecord(wait_id="b", created_at=3.0) raw = RedisInputStorage._dumps(record).encode() assert RedisInputStorage._loads(raw) == record + + +@pytest.mark.asyncio +async def test_redis_loads_invalid_payload_returns_none() -> None: + assert RedisInputStorage._loads(b"nope") is None + assert RedisInputStorage._loads('{"wait_id":""}') is None diff --git a/uv.lock b/uv.lock index 33cfe5e..be28560 100644 --- a/uv.lock +++ b/uv.lock @@ -30,7 +30,7 @@ wheels = [ [[package]] name = "aiogram-input" -version = "4.2.0" +version = "4.2.1" source = { editable = "." } dependencies = [ { name = "aiogram" },