Add max_batch_size parameter to BatchBase and related classes

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
 

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

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

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)