From 15a3cbb59afb59f3f511850f10fc893bfcbee49f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nils=20O=2E=20Sel=C3=A5sdal?= <noselasd@fiane.dyndns.org>
Date: Sun, 9 Mar 2014 23:54:48 +0100
Subject: [PATCH] Small refactoring of ratelimit code to make more sense

---
 include/ucore/rate_limit.h |  4 +-
 src/rate_limit.c           | 81 +++++++++++++++++++++-----------------
 2 files changed, 46 insertions(+), 39 deletions(-)

diff --git a/include/ucore/rate_limit.h b/include/ucore/rate_limit.h
index 7d39755..367aae9 100644
--- a/include/ucore/rate_limit.h
+++ b/include/ucore/rate_limit.h
@@ -62,7 +62,7 @@ struct RateLimit {
     //money we have to "buy" tickets
     long funds;
     //time of the previos period
-    long last_ts;
+    long last_refill_ts;
 };
 
 /** Static initializer for a struct RateLimit
@@ -91,7 +91,7 @@ void uc_ratelimit_init(struct RateLimit *r, long tickets, long period);
  * Reset a RateLimit, filling up the tickets again.
  *
  * @param r RateLimit to reset
- * @param current ts, if non-zero, re-sets the last_ts to the current_ts
+ * @param current ts, if non-zero, re-sets the last_refill_ts to the current_ts
  */
 void uc_ratelimit_reset(struct RateLimit *r, long current_ts);
 
diff --git a/src/rate_limit.c b/src/rate_limit.c
index 14299fe..2facdc1 100644
--- a/src/rate_limit.c
+++ b/src/rate_limit.c
@@ -8,11 +8,12 @@ void uc_ratelimit_init(struct RateLimit *r, long tickets, long period)
 
     r->period = period;
     r->funds = 0; //we fill it on first _allow call
-    r->last_ts = -period; //-period ensures the first tick can start 
+    r->last_refill_ts = -period; //-period ensures the first tick can start 
                           //at >= 0 and still refill the bicket
 
     assert(tickets > 0);
     assert(r->period > 0);
+    assert(r->period !=  LONG_MAX);
 }
 
 void uc_ratelimit_reset(struct RateLimit *r, long current_ts)
@@ -20,54 +21,60 @@ void uc_ratelimit_reset(struct RateLimit *r, long current_ts)
     r->funds = r->period * r->tickets;
 
     if (current_ts) {
-        r->last_ts = current_ts;
+        r->last_refill_ts = current_ts;
     }
 }
 
+static void uc_ratelimit_refill(struct RateLimit *r, long current_ts)
+{
+
+    long diff_period;
+    //Find elapsed time
+    diff_period = current_ts - r->last_refill_ts;
+
+    if (diff_period > r->period + 1) {
+        //help prevent overflow when calculating available_tickets below
+        diff_period = r->period + 1;
+    } else if (diff_period < 0) {
+        //time went backwards, skip this round
+        diff_period = 0;  
+    }
+
+    //Calculate the cost of tickets that became available 
+    r->funds += diff_period * r->tickets;
+
+    //throttle handing out tickets
+    if (r->funds > r->period * r->tickets) {
+        r->funds = r->period * r->tickets;
+    }
+
+    //save this period. 
+    r->last_refill_ts = current_ts;
+}
+
 int uc_ratelimit_allow(struct RateLimit *r, long current_ts)
 {
     int allowed;
-    long diff_period;
 
-    if (r->funds >= r->period) {
-        //we have enough
+    if (r->funds < r->period) { //not enough funds
+
+        uc_ratelimit_refill(r, current_ts);
+
+        //as we keep track of the last timestamp of refilling
+        //We don't need to refill when there is enough funds
+    }
+  
+    //If we have enough to buy atleast one ticket, we can allow 
+    if (r->funds >=  r->period) {
         r->funds -= r->period;
         allowed = 1;
-    } else { //refill
-        //Find elapsed time
-        diff_period = current_ts - r->last_ts;
-
-        if (diff_period > r->period + 1) {
-            //help prevent overflow when calculating available_tickets below
-            diff_period = r->period + 1;
-        } else if (diff_period < 0) {
-            //time went backwards, skip this round
-            diff_period = 0;  
-        }
-
-        //Calculate the cost of tickets that became available 
-        r->funds += diff_period * r->tickets;
-
-
-        //throttle handing out tickets
-        if (r->funds > r->period * r->tickets) {
-            r->funds = r->period * r->tickets;
-        }
-
-        //If we have enough to buy atleast one ticket, we can allow 
-        if (r->funds >=  r->period) {
-            r->funds -= r->period;
-            allowed = 1;
-        } else {
-            //not enough to buy a ticket
-            allowed = 0;
-        }
-
-        //save this period. 
-        r->last_ts = current_ts;
+    } else {
+        //not enough to buy a ticket
+        allowed = 0;
     }
 
 
+
     return allowed;
 }