ClickHouse/src/Common/OvercommitTracker.h

137 lines
3.4 KiB
C++
Raw Normal View History

2021-10-22 12:56:09 +00:00
#pragma once
2021-12-06 22:06:00 +00:00
#include <base/logger_useful.h>
2021-10-22 12:56:09 +00:00
#include <base/types.h>
#include <boost/core/noncopyable.hpp>
#include <cassert>
2021-10-22 15:15:33 +00:00
#include <chrono>
2021-10-22 12:56:09 +00:00
#include <condition_variable>
#include <mutex>
#include <unordered_map>
struct OvercommitRatio
{
2021-11-09 13:40:23 +00:00
OvercommitRatio(Int64 committed_, Int64 soft_limit_)
: committed(committed_)
2021-10-22 12:56:09 +00:00
, soft_limit(soft_limit_)
{}
friend bool operator<(OvercommitRatio const& lhs, OvercommitRatio const& rhs) noexcept
{
2021-10-26 13:55:57 +00:00
// (a / b < c / d) <=> (a * d < c * b)
return (lhs.committed * rhs.soft_limit) < (rhs.committed * lhs.soft_limit)
|| (lhs.soft_limit == 0 && rhs.soft_limit > 0)
|| (lhs.committed == 0 && rhs.committed == 0 && lhs.soft_limit > rhs.soft_limit);
2021-10-22 12:56:09 +00:00
}
2021-11-09 13:40:23 +00:00
Int64 committed;
2021-10-22 12:56:09 +00:00
Int64 soft_limit;
};
class MemoryTracker;
struct OvercommitTracker : boost::noncopyable
{
2021-10-22 15:15:33 +00:00
OvercommitTracker();
void setMaxWaitTime(UInt64 wait_time);
2021-10-22 12:56:09 +00:00
bool needToStopQuery(MemoryTracker * tracker);
void unsubscribe(MemoryTracker * tracker);
2021-10-22 12:56:09 +00:00
virtual ~OvercommitTracker() = default;
protected:
virtual void pickQueryToExcludeImpl() = 0;
mutable std::mutex overcommit_m;
mutable std::condition_variable cv;
2021-10-22 15:15:33 +00:00
std::chrono::microseconds max_wait_time;
2021-10-22 12:56:09 +00:00
enum class QueryCancelationState
{
NONE,
RUNNING,
};
2021-11-09 13:40:23 +00:00
// Specifies memory tracker of the chosen to stop query.
// If soft limit is not set, all the queries which reach hard limit must stop.
// This case is represented as picked tracker pointer is set to nullptr and
// overcommit tracker is in RUNNING state.
2021-10-22 12:56:09 +00:00
MemoryTracker * picked_tracker;
2021-10-22 15:15:33 +00:00
QueryCancelationState cancelation_state;
2021-10-22 12:56:09 +00:00
private:
void pickQueryToExclude()
{
if (cancelation_state != QueryCancelationState::RUNNING)
{
pickQueryToExcludeImpl();
cancelation_state = QueryCancelationState::RUNNING;
}
}
friend struct BlockQueryIfMemoryLimit;
};
namespace DB
{
class ProcessList;
struct ProcessListForUser;
}
struct UserOvercommitTracker : OvercommitTracker
{
explicit UserOvercommitTracker(DB::ProcessListForUser * user_process_list_);
~UserOvercommitTracker() override = default;
protected:
void pickQueryToExcludeImpl() override final;
private:
DB::ProcessListForUser * user_process_list;
2021-12-06 22:06:00 +00:00
Poco::Logger * logger = &Poco::Logger::get("UserOvercommitTracker");
2021-10-22 12:56:09 +00:00
};
struct GlobalOvercommitTracker : OvercommitTracker
{
explicit GlobalOvercommitTracker(DB::ProcessList * process_list_)
: process_list(process_list_)
{}
~GlobalOvercommitTracker() override = default;
protected:
void pickQueryToExcludeImpl() override final;
private:
DB::ProcessList * process_list;
2021-12-06 22:06:00 +00:00
Poco::Logger * logger = &Poco::Logger::get("GlobalOvercommitTracker");
2021-10-22 12:56:09 +00:00
};
struct BlockQueryIfMemoryLimit
{
BlockQueryIfMemoryLimit(OvercommitTracker const & overcommit_tracker)
: mutex(overcommit_tracker.overcommit_m)
, lk(mutex)
{
if (overcommit_tracker.cancelation_state == OvercommitTracker::QueryCancelationState::RUNNING)
{
2021-10-22 15:15:33 +00:00
overcommit_tracker.cv.wait_for(lk, overcommit_tracker.max_wait_time, [&overcommit_tracker]()
2021-10-22 12:56:09 +00:00
{
return overcommit_tracker.cancelation_state == OvercommitTracker::QueryCancelationState::NONE;
});
}
}
~BlockQueryIfMemoryLimit() = default;
private:
std::mutex & mutex;
std::unique_lock<std::mutex> lk;
};