Small refactoring of ratelimit code to make more sense
This commit is contained in:
@@ -62,7 +62,7 @@ struct RateLimit {
|
|||||||
//money we have to "buy" tickets
|
//money we have to "buy" tickets
|
||||||
long funds;
|
long funds;
|
||||||
//time of the previos period
|
//time of the previos period
|
||||||
long last_ts;
|
long last_refill_ts;
|
||||||
};
|
};
|
||||||
|
|
||||||
/** Static initializer for a struct RateLimit
|
/** Static initializer for a struct RateLimit
|
||||||
@@ -91,7 +91,7 @@ void uc_ratelimit_init(struct RateLimit *r, long tickets, long period);
|
|||||||
* Reset a RateLimit, filling up the tickets again.
|
* Reset a RateLimit, filling up the tickets again.
|
||||||
*
|
*
|
||||||
* @param r RateLimit to reset
|
* @param r RateLimit to reset
|
||||||
* @param current ts, if non-zero, re-sets the last_ts to the current_ts
|
* @param current ts, if non-zero, re-sets the last_refill_ts to the current_ts
|
||||||
*/
|
*/
|
||||||
void uc_ratelimit_reset(struct RateLimit *r, long current_ts);
|
void uc_ratelimit_reset(struct RateLimit *r, long current_ts);
|
||||||
|
|
||||||
|
|||||||
+44
-37
@@ -8,11 +8,12 @@ void uc_ratelimit_init(struct RateLimit *r, long tickets, long period)
|
|||||||
|
|
||||||
r->period = period;
|
r->period = period;
|
||||||
r->funds = 0; //we fill it on first _allow call
|
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
|
//at >= 0 and still refill the bicket
|
||||||
|
|
||||||
assert(tickets > 0);
|
assert(tickets > 0);
|
||||||
assert(r->period > 0);
|
assert(r->period > 0);
|
||||||
|
assert(r->period != LONG_MAX);
|
||||||
}
|
}
|
||||||
|
|
||||||
void uc_ratelimit_reset(struct RateLimit *r, long current_ts)
|
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;
|
r->funds = r->period * r->tickets;
|
||||||
|
|
||||||
if (current_ts) {
|
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 uc_ratelimit_allow(struct RateLimit *r, long current_ts)
|
||||||
{
|
{
|
||||||
int allowed;
|
int allowed;
|
||||||
long diff_period;
|
|
||||||
|
|
||||||
if (r->funds >= r->period) {
|
if (r->funds < r->period) { //not enough funds
|
||||||
//we have enough
|
|
||||||
|
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;
|
r->funds -= r->period;
|
||||||
allowed = 1;
|
allowed = 1;
|
||||||
} else { //refill
|
} else {
|
||||||
//Find elapsed time
|
//not enough to buy a ticket
|
||||||
diff_period = current_ts - r->last_ts;
|
allowed = 0;
|
||||||
|
|
||||||
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;
|
return allowed;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user