mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-12-04 05:22:17 +00:00
fdfb1eda55
And also update the test, since now you could have slightly less sleep intervals, if query spend some time in other places. But what is important is that query_duration_ms does not exceeded calculated delay. Signed-off-by: Azat Khuzhin <a.khuzhin@semrush.com>
76 lines
2.8 KiB
C++
76 lines
2.8 KiB
C++
#pragma once
|
|
|
|
#include <Common/Throttler_fwd.h>
|
|
#include <Common/ProfileEvents.h>
|
|
|
|
#include <mutex>
|
|
#include <memory>
|
|
#include <base/sleep.h>
|
|
#include <base/types.h>
|
|
#include <atomic>
|
|
|
|
namespace DB
|
|
{
|
|
|
|
/** Allows you to limit the speed of something (in tokens per second) using sleep.
|
|
* Implemented using Token Bucket Throttling algorithm.
|
|
* Also allows you to set a limit on the maximum number of tokens. If exceeded, an exception will be thrown.
|
|
*/
|
|
class Throttler
|
|
{
|
|
public:
|
|
static const size_t default_burst_seconds = 1;
|
|
|
|
Throttler(size_t max_speed_, size_t max_burst_, const std::shared_ptr<Throttler> & parent_ = nullptr)
|
|
: max_speed(max_speed_), max_burst(max_burst_), limit_exceeded_exception_message(""), tokens(max_burst), parent(parent_) {}
|
|
|
|
explicit Throttler(size_t max_speed_, const std::shared_ptr<Throttler> & parent_ = nullptr);
|
|
|
|
Throttler(size_t max_speed_, size_t max_burst_, size_t limit_, const char * limit_exceeded_exception_message_,
|
|
const std::shared_ptr<Throttler> & parent_ = nullptr)
|
|
: max_speed(max_speed_), max_burst(max_burst_), limit(limit_), limit_exceeded_exception_message(limit_exceeded_exception_message_), tokens(max_burst), parent(parent_) {}
|
|
|
|
Throttler(size_t max_speed_, size_t limit_, const char * limit_exceeded_exception_message_,
|
|
const std::shared_ptr<Throttler> & parent_ = nullptr);
|
|
|
|
/// Use `amount` tokens, sleeps if required or throws exception on limit overflow.
|
|
/// Returns duration of sleep in nanoseconds (to distinguish sleeping on different kinds of throttlers for metrics)
|
|
UInt64 add(size_t amount);
|
|
|
|
UInt64 add(size_t amount, ProfileEvents::Event event_amount, ProfileEvents::Event event_sleep_us)
|
|
{
|
|
UInt64 sleep_ns = add(amount);
|
|
ProfileEvents::increment(event_amount, amount);
|
|
ProfileEvents::increment(event_sleep_us, sleep_ns / 1000UL);
|
|
return sleep_ns;
|
|
}
|
|
|
|
/// Not thread safe
|
|
void setParent(const std::shared_ptr<Throttler> & parent_)
|
|
{
|
|
parent = parent_;
|
|
}
|
|
|
|
/// Reset all throttlers internal stats
|
|
void reset();
|
|
|
|
/// Is throttler already accumulated some sleep time and throttling.
|
|
bool isThrottling() const;
|
|
|
|
private:
|
|
size_t count{0};
|
|
const size_t max_speed{0}; /// in tokens per second.
|
|
const size_t max_burst{0}; /// in tokens.
|
|
const UInt64 limit{0}; /// 0 - not limited.
|
|
const char * limit_exceeded_exception_message = nullptr;
|
|
std::mutex mutex;
|
|
std::atomic<UInt64> accumulated_sleep{0}; // Accumulated sleep time over all waiting threads
|
|
double tokens{0}; /// Amount of tokens available in token bucket. Updated in `add` method.
|
|
UInt64 prev_ns{0}; /// Previous `add` call time (in nanoseconds).
|
|
|
|
/// Used to implement a hierarchy of throttlers
|
|
std::shared_ptr<Throttler> parent;
|
|
};
|
|
|
|
}
|