Merge lp:~sidnei/txstatsd/add-rate-limit into lp:~txstatsd-dev/txstatsd/distinct-plugin

[Sidnei da Silva]

Add a simple rate limit implementation to keep rate-limiting stats in redis.

Add a simple rate limit implementation to keep rate-limiting stats in redis.

[Lucio Torre (lucio.torre)]

review: Approve

=== added file '.bzrignore'
--- .bzrignore	1970-01-01 00:00:00 +0000
+++ .bzrignore	2013-12-03 12:47:43 +0000
@@ -0,0 +1,3 @@
+twisted/plugins/dropin.cache
+tmp
+_trial_temp

=== modified file 'Makefile'
--- Makefile	2011-12-15 20:19:11 +0000
+++ Makefile	2013-12-03 12:47:43 +0000
@@ -1,26 +1,28 @@
 
 trial:
-	trial distinctdb/
+	@trial distinctdb/
 
-start-database: 
-	./bin/start-database.sh
-	psql -h `pwd`/tmp/db1 -d distinct < bin/schema.sql
+start-database:
+	@./bin/start-database.sh
+	@psql -h `pwd`/tmp/db1 -d distinct < bin/schema.sql
 
 stop-database:
-	./bin/stop-database.sh
+	@./bin/stop-database.sh
 
 start-redis:
-	./bin/start-redis.sh
+	@./bin/start-redis.sh
 
 stop-redis:
-	./bin/stop-redis.sh
+	@./bin/stop-redis.sh
 
 start: start-database start-redis
 
 stop: stop-redis stop-database
 
 clean:
-	rm -rf ./tmp/
-test: start trial stop clean
+	@rm -rf ./tmp/
+
+test: start
+	@$(MAKE) trial stop clean || ($(MAKE) stop clean; false)
 
 .PHONY: trial test start-database stop-database start-redis stop-redis start stop

=== modified file 'bin/start-database.sh'
--- bin/start-database.sh	2012-04-26 18:16:25 +0000
+++ bin/start-database.sh	2013-12-03 12:47:43 +0000
@@ -1,4 +1,6 @@
-#! /bin/bash
+#!/usr/bin/env bash
+
+set -u
 
 ROOTDIR=${ROOTDIR:-`bzr root`}
 if [ ! -d "$ROOTDIR"  ]; then
