1 | #pragma once |
2 | |
3 | #include <map> |
4 | #include <memory> |
5 | #include <stack> |
6 | #include <mutex> |
7 | |
8 | |
9 | namespace DB |
10 | { |
11 | |
12 | |
13 | /** Pool for objects that cannot be used from different threads simultaneously. |
14 | * Allows to create an object for each thread. |
15 | * Pool has unbounded size and objects are not destroyed before destruction of pool. |
16 | * |
17 | * Use it in cases when thread local storage is not appropriate |
18 | * (when maximum number of simultaneously used objects is less |
19 | * than number of running/sleeping threads, that has ever used object, |
20 | * and creation/destruction of objects is expensive). |
21 | */ |
22 | template <typename T> |
23 | class SimpleObjectPool |
24 | { |
25 | protected: |
26 | |
27 | /// Hold all avaiable objects in stack. |
28 | std::mutex mutex; |
29 | std::stack<std::unique_ptr<T>> stack; |
30 | |
31 | /// Specialized deleter for std::unique_ptr. |
32 | /// Returns underlying pointer back to stack thus reclaiming its ownership. |
33 | struct Deleter |
34 | { |
35 | SimpleObjectPool<T> * parent; |
36 | |
37 | Deleter(SimpleObjectPool<T> * parent_ = nullptr) : parent{parent_} {} |
38 | |
39 | void operator()(T * owning_ptr) const |
40 | { |
41 | std::lock_guard lock{parent->mutex}; |
42 | parent->stack.emplace(owning_ptr); |
43 | } |
44 | }; |
45 | |
46 | public: |
47 | using Pointer = std::unique_ptr<T, Deleter>; |
48 | |
49 | /// Extracts and returns a pointer from the stack if it's not empty, |
50 | /// creates a new one by calling provided f() otherwise. |
51 | template <typename Factory> |
52 | Pointer get(Factory && f) |
53 | { |
54 | std::unique_lock lock(mutex); |
55 | |
56 | if (stack.empty()) |
57 | { |
58 | lock.unlock(); |
59 | return { f(), this }; |
60 | } |
61 | |
62 | auto object = stack.top().release(); |
63 | stack.pop(); |
64 | |
65 | return { object, this }; |
66 | } |
67 | |
68 | /// Like get(), but creates object using default constructor. |
69 | Pointer getDefault() |
70 | { |
71 | return get([] { return new T; }); |
72 | } |
73 | }; |
74 | |
75 | |
76 | /// Like SimpleObjectPool, but additionally allows store different kind of objects that are identified by Key |
77 | template <typename T, typename Key> |
78 | class ObjectPoolMap |
79 | { |
80 | private: |
81 | |
82 | using Object = SimpleObjectPool<T>; |
83 | |
84 | /// Key -> objects |
85 | using Container = std::map<Key, std::unique_ptr<Object>>; |
86 | |
87 | Container container; |
88 | std::mutex mutex; |
89 | |
90 | public: |
91 | |
92 | using Pointer = typename Object::Pointer; |
93 | |
94 | template <typename Factory> |
95 | Pointer get(const Key & key, Factory && f) |
96 | { |
97 | std::unique_lock lock(mutex); |
98 | |
99 | auto it = container.find(key); |
100 | if (container.end() == it) |
101 | it = container.emplace(key, std::make_unique<Object>()).first; |
102 | |
103 | return it->second->get(std::forward<Factory>(f)); |
104 | } |
105 | }; |
106 | |
107 | |
108 | } |
109 | |