rate limiter

could be used for API rate limiting or login lockout use cases
fathomdb · Mar 19, 2013 · 50ee2f6 · 50ee2f6
1 parent 5d4b354
commit 50ee2f6
Show file tree

Hide file tree

Showing 7 changed files with 221 additions and 0 deletions.
diff --git a/Makefile.am b/Makefile.am
@@ -51,6 +51,7 @@ nobase_include_HEADERS=			\
 	mordor/predef.h			\
 	mordor/protobuf.h		\
 	mordor/ragel.h			\
+	mordor/rate_limiter.h		\
 	mordor/scheduler.h		\
 	mordor/semaphore.h		\
 	mordor/sleep.h			\
@@ -300,6 +301,7 @@ mordor_tests_run_tests_SOURCES=				\
 	mordor/tests/notify_stream.cpp			\
 	mordor/tests/oauth.cpp				\
 	mordor/tests/pipe_stream.cpp			\
+	mordor/tests/rate_limiter.cpp			\
 	mordor/tests/scheduler.cpp			\
 	mordor/tests/socket.cpp				\
 	mordor/tests/ssl_stream.cpp			\

diff --git a/mordor/mordor.vcxproj b/mordor/mordor.vcxproj
@@ -204,6 +204,7 @@
     <ClInclude Include="main.h" />
     <ClInclude Include="daemon.h" />
     <ClInclude Include="parallel.h" />
+    <ClInclude Include="rate_limiter.h" />
     <ClInclude Include="socks.h" />
     <ClInclude Include="streams\buffer.h" />
     <ClInclude Include="streams\buffered.h" />

diff --git a/mordor/mordor.vcxproj.filters b/mordor/mordor.vcxproj.filters
@@ -499,6 +499,9 @@
     <ClInclude Include="xml\dom_parser.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="rate_limiter.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <Ragel Include="http\http_parser.rl" />