@@ -12,6 +14,7 @@
 
 function setup_database() {
     local TESTDIR=$1
+    local PGBINDIR=
 
     echo "## Starting postgres in $TESTDIR ##"
     mkdir -p "$TESTDIR/data"
@@ -19,14 +22,15 @@
 
     export PGHOST="$TESTDIR"
     export PGDATA="$TESTDIR/data"
-    if [ -d /usr/lib/postgresql/8.4 ]; then
-        export PGBINDIR=/usr/lib/postgresql/8.4/bin
-    elif [ -d /usr/lib/postgresql/8.3 ]; then
-        export PGBINDIR=/usr/lib/postgresql/8.3/bin
-    elif [ -d /usr/lib/postgresql/9.1 ]; then
-        export PGBINDIR=/usr/lib/postgresql/9.1/bin
-    else
+    for pgver in 9.1 8.4; do
+        if [ -d /usr/lib/postgresql/$pgver ]; then
+            PGBINDIR=/usr/lib/postgresql/$pgver/bin
+            break
+        fi
+    done
+    if [ -z "$PGBINDIR" ]; then
         echo "Cannot find valid parent for PGBINDIR"
+        exit 1
     fi
     $PGBINDIR/initdb -E UNICODE -D $PGDATA
     # set up the database options file

=== modified file 'bin/start-redis.sh'
--- bin/start-redis.sh	2011-12-16 00:27:45 +0000
+++ bin/start-redis.sh	2013-12-03 12:47:43 +0000
@@ -1,5 +1,4 @@
-#!/bin/bash
+#!/usr/bin/env bash
 
 mkdir -p tmp/redis
 /sbin/start-stop-daemon --start -b -m -d . -p tmp/redis.pid --exec /usr/bin/redis-server -- `pwd`/bin/redis.conf
-

=== modified file 'bin/stop-database.sh'
--- bin/stop-database.sh	2012-04-26 18:16:25 +0000
+++ bin/stop-database.sh	2013-12-03 12:47:43 +0000
@@ -1,4 +1,6 @@
-#! /bin/bash
+#!/usr/bin/env bash
+
+set -u
 
 ROOTDIR=${ROOTDIR:-`bzr root`}
 if [ ! -d "$ROOTDIR"  ]; then
@@ -6,14 +8,15 @@
     exit 1
 fi
 
-if [ -d /usr/lib/postgresql/8.4 ]; then
-    PGBINDIR=/usr/lib/postgresql/8.4/bin
-elif [ -d /usr/lib/postgresql/8.3 ]; then
-    PGBINDIR=/usr/lib/postgresql/8.3/bin
-elif [ -d /usr/lib/postgresql/9.1 ]; then
-    PGBINDIR=/usr/lib/postgresql/9.1/bin
-else
+for pgver in 9.1 8.4; do
+    if [ -d /usr/lib/postgresql/$pgver ]; then
+        PGBINDIR=/usr/lib/postgresql/$pgver/bin
+        break
+    fi
+done
+if [ -z "$PGBINDIR" ]; then
     echo "Cannot find valid parent for PGBINDIR"
+    exit 1
 fi
 
 # setting PGDATA tells pg_ctl which DB to talk to

=== modified file 'bin/stop-redis.sh'
--- bin/stop-redis.sh	2011-12-16 00:27:45 +0000
+++ bin/stop-redis.sh	2013-12-03 12:47:43 +0000
@@ -1,3 +1,3 @@
-#!/bin/bash
+#!/usr/bin/env bash
 
-/sbin/start-stop-daemon --stop -p tmp/redis.pid --exec /usr/bin/redis-server -- `pwd`/bin/redis.conf
+/sbin/start-stop-daemon --oknodo --stop -p tmp/redis.pid --exec /usr/bin/redis-server -- `pwd`/bin/redis.conf

=== modified file 'distinctdb/distinctmetric.py'
--- distinctdb/distinctmetric.py	2012-06-27 20:04:53 +0000
+++ distinctdb/distinctmetric.py	2013-12-03 12:47:43 +0000
@@ -6,6 +6,7 @@
 import psycopg2
 import redis
 
+from .ratelimit import RateLimiter
 from zope.interface import implements
 from twisted.internet import reactor, threads
 from twisted.web import server, resource
@@ -15,6 +16,10 @@
 ONE_HOUR = 60 * ONE_MINUTE
 ONE_DAY = 24 * ONE_HOUR
 
+MINUTE_LIMITER = RateLimiter(span=ONE_MINUTE, bucket_size=10)
+HOUR_LIMITER = RateLimiter(span=ONE_HOUR, bucket_size=10 * ONE_MINUTE)
+DAY_LIMITER = RateLimiter(span=ONE_DAY, bucket_size=3 * ONE_HOUR)
+
 
 class JSONMethodResource(resource.Resource):
     """Renders the result of calling C{name} on C{target} as a json result."""
@@ -51,10 +56,10 @@
     def __init__(self, reporter):
         resource.Resource.__init__(self)
         self.reporter = reporter
-        self.putChild("top",
-            JSONMethodResource(self.reporter, "_get_distinct_top_value"))
-        self.putChild("count",
-            JSONMethodResource(self.reporter, "_get_distinct_count"))
+        self.putChild("top", JSONMethodResource(self.reporter,
+                                                "_get_distinct_top_value"))
+        self.putChild("count", JSONMethodResource(self.reporter,
+                                                  "_get_distinct_count"))
 
 
 class DistinctMetricReporter(object):
@@ -64,11 +69,12 @@
     """
     implements(IMetric)
 
-    periods = [5 * ONE_MINUTE, ONE_HOUR, ONE_DAY]
+    periods = (5 * ONE_MINUTE, ONE_HOUR, ONE_DAY)
+    limiters = (HOUR_LIMITER,)
 
     def __init__(self, name, wall_time_func=time.time, prefix="",
-            bucket_size=ONE_DAY, dsn=None, redis_host=None, redis_port=None,
-            longterm_distinct_days=60):
+                 bucket_size=ONE_DAY, dsn=None, redis_host=None,
+                 redis_port=None, longterm_distinct_days=60):
         """Construct a metric we expect to be periodically updated.
 
         @param name: Indicates what is being instrumented.
@@ -85,7 +91,7 @@
         self.dsn = dsn
         self.redis_host = redis_host
         if redis_port is None:
-            redis.port = 6379
+            redis_port = 6379
         self.redis_port = redis_port
         self.metric_id = None
         self.build_bucket()
@@ -94,7 +100,7 @@
         self.longterm_distinct_count = None
         self.longterm_distinct_days = longterm_distinct_days
 
-        if redis_host != None:
+        if redis_host is not None:
             self.redis = redis.client.Redis(host=redis_host, port=redis_port)
 
         if self.dsn is not None:
@@ -143,6 +149,9 @@
                 bucket = self.bucket_name_for(period)
                 self.redis.zremrangebyscore(bucket, 0, now - period)
                 self.redis_count[bucket] = self.redis.zcard(bucket)
+            for item, count in self.bucket.iteritems():
+                for limiter in self.limiters:
+                    limiter.add(self.redis, "%s:%s" % (self.name, item), count)
         finally:
             self.redis_flush_lock.release()
 
@@ -179,7 +188,7 @@
                  ".count_1day": self.count_1day(),
                  }
 
