Small refactoring of ratelimit code to make more sense

This commit is contained in:
Nils O. Selåsdal
2014-03-09 23:54:48 +01:00
parent a112b44c6e
commit d831ab9510
2 changed files with 46 additions and 39 deletions
+44 -37
View File
@@ -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;
}