From 941b9b5350cba66e32ec960da8e9af9a5d283f22 Mon Sep 17 00:00:00 2001
From: Ethan Dalool <ethan@voussoir.net>
Date: Mon, 17 Feb 2020 23:39:43 -0800
Subject: [PATCH] Add a lock to make Ratelimiter thread-safe.

---
 voussoirkit/ratelimiter.py | 21 +++++++++++++--------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/voussoirkit/ratelimiter.py b/voussoirkit/ratelimiter.py
index 02c2ad2..26629af 100644
--- a/voussoirkit/ratelimiter.py
+++ b/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)