-        if self.longterm_distinct_count != None:
+        if self.longterm_distinct_count is not None:
             items[".count_%dday" % (self.longterm_distinct_days,)] = \
                 self.longterm_distinct_count
 
@@ -197,7 +206,7 @@
             row = cr.fetchone()
             if row is None:
                 cr.execute("INSERT INTO paths (path) VALUES (%s) "
-                            "RETURNING (id)", (path,))
+                           "RETURNING (id)", (path,))
                 row = cr.fetchone()
                 cr.execute("commit")
 
@@ -205,8 +214,8 @@
 
         for i, (k, v) in enumerate(bucket.iteritems()):
             cr.execute("INSERT INTO points (path_id, bucket, value, count) "
-                "VALUES (%s, %s, %s, %s)", (self.metric_id, bucket_no,
-                                            k, v))
+                       "VALUES (%s, %s, %s, %s)", (self.metric_id, bucket_no,
+                                                   k, v))
             if i % 1000 == 0:
                 cr.execute("commit")
         cr.execute("commit")
@@ -241,7 +250,7 @@
         cr.execute("SELECT value, SUM(count) AS cnt FROM points "
                    "INNER JOIN paths ON (paths.id = points.path_id) "
                    "WHERE paths.path = %s AND bucket BETWEEN %s AND %s"
-                   "GROUP BY value ORDER BY cnt DESC LIMIT %s", (
+                   "GROUP BY value ORDER BY cnt DESC, value ASC LIMIT %s", (
                        path, since_bucket, until_bucket, how_many,))
         rows = cr.fetchall()
         return rows

=== added file 'distinctdb/ratelimit.py'
--- distinctdb/ratelimit.py	1970-01-01 00:00:00 +0000
+++ distinctdb/ratelimit.py	2013-12-03 12:47:43 +0000
@@ -0,0 +1,49 @@
+"""
+A rate limiter, heavily based on https://chris6f.com/rate-limiting-with-redis
+"""
+
+import time
+import math
+
+
+class RateLimiter(object):
+
+    def __init__(self, span, bucket_size, expires=None, tick=time.time):
+        self.bucket_size = bucket_size
+        self.buckets = int(math.floor(float(span) / bucket_size)) + 2
+        self.span = span
+        self.expires = span if expires is None else expires
+        self.tick = tick
+
+    def add(self, client, key, count=1):
+        p = client.pipeline()
+        bucket = self.bucket()
+        key = "rl:%s:%s" % (self.span, key)
+        p.hincrby(key, bucket, count)
+        p.expire(key, self.expires)
+        p.hdel(key, (bucket + 1) % self.buckets)
+        p.hdel(key, (bucket + 2) % self.buckets)
+        p.execute()
+
+    def count(self, client, key, interval):
+        p = client.pipeline()
+        bucket = self.bucket()
+
+        key = "rl:%s:%s" % (self.span, key)
+
+        # Asking for an interval smaller than the bucket size rounds it up to
+        # one bucket. We do this to avoid getting false negatives when
+        # accidentally asking for too small an interval.
+        count = max(int(math.floor(float(interval) / self.bucket_size)), 1)
+        for i in xrange(0, count):
+            p.hget(key, (bucket - i) % self.buckets)
+
+        results = p.execute()
+        total = 0
+        for value in results:
+            if value is not None:
+                total += int(value)
+        return total
+
+    def bucket(self):
+        return int((self.tick() / self.bucket_size) % self.buckets)

=== modified file 'distinctdb/tests/test_distinct.py'
--- distinctdb/tests/test_distinct.py	2012-06-27 20:04:53 +0000
+++ distinctdb/tests/test_distinct.py	2013-12-03 12:47:43 +0000
@@ -13,8 +13,7 @@
     import subprocess
 
     def check_output(args):
-        return subprocess.Popen(args,
-            stdout=subprocess.PIPE).communicate()[0]
+        return subprocess.Popen(args, stdout=subprocess.PIPE).communicate()[0]
 
 import psycopg2
 import redis
@@ -49,7 +48,8 @@
         return 42
 
     def _get_distinct_top_value(self, since, until, how_many=20):
-        self.called.append(("get_distinct_top_value", (since, until, how_many)))
+        self.called.append(("get_distinct_top_value",
+                            (since, until, how_many)))
         return [("one", 1), ("two", 1)]
 
 
@@ -105,7 +105,7 @@
         dmr.flush(1, day)
         dmr.update("three")
         self.assertEquals(result,
-            {"bucket": {"one": 2, "two": 1}, "bucket_no": 0})
+                          {"bucket": {"one": 2, "two": 1}, "bucket_no": 0})
         result = dmr.flush(1, day)
 
         self.assertTrue(("test.max", 1, day) in result)
