/
ratelimiter.py
73 lines (61 loc) · 2.28 KB
/
ratelimiter.py
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
'''
Created on 2018/4/19
Preventing many time-consuming operations to be done in the same loop
:author: hubo
'''
from vlcp.event.event import withIndices, Event
@withIndices("limiter", "index")
class RateLimitingEvent(Event):
pass
class RateLimiter(object):
"""
Limit operations executed in current loop, ensure sockets are
still processed in time-consuming operations
"""
def __init__(self, limit, container):
"""
:param limit: "resources" limited in a single loop. "resources"
can be any countable things like operations executed
or bytes sent
:param container: a `RoutineContainer`
"""
self._container = container
self._limit = limit
if self._limit <= 0:
raise ValueError("Limit must be greater than 0")
self._counter = 0
self._task = None
self._bottom_line = limit
async def _limiter_task(self):
current_index = 0
while True:
await self._container.do_events()
current_index += 1
if current_index * self._limit >= self._counter:
# Last event covers all (NOTICE: self._counter - 1 is the last limited)
break
else:
# This will release from current_index * limit to (current_index + 1) * limit - 1
self._container.scheduler.emergesend(RateLimitingEvent(self, current_index))
self._bottom_line += self._limit
# Reset counter
self._counter = 0
self._task = None
self._bottom_line = self._limit
async def limit(self, use = 1):
"""
Acquire "resources", wait until enough "resources" are acquired. For each loop,
`limit` number of "resources" are permitted.
:param use: number of "resouces" to be used.
:return: True if is limited
"""
c = self._counter
self._counter = c + use
if self._task is None:
self._task = self._container.subroutine(self._limiter_task(), False)
if c >= self._bottom_line:
# Limited
await RateLimitingEvent.createMatcher(self, c // self._limit)
return True
else:
return False