Add rate limit API
This commit is contained in:
@@ -0,0 +1,70 @@
|
|||||||
|
#ifndef UC_RATE_LIMIT_H_
|
||||||
|
#define UC_RATE_LIMIT_H_
|
||||||
|
|
||||||
|
/** This is a rate limiter, based on the token bucket principle.
|
||||||
|
* We have X amount of tickets(=our rate) per Y amount of time,
|
||||||
|
*
|
||||||
|
* The allowed tickets increases with X per Y amount of time.
|
||||||
|
*
|
||||||
|
* Bursts will be allowed up till the number of tickets initially
|
||||||
|
* specified.
|
||||||
|
* e.g. if we are limited to 50 tickets * per 2 second, 50 tickets
|
||||||
|
* can be immedietely handed out in the
|
||||||
|
* first second. At the 2. second, 25 more tickets become available
|
||||||
|
* handed out, at the 3. second 25 more, and so on.
|
||||||
|
* (So this is throttleed, so no more than the initial tickets, 50 in
|
||||||
|
* this case, can be handed out in any given period)
|
||||||
|
*
|
||||||
|
* Another example is if we rate 1000 tickets per 60 seconds,
|
||||||
|
* 1000 tickets can be handed out the 1. second, 16/17 more the
|
||||||
|
* next second, 16/17 more the 3. second and so on.
|
||||||
|
*
|
||||||
|
* Specify how many tickes we have available per period of time.
|
||||||
|
* when work is performed, call uc_ratelimit_cllow() with the current
|
||||||
|
* time.
|
||||||
|
* The unit of time passed to uc_ratelimit_allow must be the same as
|
||||||
|
* for uc_ratelimit_init.
|
||||||
|
*
|
||||||
|
* Example for rate limiting accepting connections:
|
||||||
|
* struct RateLimit r;
|
||||||
|
* //Limit to 20 per 1 second, which also limits us to max 20
|
||||||
|
* per second. (1 second since we will pass in time that have 1 second resolution)
|
||||||
|
* 40 per 2 second would even out to the same rate, but allow a burst of 40
|
||||||
|
* (if there's 40 tickets available)
|
||||||
|
*
|
||||||
|
* uc_ratelimit_init(&r,20, 1, time(NULL));
|
||||||
|
*
|
||||||
|
* for (;;) {
|
||||||
|
* int fd = accept(..);
|
||||||
|
* if (uc_ratelimit_allow(&r, time(NULL)) {
|
||||||
|
* //handle new connection
|
||||||
|
* } else {
|
||||||
|
* //exceeded limit
|
||||||
|
* close(fd);
|
||||||
|
* }
|
||||||
|
* }
|
||||||
|
*/
|
||||||
|
struct RateLimit {
|
||||||
|
//how many tickets we want
|
||||||
|
long tickets;
|
||||||
|
//per this period of time
|
||||||
|
long period;
|
||||||
|
|
||||||
|
//internal fields:
|
||||||
|
|
||||||
|
//scaled up number of tickets, to avoid floating point
|
||||||
|
//math
|
||||||
|
long scaled_tickets;
|
||||||
|
|
||||||
|
//number of tickets we have (scaled up)
|
||||||
|
long available_tickets;
|
||||||
|
//time of the previos period
|
||||||
|
long last_ts;
|
||||||
|
};
|
||||||
|
|
||||||
|
void uc_ratelimit_init(struct RateLimit *r, long tickets, long period, long current_ts);
|
||||||
|
void uc_ratelimit_reset(struct RateLimit *r, long current_ts);
|
||||||
|
int uc_ratelimit_allow(struct RateLimit *r, long current_ts);
|
||||||
|
|
||||||
|
#endif
|
||||||
|
|
||||||
@@ -0,0 +1,81 @@
|
|||||||
|
#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
|
||||||
Reference in New Issue
Block a user