Merge pull request #236 from draios/expose-tokens

Add ability to get number of tokens.

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;
+}

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:
 
 	//