Claude-coauthored rate limiting put one customer's writes on one DynamoDB partition, so heavy traffic could throttle that customer and nearby tenants.

++import asyncio+import time+from contextlib import asynccontextmanager, contextmanager+from enum import Enum+from typing import Any, AsyncIterator, Iterator, Optional+

+    EntityNotFoundError,+    RateLimitExceeded,+    RateLimiterUnavailable,+)+from .lease import Lease, LeaseEntry, SyncLease+from .models import (+    BucketState,+    Entity,+    EntityCapacity,+    Limit,+    LimitStatus,+    ResourceCapacity,+)+from .repository import Repository+from .schema import DEFAULT_RESOURCE+++class FailureMode(Enum):+    """Behavior when DynamoDB is unavailable."""++    FAIL_OPEN = "open"  # Allow requests+    FAIL_CLOSED = "closed"  # Reject requests+++class RateLimiter:+    """+    Async rate limiter backed by DynamoDB.++    Implements token bucket algorithm with support for:+    - Multiple limits per entity/resource+    - Two-level hierarchy (parent/child entities)+    - Cascade mode (consume from entity + parent)+    - Stored limit configs+    - Usage analytics+    """++    def __init__(+        self,+        table_name: str,+        region: Optional[str] = None,+        endpoint_url: Optional[str] = None,+        create_table: bool = False,+        failure_mode: FailureMode = FailureMode.FAIL_CLOSED,+    ) -> None:+        self.table_name = table_name+        self.failure_mode = failure_mode+        self._create_table = create_table+        self._repository = Repository(+            table_name=table_name,+            region=region,+            endpoint_url=endpoint_url,+        )+        self._initialized = False++    async def _ensure_initialized(self) -> None:+        """Ensure table exists if create_table=True."""+        if self._initialized:+            return+        if self._create_table:+            await self._repository.create_table()+        self._initialized = True++    async def close(self) -> None:+        """Close the underlying connections."""+        await self._repository.close()++    async def __aenter__(self) -> "RateLimiter":+        await self._ensure_initialized()+        return self++    async def __aexit__(self, *args: Any) -> None:+        await self.close()++    # -------------------------------------------------------------------------+    # Entity management+    # -------------------------------------------------------------------------++    async def create_entity(+        self,+        entity_id: str,+        name: Optional[str] = None,+        parent_id: Optional[str] = None,+        metadata: Optional[dict] = None,+    ) -> Entity:+        """+        Create a new entity.++        Args:+            entity_id: Unique identifier for the entity+            name: Human-readable name (defaults to entity_id)+            parent_id: Parent entity ID (None for root/project entities)+            metadata: Additional metadata to store++        Returns:+            The created Entity++        Raises:+            EntityExistsError: If entity already exists+        """+        await self._ensure_initialized()+        return await self._repository.create_entity(+            entity_id=entity_id,

++import time+from typing import Any, Optional++import aioboto3+from botocore.exceptions import ClientError++from . import schema+from .exceptions import EntityExistsError, EntityNotFoundError+from .models import BucketState, Entity, Limit+++class Repository:+    """+    Async DynamoDB repository for rate limiter data.++    Handles all DynamoDB operations including entities, buckets,+    limit configs, and transactions.+    """++    def __init__(+        self,+        table_name: str,+        region: Optional[str] = None,+        endpoint_url: Optional[str] = None,+    ) -> None:+        self.table_name = table_name+        self.region = region+        self.endpoint_url = endpoint_url+        self._session: Optional[aioboto3.Session] = None+        self._client: Any = None++    async def _get_client(self) -> Any:+        """Get or create the DynamoDB client."""+        if self._client is None:+            self._session = aioboto3.Session()+            self._client = await self._session.client(+                "dynamodb",+                region_name=self.region,+                endpoint_url=self.endpoint_url,+            ).__aenter__()+        return self._client++    async def close(self) -> None:+        """Close the DynamoDB client."""+        if self._client is not None:+            await self._client.__aexit__(None, None, None)+            self._client = None+            self._session = None++    def _now_ms(self) -> int:+        """Current time in milliseconds."""+        return int(time.time() * 1000)++    # -------------------------------------------------------------------------+    # Table operations+    # -------------------------------------------------------------------------++    async def create_table(self) -> None:+        """Create the DynamoDB table if it doesn't exist."""+        client = await self._get_client()+        definition = schema.get_table_definition(self.table_name)++        try:+            await client.create_table(**definition)+            # Wait for table to be active+            waiter = client.get_waiter("table_exists")+            await waiter.wait(TableName=self.table_name)+        except ClientError as e:+            if e.response["Error"]["Code"] != "ResourceInUseException":+                raise++    async def delete_table(self) -> None:+        """Delete the DynamoDB table."""+        client = await self._get_client()

