1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
|
# Original work Copyright 2013 Arnaud Porterie
# Modified work Copyright 2016 Frazer McLean
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import collections
import functools
import threading
import time
import sys
from textwrap import dedent
try:
import asyncio
except ImportError:
asyncio = None
__author__ = 'Frazer McLean <frazer@frazermclean.co.uk>'
__version__ = '1.2.0.post0'
__license__ = 'Apache'
__description__ = 'Simple python rate limiting object'
class RateLimiter(object):
"""Provides rate limiting for an operation with a configurable number of
requests for a time period.
"""
def __init__(self, max_calls, period=1.0, callback=None):
"""Initialize a RateLimiter object which enforces as much as max_calls
operations on period (eventually floating) number of seconds.
"""
if period <= 0:
raise ValueError('Rate limiting period should be > 0')
if max_calls <= 0:
raise ValueError('Rate limiting number of calls should be > 0')
# We're using a deque to store the last execution timestamps, not for
# its maxlen attribute, but to allow constant time front removal.
self.calls = collections.deque()
self.period = period
self.max_calls = max_calls
self.callback = callback
self._lock = threading.Lock()
self._alock = None
# Lock to protect creation of self._alock
self._init_lock = threading.Lock()
def _init_async_lock(self):
with self._init_lock:
if self._alock is None:
self._alock = asyncio.Lock()
def __call__(self, f):
"""The __call__ function allows the RateLimiter object to be used as a
regular function decorator.
"""
@functools.wraps(f)
def wrapped(*args, **kwargs):
with self:
return f(*args, **kwargs)
return wrapped
def __enter__(self):
with self._lock:
# We want to ensure that no more than max_calls were run in the allowed
# period. For this, we store the last timestamps of each call and run
# the rate verification upon each __enter__ call.
if len(self.calls) >= self.max_calls:
until = time.time() + self.period - self._timespan
if self.callback:
t = threading.Thread(target=self.callback, args=(until,))
t.daemon = True
t.start()
sleeptime = until - time.time()
if sleeptime > 0:
time.sleep(sleeptime)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
with self._lock:
# Store the last operation timestamp.
self.calls.append(time.time())
# Pop the timestamp list front (ie: the older calls) until the sum goes
# back below the period. This is our 'sliding period' window.
while self._timespan >= self.period:
self.calls.popleft()
async def __aenter__(self):
if self._alock is None:
self._init_async_lock()
async with self._alock:
# We want to ensure that no more than max_calls were run in the allowed
# period. For this, we store the last timestamps of each call and run
# the rate verification upon each __enter__ call.
if len(self.calls) >= self.max_calls:
until = time.time() + self.period - self._timespan
if self.callback:
asyncio.ensure_future(self.callback(until))
sleeptime = until - time.time()
if sleeptime > 0:
await asyncio.sleep(sleeptime)
return self
async def __aexit__(self, exc_type, exc_value, traceback):
return self.__exit__(exc_type, exc_value, traceback)
@property
def _timespan(self):
return self.calls[-1] - self.calls[0]
|
