Add ability to claim multiple tokens.

This way you can use it as a form of bandwidth throttling.
This commit is contained in:
Mark Stemm 2017-05-02 11:46:20 -07:00
parent a3c83e7f6e
commit 4d148ce28f
2 changed files with 18 additions and 13 deletions

View File

@ -45,15 +45,15 @@ void token_bucket::init(double rate, double max_tokens, uint64_t now)
m_last_seen = now; m_last_seen = now;
} }
bool token_bucket::claim(uint64_t now) bool token_bucket::claim()
{ {
// Determine the number of tokens gained. Delta between uint64_t now = sinsp_utils::get_current_time_ns();
// last_seen and now, divided by the rate.
if(now == 0)
{
now = sinsp_utils::get_current_time_ns();
}
return claim(1, now);
}
bool token_bucket::claim(double tokens, uint64_t now)
{
double tokens_gained = m_rate * ((now - m_last_seen) / (1000000000.0)); double tokens_gained = m_rate * ((now - m_last_seen) / (1000000000.0));
m_last_seen = now; m_last_seen = now;
@ -68,14 +68,14 @@ bool token_bucket::claim(uint64_t now)
} }
// //
// If tokens is < 1, can't claim. // If m_tokens is < tokens, can't claim.
// //
if(m_tokens < 1) if(m_tokens < tokens)
{ {
return false; return false;
} }
m_tokens--; m_tokens -= tokens;
return true; return true;
} }

View File

@ -34,10 +34,15 @@ public:
void init(double rate, double max_tokens, uint64_t now = 0); void init(double rate, double max_tokens, uint64_t now = 0);
// //
// Returns true if a token can be claimed. Also updates // Try to claim tokens tokens from the token bucket, using a
// internal metrics. // timestamp of now. Returns true if the tokens could be
// claimed. Also updates internal metrics.
// //
bool claim(uint64_t now = 0); bool claim(double tokens, uint64_t now);
// Simpler version of claim that claims a single token and
// uses the current time for now
bool claim();
// Return the current number of tokens available // Return the current number of tokens available
double get_tokens(); double get_tokens();