🐛 Fixed deadlock in rate limiter

This commit is contained in:
GeoffreyCoulaud
2023-06-03 17:50:00 +02:00
parent dc47c850ce
commit 216c3f5dae

View File

@@ -27,7 +27,7 @@ class TokenBucketRateLimiter(AbstractContextManager):
@n_tokens.setter @n_tokens.setter
def n_tokens(self, value: int): def n_tokens(self, value: int):
with self.__n_tokens_lock: with self.__n_tokens_lock:
self.n_tokens = value self.__n_tokens = value
def __init__( def __init__(
self, self,
@@ -35,22 +35,24 @@ class TokenBucketRateLimiter(AbstractContextManager):
max_tokens: Optional[int] = None, max_tokens: Optional[int] = None,
initial_tokens: Optional[int] = None, initial_tokens: Optional[int] = None,
) -> None: ) -> None:
"""Initialize the limiter"""
# Initialize default values # Initialize default values
self.queue_lock = Lock()
if max_tokens is not None: if max_tokens is not None:
self.MAX_TOKENS = max_tokens self.MAX_TOKENS = max_tokens
if refill_spacing_seconds is not None: if refill_spacing_seconds is not None:
self.REFILL_SPACING_SECONDS = refill_spacing_seconds self.REFILL_SPACING_SECONDS = refill_spacing_seconds
# Initialize the bucket # Create locks
self.__n_tokens_lock = Lock()
self.queue_lock = Lock()
# Initialize the number of tokens in the bucket
self.bucket = BoundedSemaphore(self.MAX_TOKENS) self.bucket = BoundedSemaphore(self.MAX_TOKENS)
missing = 0 if initial_tokens is None else self.MAX_TOKENS - initial_tokens missing = 0 if initial_tokens is None else self.MAX_TOKENS - initial_tokens
missing = max(0, min(missing, self.MAX_TOKENS)) missing = max(0, min(missing, self.MAX_TOKENS))
for _ in range(missing): for _ in range(missing):
self.bucket.acquire() self.bucket.acquire()
# Initialize the counter
self.__n_tokens_lock = Lock()
self.n_tokens = self.MAX_TOKENS - missing self.n_tokens = self.MAX_TOKENS - missing
# Spawn daemon thread that refills the bucket # Spawn daemon thread that refills the bucket