+    # Entity operations+    # -------------------------------------------------------------------------++    async def create_entity(+        self,+        entity_id: str,+        name: Optional[str] = None,+        parent_id: Optional[str] = None,+        metadata: Optional[dict] = None,+    ) -> Entity:+        """Create a new entity."""+        client = await self._get_client()+        now = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())++        item: dict[str, Any] = {+            "PK": {"S": schema.pk_entity(entity_id)},+            "SK": {"S": schema.sk_meta()},+            "entity_id": {"S": entity_id},+            "data": {+                "M": {+                    "name": {"S": name or entity_id},+                    "parent_id": {"S": parent_id} if parent_id else {"NULL": True},+                    "metadata": {"M": self._serialize_map(metadata or {})},+                    "created_at": {"S": now},+                }+            },+        }++        # Add GSI1 keys for parent lookup if this is a child+        if parent_id:+            item["GSI1PK"] = {"S": schema.gsi1_pk_parent(parent_id)}+            item["GSI1SK"] = {"S": schema.gsi1_sk_child(entity_id)}++        try:

 ) from .repository import Repository-from .schema import DEFAULT_RESOURCE+from .schema import DEFAULT_RESOURCE, WCU_LIMIT_NAME  _UNSET: Any = object()  # sentinel for detecting explicitly-passed deprecated params @@ -734,6 +734,21 @@ class RateLimiter:             if result.parent_result is not None and result.parent_result.success:                 assert result.parent_id is not None  # set by repository cache path                 await self._compensate_speculative(result.parent_id, resource, consume)++            # Shard doubling: if wcu exhausted, double shard_count and update cache+            if self._is_wcu_exhausted(result.old_buckets):+                await self._repository.bump_shard_count(entity_id, resource, result.shard_count)+                # Fall through to slow path (new shard bucket will be created)+                return None++            # Shard retry: if multi-shard, try another shard+            if result.shard_count > 1:+                retry_result = await self._retry_on_other_shard(+                    entity_id, resource, consume, ttl_seconds=None, result=result+                )+                if retry_result is not None:+                    return retry_result+             self._check_speculative_failure(result, consume, now_ms)             return None @@ -986,6 +1001,80 @@ class RateLimiter:         if not would_help:             raise RateLimitExceeded(statuses) +    _MAX_SHARD_RETRIES = 2++    @staticmethod+    def _is_wcu_exhausted(old_buckets: "list[BucketState] | None") -> bool:+        """Check if wcu infrastructure limit is exhausted in failed result."""+        if old_buckets is None:+            return False+        for b in old_buckets:+            if b.limit_name == WCU_LIMIT_NAME and b.tokens_milli < 1000:+                return True+        return False++    async def _retry_on_other_shard(+        self,+        entity_id: str,+        resource: str,+        consume: dict[str, int],+        ttl_seconds: int | None,+        result: "SpeculativeResult",+    ) -> "Lease | None":+        """Retry speculative consume on untried shards.++        Returns a Lease if a retry succeeds, None otherwise.+        """+        import random++        tried_shards = {result.shard_id}+        shard_count = result.shard_count++        for _ in range(self._MAX_SHARD_RETRIES):+            untried = [s for s in range(shard_count) if s not in tried_shards]+            if not untried:+                break+            new_shard = random.choice(untried)+            tried_shards.add(new_shard)++            retry = await self._repository.speculative_consume(

