|
@@ -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,
|
|
|
)
|
|
|
|