summaryrefslogtreecommitdiff
path: root/src/utils/tokens_bucket.hpp
diff options
context:
space:
mode:
Diffstat (limited to 'src/utils/tokens_bucket.hpp')
-rw-r--r--src/utils/tokens_bucket.hpp58
1 files changed, 58 insertions, 0 deletions
diff --git a/src/utils/tokens_bucket.hpp b/src/utils/tokens_bucket.hpp
new file mode 100644
index 0000000..d44eb06
--- /dev/null
+++ b/src/utils/tokens_bucket.hpp
@@ -0,0 +1,58 @@
+/**
+ * Implementation of the token bucket algorithm.
+ *
+ * It uses a repetitive TimedEvent, started at construction, to fill the
+ * bucket.
+ *
+ * Every n seconds, it executes the given callback. If the callback
+ * returns true, we add a token (if the limit is not yet reached).
+ *
+ */
+
+#pragma once
+
+#include <utils/timed_events.hpp>
+#include <logger/logger.hpp>
+
+class TokensBucket
+{
+public:
+ TokensBucket(std::size_t max_size, std::chrono::milliseconds fill_duration, std::function<bool()> callback, std::string name):
+ limit(max_size),
+ tokens(limit),
+ fill_duration(fill_duration),
+ callback(std::move(callback))
+ {
+ log_debug("creating TokensBucket with max size: ", max_size);
+ TimedEvent event(std::move(fill_duration), [this]() { this->add_token(); }, std::move(name));
+ TimedEventsManager::instance().add_event(std::move(event));
+ }
+
+ bool use_token()
+ {
+ if (this->tokens > 0)
+ {
+ this->tokens--;
+ return true;
+ }
+ else
+ return false;
+ }
+
+ void set_limit(std::size_t limit)
+ {
+ this->limit = limit;
+ }
+
+private:
+ std::size_t limit;
+ std::size_t tokens;
+ std::chrono::milliseconds fill_duration;
+ std::function<bool()> callback;
+
+ void add_token()
+ {
+ if (this->callback() && this->tokens != limit)
+ this->tokens++;
+ }
+};