ClickHouse/dbms/include/DB/Functions/ObjectPool.h

108 lines
2.4 KiB
C++

#pragma once
#include <map>
#include <memory>
#include <stack>
#include <mutex>
namespace DB
{
/** Pool for objects that cannot be used from different threads simultaneously.
* Allows to create an object for each thread.
* Pool has unbounded size and objects are not destroyed before destruction of pool.
*
* Use it in cases when thread local storage is not appropriate
* (when maximum number of simultaneously used objects is less
* than number of running/sleeping threads, that has ever used object,
* and creation/destruction of objects is expensive).
*/
template <typename T>
class SimpleObjectPool
{
protected:
/// Hold all avaiable objects in stack.
std::mutex mutex;
std::stack<std::unique_ptr<T>> stack;
/// Specialized deleter for std::unique_ptr.
/// Returns underlying pointer back to stack thus reclaiming its ownership.
struct Deleter
{
SimpleObjectPool<T> * parent;
Deleter(SimpleObjectPool<T> * parent_ = nullptr) : parent{parent_} {}
void operator()(T * owning_ptr) const
{
std::lock_guard<std::mutex> lock{parent->mutex};
parent->stack.emplace(owning_ptr);
}
};
public:
using Pointer = std::unique_ptr<T, Deleter>;
/// Extracts and returns a pointer from the stack if it's not empty,
/// creates a new one by calling provided f() otherwise.
template <typename Factory>
Pointer get(Factory && f)
{
std::unique_lock<std::mutex> lock(mutex);
if (stack.empty())
{
lock.unlock();
return { f(), this };
}
auto object = stack.top().release();
stack.pop();
return { object, this };
}
/// Like get(), but creates object using default constructor.
Pointer getDefault()
{
return get([] { return new T; });
}
};
/// Like SimpleObjectPool, but additionally allows store different kind of objects that are identified by Key
template <typename T, typename Key>
class ObjectPoolMap
{
private:
using Object = SimpleObjectPool<T>;
/// Key -> objects
using Container = std::map<Key, std::unique_ptr<Object>>;
Container container;
std::mutex mutex;
public:
using Pointer = typename Object::Pointer;
template <typename Factory>
Pointer get(const Key & key, Factory && f)
{
std::unique_lock<std::mutex> lock(mutex);
auto it = container.find(key);
if (container.end() == it)
it = container.emplace(key, std::make_unique<Object>()).first;
return it->second->get(std::forward<Factory>(f));
}
};
}