def test_parsing_rates(self) -> None:
qa_pairs = [
("1/s", (1, 1)),
("10/10s", (10, 10)),
("1/m", (1, 60)),
("1/5m", (1, 300)),
]
for q, a in qa_pairs:
with self.subTest("Parsing rates...", rate=q):
self.assertEqual(parse_rate(q), a)
def is_rate_okay(task: Task, rate: str = "1/s", key=None) -> bool:
"""Keep a global throttle for tasks
Can be used via the `throttle_task` decorator above.
This implements the timestamp-based algorithm detailed here:
https://www.figma.com/blog/an-alternative-approach-to-rate-limiting/
Basically, you keep track of the number of requests and use the key
expiration as a reset of the counter.
So you have a rate of 5/m, and your first task comes in. You create a key:
celery_throttle:task_name = 1
celery_throttle:task_name.expires = 60
Another task comes in a few seconds later:
celery_throttle:task_name = 2
Do not update the ttl, it now has 58s remaining
And so forth, until:
celery_throttle:task_name = 6
(10s remaining)
We're over the threshold. Re-queue the task for later. 10s later:
Key expires b/c no more ttl.
Another task comes in:
celery_throttle:task_name = 1
celery_throttle:task_name.expires = 60
And so forth.
:param task: The task that is being checked
:param rate: How many times the task can be run during the time period.
Something like, 1/s, 2/h or similar.
:param key: If given, add this to the key placed in Redis for the item.
Typically, this will correspond to the value of an argument passed to the
throttled task.
:return: Whether the task should be throttled or not.
"""
key = f"celery_throttle:{task.name}{':' + str(key) if key else ''}"
r = make_redis_interface("CACHE")
allowed_task_count, duration = parse_rate(rate)
# Check the count in redis
actual_task_count = r.get(key)
if actual_task_count is None:
# No key. Set the value to 1 and set the ttl of the key.
r.set(key, 1, ex=duration)
return True
else:
# Key found. Check it.
if int(actual_task_count) <= allowed_task_count:
# We're OK to run the task. Increment our counter, and say things
# are OK by returning True.
new_count = r.incr(key, 1)
if new_count == 1:
# Safety check. If the count is 1 after incrementing, that
# means we created the key via the incr command. This can
# happen when it expires between when we `get` its value up
# above and when we increment it here. If that happens, it
# lacks a ttl! Set one.
#
# N.B. There's no need to worry about a race condition between
# our incr above, and the `expire` line here b/c without a ttl
# on this key, it can't expire between these two commands.
r.expire(key, duration)
return True
else:
# Over the threshold.
return False
def is_rate_okay(task: Task, rate: str = "1/s", key=None) -> bool:
"""Keep a global throttle for tasks
Can be used via the `throttle_task` decorator above.
This implements the timestamp-based algorithm detailed here:
https://www.figma.com/blog/an-alternative-approach-to-rate-limiting/
Basically, you keep track of the number of requests and use the key
expiration as a reset of the counter.
So you have a rate of 5/m, and your first task comes in. You create a key:
celery_throttle:task_name = 1
celery_throttle:task_name.expires = 60
Another task comes in a few seconds later:
celery_throttle:task_name = 2
Do not update the ttl, it now has 58s remaining
And so forth, until:
celery_throttle:task_name = 6
(10s remaining)
We're over the threshold. Re-queue the task for later. 10s later:
Key expires b/c no more ttl.
Another task comes in:
celery_throttle:task_name = 1
celery_throttle:task_name.expires = 60
And so forth.
:param task: The task that is being checked
:param rate: How many times the task can be run during the time period.
Something like, 1/s, 2/h or similar.
:param key: If given, add this to the key placed in Redis for the item.
Typically, this will correspond to the value of an argument passed to the
throttled task.
:return: Whether the task should be throttled or not.
"""
key = f"celery_throttle:{task.name}{':' + str(key) if key else ''}"
r = make_redis_interface("CACHE")
num_tasks, duration = parse_rate(rate)
# Check the count in redis
count = r.get(key)
if count is None:
# No key. Set the value to 1 and set the ttl of the key.
r.set(key, 1)
r.expire(key, duration)
return True
else:
# Key found. Check it.
if int(count) <= num_tasks:
# We're OK, run it.
r.incr(key, 1)
return True
else:
return False
def get_task_wait(
task: Task,
rate: str = "1/s",
key: str = None,
) -> float:
"""Keep a global throttle for tasks
Can be used via the `throttle_task` decorator above.
This implements the timestamp-based algorithm detailed here:
https://www.figma.com/blog/an-alternative-approach-to-rate-limiting/
Basically, you keep track of the number of requests and use the key
expiration as a reset of the counter.
So you have a rate of 5/m, and your first task comes in. You create a key:
celery_throttle:task_name = 1
celery_throttle:task_name.expires = 60
Another task comes in a few seconds later:
celery_throttle:task_name = 2
Do not update the ttl, it now has 58s remaining
And so forth, until:
celery_throttle:task_name = 6
(10s remaining)
We're over the threshold. Re-queue the task for later. 10s later:
Key expires b/c no more ttl.
Another task comes in:
celery_throttle:task_name = 1
celery_throttle:task_name.expires = 60
And so forth.
---
There is also a scheduler that figures out when to re-queue tasks. The idea
of the scheduler is simple: If you know the rate the tasks can be
processed, and if you're getting tasks faster than that rate, you can
schedule each one to take its turn at a reasonable specified time. This is
implemented by keeping a timestamp in redis indicating when the throttle
will no longer be clogged up.
Say you have a rate of 1/5s, and you get tasks as follows:
Elapsed Time | Task Number
-------------+------------
1s | 1
2s | 2
3s | 3
Task number 1 runs immediately, but sets a throttle for five seconds until
more work can be done. The second comes in and sees that the throttle has a
ttl of three remaining seconds, so it waits that long. Next, task number 3
comes in. It sees that the current window is full, and that the next one is
too — only one task every five seconds, right? It has to wait seven
seconds: two seconds (for the current window) *plus* 5 seconds (for the
next one, which is occupied by task two).
And so forth.
:param task: The task that is being checked
:param rate: How many times the task can be run during the time period.
Something like, 1/s, 2/h or similar.
:param key: If given, add this to the key placed in Redis for the item.
Typically, this will correspond to the value of an argument passed to the
throttled task.
:return: If throttled returns a float of how many seconds the task should
wait until the next open window for processing. If not throttled, returns
zero (i.e., don't wait).
"""
task_sub_key = f"{task.name}{':' + str(key) if key else ''}"
throttle_key = f"celery_throttle:{task_sub_key}"
r = make_redis_interface("CACHE")
allowed_task_count, duration = parse_rate(rate)
# Check the count in redis
actual_task_count = r.get(throttle_key)
if actual_task_count is None:
# No key. Set the value to 1 and set the ttl of the key.
r.set(throttle_key, 1, ex=duration)
return 0
# Key found. Check if we should throttle.
if int(actual_task_count) < allowed_task_count:
# We're OK to run the task. Increment our counter, and say things are
# OK by returning 0.
new_count = r.incr(throttle_key, 1)
if new_count == 1:
# Safety check. If the count is 1 after incrementing, that means we
# created the key via the incr command. This can happen when it
# expires between when we `get` its value up above and when we
# increment it here. If that happens, it lacks a ttl! Set one.
#
# N.B. There's no need to worry about a race condition between our
# incr above, and the `expire` line here b/c without a ttl on this
# key, it can't expire between these two commands.
r.expire(throttle_key, duration)
return 0
# Over the threshold. Find the next window and schedule the task.
schedule_key = f"celery_throttle:schedule:{task_sub_key}"
n = now()
delay = r.get(schedule_key)
if delay is None:
# No schedule yet. Run the task when the current throttle expires.
return set_for_next_window(r, throttle_key, schedule_key, n)
# We have a delay, so use it if it's in the future
delay = parser.parse(delay)
if delay < n:
# Delay is in the past. Run the task when the current throttle expires.
return set_for_next_window(r, throttle_key, schedule_key, n)
# Delay is in the future; use it and supplement it
new_time = delay + timedelta(seconds=duration / allowed_task_count)
r.set(schedule_key, str(new_time))
return (new_time - n).total_seconds()