Skip to content

Commit

Permalink
fix: Using sliding window rate limiter
Browse files Browse the repository at this point in the history
  • Loading branch information
Bhargav Dodla committed Feb 5, 2025
1 parent 5710727 commit d578328
Showing 1 changed file with 27 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -242,6 +242,24 @@ def __exit__(self, exc_type, exc_val, exc_tb):
pass


class SlidingWindowRateLimiter:
def __init__(self, max_calls, period):
self.max_calls = max_calls
self.period = period
self.timestamps = [0] * max_calls
self.index = 0

def acquire(self):
now = time.time()
window_start = now - self.period

if self.timestamps[self.index] <= window_start:
self.timestamps[self.index] = now
self.index = (self.index + 1) % self.max_calls
return True
return False


class CassandraOnlineStore(OnlineStore):
"""
Cassandra/Astra DB online store implementation for Feast.
Expand Down Expand Up @@ -410,7 +428,8 @@ def on_failure(exc, concurrent_queue):
raise Exception("Error writing a batch") from exc

write_concurrency = config.online_store.write_concurrency
concurrent_queue: queue.Queue = queue.Queue(maxsize=write_concurrency)
rate_limiter = SlidingWindowRateLimiter(write_concurrency, 1)
concurrent_queue: queue.Queue = queue.Queue(maxsize=write_concurrency / 5)

project = config.project
ttl = (
Expand Down Expand Up @@ -444,6 +463,12 @@ def on_failure(exc, concurrent_queue):
timestamp,
)
batch.add(insert_cql, params)

# Wait until the rate limiter allows
if not rate_limiter.acquire():
while not rate_limiter.acquire():
time.sleep(0.001)

future = session.execute_async(batch)
concurrent_queue.put(future)
future.add_callbacks(
Expand All @@ -460,6 +485,7 @@ def on_failure(exc, concurrent_queue):
# this happens N-1 times, will be corrected outside:
if progress:
progress(1)
# Wait for all tasks to complete
while not concurrent_queue.empty():
time.sleep(0.001)

Expand Down

0 comments on commit d578328

Please sign in to comment.