Add ratelimiting package into Nova. After Austin it'll be pulled out into PyPI.

2 files changed, 163 insertions, 0 deletions

diff --git a/nova/api/rackspace/ratelimiting/__init__.py b/nova/api/rackspace/ratelimiting/__init__.py
new file mode 100644
index 000000000..176e7d66e
--- /dev/null
+++ b/nova/api/rackspace/ratelimiting/__init__.py
@@ -0,0 +1,103 @@
+"""Rate limiting of arbitrary actions."""
+
+import time
+import urllib
+import webob.dec
+import webob.exc
+
+
+# Convenience constants for the limits dictionary passed to Limiter().
+PER_SECOND = 1
+PER_MINUTE = 60
+PER_HOUR = 60 * 60
+PER_DAY = 60 * 60 * 24
+
+class Limiter(object):
+
+    """Class providing rate limiting of arbitrary actions."""
+
+    def __init__(self, limits):
+        """Create a rate limiter.
+
+        limits: a dict mapping from action name to a tuple.  The tuple contains
+        the number of times the action may be performed, and the time period
+        (in seconds) during which the number must not be exceeded for this
+        action.  Example: dict(reboot=(10, ratelimiting.PER_MINUTE)) would
+        allow 10 'reboot' actions per minute.
+        """
+        self.limits = limits
+        self._levels = {}
+
+    def perform(self, action_name, username='nobody'):
+        """Attempt to perform an action by the given username.
+
+        action_name: the string name of the action to perform.  This must
+        be a key in the limits dict passed to the ctor.
+
+        username: an optional string name of the user performing the action.
+        Each user has her own set of rate limiting counters.  Defaults to
+        'nobody' (so that if you never specify a username when calling
+        perform(), a single set of counters will be used.)
+
+        Return None if the action may proceed.  If the action may not proceed
+        because it has been rate limited, return the float number of seconds
+        until the action would succeed.
+        """
+        # Think of rate limiting as a bucket leaking water at 1cc/second.  The
+        # bucket can hold as many ccs as there are seconds in the rate
+        # limiting period (e.g. 3600 for per-hour ratelimits), and if you can
+        # perform N actions in that time, each action fills the bucket by
+        # 1/Nth of its volume.  You may only perform an action if the bucket
+        # would not overflow.
+        now = time.time()
+        key = '%s:%s' % (username, action_name)
+        last_time_performed, water_level = self._levels.get(key, (now, 0))
+        # The bucket leaks 1cc/second.
+        water_level -= (now - last_time_performed)
+        if water_level < 0:
+            water_level = 0
+        num_allowed_per_period, period_in_secs = self.limits[action_name]
+        # Fill the bucket by 1/Nth its capacity, and hope it doesn't overflow.
+        capacity = period_in_secs
+        new_level = water_level + (capacity * 1.0 / num_allowed_per_period)
+        if new_level > capacity:
+            # Delay this many seconds.
+            return new_level - capacity
+        self._levels[key] = (now, new_level)
+        return None
+
+
+# If one instance of this WSGIApps is unable to handle your load, put a
+# sharding app in front that shards by username to one of many backends.
+
+class WSGIApp(object):
+
+    """Application that tracks rate limits in memory.  Send requests to it of
+    this form:
+
+    POST /limiter/<username>/<urlencoded action>
+
+    and receive a 200 OK, or a 403 Forbidden with an X-Wait-Seconds header
+    containing the number of seconds to wait before the action would succeed.
+    """
+
+    def __init__(self, limiter):
+        """Create the WSGI application using the given Limiter instance."""
+        self.limiter = limiter
+
+    @webob.dec.wsgify
+    def __call__(req):
+        parts = req.path_info.split('/')
+        # format: /limiter/<username>/<urlencoded action>
+        if req.method != 'POST':
+            raise webob.exc.HTTPMethodNotAllowed()
+        if len(parts) != 4 or parts[1] != 'limiter':
+            raise webob.exc.HTTPNotFound()
+        username = parts[2]
+        action_name = urllib.unquote(parts[3])
+        delay = self.limiter.perform(action_name, username)
+        if delay:
+            return webob.exc.HTTPForbidden(
+                    headers={'X-Wait-Seconds': delay})
+        else:
+            return '' # 200 OK
diff --git a/nova/api/rackspace/ratelimiting/tests.py b/nova/api/rackspace/ratelimiting/tests.py
new file mode 100644
index 000000000..1983cdea8
--- /dev/null
+++ b/nova/api/rackspace/ratelimiting/tests.py
@@ -0,0 +1,60 @@
+import ratelimiting
+import time
+import unittest
+
+class Test(unittest.TestCase):
+
+    def setUp(self):
+        self.limits = {
+                'a': (5, ratelimiting.PER_SECOND),
+                'b': (5, ratelimiting.PER_MINUTE),
+                'c': (5, ratelimiting.PER_HOUR),
+                'd': (1, ratelimiting.PER_SECOND),
+                'e': (100, ratelimiting.PER_SECOND)}
+        self.rl = ratelimiting.Limiter(self.limits)
+
+    def exhaust(self, action, times_until_exhausted, **kwargs):
+        for i in range(times_until_exhausted):
+            when = self.rl.perform(action, **kwargs)
+            self.assertEqual(when, None)
+        num, period = self.limits[action]
+        delay = period * 1.0 / num
+        # Verify that we are now thoroughly delayed
+        for i in range(10):
+            when = self.rl.perform(action, **kwargs)
+            self.assertAlmostEqual(when, delay, 2)
+
+    def test_second(self):
+        self.exhaust('a', 5)
+        time.sleep(0.2)
+        self.exhaust('a', 1)
+        time.sleep(1)
+        self.exhaust('a', 5)
+
+    def test_minute(self):
+        self.exhaust('b', 5)
+
+    def test_one_per_period(self):
+        def allow_once_and_deny_once():
+            when = self.rl.perform('d')
+            self.assertEqual(when, None)
+            when = self.rl.perform('d')
+            self.assertAlmostEqual(when, 1, 2)
+            return when
+        time.sleep(allow_once_and_deny_once())
+        time.sleep(allow_once_and_deny_once())
+        allow_once_and_deny_once()
+
+    def test_we_can_go_indefinitely_if_we_spread_out_requests(self):
+        for i in range(200):
+            when = self.rl.perform('e')
+            self.assertEqual(when, None)
+            time.sleep(0.01)
+
+    def test_users_get_separate_buckets(self):
+        self.exhaust('c', 5, username='alice')
+        self.exhaust('c', 5, username='bob')
+        self.exhaust('c', 5, username='chuck')
+        self.exhaust('c', 0, username='chuck')
+        self.exhaust('c', 0, username='bob')
+        self.exhaust('c', 0, username='alice')