From 32085b4a6768daa7223002893a7f1de83fa77e2f Mon Sep 17 00:00:00 2001
From: Andrei <b1naryth1ef@gmail.com>
Date: Tue, 11 Jul 2017 13:50:30 -0700
Subject: [PATCH] [util] fix SimpleLimiter not limiting correctly

This would cause our GatewayClient to get rate limited when guild member
syncing was enabled for clients with a large number of guilds. The bug
had to do with the limiter releaseing all the requests after the initial
period it blocked for. This was resolved by rewriting the SimpleLimiter
to use a semaphore.
---
 disco/util/limiter.py      | 28 +++++------------------
 tests/test_util_limiter.py | 46 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 22 deletions(-)
 create mode 100644 tests/test_util_limiter.py

diff --git a/disco/util/limiter.py b/disco/util/limiter.py
index ccb7622..493d4ae 100644
--- a/disco/util/limiter.py
+++ b/disco/util/limiter.py
@@ -6,33 +6,17 @@ class SimpleLimiter(object):
     def __init__(self, total, per):
         self.total = total
         self.per = per
+        self._lock = gevent.lock.Semaphore(total)
 
         self.count = 0
         self.reset_at = 0
-
-        self.event = None
-
-    def backoff(self):
-        self.event = gevent.event.Event()
-        gevent.sleep(self.reset_at - time.time())
-        self.count = 0
-        self.reset_at = 0
-        if self.event:
-            self.event.set()
         self.event = None
 
     def check(self):
-        if self.event:
-            self.event.wait()
-
-        self.count += 1
+        self._lock.acquire()
 
-        if not self.reset_at:
-            self.reset_at = time.time() + self.per
-            return
-        elif self.reset_at < time.time():
-            self.count = 1
-            self.reset_at = time.time()
+        def _release_lock():
+            gevent.sleep(self.per)
+            self._lock.release()
 
-        if self.count > self.total and self.reset_at > time.time():
-            self.backoff()
+        gevent.spawn(_release_lock)
diff --git a/tests/test_util_limiter.py b/tests/test_util_limiter.py
new file mode 100644
index 0000000..0c1fabc
--- /dev/null
+++ b/tests/test_util_limiter.py
@@ -0,0 +1,46 @@
+import time
+import gevent
+from unittest import TestCase
+
+from disco.util.limiter import SimpleLimiter
+
+
+class TestSimpleLimiter(TestCase):
+    def test_many_wait_ratelimiter(self):
+        limit = SimpleLimiter(5, 1)
+        many = []
+
+        def check(lock):
+            limit.check()
+            lock.release()
+
+        start = time.time()
+        for _ in range(16):
+            lock = gevent.lock.Semaphore()
+            lock.acquire()
+            many.append(lock)
+            gevent.spawn(check, lock)
+
+        for item in many:
+            item.acquire()
+
+        self.assertGreater(time.time() - start, 3)
+
+    def test_nowait_ratelimiter(self):
+        limit = SimpleLimiter(5, 1)
+
+        start = time.time()
+        for _ in range(5):
+            limit.check()
+
+        self.assertLess(time.time() - start, 1)
+
+    def test_single_wait_ratelimiter(self):
+        limit = SimpleLimiter(5, 1)
+
+        start = time.time()
+        for _ in range(10):
+            limit.check()
+
+
+        self.assertEqual(int(time.time() - start), 1)