From 4d148ce28fe8841ee3cc31783cbd59cf7d490487 Mon Sep 17 00:00:00 2001
From: Mark Stemm <mark.stemm@sysdig.com>
Date: Tue, 2 May 2017 11:46:20 -0700
Subject: [PATCH] Add ability to claim multiple tokens.

This way you can use it as a form of bandwidth throttling.
---
 userspace/engine/token_bucket.cpp | 20 ++++++++++----------
 userspace/engine/token_bucket.h   | 11 ++++++++---
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/userspace/engine/token_bucket.cpp b/userspace/engine/token_bucket.cpp
index e9f28ce0..1edb0895 100644
--- a/userspace/engine/token_bucket.cpp
+++ b/userspace/engine/token_bucket.cpp
@@ -45,15 +45,15 @@ void token_bucket::init(double rate, double max_tokens, uint64_t now)
 	m_last_seen = now;
 }
 
-bool token_bucket::claim(uint64_t now)
+bool token_bucket::claim()
 {
-	// Determine the number of tokens gained. Delta between
-	// last_seen and now, divided by the rate.
-	if(now == 0)
-	{
-		now = sinsp_utils::get_current_time_ns();
-	}
+	uint64_t now = sinsp_utils::get_current_time_ns();
 
+	return claim(1, now);
+}
+
+bool token_bucket::claim(double tokens, uint64_t now)
+{
 	double tokens_gained = m_rate * ((now - m_last_seen) / (1000000000.0));
 	m_last_seen = now;
 
@@ -68,14 +68,14 @@ bool token_bucket::claim(uint64_t now)
 	}
 
 	//
-	// If tokens is < 1, can't claim.
+	// If m_tokens is < tokens, can't claim.
 	//
-	if(m_tokens < 1)
+	if(m_tokens < tokens)
 	{
 		return false;
 	}
 
-	m_tokens--;
+	m_tokens -= tokens;
 
 	return true;
 }
diff --git a/userspace/engine/token_bucket.h b/userspace/engine/token_bucket.h
index 076ba07b..5056192c 100644
--- a/userspace/engine/token_bucket.h
+++ b/userspace/engine/token_bucket.h
@@ -34,10 +34,15 @@ public:
 	void init(double rate, double max_tokens, uint64_t now = 0);
 
 	//
-	// Returns true if a token can be claimed. Also updates
-	// internal metrics.
+	// 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(uint64_t now = 0);
+	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();
 
 	// Return the current number of tokens available
 	double get_tokens();