github/falco
1
0
Fork 0
mirror of https://github.com/falcosecurity/falco.git synced 2025-07-06 11:26:44 +00:00
Code Issues Projects Releases Wiki Activity

new: dependency inject the timer for token bucket

Browse Source 
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

64
userspace/engine/token_bucket.cpp
View File

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

84
userspace/engine/token_bucket.h
View File

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