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>

userspace/engine/token_bucket.cpp

@@ -18,13 +18,19 @@ 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)
 {
+}
+
+token_bucket::token_bucket(std::function<uint64_t()> timer)
+{
+    m_timer = timer;
     init(1, 1);
 }
 
@@ -37,20 +43,12 @@ 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_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)

userspace/engine/token_bucket.h

@@ -20,13 +20,17 @@ 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:
+  public:
     token_bucket();
+    token_bucket(std::function<uint64_t()> timer);
     virtual ~token_bucket();
 
     //
@@ -51,7 +55,8 @@ public:
     // 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.
@@ -75,4 +80,3 @@ private:
     //
     uint64_t m_last_seen;
 };
-