mirror of
https://github.com/ClickHouse/ClickHouse.git
synced 2024-12-12 09:22:05 +00:00
87 lines
2.5 KiB
C++
87 lines
2.5 KiB
C++
#pragma once
|
||
|
||
#include <mutex>
|
||
#include <memory>
|
||
#include <DB/Common/Stopwatch.h>
|
||
#include <DB/Common/Exception.h>
|
||
#include <DB/IO/WriteHelpers.h>
|
||
|
||
|
||
namespace DB
|
||
{
|
||
|
||
namespace ErrorCodes
|
||
{
|
||
extern const int LIMIT_EXCEEDED;
|
||
}
|
||
|
||
|
||
/** Позволяет ограничить скорость чего либо (в штуках в секунду) с помощью sleep.
|
||
* Особенности работы:
|
||
* - считается только средняя скорость, от момента первого вызова функции add;
|
||
* если были периоды с низкой скоростью, то в течение промежутка времени после них, скорость будет выше;
|
||
*
|
||
* Также позволяет задать ограничение на максимальное количество в штуках. При превышении кидается исключение.
|
||
*/
|
||
class Throttler
|
||
{
|
||
public:
|
||
Throttler(size_t max_speed_, size_t limit_, const char * limit_exceeded_exception_message_)
|
||
: max_speed(max_speed_), limit(limit_), limit_exceeded_exception_message(limit_exceeded_exception_message_) {}
|
||
|
||
void add(size_t amount)
|
||
{
|
||
size_t new_count;
|
||
UInt64 elapsed_ns = 0;
|
||
|
||
{
|
||
std::lock_guard<std::mutex> lock(mutex);
|
||
|
||
if (max_speed)
|
||
{
|
||
if (0 == count)
|
||
{
|
||
watch.start();
|
||
elapsed_ns = 0;
|
||
}
|
||
else
|
||
elapsed_ns = watch.elapsed();
|
||
}
|
||
|
||
count += amount;
|
||
new_count = count;
|
||
}
|
||
|
||
if (limit && new_count > limit)
|
||
throw Exception(limit_exceeded_exception_message + std::string(" Maximum: ") + toString(limit), ErrorCodes::LIMIT_EXCEEDED);
|
||
|
||
if (max_speed)
|
||
{
|
||
/// Сколько должно было бы пройти времени, если бы скорость была равна max_speed.
|
||
UInt64 desired_ns = new_count * 1000000000 / max_speed;
|
||
|
||
if (desired_ns > elapsed_ns)
|
||
{
|
||
UInt64 sleep_ns = desired_ns - elapsed_ns;
|
||
timespec sleep_ts;
|
||
sleep_ts.tv_sec = sleep_ns / 1000000000;
|
||
sleep_ts.tv_nsec = sleep_ns % 1000000000;
|
||
nanosleep(&sleep_ts, nullptr); /// NOTE Завершается раньше в случае сигнала. Это считается нормальным.
|
||
}
|
||
}
|
||
}
|
||
|
||
private:
|
||
size_t max_speed = 0;
|
||
size_t count = 0;
|
||
size_t limit = 0; /// 0 - не ограничено.
|
||
const char * limit_exceeded_exception_message = nullptr;
|
||
Stopwatch watch {CLOCK_MONOTONIC_COARSE};
|
||
std::mutex mutex;
|
||
};
|
||
|
||
|
||
using ThrottlerPtr = std::shared_ptr<Throttler>;
|
||
|
||
}
|