diff --git a/lib/tfw/decorators/rate_limiter.py b/lib/tfw/decorators/rate_limiter.py index d1972b8..abe6453 100644 --- a/lib/tfw/decorators/rate_limiter.py +++ b/lib/tfw/decorators/rate_limiter.py @@ -1,25 +1,100 @@ # Copyright (C) 2018 Avatao.com Innovative Learning Kft. # All Rights Reserved. See LICENSE file for details. -from functools import wraps +from functools import wraps, partial from time import time, sleep +from tfw.decorators.lazy_property import lazy_property + class RateLimiter: + """ + Decorator class for rate limiting, blocking. + + When applied to a function this decorator will apply rate limiting + if the function is invoked more frequently than rate_per_seconds. + + By default rate limiting means sleeping until the next invocation time + as per __init__ parameter rate_per_seconds. + + Note that this decorator BLOCKS THE THREAD it is being executed on, + so it is only acceptable for stuff running on a separate thread. + + If this is no good for you please refer to AsyncRateLimiter in this module, + which is designed not to block and use the IOLoop it is being called from. + """ def __init__(self, rate_per_second): + """ + :param rate_per_second: max frequency the decorated method should be + invoked with + """ self.min_interval = 1 / float(rate_per_second) + self.fun = None self.last_call = time() + def action(self, seconds_to_next_call): + if seconds_to_next_call: + sleep(seconds_to_next_call) + self.fun() + def __call__(self, fun): @wraps(fun) def wrapper(*args, **kwargs): - self._limit_rate() - fun(*args, **kwargs) + self.fun = partial(fun, *args, **kwargs) + limit_seconds = self._limit_rate() + self.action(limit_seconds) return wrapper def _limit_rate(self): - since_last_call = time() - self.last_call - to_next_call = self.min_interval - since_last_call + seconds_since_last_call = time() - self.last_call + seconds_to_next_call = self.min_interval - seconds_since_last_call + + if seconds_to_next_call > 0: + return seconds_to_next_call self.last_call = time() - if to_next_call > 0: - sleep(to_next_call) + return 0 + + +class AsyncRateLimiter(RateLimiter): + """ + Decorator class for rate limiting, non-blocking. + + The semantics of the rate limiting: + - unlike RateLimiter this decorator never blocks, instead it adds an async + callback version of the decorated function to the IOLoop + (to be executed after the rate limiting has expired). + - the timing works similarly to RateLimiter + """ + def __init__(self, rate_per_second, ioloop_factory): + """ + :param rate_per_second: max frequency the decorated method should be + invoked with + :param ioloop_factory: callable that should return an instance of the + IOLoop of the application + """ + self._ioloop_factory = ioloop_factory + self._ioloop = None + self._last_callback = None + + self._make_action_thread_safe() + super().__init__(rate_per_second=rate_per_second) + + def _make_action_thread_safe(self): + self.action = partial(self.ioloop.add_callback, self.action) + + @lazy_property + def ioloop(self): + return self._ioloop_factory() + + def action(self, seconds_to_next_call): + if self._last_callback: + self.ioloop.remove_timeout(self._last_callback) + + self._last_callback = self.ioloop.call_later( + seconds_to_next_call, + self.fun_with_debounce + ) + + def fun_with_debounce(self): + self.last_call = time() + self.fun()