feat: dynamic batch size recovery for storage operations (#181)

* feat(providers): add dynamic batch size recovery for storage operations

After a timeout reduces batch size, subsequent successful requests with smaller payloads automatically double the batch size back toward the original. This prevents one large-storage account from permanently slowing all subsequent operations.

Changes:
- get_account_storage_with_fallback: tracks last_successful_avg_bytes and failure_avg_bytes_ref, doubles batch on smaller payload (measured as avg bytes/key)
- set_state_with_batching: tracks failure_avg_bytes_ref from failed batch, doubles batch on smaller payload (measured as avg bytes/pair)
- Both operations reset batch size per account (per-operation isolation, no cross-account state)
- Increase is capped at original configured batch size, matching the symmetric ÷2 reduction rule
- INFO log on each batch size increase with old/new values and payload comparison reason

* test(providers): add tests for dynamic batch size increase and per-account reset

Add 14 new tests across 3 test classes:

TestGetAccountStorageWithFallbackBatchIncrease (6 tests):
- Batch increases on smaller payload (avg bytes/key strictly less)
- No increase on same or larger payload
- No increase when first request times out (no prior success reference)
- Recovery tracking resets on re-failure
- Increase capped at original batch size
- INFO log emitted with old/new sizes and payload comparison

TestSetStateWithBatchingBatchIncrease (5 tests):
- Batch increases on smaller avg bytes/pair
- No increase when avg bytes/pair same or larger
- Recovery reset on re-failure
- Increase capped at original
- INFO log emitted

TestBatchSizeResetBetweenAccounts (3 tests):
- Fetch batch size resets between account calls
- Push batch size resets between account calls
- Fetch fallback to standard endpoint doesn't affect next account's paginated batch

* docs(changelog): record dynamic batch size recovery and batch reset tests

* Bump version: 3.1.1-dev1 → 3.1.1-dev2

---------

Co-authored-by: github-actions <41898282+github-actions@users.noreply.github.com>

Author: EtienneWallet

docs/source/dev_documentation/changelog.md

@@ -5,6 +5,8 @@
 ### Added
 
 - Explicit test for variadic counted values
+- Dynamic batch size recovery: after a timeout reduces the batch size, subsequent successful requests with smaller payloads automatically double the batch size back toward the original (for both storage fetch and push operations)
+- Tests verifying batch size resets between accounts
 
 ### Fixed
 

mxops/common/providers.py

@@ -86,6 +86,8 @@ def set_state_with_batching(
     # for a given input dict
     pairs_list = list(pairs.items())
     current_batch_size = batch_size
+    original_batch_size = batch_size
+    failure_avg_bytes_ref: float | None = None
     pairs_sent = 0
     batches_sent = 0
     is_first_request = True
@@ -119,6 +121,11 @@ def set_state_with_batching(
 
             except Exception as e:
                 if _is_retryable_error(e):
+                    if batch_pairs:
+                        failed_total = sum(
+                            len(k) + len(v) for k, v in batch_pairs.items()
+                        )
+                        failure_avg_bytes_ref = failed_total / len(batch_pairs)
                     new_batch_size = current_batch_size // 2
                     logger.info(
                         f"Set state for {address} failed with batch size "
@@ -151,6 +158,25 @@ def set_state_with_batching(
                 f"Consider retrying or increasing min_batch_size."
             )
 
+        # Dynamic batch increase after successful push
+        if (
+            batch_pairs
+            and failure_avg_bytes_ref is not None
+            and current_batch_size < original_batch_size
+        ):
+            total_bytes = sum(len(k) + len(v) for k, v in batch_pairs.items())
+            current_avg = total_bytes / len(batch_pairs)
+            if current_avg < failure_avg_bytes_ref:
+                old_batch_size = current_batch_size
+                current_batch_size = min(current_batch_size * 2, original_batch_size)
+                logger.info(
+                    f"Set state for {address}: avg bytes/pair "
+                    f"({current_avg:.0f}) < failure reference "
+                    f"({failure_avg_bytes_ref:.0f}), increasing "
+                    f"batch size from {old_batch_size} to "
+                    f"{current_batch_size}"
+                )
+
         # Delay between successful batches
         if pairs_sent < len(pairs_list) and request_delay > 0:
             sleep(request_delay)
@@ -231,6 +257,9 @@ def get_account_storage_with_fallback(
     iterator_state: list[list[int]] = []
     iteration_succeeded = False
     current_batch_size = num_keys
+    original_batch_size = num_keys
+    last_successful_avg_bytes: float | None = None
+    failure_avg_bytes_ref: float | None = None
     bech32_address = address.to_bech32()
 
     try:
@@ -276,6 +305,7 @@ def get_account_storage_with_fallback(
 
                 except Exception as e:
                     if _is_retryable_error(e):
+                        failure_avg_bytes_ref = last_successful_avg_bytes
                         new_batch_size = current_batch_size // 2
                         logger.info(
                             f"Storage iteration for {bech32_address} failed with "
@@ -296,6 +326,28 @@ def get_account_storage_with_fallback(
                 )
                 break
 
+            # Dynamic batch increase after successful fetch
+            if pairs:
+                total_bytes = sum(len(k) + len(v) for k, v in pairs.items())
+                current_avg = total_bytes / len(pairs)
+                if (
+                    failure_avg_bytes_ref is not None
+                    and current_avg < failure_avg_bytes_ref
+                    and current_batch_size < original_batch_size
+                ):
+                    old_batch_size = current_batch_size
+                    current_batch_size = min(
+                        current_batch_size * 2, original_batch_size
+                    )
+                    logger.info(
+                        f"Storage iteration for {bech32_address}: "
+                        f"avg bytes/key ({current_avg:.0f}) < failure "
+                        f"reference ({failure_avg_bytes_ref:.0f}), "
+                        f"increasing batch size from {old_batch_size} "
+                        f"to {current_batch_size}"
+                    )
+                last_successful_avg_bytes = current_avg
+
             if iteration_succeeded:
                 break
 

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mxops"
-version = "3.1.1-dev1"
+version = "3.1.1-dev2"
 authors = [
   {name="Etienne Wallet"},
 ]

setup.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 3.1.1-dev1
+current_version = 3.1.1-dev2
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(-(?P<release>[^-0-9]+)(?P<build>\d+))?
 serialize = 
 	{major}.{minor}.{patch}-{release}{build}