Nils O. Selåsdal
74 lines
1.8 KiB
C
74 lines
1.8 KiB
C
#include <assert.h>
|
|
#include <limits.h>
|
|
#include "ucore/rate_limit.h"
|
|
|
|
void uc_ratelimit_init(struct RateLimit *r, long tickets, long period)
|
|
{
|
|
r->tickets = tickets;
|
|
|
|
r->period = period;
|
|
r->funds = 0; //we fill it on first _allow call
|
|
r->last_ts = -period; //-period ensures the first tick can start
|
|
//at >= 0 and still refill the bicket
|
|
|
|
assert(tickets > 0);
|
|
assert(r->period > 0);
|
|
}
|
|
|
|
void uc_ratelimit_reset(struct RateLimit *r, long current_ts)
|
|
{
|
|
r->funds = r->period * r->tickets;
|
|
|
|
if (current_ts) {
|
|
r->last_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
|
|
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;
|
|
}
|
|
|
|
|
|
return allowed;
|
|
}
|
|
|