mirror of
https://github.com/falcosecurity/falco.git
synced 2025-07-05 10:56:47 +00:00
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:
parent
1a0cf69b03
commit
a09f71b457
@ -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;
|
||||
}
|
||||
|
@ -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;
|
||||
};
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user