Rework rate limit to only refill when needed

include/ucore/rate_limit.h

@@ -65,8 +65,8 @@ struct RateLimit {
  * @param tickets number of tickets (bucket depth)
  * @param period per this period
  */
-#define UC_RATELIMIT_INITIALIZER(tickets, period)\
+#define UC_RATELIMIT_INITIALIZER(tickets)\
-{ tickets, period, tickets * period, 0}
+{ tickets, period, 0, -period}
 
   
 
@@ -79,9 +79,8 @@ struct RateLimit {
  * @param r RateLimit to initialize
  * @param tickets number of tickets (bucket depth)
  * @param period per this period
- * @param current_ts the current timestamp
  */
-void uc_ratelimit_init(struct RateLimit *r, long tickets, long period, long current_ts);
+void uc_ratelimit_init(struct RateLimit *r, long tickets, long period);
 
 /**
  * Reset a RateLimit, filling up the tickets again.

src/rate_limit.c

@@ -1,18 +1,18 @@
 #include <assert.h>
+#include <limits.h>
 #include "ucore/rate_limit.h"
 
-void uc_ratelimit_init(struct RateLimit *r, long tickets, long period, long current_ts)
+void uc_ratelimit_init(struct RateLimit *r, long tickets, long period)
 {
     r->tickets = tickets;
 
     r->period = period;
-    r->funds = r->period * tickets;
+    r->funds = 0; //we fill it on first _allow call
-    r->last_ts = current_ts;
+    r->last_ts = -period; //-period ensures the first tick can start 
+                          //at >= 0 and still refill the bicket
 
     assert(tickets > 0);
-    assert(period > 0);
     assert(r->period > 0);
-    assert(r->funds > 0);
 }
 
 void uc_ratelimit_reset(struct RateLimit *r, long current_ts)
@@ -29,6 +29,11 @@ int uc_ratelimit_allow(struct RateLimit *r, long current_ts)
     int allowed;
     long diff_period;
 
+    if (r->funds >= r->period) {
+        //we have enough
+        r->funds -= r->period;
+        allowed = 1;
+    } else { //refill
         //Find elapsed time
         diff_period = current_ts - r->last_ts;
 
@@ -36,19 +41,13 @@ int uc_ratelimit_allow(struct RateLimit *r, long current_ts)
             //help prevent overflow when calculating available_tickets below
             diff_period = r->period + 1;
         } else if (diff_period < 0) {
-        //If time went backwards, we halt the time for this round
+            //time went backwards, skip this round
             diff_period = 0;  
         }
 
-    //Calculate the cost of tickets that became available since 
+        //Calculate the cost of tickets that became available 
-    //the last time. 
         r->funds += diff_period * r->tickets;
 
-    //the below is the real algorithm, with period = 
-    //tickets * period. 
-    //we have optimized this to the expression used above
-    //r->funds += diff_period  * 
-    //                        ((r->tickets * r->period) / r->period);
 
         //throttle handing out tickets
         if (r->funds > r->period * r->tickets) {
@@ -66,6 +65,8 @@ int uc_ratelimit_allow(struct RateLimit *r, long current_ts)
 
         //save this period. 
         r->last_ts = current_ts;
+    }
+
 
     return allowed;
 }

test/test_ratelimit.c

@@ -7,9 +7,9 @@ START_TEST (test_ratelimit_1)
     long now = 10;
     int i;
 
-    uc_ratelimit_init(&r, 30, 2, now);
+    uc_ratelimit_init(&r, 30, 2);
     for (i = 0; i < 30; i++) {
-        fail_if(!uc_ratelimit_allow(&r, now));
+        fail_if(!uc_ratelimit_allow(&r, now), "i = %d", i);
     }
     fail_if(uc_ratelimit_allow(&r, now));
 END_TEST
@@ -18,7 +18,7 @@ START_TEST (test_ratelimit_2)
     struct RateLimit r;
     long now = 10;
 
-    uc_ratelimit_init(&r, 1, 1, now);
+    uc_ratelimit_init(&r, 1, 1);
     fail_if(!uc_ratelimit_allow(&r, now));
     fail_if(uc_ratelimit_allow(&r, now));
 
@@ -34,7 +34,7 @@ START_TEST (test_ratelimit_3)
     long now = 10;
     int i;
 
-    uc_ratelimit_init(&r, 10, 5, now);
+    uc_ratelimit_init(&r, 10, 5);
 
     for (i = 0; i < 10; i++) {
         fail_if(!uc_ratelimit_allow(&r, now));
@@ -70,7 +70,7 @@ START_TEST (test_ratelimit_4)
     long now = 10;
     int i;
 
-    uc_ratelimit_init(&r, 10, 5, now);
+    uc_ratelimit_init(&r, 10, 5);
 
     for (i = 0; i < 10; i++) {
         fail_if(!uc_ratelimit_allow(&r, now));
@@ -92,7 +92,7 @@ START_TEST (test_ratelimit_5)
     struct RateLimit r;
     long now = 1400000000;
 
-    uc_ratelimit_init(&r, 1, 10, now);
+    uc_ratelimit_init(&r, 1, 10);
 
     fail_if(!uc_ratelimit_allow(&r, now));
     fail_if(uc_ratelimit_allow(&r, now));
@@ -117,7 +117,7 @@ START_TEST (test_ratelimit_time_backward)
     struct RateLimit r;
     long now = 1400000000;
 
-    uc_ratelimit_init(&r, 1, 10, now);
+    uc_ratelimit_init(&r, 1, 10);
 
     fail_if(!uc_ratelimit_allow(&r, now));
     fail_if(uc_ratelimit_allow(&r, now));