diff --git a/userspace/engine/token_bucket.cpp b/userspace/engine/token_bucket.cpp
index 93f04924..b72ae399 100644
--- a/userspace/engine/token_bucket.cpp
+++ b/userspace/engine/token_bucket.cpp
@@ -31,7 +31,7 @@ token_bucket::~token_bucket()
 {
 }
 
-void token_bucket::init(uint32_t rate, uint32_t max_tokens)
+void token_bucket::init(double rate, double max_tokens)
 {
 	m_rate = rate;
 	m_max_tokens = max_tokens;
@@ -39,12 +39,16 @@ void token_bucket::init(uint32_t rate, uint32_t max_tokens)
 	m_last_seen = sinsp_utils::get_current_time_ns();
 }
 
-bool token_bucket::claim()
+bool token_bucket::claim(uint64_t now)
 {
 	// Determine the number of tokens gained. Delta between
 	// last_seen and now, divided by the rate.
-	uint64_t now = sinsp_utils::get_current_time_ns();
-	uint64_t tokens_gained = (now - m_last_seen) / (m_rate * 1000000000);
+	if(now == 0)
+	{
+		now = sinsp_utils::get_current_time_ns();
+	}
+
+	double tokens_gained = (now - m_last_seen) / (m_rate * 1000000000);
 	m_last_seen = now;
 
 	m_tokens += tokens_gained;
diff --git a/userspace/engine/token_bucket.h b/userspace/engine/token_bucket.h
index b9c0d8b9..b60b5319 100644
--- a/userspace/engine/token_bucket.h
+++ b/userspace/engine/token_bucket.h
@@ -31,30 +31,30 @@ public:
 	//
 	// Initialize the token bucket and start accumulating tokens
 	//
-	void init(uint32_t rate, uint32_t max_tokens);
+	void init(double rate, double max_tokens);
 
 	//
 	// Returns true if a token can be claimed. Also updates
 	// internal metrics.
 	//
-	bool claim();
+	bool claim(uint64_t now = 0);
 private:
 
 	//
 	// The number of tokens generated per second.
 	//
-	uint64_t m_rate;
+	double m_rate;
 
 	//
 	// The maximum number of tokens that can be banked for future
 	// claim()s.
 	//
-	uint64_t m_max_tokens;
+	double m_max_tokens;
 
 	//
 	// The current number of tokens
 	//
-	uint64_t m_tokens;
+	double m_tokens;
 
 	//
 	// The last time claim() was called (or the object was created).