From 06083a054e58b511eb2fd2ffac91e99024a99aa7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nils=20O=2E=20Sel=C3=A5sdal?= <noselasd@fiane.dyndns.org>
Date: Wed, 9 Oct 2013 22:55:29 +0200
Subject: [PATCH] Add rate limit API

---
 include/ucore/rate_limit.h | 70 ++++++++++++++++++++++++++++++++
 src/rate_limit.c           | 81 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 151 insertions(+)
 create mode 100644 include/ucore/rate_limit.h
 create mode 100644 src/rate_limit.c

diff --git a/include/ucore/rate_limit.h b/include/ucore/rate_limit.h
new file mode 100644
index 0000000..fa7ec01
--- /dev/null
+++ b/include/ucore/rate_limit.h
@@ -0,0 +1,70 @@
+#ifndef UC_RATE_LIMIT_H_
+#define UC_RATE_LIMIT_H_
+
+/** This is a rate limiter, based on the token bucket principle.
+ * We have X amount of tickets(=our rate) per Y amount of time,
+ *  
+ * The allowed tickets increases with X per Y amount of time.
+ * 
+ * Bursts will be allowed up till the number of tickets initially
+ * specified.
+ * e.g. if we are limited to 50 tickets * per 2 second, 50 tickets 
+ * can be immedietely handed out in the
+ * first second. At the 2. second, 25 more tickets become available
+ * handed out, at the 3. second 25 more, and so on. 
+ * (So this is throttleed, so no more than the initial tickets, 50 in 
+ * this case, can be handed out in any given period)
+ *
+ * Another example is if we rate 1000 tickets per 60 seconds,
+ * 1000 tickets can be handed out the 1. second, 16/17 more the 
+ * next second, 16/17 more the 3. second and so on.
+ *
+ * Specify how many tickes we have available per period of time.
+ * when work is performed, call uc_ratelimit_cllow() with the current
+ * time. 
+ * The unit of time passed to uc_ratelimit_allow must be the same as 
+ * for uc_ratelimit_init.
+ *
+ * Example for rate limiting accepting connections:
+ * struct RateLimit r;
+ * //Limit to 20 per 1 second, which also limits us to max 20 
+ * per second. (1 second since we will pass in time that have 1 second resolution)
+ * 40 per 2 second would even out to the same rate, but allow a burst of 40
+ * (if there's 40 tickets available)
+ *
+ * uc_ratelimit_init(&r,20, 1, time(NULL));
+ *
+ * for (;;) {
+ *    int fd = accept(..);
+ *    if (uc_ratelimit_allow(&r, time(NULL)) {
+ *       //handle new connection
+ *    } else {
+ *      //exceeded limit
+ *      close(fd);
+ *    }
+ *  }
+ */
+struct RateLimit {
+    //how many tickets we want
+    long tickets;
+    //per this period of time
+    long period;
+
+    //internal fields:
+
+    //scaled up number of tickets, to avoid floating point
+    //math
+    long scaled_tickets;
+
+    //number of tickets we have (scaled up)
+    long available_tickets;
+    //time of the previos period
+    long last_ts;
+};
+
+void uc_ratelimit_init(struct RateLimit *r, long tickets, long period, long current_ts);
+void uc_ratelimit_reset(struct RateLimit *r, long current_ts);
+int uc_ratelimit_allow(struct RateLimit *r, long current_ts);
+
+#endif
+
diff --git a/src/rate_limit.c b/src/rate_limit.c
new file mode 100644
index 0000000..f7c33b4
--- /dev/null
+++ b/src/rate_limit.c
@@ -0,0 +1,81 @@
+#include "ucore/rate_limit.h"
+
+#define SCALE_FACTOR 1000
+
+
+void uc_ratelimit_init(struct RateLimit *r, long tickets, long period, long current_ts)
+{
+    r->tickets = tickets;
+    r->scaled_tickets = tickets * SCALE_FACTOR;
+    r->period = period;
+    r->available_tickets = r->scaled_tickets;
+    r->last_ts = current_ts;
+}
+
+void uc_ratelimit_reset(struct RateLimit *r, long current_ts)
+{
+    r->available_tickets = r->scaled_tickets;
+
+    if (current_ts)
+        r->last_ts = current_ts;
+}
+
+int uc_ratelimit_allow(struct RateLimit *r, long current_ts)
+{
+    int allowed;
+    long diff_period;
+
+    //Find elapsed time
+    diff_period = current_ts - r->last_ts;
+
+    //If time went backwards, we halt the time for this round
+    if (diff_period < 0) 
+        diff_period = 0;
+
+    //Calculate the number of tickets that became available since 
+    //the last time. (+ ScALE_FACTOR/2 is to round up at halfway points for
+    //integer arithmetics)
+    r->available_tickets += diff_period  * 
+                            ((r->scaled_tickets + SCALE_FACTOR/2) / r->period);
+
+    //throttle handing out tickets
+    //
+    if (r->available_tickets > r->scaled_tickets) 
+        r->available_tickets = r->scaled_tickets;
+
+    //If we have at least onen ticket, we can allow 
+    if (r->available_tickets >=  1 * SCALE_FACTOR) {
+        r->available_tickets -= 1 * SCALE_FACTOR;
+        allowed = 1;
+    } else {
+        //no more tickets.
+        allowed = 0;
+    }
+
+    //save this period. 
+    r->last_ts = current_ts;
+
+    return allowed;
+}
+
+#if 0
+int main()
+{
+    int i;
+    struct RateLimit r;
+    time_t start;
+
+    uc_ratelimit_init(&r, 10, 60, time(NULL));
+
+    for(i = 0; i < 120; i++) {
+        int k;
+        long allowed = 0;
+        sleep(1);
+        for (k = 0; k < 500; k++) 
+            allowed += uc_ratelimit_allow(&r, time(NULL));
+        printf("%ld of 500 can pass\n", allowed);
+    }
+
+    return 0;
+}
+#endif