From a09f71b457c0e2065ded72ed959e4c9f169634da Mon Sep 17 00:00:00 2001 From: Leonardo Di Donato Date: Thu, 27 Jun 2019 10:34:47 +0000 Subject: [PATCH] new: dependency inject the timer for token bucket Co-Authored-By: Lorenzo Fontana Signed-off-by: Leonardo Di Donato --- userspace/engine/token_bucket.cpp | 64 +++++++++++------------ userspace/engine/token_bucket.h | 86 ++++++++++++++++--------------- 2 files changed, 76 insertions(+), 74 deletions(-) diff --git a/userspace/engine/token_bucket.cpp b/userspace/engine/token_bucket.cpp index ec6025de..6dd7ea00 100644 --- a/userspace/engine/token_bucket.cpp +++ b/userspace/engine/token_bucket.cpp @@ -18,14 +18,20 @@ limitations under the License. */ #include +#include #include -#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 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; } diff --git a/userspace/engine/token_bucket.h b/userspace/engine/token_bucket.h index bb36b3f6..03c4790e 100644 --- a/userspace/engine/token_bucket.h +++ b/userspace/engine/token_bucket.h @@ -20,59 +20,63 @@ limitations under the License. #pragma once #include +#include + +using token_timer = std::function; // 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 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 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 last time claim() was called (or the object was created). - // Nanoseconds since the epoch. - // - uint64_t m_last_seen; -}; + // + // 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; +}; \ No newline at end of file