|
|
|
@ -57,6 +57,7 @@ class Ratelimiter: |
|
|
|
|
rate_hz: Optional[float] = None, |
|
|
|
|
burst_count: Optional[int] = None, |
|
|
|
|
update: bool = True, |
|
|
|
|
n_actions: int = 1, |
|
|
|
|
_time_now_s: Optional[int] = None, |
|
|
|
|
) -> Tuple[bool, float]: |
|
|
|
|
"""Can the entity (e.g. user or IP address) perform the action? |
|
|
|
@ -76,6 +77,9 @@ class Ratelimiter: |
|
|
|
|
burst_count: How many actions that can be performed before being limited. |
|
|
|
|
Overrides the value set during instantiation if set. |
|
|
|
|
update: Whether to count this check as performing the action |
|
|
|
|
n_actions: The number of times the user wants to do this action. If the user |
|
|
|
|
cannot do all of the actions, the user's action count is not incremented |
|
|
|
|
at all. |
|
|
|
|
_time_now_s: The current time. Optional, defaults to the current time according |
|
|
|
|
to self.clock. Only used by tests. |
|
|
|
|
|
|
|
|
@ -124,17 +128,20 @@ class Ratelimiter: |
|
|
|
|
time_delta = time_now_s - time_start |
|
|
|
|
performed_count = action_count - time_delta * rate_hz |
|
|
|
|
if performed_count < 0: |
|
|
|
|
# Allow, reset back to count 1 |
|
|
|
|
allowed = True |
|
|
|
|
performed_count = 0 |
|
|
|
|
time_start = time_now_s |
|
|
|
|
action_count = 1.0 |
|
|
|
|
elif performed_count > burst_count - 1.0: |
|
|
|
|
|
|
|
|
|
# This check would be easier read as performed_count + n_actions > burst_count, |
|
|
|
|
# but performed_count might be a very precise float (with lots of numbers |
|
|
|
|
# following the point) in which case Python might round it up when adding it to |
|
|
|
|
# n_actions. Writing it this way ensures it doesn't happen. |
|
|
|
|
if performed_count > burst_count - n_actions: |
|
|
|
|
# Deny, we have exceeded our burst count |
|
|
|
|
allowed = False |
|
|
|
|
else: |
|
|
|
|
# We haven't reached our limit yet |
|
|
|
|
allowed = True |
|
|
|
|
action_count += 1.0 |
|
|
|
|
action_count = performed_count + n_actions |
|
|
|
|
|
|
|
|
|
if update: |
|
|
|
|
self.actions[key] = (action_count, time_start, rate_hz) |
|
|
|
@ -182,6 +189,7 @@ class Ratelimiter: |
|
|
|
|
rate_hz: Optional[float] = None, |
|
|
|
|
burst_count: Optional[int] = None, |
|
|
|
|
update: bool = True, |
|
|
|
|
n_actions: int = 1, |
|
|
|
|
_time_now_s: Optional[int] = None, |
|
|
|
|
): |
|
|
|
|
"""Checks if an action can be performed. If not, raises a LimitExceededError |
|
|
|
@ -201,6 +209,9 @@ class Ratelimiter: |
|
|
|
|
burst_count: How many actions that can be performed before being limited. |
|
|
|
|
Overrides the value set during instantiation if set. |
|
|
|
|
update: Whether to count this check as performing the action |
|
|
|
|
n_actions: The number of times the user wants to do this action. If the user |
|
|
|
|
cannot do all of the actions, the user's action count is not incremented |
|
|
|
|
at all. |
|
|
|
|
_time_now_s: The current time. Optional, defaults to the current time according |
|
|
|
|
to self.clock. Only used by tests. |
|
|
|
|
|
|
|
|
@ -216,6 +227,7 @@ class Ratelimiter: |
|
|
|
|
rate_hz=rate_hz, |
|
|
|
|
burst_count=burst_count, |
|
|
|
|
update=update, |
|
|
|
|
n_actions=n_actions, |
|
|
|
|
_time_now_s=time_now_s, |
|
|
|
|
) |
|
|
|
|
|
|
|
|
|