Add a lock to make Ratelimiter thread-safe.

voussoirkit/ratelimiter.py

@@ -1,6 +1,6 @@
+import threading
 import time
 
-
 class Ratelimiter:
     def __init__(self, allowance, period=1, operation_cost=1, mode='sleep'):
         '''
@@ -30,6 +30,7 @@ class Ratelimiter:
         self.period = period
         self.operation_cost = operation_cost
         self.mode = mode
+        self.lock = threading.Lock()
 
         self.last_operation = time.time()
         self.balance = 0
@@ -38,13 +39,7 @@ class Ratelimiter:
     def gain_rate(self):
         return self.allowance / self.period
 
-    def limit(self, cost=None):
-        '''
-        See the main class docstring for info about cost and mode behavior.
-        '''
-        if cost is None:
-            cost = self.operation_cost
-
+    def _limit(self, cost):
         time_diff = time.time() - self.last_operation
         self.balance += time_diff * self.gain_rate
         self.balance = min(self.balance, self.allowance)
@@ -65,3 +60,13 @@ class Ratelimiter:
 
         self.last_operation = time.time()
         return successful
+
+    def limit(self, cost=None):
+        '''
+        See the main class docstring for info about cost and mode behavior.
+        '''
+        if cost is None:
+            cost = self.operation_cost
+
+        with self.lock:
+            return self._limit(cost)