diff --git a/mordor/rate_limiter.h b/mordor/rate_limiter.h
@@ -0,0 +1,121 @@
+#ifndef __MORDOR_RATE_LIMITER_H__
+#define __MODOR_RATE_LIMITER_H__
+// Copyright (c) 2013 - Cody Cutrer
+
+#include <list>
+#include <map>
+
+#include <boost/thread/mutex.hpp>
+
+#include "assert.h"
+#include "config.h"
+#include "log.h"
+#include "timer.h"
+
+namespace Mordor {
+
+template <class T>
+class RateLimiter
+{
+private:
+    struct Bucket
+    {
+        Bucket()
+            : m_count(0u)
+        {}
+
+        std::list<unsigned long long> m_timestamps;
+        size_t m_count;
+        Timer::ptr m_timer;
+    };
+
+public:
+    RateLimiter(TimerManager &timerManager, ConfigVar<size_t>::ptr countLimit,
+        ConfigVar<unsigned long long>::ptr timeLimit)
+        : m_timerManager(timerManager),
+          m_countLimit(countLimit),
+          m_timeLimit(timeLimit)
+    {}
+
+    bool allowed(const T &key)
+    {
+        boost::mutex::scoped_lock lock(m_mutex);
+        unsigned long long now = m_timerManager.now();
+        Bucket &bucket = m_buckets[key];
+        size_t countLimit = m_countLimit->val();
+        trim(bucket, now, countLimit);
+        if (bucket.m_count >= countLimit) {
+            startTimer(key, bucket);
+            return false;
+        }
+        bucket.m_timestamps.push_back(now);
+        ++bucket.m_count;
+        startTimer(key, bucket);
+        MORDOR_ASSERT(bucket.m_count == bucket.m_timestamps.size());
+        return true;
+    }
+
+    void reset(const T &key)
+    {
+        boost::mutex::scoped_lock lock(m_mutex);
+        std::map<T, Bucket>::iterator it = m_buckets.find(key);
+        if (it != m_buckets.end()) {
+            if (it->second.m_timer)
+                it->second.m_timer->cancel();
+            m_buckets.erase(key);
+        }
+    }
+
+private:
+    void trimKey(const T& key)
+    {
+        boost::mutex::scoped_lock lock(m_mutex);
+        Bucket &bucket = m_buckets[key];
+        trim(bucket, m_timerManager.now());
+        startTimer(key, bucket, m_countLimit->val());
+    }
+
+    void trim(Bucket &bucket, unsigned long long now, size_t countLimit)
+    {
+        unsigned long long timeLimit = m_timeLimit->val();
+        MORDOR_ASSERT(bucket.m_count == bucket.m_timestamps.size());
+        while(!bucket.m_timestamps.empty() && (bucket.m_timestamps.front() < now - timeLimit || bucket.m_count > countLimit))
+        {
+            drop(bucket);
+        }
+        MORDOR_ASSERT(bucket.m_count == bucket.m_timestamps.size());
+    }
+
+    void drop(Bucket &bucket)
+    {
+        bucket.m_timestamps.pop_front();
+        --bucket.m_count;
+        if (bucket.m_timer) {
+            bucket.m_timer->cancel();
+            bucket.m_timer.reset();
+        }
+    }
+
+    void startTimer(const T &key, Bucket &bucket)
+    {
+        // If there are still timestamps, set a timer to clear it
+        if(!bucket.m_timestamps.empty()) {
+            if (!bucket.m_timer) {
+                //bucket.m_timer = m_timerManager.registerTimer(bucket.m_timestamps.front() + timeLimit - now,
+                //    boost::bind(&RateLimiter::trimKey, this, key));
+            }
+        } else {
+            m_buckets.erase(key);
+        }
+    }
+private:
+    TimerManager &m_timerManager;
+    ConfigVar<size_t>::ptr m_countLimit;
+    ConfigVar<unsigned long long>::ptr m_timeLimit;
+    std::map<T, Bucket> m_buckets;
+    boost::mutex m_mutex;
+};
+
+}
+
+#endif
diff --git a/mordor/tests/rate_limiter.cpp b/mordor/tests/rate_limiter.cpp
@@ -0,0 +1,90 @@
+// Copyright (c) 2013 - Cody Cutrer
+
+#include <boost/bind.hpp>
+
+#include "mordor/rate_limiter.h"
+#include "mordor/iomanager.h"
+#include "mordor/sleep.h"
+#include "mordor/test/test.h"
+
+using namespace Mordor;
+using namespace Mordor::Test;
+
+static ConfigVar<size_t>::ptr g_countLimit = Config::lookup(
+    "ratelimiter.count", (size_t)0u, "Config var used by unit test");
+static ConfigVar<unsigned long long>::ptr g_timeLimit = Config::lookup(
+    "ratelimiter.time", 0ull, "Config var used by unit test");
+
+MORDOR_UNITTEST(RateLimiter, countLimits)
+{
+    IOManager ioManager;
+    RateLimiter<int> limiter(ioManager, g_countLimit, g_timeLimit);
+    // max of 3, tenth of a second
+    g_countLimit->val(3);
+    g_timeLimit->val(100000ull);
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(!limiter.allowed(1));
+    MORDOR_TEST_ASSERT(!limiter.allowed(1));
+    // sleep .2s; should allow three again
+    Mordor::sleep(ioManager, 200000ull);
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(!limiter.allowed(1));
+}
+
+MORDOR_UNITTEST(RateLimiter, countLimitsSlidingTime)
+{
+    IOManager ioManager;
+    RateLimiter<int> limiter(ioManager, g_countLimit, g_timeLimit);
+    // max of 3, half a second
+    g_countLimit->val(3);
+    g_timeLimit->val(500000ull);
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    Mordor::sleep(ioManager, 250000ull);
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(!limiter.allowed(1));
+    // sleep .3s; should only allow 1 more as the first slid off
+    Mordor::sleep(ioManager, 350000ull);
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(!limiter.allowed(1));
+}
+
+MORDOR_UNITTEST(RateLimiter, reset)
+{
+    IOManager ioManager;
+    RateLimiter<int> limiter(ioManager, g_countLimit, g_timeLimit);
+    // max of 3, 5 seconds
+    g_countLimit->val(3);
+    g_timeLimit->val(5000000ull);
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(!limiter.allowed(1));
+    // Resetting (i.e. successful login) should allow a full new batch
+    limiter.reset(1);
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(!limiter.allowed(1));
+}
+
+MORDOR_UNITTEST(RateLimiter, uniqueKeys)
+{
+    IOManager ioManager;
+    RateLimiter<int> limiter(ioManager, g_countLimit, g_timeLimit);
+    // max of 1, 5 seconds
+    g_countLimit->val(1);
+    g_timeLimit->val(5000000ull);
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(!limiter.allowed(1));
+    MORDOR_TEST_ASSERT(limiter.allowed(2));
+    MORDOR_TEST_ASSERT(!limiter.allowed(2));
+    limiter.reset(1);
+    MORDOR_TEST_ASSERT(!limiter.allowed(2));
+    MORDOR_TEST_ASSERT(limiter.allowed(1));
+    MORDOR_TEST_ASSERT(!limiter.allowed(1));
+}
diff --git a/mordor/tests/tests.vcxproj b/mordor/tests/tests.vcxproj
@@ -225,6 +225,7 @@
       <PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Release|x64'">Create</PrecompiledHeader>
     </ClCompile>
     <ClCompile Include="pipe_stream.cpp" />
+    <ClCompile Include="rate_limiter.cpp" />
     <ClCompile Include="run_tests.cpp" />
     <ClCompile Include="scheduler.cpp" />
     <ClCompile Include="socket.cpp" />

diff --git a/mordor/tests/tests.vcxproj.filters b/mordor/tests/tests.vcxproj.filters
@@ -141,5 +141,8 @@
     <ClCompile Include="util.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="rate_limiter.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
 </Project>