fix(rate-limiting): fold policy scoping fields into the partition key

_partition_key only distinguished entries by tag_id/name/limit/period_seconds/
scope_by_key_hash/max_in_memory_cache_size, but two entries can share all of
those while disagreeing on enabled_for/disabled_for/apply_to_key_alias (the
same class of collision _DedupSignature and _policy_fingerprint already
guard against for dedup and bucket keys). Two such entries setting the same
max_in_memory_cache_size to get their own dedicated partition collided onto
one shared partition instead, letting one entry's high-cardinality traffic
evict the other's active counters. Folds _policy_fingerprint into the
partition key, keeping max_in_memory_cache_size as the trailing element
since _partition_for reads it via partition_key[-1].
This commit is contained in:
Deepanshu 2026-08-25 12:31:04 -04:00
parent 93b5dd7e9f
commit 311bf593f4
2 changed files with 43 additions and 4 deletions

View file

@ -883,7 +883,17 @@ def _resolve_max_in_memory_cache_size() -> int | None:
# always resolves to the same signature across index rebuilds, which is what
# keeps _PROXY_ModelBasedTagRateLimitsHook._partitions from leaking a fresh partition
# every time _TagRateLimitIndex rebuilds and reconstructs `_ConfiguredLimit`s.
_PartitionKey: TypeAlias = tuple[str, str, float, int, bool, int] | None
# The str before the trailing int is `_policy_fingerprint(entry)`: two
# entries can share tag_id/name (see
# test_bucket_key_differs_for_same_named_entries_with_different_scoping_only)
# while disagreeing on limit/period_seconds/scope_by_key_hash/enabled_for/
# disabled_for/apply_to_key_alias -- _DedupSignature already treats that as
# two distinct policies, so a shared max_in_memory_cache_size must not route
# them onto the same partition either, or one entry's high-cardinality
# traffic can evict the other's active counters from a cache neither entry
# asked to share. max_in_memory_cache_size stays the trailing element:
# `partition_key[-1]` reads it directly to size the partition's cache.
_PartitionKey: TypeAlias = tuple[str, str, str, int] | None
# Grouping type for async_log_success_event's per-partition tokens/dollars
# pipeline dispatch -- named only so the declaration fits on one line; see
# that method for why the grouping is needed.
@ -896,9 +906,7 @@ def _partition_key(entry: TagRateLimitEntry) -> _PartitionKey:
return (
entry.tag_id,
entry.name,
entry.limit,
entry.period_seconds,
entry.scope_by_key_hash,
_policy_fingerprint(entry),
entry.max_in_memory_cache_size,
)

View file

@ -3748,6 +3748,37 @@ def test_partition_key_distinguishes_entries_that_differ_only_by_scope_by_key_ha
assert _partition_key(unscoped) != _partition_key(scoped)
def test_partition_key_distinguishes_entries_that_differ_only_by_scoping_fields():
"""
A plain, unscoped entry and a scoped override can legitimately share
name/tag_id/limit/period_seconds/scope_by_key_hash (see
test_bucket_key_differs_for_same_named_entries_with_different_scoping_only)
while disagreeing on enabled_for/disabled_for/apply_to_key_alias --
_DedupSignature and _policy_fingerprint already treat that as two
distinct policies, so a shared max_in_memory_cache_size must not route
them onto the same in-memory partition either, or one entry's
high-cardinality traffic can evict the other's active counters from a
cache neither entry asked to share.
"""
base_kwargs = {"name": "daily", "tag_id": "end_user_id", "limit": 100, "period_seconds": 86400, "max_in_memory_cache_size": 50}
unscoped = TagRateLimitEntry(**base_kwargs)
enabled_for_scoped = TagRateLimitEntry(
**base_kwargs, enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",))
)
disabled_for_scoped = TagRateLimitEntry(
**base_kwargs, disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",))
)
alias_scoped = TagRateLimitEntry(**base_kwargs, apply_to_key_alias=("premium-key",))
keys = {
_partition_key(unscoped),
_partition_key(enabled_for_scoped),
_partition_key(disabled_for_scoped),
_partition_key(alias_scoped),
}
assert len(keys) == 4
@pytest.mark.asyncio
async def test_scope_by_key_hash_composes_with_max_in_memory_cache_size_and_key_ttl_seconds_overrides(time_controller):
"""