ClickHouse/dbms/include/DB/Storages/MergeTree/BackgroundProcessingPool.h
Michael Kolupaev 7ec176038c Merge
2014-04-11 17:05:20 +04:00

259 lines
5.8 KiB
C++
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#pragma once
#include <thread>
#include <set>
#include <map>
#include <Poco/Mutex.h>
#include <Poco/RWLock.h>
#include <DB/Core/Types.h>
#include <DB/Core/Exception.h>
#include <DB/Core/ErrorCodes.h>
namespace DB
{
/** В нескольких потоках в бесконечном цикле выполняет указанные функции.
*/
class BackgroundProcessingPool
{
public:
typedef std::map<String, int> Counters;
/** Используется изнутри таски. Позволяет инкрементировать какие-нибудь счетчики.
* После завершения таски, все изменения откатятся.
* Например, чтобы можно было узнавать количество потоков, выполняющих большое слияние,
* можно в таске, выполняющей большое слияние, инкрементировать счетчик. Декрементировать обратно его не нужно.
*/
class Context
{
public:
void incrementCounter(const String & name, int value = 1)
{
Poco::ScopedLock<Poco::FastMutex> lock(pool.mutex);
local_counters[name] += value;
pool.counters[name] += value;
}
private:
friend class BackgroundProcessingPool;
Context(BackgroundProcessingPool & pool_, Counters & local_counters_) : pool(pool_), local_counters(local_counters_) {}
BackgroundProcessingPool & pool;
Counters & local_counters;
};
/// Возвращает true, если что-то получилось сделать. В таком случае поток не будет спать перед следующим вызовом.
typedef std::function<bool (Context & context)> Task;
typedef std::shared_ptr<void> TaskHandle;
BackgroundProcessingPool() : size(1), sleep_seconds(1), shutdown(false) {}
void setNumberOfThreads(int size_)
{
Poco::ScopedLock<Poco::FastMutex> tlock(threads_mutex);
Poco::ScopedLock<Poco::FastMutex> lock(mutex);
if (size_ == size)
return;
if (threads.empty())
{
size = size_;
return;
}
throw Exception("setNumberOfThreads is not implemented for non-empty pool", ErrorCodes::NOT_IMPLEMENTED);
}
int getNumberOfThreads()
{
Poco::ScopedLock<Poco::FastMutex> lock(mutex);
return size;
}
void setSleepTime(double seconds)
{
Poco::ScopedLock<Poco::FastMutex> lock(mutex);
sleep_seconds = seconds;
}
int getCounter(const String & name)
{
Poco::ScopedLock<Poco::FastMutex> lock(mutex);
return counters[name];
}
TaskHandle addTask(const Task & task)
{
Poco::ScopedLock<Poco::FastMutex> lock(threads_mutex);
TaskInfoPtr res = std::make_shared<TaskInfo>(task);
{
Poco::ScopedLock<Poco::FastMutex> lock(mutex);
tasks.push_back(res);
}
if (threads.empty())
{
shutdown = false;
counters.clear();
threads.resize(size);
for (std::thread & thread : threads)
thread = std::thread(std::bind(&BackgroundProcessingPool::threadFunction, this));
}
return res;
}
void removeTask(const TaskHandle & handle)
{
Poco::ScopedLock<Poco::FastMutex> tlock(threads_mutex);
TaskInfoPtr task = std::static_pointer_cast<TaskInfo>(handle);
/// Дождемся завершения всех выполнений этой задачи.
{
Poco::ScopedWriteRWLock wlock(task->lock);
task->removed = true;
}
{
Poco::ScopedLock<Poco::FastMutex> lock(mutex);
auto it = std::find(tasks.begin(), tasks.end(), task);
if (it == tasks.end())
throw Exception("Task not found", ErrorCodes::LOGICAL_ERROR);
tasks.erase(it);
}
if (tasks.empty())
{
shutdown = true;
for (std::thread & thread : threads)
thread.join();
threads.clear();
counters.clear();
}
}
~BackgroundProcessingPool()
{
try
{
Poco::ScopedLock<Poco::FastMutex> lock(threads_mutex);
if (!threads.empty())
{
LOG_ERROR(&Logger::get("~BackgroundProcessingPool"), "Destroying non-empty BackgroundProcessingPool");
shutdown = true;
for (std::thread & thread : threads)
thread.join();
}
}
catch (...)
{
tryLogCurrentException(__PRETTY_FUNCTION__);
}
}
private:
struct TaskInfo
{
Task function;
Poco::RWLock lock;
volatile bool removed;
TaskInfo(const Task & function_) : function(function_), removed(false) {}
};
typedef std::shared_ptr<TaskInfo> TaskInfoPtr;
typedef std::vector<TaskInfoPtr> Tasks;
typedef std::vector<std::thread> Threads;
Poco::FastMutex threads_mutex;
Poco::FastMutex mutex;
int size;
Tasks tasks;
Threads threads;
Counters counters;
double sleep_seconds;
bool shutdown;
void threadFunction()
{
size_t i = static_cast<size_t>(rand());
while (!shutdown)
{
Counters counters_diff;
bool need_sleep = false;
size_t tasks_count = 1;
try
{
TaskInfoPtr task;
{
Poco::ScopedLock<Poco::FastMutex> lock(mutex);
tasks_count = tasks.size();
if (!tasks.empty())
{
need_sleep = true;
i %= tasks_count;
task = tasks[i];
++i;
}
}
if (shutdown)
break;
if (!task)
{
std::this_thread::sleep_for(std::chrono::duration<double>(sleep_seconds));
continue;
}
Poco::ScopedReadRWLock rlock(task->lock);
if (task->removed)
continue;
Context context(*this, counters_diff);
if (task->function(context))
{
--i;
need_sleep = false;
}
}
catch (...)
{
need_sleep = true;
tryLogCurrentException(__PRETTY_FUNCTION__);
}
/// Вычтем все счетчики обратно.
if (!counters_diff.empty())
{
Poco::ScopedLock<Poco::FastMutex> lock(mutex);
for (const auto & it : counters_diff)
{
counters[it.first] -= it.second;
}
}
if (shutdown)
break;
if (need_sleep)
{
std::this_thread::sleep_for(std::chrono::duration<double>(sleep_seconds / tasks_count));
}
}
}
};
}