Drafting ThreadsafeBurstyRateLimiterState

Related to #133

Author: Yomguithereal

minet/crawl/crawler.py

@@ -403,7 +403,7 @@ def __init__(
             "spoof_ua": spoof_ua,
             "use_pycurl": use_pycurl,
             "compressed": compressed,
-            "known_encoding": known_encoding
+            "known_encoding": known_encoding,
         }
 
     def __repr__(self):

minet/rate_limiting.py

@@ -1,8 +1,11 @@
+from typing import Optional
+
 import time
 import functools
+from threading import Event, Lock
 
 
-class RateLimiter(object):
+class RateLimiter:
     """
     Naive rate limiter context manager with smooth output.
 
@@ -70,7 +73,7 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
         return self.exit()
 
 
-class RetryableIterator(object):
+class RetryableIterator:
     """
     Iterator exposing a #.retry method that will make sure the next item
     is the same as the current one.
@@ -99,7 +102,7 @@ def retry(self):
         self.retried = True
 
 
-class RateLimitedIterator(object):
+class RateLimitedIterator:
     """
     Handy iterator wrapper that will yield its items while respecting a given
     rate limit and that will not sleep needlessly when the iterator is
@@ -143,7 +146,7 @@ def __iter__(self):
             self.rate_limiter.exit()
 
 
-class RateLimiterState(object):
+class RateLimiterState:
     def __init__(self, max_per_period: int, period: float = 1.0):
         max_per_second = max_per_period / period
         self.min_interval = 1.0 / max_per_second
@@ -163,6 +166,47 @@ def update(self):
         self.last_entry = time.perf_counter()
 
 
+class ThreadsafeBurstyRateLimiterState:
+    def __init__(self, max_per_period: int, period: float = 1.0):
+        self.max_per_period = max_per_period
+        self.period = period
+
+        self.current_burst = 0
+        self.time_of_next_burst: Optional[float] = None
+
+        self.event = Event()
+        self.lock = Lock()
+
+        self.event.set()
+
+    def wait_if_needed(self):
+        self.event.wait()
+        self.lock.acquire()
+
+        if self.current_burst < self.max_per_period:
+            if self.time_of_next_burst is None:
+                self.time_of_next_burst = time.perf_counter() + self.period
+
+            self.current_burst += 1
+            self.lock.release()
+            return
+
+        assert self.time_of_next_burst is not None
+
+        delta = time.perf_counter() - self.time_of_next_burst
+
+        self.time_of_next_burst = None
+        self.current_burst = 0
+
+        if delta > 0:
+            self.event.clear()
+            self.lock.release()
+            time.sleep(delta)
+            self.event.set()
+        else:
+            self.lock.release()
+
+
 def rate_limited(max_per_period, period=1.0):
     state = RateLimiterState(max_per_period, period)
 

minet/twitter/api_scraper.py

@@ -481,7 +481,7 @@ def request(self, url, headers=None, method="GET"):
             pool_manager=self.pool_manager,
             spoof_ua=True,
             method=method,
-            headers=headers
+            headers=headers,
         )
 
     # def acquire_guest_token(self):

minet/web.py

@@ -1127,7 +1127,7 @@ def request_jsonrpc(
         url,
         pool_manager=pool_manager,
         method="POST",
-        json_body={"method": method, "params": params}
+        json_body={"method": method, "params": params},
     )