1 /**
2  * Implementation of the token bucket algorithm.
3  *
4  * It uses a repetitive TimedEvent, started at construction, to fill the
5  * bucket.
6  *
7  * Every n seconds, it executes the given callback. If the callback
8  * returns true, we add a token (if the limit is not yet reached).
9  *
10  */
11 
12 #pragma once
13 
14 #include <utils/timed_events.hpp>
15 #include <logger/logger.hpp>
16 
17 class TokensBucket
18 {
19 public:
TokensBucket(long int max_size,std::chrono::milliseconds fill_duration,std::function<bool ()> callback,std::string name)20   TokensBucket(long int max_size, std::chrono::milliseconds fill_duration, std::function<bool()> callback, std::string name):
21       limit(max_size),
22       tokens(static_cast<std::size_t>(limit)),
23       callback(std::move(callback))
24   {
25     log_debug("creating TokensBucket with max size: ", max_size);
26     TimedEvent event(std::move(fill_duration), [this]() { this->add_token(); }, std::move(name));
27     TimedEventsManager::instance().add_event(std::move(event));
28   }
29 
use_token()30   bool use_token()
31   {
32     if (this->limit < 0)
33       return true;
34     if (this->tokens > 0)
35       {
36         this->tokens--;
37         return true;
38       }
39     else
40       return false;
41   }
42 
set_limit(long int limit)43   void set_limit(long int limit)
44   {
45     this->limit = limit;
46   }
47 
48 private:
49   long int limit;
50   std::size_t tokens;
51   std::function<bool()> callback;
52 
add_token()53   void add_token()
54   {
55     if (this->limit < 0)
56       return;
57     if (this->callback() && this->tokens != static_cast<decltype(this->tokens)>(this->limit))
58       this->tokens++;
59   }
60 };
61