@@ -119,7 +119,8 @@
         o = TestOptions()
         config_file = ConfigParser.RawConfigParser()
         config_file.readfp(StringIO("[statsd]\n\n[plugin_distinctdb]\n"
-            "dsn = dbdsn\nbucket_size = 100\nlongterm_distinct_days = 30"))
+                                    "dsn = dbdsn\nbucket_size = 100\n"
+                                    "longterm_distinct_days = 30"))
         o.configure(config_file)
         dmf = distinctdbplugin.DistinctMetricFactory()
         dmf.configure(o)
@@ -184,7 +185,7 @@
 class TestPlugin(TestCase):
 
     def test_factory(self):
-        self.assertTrue(distinctdbplugin.distinct_metric_factory in \
+        self.assertTrue(distinctdbplugin.distinct_metric_factory in
                         list(getPlugins(IMetricFactory)))
 
 
@@ -473,8 +474,9 @@
 
         o = TestOptions()
         config_file = ConfigParser.RawConfigParser()
-        config_file.readfp(StringIO("[statsd]\n\n[plugin_distinctdb]\n"
-            "redis_host = localhost\nredis_port = 16379"))
+        config_file.readfp(
+            StringIO("[statsd]\n\n[plugin_distinctdb]\n"
+                     "redis_host = localhost\nredis_port = 16379"))
         o.configure(config_file)
         dmf = distinctdbplugin.DistinctMetricFactory()
         dmf.configure(o)
@@ -483,8 +485,8 @@
         self.assertEquals(dmr.redis_port, 16379)
 
     def test_usage(self):
-        dmr = distinct.DistinctMetricReporter("test",
-            redis_host="localhost", redis_port=16379)
+        dmr = distinct.DistinctMetricReporter(
+            "test", redis_host="localhost", redis_port=16379)
 
         self.assertEquals(dmr.count_1hour(), 0)
         dmr._update_count("one", 0)
@@ -500,8 +502,8 @@
         self.assertEquals(dmr.count_1hour(), 1)
 
     def test_load(self):
-        dmr = distinct.DistinctMetricReporter("test",
-            redis_host="localhost", redis_port=16379)
+        dmr = distinct.DistinctMetricReporter(
+            "test", redis_host="localhost", redis_port=16379)
         start = time.time()
         for i in range(10000):
             dmr.update(str(i % 1000))
@@ -511,10 +513,17 @@
             if w == 0:
                 break
             time.sleep(0.1)
+
         dmr._flush_redis(time.time())
         duration = time.time() - start
         self.assertEquals(dmr.count_1hour(), 1000)
         self.assertTrue(duration < 10)
+        # self.assertEquals(10, distinct.MINUTE_LIMITER.count(
+        #     dmr.redis, "test:42", 10))
+        self.assertEquals(10, distinct.HOUR_LIMITER.count(
+            dmr.redis, "test:42", 10 * distinct.ONE_MINUTE))
+        # self.assertEquals(10, distinct.DAY_LIMITER.count(
+        #     dmr.redis, "test:42", 3 * distinct.ONE_HOUR))
 
     def test_bucket_name(self):
         dmr1 = distinct.DistinctMetricReporter("somename")

=== added file 'distinctdb/tests/test_ratelimit.py'
--- distinctdb/tests/test_ratelimit.py	1970-01-01 00:00:00 +0000
+++ distinctdb/tests/test_ratelimit.py	2013-12-03 12:47:43 +0000
@@ -0,0 +1,79 @@
+import time
+import redis
+
+from distinctdb.ratelimit import RateLimiter
+from twisted.trial.unittest import TestCase
+
+
+class Clock(object):
+
+    def __init__(self):
+        self.value = 0
+
+    def set(self, value):
+        self.value = value
+
+    def __call__(self):
+        return self.value
+
+
+class RateLimitTest(TestCase):
+
+    def setUp(self):
+        self.client = redis.client.Redis(host="localhost", port=16379)
+        self.clock = Clock()
+        self.limiter = RateLimiter(span=60, bucket_size=10,
+                                   tick=self.clock)
+        return super(RateLimitTest, self).setUp()
+
+    def tearDown(self):
+        self.client.flushdb()
+        return super(RateLimitTest, self).tearDown()
+
+    def test_count(self):
+        for i in xrange(10):
+            self.clock.set(i * 10)
+            self.limiter.add(self.client, 'test')
+
+        self.assertEqual(6, self.limiter.count(self.client, 'test', 80))
+        self.assertEqual(6, self.limiter.count(self.client, 'test', 60))
+        self.assertEqual(4, self.limiter.count(self.client, 'test', 40))
+        self.assertEqual(2, self.limiter.count(self.client, 'test', 20))
+        self.assertEqual(1, self.limiter.count(self.client, 'test', 10))
+
+        # Asking for an interval smaller than the bucket size rounds it up to
+        # one bucket. We do this to avoid getting false negatives when
+        # accidentally asking for too small an interval.
+        self.assertEqual(1, self.limiter.count(self.client, 'test', 1))
+
+        self.clock.set(100)
+        self.assertEqual(3, self.limiter.count(self.client, 'test', 40))
+        self.assertEqual(1, self.limiter.count(self.client, 'test', 20))
+        self.assertEqual(0, self.limiter.count(self.client, 'test', 10))
+
+        self.clock.set(110)
+        self.assertEqual(2, self.limiter.count(self.client, 'test', 40))
+        self.assertEqual(0, self.limiter.count(self.client, 'test', 20))
+        self.assertEqual(0, self.limiter.count(self.client, 'test', 10))
+
+
+class RateLimitExpireTest(TestCase):
+
+    def setUp(self):
+        self.client = redis.client.Redis(host="localhost", port=16379)
+        self.clock = Clock()
+        self.limiter = RateLimiter(span=1, bucket_size=1,
+                                   tick=self.clock)
+        return super(RateLimitExpireTest, self).setUp()
+
+    def tearDown(self):
+        self.client.flushdb()
+        return super(RateLimitExpireTest, self).tearDown()
+
+    def test_expire(self):
+        for i in xrange(2):
+            self.limiter.add(self.client, 'test')
+
+        self.assertEqual(2, self.limiter.count(self.client, 'test', 1))
+        time.sleep(2)
+        self.assertEqual(0, self.limiter.count(self.client, 'test', 1))