diff --git a/messaging/limiter.py b/messaging/limiter.py index 37db3ce5..55e72c97 100644 --- a/messaging/limiter.py +++ b/messaging/limiter.py @@ -8,7 +8,7 @@ using a leaky bucket algorithm (aiolimiter) and a task queue. import asyncio import logging import os -from typing import Awaitable, Callable, Any, Optional +from typing import Awaitable, Callable, Any, Optional, List, Dict from aiolimiter import AsyncLimiter logger = logging.getLogger(__name__) @@ -52,7 +52,7 @@ class GlobalRateLimiter: # Custom queue state self._queue_list: List[str] = [] # List of dedup_keys in order self._queue_map: Dict[ - str, tuple[Callable[[], Awaitable[Any]], asyncio.Future] + str, tuple[Callable[[], Awaitable[Any]], List[asyncio.Future]] ] = {} self._condition = asyncio.Condition() @@ -74,7 +74,7 @@ class GlobalRateLimiter: await self._condition.wait() dedup_key = self._queue_list.pop(0) - func, future = self._queue_map.pop(dedup_key) + func, futures = self._queue_map.pop(dedup_key) # Check for manual pause (FloodWait) now = asyncio.get_event_loop().time() @@ -89,8 +89,9 @@ class GlobalRateLimiter: async with self.limiter: try: result = await func() - if not future.done(): - future.set_result(result) + for f in futures: + if not f.done(): + f.set_result(result) except Exception as e: # Handle Telegram FloodWaitError specifically error_msg = str(e).lower() @@ -107,14 +108,15 @@ class GlobalRateLimiter: asyncio.get_event_loop().time() + seconds ) - # Re-queue the task at the front (as a high priority update) - await self._enqueue_internal( - func, future, dedup_key, front=True + # Re-queue the tasks at the front (as a high priority update) + await self._enqueue_internal_multi( + func, futures, dedup_key, front=True ) await asyncio.sleep(seconds) else: - if not future.done(): - future.set_exception(e) + for f in futures: + if not f.done(): + f.set_exception(e) except asyncio.CancelledError: break except Exception as e: @@ -122,18 +124,24 @@ class GlobalRateLimiter: await asyncio.sleep(1) async def _enqueue_internal(self, func, future, dedup_key, front=False): + async def callback(f): + # This is just a placeholder to use the same internal logic + pass + + await self._enqueue_internal_multi(func, [future], dedup_key, front) + + async def _enqueue_internal_multi(self, func, futures, dedup_key, front=False): async with self._condition: if dedup_key in self._queue_map: - # Compaction: Update existing task with new func, keep old future - old_func, old_future = self._queue_map[dedup_key] - self._queue_map[dedup_key] = (func, old_future) - # Chain them so the new one resolves the old caller's future - # but we actually just use the same future object for simplicity - # The user just wants to know when "their" request is done, - # which is now the new request's completion. - logger.debug(f"Compacted task for key: {dedup_key}") + # Compaction: Update existing task with new func, append new futures + old_func, old_futures = self._queue_map[dedup_key] + old_futures.extend(futures) + self._queue_map[dedup_key] = (func, old_futures) + logger.debug( + f"Compacted task for key: {dedup_key} (now {len(old_futures)} futures)" + ) else: - self._queue_map[dedup_key] = (func, future) + self._queue_map[dedup_key] = (func, futures) if front: self._queue_list.insert(0, dedup_key) else: diff --git a/tests/verify_limiter_v2.py b/tests/verify_limiter_v2.py new file mode 100644 index 00000000..736e6063 --- /dev/null +++ b/tests/verify_limiter_v2.py @@ -0,0 +1,77 @@ +import asyncio +import time +import os +from messaging.limiter import GlobalRateLimiter + +import logging + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +async def test_compaction_and_hang(): + # Set small rate for testing + os.environ["MESSAGING_RATE_LIMIT"] = "1" + os.environ["MESSAGING_RATE_WINDOW"] = "0.5" # Fast for testing + + limiter = await GlobalRateLimiter.get_instance() + + call_counts = {} + + async def mock_edit(msg_id, content): + call_counts[msg_id] = call_counts.get(msg_id, 0) + 1 + logger.info(f"Executing actual Telegram edit for {msg_id}: {content}") + await asyncio.sleep(0.1) # Simulate network lag + return f"result_{content}" + + print("\n--- Starting Hang/Multi-Future Test ---") + + msg_id = "test_msg_123" + + # We will enqueue 3 edits and await all of them. + # Previously, the 2nd and 3rd would HANG. + + async def task(i): + logger.info(f"Task {i} started, enqueuing edit...") + res = await limiter.enqueue( + lambda i=i: mock_edit(msg_id, f"v{i}"), dedup_key=f"edit:{msg_id}" + ) + logger.info(f"Task {i} completed with: {res}") + return res + + start_time = time.time() + + # Run tasks concurrently + results = await asyncio.gather(task(1), task(2), task(3)) + + end_time = time.time() + duration = end_time - start_time + + print(f"\nAll tasks finished in {duration:.2f}s") + print(f"Results: {results}") + print(f"Call counts: {call_counts}") + + # Check that they all got the LAST result + for res in results: + assert res == "result_v3", f"Expected result_v3, got {res}" + + # Check that they didn't hang (should be < 1s given the compaction) + assert duration < 2.0, "Tasks took too long, might have hung or not compacted" + + # Check call counts: + # T1 might go through immediately or be compacted if T2/T3 arrive fast enough. + # Given the loop speed, we expect 1-2 calls. + assert call_counts[msg_id] <= 2, f"Too many calls: {call_counts[msg_id]}" + + print("\nPASSED: All futures resolved, no hang detected.") + + +if __name__ == "__main__": + # Fix encoding for windows terminal to avoid Unicode print crash + import sys + import io + + if sys.platform == "win32": + sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8") + + asyncio.run(test_compaction_and_hang())