blob: 263359a8d501cd9e8a4dd9fc1adbd3026b6bbd09 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
|
/**
* 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(long int max_size, std::chrono::milliseconds fill_duration, std::function<bool()> callback, std::string name):
limit(max_size),
tokens(static_cast<std::size_t>(limit)),
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->limit < 0)
return true;
if (this->tokens > 0)
{
this->tokens--;
return true;
}
else
return false;
}
void set_limit(long int limit)
{
this->limit = limit;
}
private:
long int limit;
std::size_t tokens;
std::function<bool()> callback;
void add_token()
{
if (this->limit < 0)
return;
if (this->callback() && this->tokens != static_cast<decltype(this->tokens)>(this->limit))
this->tokens++;
}
};
|
