new: dependency inject the timer for token bucket

Co-Authored-By: Lorenzo Fontana <lo@linux.com>
Signed-off-by: Leonardo Di Donato <leodidonato@gmail.com>
This commit is contained in:
Leonardo Di Donato 2019-06-27 10:34:47 +00:00 committed by Lorenzo Fontana
parent 1a0cf69b03
commit a09f71b457
2 changed files with 76 additions and 74 deletions

View File

@ -18,14 +18,20 @@ limitations under the License.
*/
#include <cstddef>
#include <functional>
#include <sys/time.h>
#include "utils.h"
#include "token_bucket.h"
#include "utils.h"
token_bucket::token_bucket()
token_bucket::token_bucket() : token_bucket(sinsp_utils::get_current_time_ns)
{
init(1, 1);
}
token_bucket::token_bucket(std::function<uint64_t()> timer)
{
m_timer = timer;
init(1, 1);
}
token_bucket::~token_bucket()
@ -34,59 +40,51 @@ token_bucket::~token_bucket()
void token_bucket::init(double rate, double max_tokens, uint64_t now)
{
m_rate = rate;
m_max_tokens = max_tokens;
m_tokens = max_tokens;
if(now == 0)
{
now = sinsp_utils::get_current_time_ns();
}
m_last_seen = now;
m_rate = rate;
m_max_tokens = max_tokens;
m_tokens = max_tokens;
m_last_seen = now == 0 ? m_timer() : now;
}
bool token_bucket::claim()
{
uint64_t now = sinsp_utils::get_current_time_ns();
return claim(1, now);
return claim(1, m_timer());
}
bool token_bucket::claim(double tokens, uint64_t now)
{
double tokens_gained = m_rate * ((now - m_last_seen) / (1000000000.0));
m_last_seen = now;
double tokens_gained = m_rate * ((now - m_last_seen) / (1000000000.0));
m_last_seen = now;
m_tokens += tokens_gained;
m_tokens += tokens_gained;
//
// Cap at max_tokens
//
if(m_tokens > m_max_tokens)
//
// Cap at max_tokens
//
if(m_tokens > m_max_tokens)
{
m_tokens = m_max_tokens;
m_tokens = m_max_tokens;
}
//
// If m_tokens is < tokens, can't claim.
//
if(m_tokens < tokens)
//
// If m_tokens is < tokens, can't claim.
//
if(m_tokens < tokens)
{
return false;
return false;
}
m_tokens -= tokens;
m_tokens -= tokens;
return true;
return true;
}
double token_bucket::get_tokens()
{
return m_tokens;
return m_tokens;
}
uint64_t token_bucket::get_last_seen()
{
return m_last_seen;
return m_last_seen;
}

View File

@ -20,59 +20,63 @@ limitations under the License.
#pragma once
#include <cstdint>
#include <functional>
using token_timer = std::function<uint64_t()>;
// A simple token bucket that accumulates tokens at a fixed rate and allows
// for limited bursting in the form of "banked" tokens.
class token_bucket
{
public:
token_bucket();
virtual ~token_bucket();
public:
token_bucket();
token_bucket(std::function<uint64_t()> timer);
virtual ~token_bucket();
//
// Initialize the token bucket and start accumulating tokens
//
void init(double rate, double max_tokens, uint64_t now = 0);
//
// Initialize the token bucket and start accumulating tokens
//
void init(double rate, double max_tokens, uint64_t now = 0);
//
// Try to claim tokens tokens from the token bucket, using a
// timestamp of now. Returns true if the tokens could be
// claimed. Also updates internal metrics.
//
bool claim(double tokens, uint64_t now);
//
// Try to claim tokens tokens from the token bucket, using a
// timestamp of now. Returns true if the tokens could be
// claimed. Also updates internal metrics.
//
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();
// 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
double get_tokens();
// Return the current number of tokens available
double get_tokens();
// Return the last time someone tried to claim a token.
uint64_t get_last_seen();
// Return the last time someone tried to claim a token.
uint64_t get_last_seen();
private:
private:
std::function<uint64_t()> m_timer;
//
// The number of tokens generated per second.
//
double m_rate;
//
// The number of tokens generated per second.
//
double m_rate;
//
// The maximum number of tokens that can be banked for future
// claim()s.
//
double m_max_tokens;
//
// The maximum number of tokens that can be banked for future
// claim()s.
//
double m_max_tokens;
//
// The current number of tokens
//
double m_tokens;
//
// The current number of tokens
//
double m_tokens;
//
// The last time claim() was called (or the object was created).
// Nanoseconds since the epoch.
//
uint64_t m_last_seen;
//
// The last time claim() was called (or the object was created).
// Nanoseconds since the epoch.
//
uint64_t m_last_seen;
};