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