Nils O. Selåsdal
82 lines
1.9 KiB
C
82 lines
1.9 KiB
C
#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
|