+        consume: dict[str, int],+        result: "SpeculativeResult",+    ) -> "Lease":+        """Build a Lease from a successful speculative result."""+        entries: list[LeaseEntry] = []+        for state in result.buckets:+            amount = consume.get(state.limit_name, 0)+            if amount == 0:+                continue+            limit = Limit.from_bucket_state(state)+            entries.append(+                LeaseEntry(+                    entity_id=state.entity_id,+                    resource=state.resource,+                    limit=limit,+                    state=state,+                    consumed=amount,+                    _cascade=result.cascade,+                    _parent_id=result.parent_id,+                )+            )+        return Lease(+            entries=entries,+            repository=self._repository,+            _committed=True,+        )+     async def _try_parent_only_acquire(         self,         parent_id: str,

 import asyncio import logging+import random import time import warnings from typing import TYPE_CHECKING, Any, cast@@ -2171,6 +2172,7 @@ class Repository:         resource: str,         consume: dict[str, int],         ttl_seconds: int | None = None,+        shard_id: int | None = None,     ) -> SpeculativeResult:         """Attempt speculative UpdateItem with condition check. @@ -2178,11 +2180,16 @@ class Repository:         issues child+parent UpdateItems concurrently via asyncio.gather         and returns nested parent_result. +        When ``shard_id`` is explicitly provided, targets that shard directly+        without cascade logic (used for shard retry).+         Args:             entity_id: Entity owning the bucket             resource: Resource name             consume: Amount per limit (tokens, not milli)             ttl_seconds: TTL in seconds from now, or None for no TTL change+            shard_id: Explicit shard to target (skips random selection and+                cascade logic). None means auto-select from entity cache.          Returns:             SpeculativeResult with:@@ -2190,17 +2197,35 @@ class Repository:             - On cache miss or non-cascade: parent_result is None             - On failure: old_buckets from ALL_OLD (or None if bucket missing)         """+        # Explicit shard_id: direct single-shard consume (shard retry path)+        if shard_id is not None:+            return await self._speculative_consume_single(+                entity_id, resource, consume, ttl_seconds, shard_id=shard_id+            )+         # Check entity cache for parallel cascade opportunity (issue #318)         cache_key = (self._namespace_id, entity_id)         cache_entry = self._entity_cache.get(cache_key) +        # Determine shard from entity cache+        shard_count = 1+        if cache_entry is not None:+            shard_count = cache_entry[2].get(resource, 1)+        effective_shard_id = random.randrange(shard_count) if shard_count > 1 else 0+         if cache_entry is not None:             cascade_cached, parent_id_cached, shards_cached = cache_entry             if cascade_cached and parent_id_cached:                 child_result: SpeculativeResult                 parent_result: SpeculativeResult                 child_result, parent_result = await asyncio.gather(-                    self._speculative_consume_single(entity_id, resource, consume, ttl_seconds),+                    self._speculative_consume_single(+                        entity_id,+                        resource,+                        consume,+                        ttl_seconds,+                        shard_id=effective_shard_id,+                    ),                     self._speculative_consume_single(                         parent_id_cached, resource, consume, ttl_seconds                     ),@@ -2222,7 +2247,9 @@ class Repository:                 return child_result          # Cache miss or non-cascade: single UpdateItem-        result = await self._speculative_consume_single(entity_id, resource, consume, ttl_seconds)+        result = await self._speculative_consume_single(+            entity_id, resource, consume, ttl_seconds, shard_id=effective_shard_id+        )         if result.success:             existing_shards = self._entity_cache.get(cache_key, (False, None, {}))[2]             existing_shards = {**existing_shards, resource: result.shard_count}

                         old_buckets=old_buckets,+                        shard_id=shard_id,+                        shard_count=old_shard_count,                     )                 else:-                    return SpeculativeResult(success=False)+                    return SpeculativeResult(success=False, shard_id=shard_id)             raise +    async def bump_shard_count(self, entity_id: str, resource: str, current_count: int) -> int:+        """Double shard_count on shard 0 (conditional write).++        Also updates the entity cache with the new shard_count.++        Returns the new shard_count, or the current if another client already doubled.+        """+        new_count = current_count * 2+        client = await self._get_client()+        try:+            await client.update_item(+                TableName=self.table_name,+                Key={+                    "PK": {"S": schema.pk_bucket(self._namespace_id, entity_id, resource, 0)},+                    "SK": {"S": schema.sk_state()},+                },+                UpdateExpression="SET shard_count = :new",+                ConditionExpression="shard_count = :old",+                ExpressionAttributeValues={+                    ":old": {"N": str(current_count)},+                    ":new": {"N": str(new_count)},