From aa1ec158e3999e06fd11ef277152bfe0889eb684 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E0=AE=AE=E0=AE=A9=E0=AF=8B=E0=AE=9C=E0=AF=8D=E0=AE=95?= =?UTF-8?q?=E0=AF=81=E0=AE=AE=E0=AE=BE=E0=AE=B0=E0=AF=8D=20=E0=AE=AA?= =?UTF-8?q?=E0=AE=B4=E0=AE=A9=E0=AE=BF=E0=AE=9A=E0=AF=8D=E0=AE=9A=E0=AE=BE?= =?UTF-8?q?=E0=AE=AE=E0=AE=BF?= Date: Sun, 22 Jun 2025 11:22:51 +0530 Subject: [PATCH] Add max_batch_size parameter to BatchBase and related classes --- asynq/batching.py | 11 +++-- asynq/batching.pyi | 6 ++- asynq/tests/test_batching.py | 94 +++++++++++++++++++++++++++++++++++- examples/batching.py | 18 +++++-- 4 files changed, 118 insertions(+), 11 deletions(-) diff --git a/asynq/batching.py b/asynq/batching.py index 25d76fe..420ca39 100644 --- a/asynq/batching.py +++ b/asynq/batching.py @@ -37,9 +37,10 @@ class BatchCancelledError(BatchingError): class BatchBase(futures.FutureBase): """Abstract base class describing a batch of operations.""" - def __init__(self): + def __init__(self, max_batch_size=None): futures.FutureBase.__init__(self) # Cython doesn't support super(...) self.items = [] + self.max_batch_size = max_batch_size def is_flushed(self): return self.is_computed() @@ -50,6 +51,10 @@ def is_cancelled(self): def is_empty(self): return len(self.items) == 0 + def is_full(self): + """Returns True if the batch has reached its maximum size limit.""" + return self.max_batch_size is not None and len(self.items) >= self.max_batch_size + def get_priority(self): """Returns batch flush priority. The higher priority, the earlier batch is flushed. @@ -247,8 +252,8 @@ def __init__(self, batch_name="default", result=None): class DebugBatch(BatchBase): """Debug batch used to sync async execution.""" - def __init__(self, name="default", index=0): - super(DebugBatch, self).__init__() + def __init__(self, name="default", index=0, max_batch_size=None): + super(DebugBatch, self).__init__(max_batch_size=max_batch_size) self.name = name self.index = index diff --git a/asynq/batching.pyi b/asynq/batching.pyi index 61099d4..745fb52 100644 --- a/asynq/batching.pyi +++ b/asynq/batching.pyi @@ -10,10 +10,12 @@ class BatchCancelledError(BatchingError): ... class BatchBase(futures.FutureBase, Generic[_T]): items: List[_T] - def __init__(self) -> None: ... + max_batch_size: Optional[int] + def __init__(self, max_batch_size: Optional[int] = ...) -> None: ... def is_flushed(self) -> bool: ... def is_cancelled(self) -> bool: ... def is_empty(self) -> bool: ... + def is_full(self) -> bool: ... def get_priority(self) -> Tuple[int, int]: ... def flush(self) -> None: ... def cancel(self, error: Optional[Exception] = ...) -> None: ... @@ -28,7 +30,7 @@ class DebugBatchItem(BatchItemBase): def __init__(self, batch_name: str = ..., result: Any = ...) -> None: ... class DebugBatch(BatchBase): - def __init__(self, name: str = ..., index: int = 0) -> None: ... + def __init__(self, name: str = ..., index: int = 0, max_batch_size: Optional[int] = ...) -> None: ... class LocalDebugBatchState(threading.local): ... diff --git a/asynq/tests/test_batching.py b/asynq/tests/test_batching.py index 00605ac..14568ce 100644 --- a/asynq/tests/test_batching.py +++ b/asynq/tests/test_batching.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from asynq import asynq +from asynq import asynq, BatchBase, BatchItemBase from .debug_cache import reset_caches, mc from .caching import ExternalCacheBatchItem @@ -41,3 +41,95 @@ def foo(depth): reset_caches() foo(5) assert mc._batch.index == 2 + + +def test_batch_max_size(): + """Test that BatchBase respects max_batch_size parameter.""" + # Test basic functionality without size limit + batch = BatchBase() + assert batch.max_batch_size is None + assert not batch.is_full() + + # Test with size limit + batch_with_limit = BatchBase(max_batch_size=3) + assert batch_with_limit.max_batch_size == 3 + assert not batch_with_limit.is_full() + + # Test is_full functionality + class MockBatchItem(BatchItemBase): + def __init__(self, batch): + super().__init__(batch) + + # Add items one by one and check is_full status + MockBatchItem(batch_with_limit) + assert not batch_with_limit.is_full() + assert len(batch_with_limit.items) == 1 + + MockBatchItem(batch_with_limit) + assert not batch_with_limit.is_full() + assert len(batch_with_limit.items) == 2 + + MockBatchItem(batch_with_limit) + assert batch_with_limit.is_full() + assert len(batch_with_limit.items) == 3 + + +def test_batch_auto_switching(): + """Test automatic batch switching when max size is reached.""" + + class TestBatch(BatchBase): + switched_count = 0 + current_batch = None + + def __init__(self, max_batch_size=None): + super().__init__(max_batch_size=max_batch_size) + TestBatch.current_batch = self + + def _try_switch_active_batch(self): + TestBatch.switched_count += 1 + TestBatch.current_batch = TestBatch(max_batch_size=self.max_batch_size) + + def _flush(self): + # Mock implementation + for item in self.items: + item.set_value(f"result_{item.index}") + + class TestBatchItem(BatchItemBase): + def __init__(self, max_batch_size=None): + # Get or create current batch + if TestBatch.current_batch is None or TestBatch.current_batch.is_full(): + if TestBatch.current_batch is not None: + TestBatch.current_batch._try_switch_active_batch() + super().__init__(TestBatch.current_batch) + + # Reset state + TestBatch.switched_count = 0 + TestBatch.current_batch = TestBatch(max_batch_size=2) + + # Create items that should trigger batch switching + item1 = TestBatchItem() + assert TestBatch.switched_count == 0 + assert len(TestBatch.current_batch.items) == 1 + + item2 = TestBatchItem() + assert TestBatch.switched_count == 0 + assert len(TestBatch.current_batch.items) == 2 + assert TestBatch.current_batch.is_full() + + # This should trigger batch switching + item3 = TestBatchItem() + assert TestBatch.switched_count == 1 + assert len(TestBatch.current_batch.items) == 1 # New batch with one item + + +def test_batch_empty_check(): + """Test that is_empty works correctly.""" + batch = BatchBase() + assert batch.is_empty() + + class MockBatchItem(BatchItemBase): + def __init__(self, batch): + super().__init__(batch) + + MockBatchItem(batch) + assert not batch.is_empty() diff --git a/examples/batching.py b/examples/batching.py index 9dbffad..b95f9f2 100644 --- a/examples/batching.py +++ b/examples/batching.py @@ -33,7 +33,11 @@ class Client(object): Example usage: + # Basic client without batch size limit client = Client(['127.0.0.1:11211']) + + # Client with batch size limit of 1000 items to prevent timeouts + client = Client(['127.0.0.1:11211'], max_batch_size=1000) @client.cached('user_name') def name_of_user(uid): @@ -52,12 +56,16 @@ def all_author_names(aids): """ - def __init__(self, servers): + def __init__(self, servers, max_batch_size=None): # keep a reference to the current batch and to the underlying library's client object - self.batch = _MCBatch(self) + self.max_batch_size = max_batch_size + self.batch = _MCBatch(self, max_batch_size=max_batch_size) self._mc_client = memcache.Client(servers) def _make_batch_item(self, command, args): + # If the current batch is full, switch to a new one + if self.batch.is_full(): + self.batch._try_switch_active_batch() return _MCBatchItem(self.batch, command, args) @async_proxy() @@ -112,13 +120,13 @@ def wrapped(*args): class _MCBatch(BatchBase): - def __init__(self, client): + def __init__(self, client, max_batch_size=None): self.client = client - super().__init__() + super().__init__(max_batch_size=max_batch_size) def _try_switch_active_batch(self): if self.client.batch is self: - self.client.batch = _MCBatch(self.client) + self.client.batch = _MCBatch(self.client, max_batch_size=self.max_batch_size) def _flush(self): item_groups = itertools.groupby(self.items, lambda i: i.cmd)