diff --git a/userspace/engine/token_bucket.cpp b/userspace/engine/token_bucket.cpp index f40112e7..e9f28ce0 100644 --- a/userspace/engine/token_bucket.cpp +++ b/userspace/engine/token_bucket.cpp @@ -31,12 +31,18 @@ token_bucket::~token_bucket() { } -void token_bucket::init(double rate, double max_tokens) +void token_bucket::init(double rate, double max_tokens, uint64_t now) { m_rate = rate; m_max_tokens = max_tokens; m_tokens = max_tokens; - m_last_seen = sinsp_utils::get_current_time_ns(); + + if(now == 0) + { + now = sinsp_utils::get_current_time_ns(); + } + + m_last_seen = now; } bool token_bucket::claim(uint64_t now) @@ -73,3 +79,13 @@ bool token_bucket::claim(uint64_t now) return true; } + +double token_bucket::get_tokens() +{ + return m_tokens; +} + +uint64_t token_bucket::get_last_seen() +{ + return m_last_seen; +} diff --git a/userspace/engine/token_bucket.h b/userspace/engine/token_bucket.h index b60b5319..076ba07b 100644 --- a/userspace/engine/token_bucket.h +++ b/userspace/engine/token_bucket.h @@ -31,13 +31,20 @@ public: // // Initialize the token bucket and start accumulating tokens // - void init(double rate, double max_tokens); + void init(double rate, double max_tokens, uint64_t now = 0); // // Returns true if a token can be claimed. Also updates // internal metrics. // bool claim(uint64_t now = 0); + + // 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(); + private: //