Skip to content

fix(rate_limiter): use a threading.Lock for the synchronous methods (asyncio.Lock can't use plain with) #923

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Osamaali313 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(rate_limiter): use a threading.Lock for the synchronous methods (asyncio.Lock can't use plain with) #923

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0915ad0
fix(rate_limiter): use a threading.Lock for the synchronous methods
Osamaali313 Jul 15, 2026
18e97f5
fix(rate_limiter): close state-lock gap in _wait_if_needed; resize ad…
Osamaali313 Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 57 additions & 28 deletions ms_agent/utils/rate_limiter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# Copyright (c) ModelScope Contributors. All rights reserved.
import asyncio
import threading
import time
from collections import deque
from typing import Callable
Expand Down Expand Up @@ -47,6 +48,10 @@ def __init__(
# Concurrency control
self._semaphore = asyncio.Semaphore(max_concurrent)
self._lock = asyncio.Lock()
# Separate synchronous lock for the non-async bookkeeping methods
# (get_stats/reset/record_success/record_error); an asyncio.Lock only
# supports `async with`, so using it under a plain `with` raises.
self._state_lock = threading.Lock()

@gemini-code-assist gemini-code-assist Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Synchronization Gap between asyncio.Lock and threading.Lock

While introducing self._state_lock (a threading.Lock) prevents the TypeError in synchronous methods, it introduces a critical thread-safety issue.

The asynchronous _wait_if_needed method still accesses and mutates self._request_times and self._last_request_time under self._lock (an asyncio.Lock), which does not synchronize with self._state_lock.

Since synchronous methods like get_stats, reset, record_success, and record_error can be called from other threads (e.g., when running synchronous tools in a thread pool via run_in_executor), they can access or mutate self._request_times concurrently with _wait_if_needed. This will lead to race conditions and potential RuntimeError: deque mutated during iteration in get_stats.

Recommended Fix

To ensure thread safety without blocking the asyncio event loop during await points, use self._state_lock inside _wait_if_needed to protect all synchronous state accesses and mutations, but release it before any await calls.

Here is how _wait_if_needed should be refactored:

async def _wait_if_needed(self):
    now = time.time()

    # 1. Check minimum request interval safely
    with self._state_lock:
        last_request_time = self._last_request_time
        min_request_interval = self.min_request_interval

    if last_request_time > 0:
        elapsed = now - last_request_time
        if elapsed < min_request_interval:
            wait_time = min_request_interval - elapsed
            logger.debug(f'Enforcing min interval: waiting {wait_time:.3f}s')
            await asyncio.sleep(wait_time)
            now = time.time()

    # 2. Check requests per second limit
    async with self._lock:
        with self._state_lock:
            # Clean up expired request records
            cutoff_time = now - 1.0
            while self._request_times and self._request_times[0] < cutoff_time:
                self._request_times.popleft()

            # Check if rate limit reached
            has_reached_limit = len(self._request_times) >= self.max_requests_per_second
            wait_time = 1.0 - (now - self._request_times[0]) + 0.01 if has_reached_limit else 0

        if has_reached_limit and wait_time > 0:
            logger.debug(f'Rate limit reached: waiting {wait_time:.3f}s')
            await asyncio.sleep(wait_time)
            now = time.time()
            
            with self._state_lock:
                cutoff_time = now - 1.0
                while self._request_times and self._request_times[0] < cutoff_time:
                    self._request_times.popleft()

        with self._state_lock:
            self._request_times.append(now)
            self._last_request_time = now


logger.info(
f'RateLimiter initialized: {max_requests_per_second} req/s, '
Expand All @@ -69,42 +74,54 @@ async def _wait_if_needed(self):
now = time.time()

# 1. Check minimum request interval
if self._last_request_time > 0:
elapsed = now - self._last_request_time
with self._state_lock:
last_request_time = self._last_request_time
if last_request_time > 0:
elapsed = now - last_request_time
if elapsed < self.min_request_interval:
wait_time = self.min_request_interval - elapsed
logger.debug(
f'Enforcing min interval: waiting {wait_time:.3f}s')
await asyncio.sleep(wait_time)
now = time.time()

# 2. Check requests per second limit
# 2. Check requests per second limit. _lock serializes coroutines;
# _state_lock (a threading.Lock shared with the synchronous bookkeeping
# methods) guards the actual shared state and is never held across an
# await.
async with self._lock:
# Clean up expired request records (older than 1 second)
cutoff_time = now - 1.0
while self._request_times and self._request_times[0] < cutoff_time:
self._request_times.popleft()

# If rate limit reached, wait until oldest request expires
if len(self._request_times) >= self.max_requests_per_second:
oldest_request = self._request_times[0]
wait_time = 1.0 - (now
- oldest_request) + 0.01 # Add 10ms margin
if wait_time > 0:
logger.debug(
f'Rate limit reached ({self.max_requests_per_second} req/s): '
f'waiting {wait_time:.3f}s')
await asyncio.sleep(wait_time)
now = time.time()
# Clean up expired records
with self._state_lock:
# Clean up expired request records (older than 1 second)
cutoff_time = now - 1.0
while self._request_times and self._request_times[
0] < cutoff_time:
self._request_times.popleft()

# If rate limit reached, wait until oldest request expires
if len(self._request_times) >= self.max_requests_per_second:
oldest_request = self._request_times[0]
wait_time = 1.0 - (now - oldest_request
) + 0.01 # Add 10ms margin
else:
wait_time = 0.0

if wait_time > 0:
logger.debug(
f'Rate limit reached ({self.max_requests_per_second} req/s): '
f'waiting {wait_time:.3f}s')
await asyncio.sleep(wait_time)
now = time.time()
with self._state_lock:
# Clean up expired records after waiting
cutoff_time = now - 1.0
while self._request_times and self._request_times[
0] < cutoff_time:
self._request_times.popleft()

# Record this request time
self._request_times.append(now)
self._last_request_time = now
with self._state_lock:
# Record this request time
self._request_times.append(now)
self._last_request_time = now

async def execute(self, func: Callable, *args, **kwargs):
"""
Expand Down Expand Up @@ -137,7 +154,7 @@ def get_stats(self) -> dict:
Returns:
Dictionary containing current state
"""
with self._lock:
with self._state_lock:
now = time.time()
cutoff_time = now - 1.0
recent_requests = sum(1 for t in self._request_times
Expand All @@ -161,7 +178,7 @@ def get_stats(self) -> dict:

def reset(self):
"""Reset rate limiter state"""
with self._lock:
with self._state_lock:
self._request_times.clear()
self._last_request_time = 0.0
logger.info('RateLimiter reset')
Expand Down Expand Up @@ -219,7 +236,7 @@ def __init__(

def record_success(self):
"""Record successful request"""
with self._lock:
with self._state_lock:

@gemini-code-assist gemini-code-assist Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Immutable deque maxlen Bug in AdaptiveRateLimiter

In AdaptiveRateLimiter, the rate limit self.max_requests_per_second is dynamically adjusted in record_success and record_error. However, self._request_times is initialized in the parent class constructor as a deque with a fixed maxlen:

self._request_times = deque(maxlen=max_requests_per_second)

Since a deque's maxlen is immutable after instantiation, changing self.max_requests_per_second does not update the capacity of self._request_times.

  • If the rate limit increases (e.g., from 2 to 10), self._request_times remains capped at 2 elements. Consequently, len(self._request_times) will never reach 10, and the rate limiter will never trigger a wait, effectively disabling rate limiting.
  • If the rate limit decreases, the deque will still allow more elements than the new limit before discarding.

Recommended Fix

Whenever self.max_requests_per_second is updated, recreate self._request_times with the new maxlen to preserve the existing history up to the new limit:

self._request_times = deque(self._request_times, maxlen=new_rps)

For example, in record_success:

if new_rps > old_rps:
    self.max_requests_per_second = new_rps
    self._request_times = deque(self._request_times, maxlen=new_rps)
    logger.info(...)

And similarly in record_error whenever self.max_requests_per_second is decreased.

self._total_requests += 1
self._consecutive_successes += 1
self._consecutive_errors = 0
Expand All @@ -231,6 +248,10 @@ def record_success(self):
round(old_rps * self._recovery_factor), self._max_rps)
if new_rps > old_rps:
self.max_requests_per_second = new_rps
# deque maxlen is immutable, so resize it to the new limit
# (preserving recent history) or the rps change has no effect.
self._request_times = deque(
self._request_times, maxlen=new_rps)
logger.info(
f'Rate limit increased: {old_rps} → {new_rps} req/s '
f'(after {self._consecutive_successes} successes)')
Expand All @@ -243,7 +264,7 @@ def record_error(self, is_rate_limit_error: bool = False):
Args:
is_rate_limit_error: Whether this is a rate limit error
"""
with self._lock:
with self._state_lock:
self._total_requests += 1
self._total_errors += 1
self._consecutive_errors += 1
Expand All @@ -256,6 +277,10 @@ def record_error(self, is_rate_limit_error: bool = False):
int(old_rps * self._backoff_factor), self._min_rps)
if new_rps < old_rps:
self.max_requests_per_second = new_rps
# deque maxlen is immutable, so resize it to the new limit
# (preserving recent history) or the rps change has no effect.
self._request_times = deque(
self._request_times, maxlen=new_rps)
# Also increase minimum request interval
self.min_request_interval = min(
self.min_request_interval * 1.5, 2.0)
Expand All @@ -271,6 +296,10 @@ def record_error(self, is_rate_limit_error: bool = False):
int(old_rps * self._backoff_factor), self._min_rps)
if new_rps < old_rps:
self.max_requests_per_second = new_rps
# deque maxlen is immutable, so resize it to the new limit
# (preserving recent history) or the rps change has no effect.
self._request_times = deque(
self._request_times, maxlen=new_rps)
logger.warning(
f'Multiple errors detected! Reducing rate: {old_rps} → {new_rps} req/s '
f'(after {self._consecutive_errors} errors)')
Expand All @@ -279,7 +308,7 @@ def record_error(self, is_rate_limit_error: bool = False):
def get_stats(self) -> dict:
"""Get extended statistics"""
stats = super().get_stats()
with self._lock:
with self._state_lock:
stats.update({
'total_requests':
self._total_requests,
Expand Down