1 | // Copyright 2019 The Marl Authors. |
2 | // |
3 | // Licensed under the Apache License, Version 2.0 (the "License"); |
4 | // you may not use this file except in compliance with the License. |
5 | // You may obtain a copy of the License at |
6 | // |
7 | // https://www.apache.org/licenses/LICENSE-2.0 |
8 | // |
9 | // Unless required by applicable law or agreed to in writing, software |
10 | // distributed under the License is distributed on an "AS IS" BASIS, |
11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
12 | // See the License for the specific language governing permissions and |
13 | // limitations under the License. |
14 | |
15 | #ifndef marl_pool_h |
16 | #define marl_pool_h |
17 | |
18 | #include "conditionvariable.h" |
19 | #include "memory.h" |
20 | #include "mutex.h" |
21 | |
22 | #include <atomic> |
23 | |
24 | namespace marl { |
25 | |
26 | // PoolPolicy controls whether pool items are constructed and destructed each |
27 | // time they are borrowed from and returned to a pool, or whether they persist |
28 | // constructed for the lifetime of the pool. |
29 | enum class PoolPolicy { |
30 | // Call the Pool items constructor on borrow(), and destruct the item |
31 | // when the item is returned. |
32 | Reconstruct, |
33 | |
34 | // Construct and destruct all items once for the lifetime of the Pool. |
35 | // Items will keep their state between loans. |
36 | Preserve, |
37 | }; |
38 | |
39 | //////////////////////////////////////////////////////////////////////////////// |
40 | // Pool<T> |
41 | //////////////////////////////////////////////////////////////////////////////// |
42 | |
43 | // Pool is the abstract base class for BoundedPool<> and UnboundedPool<>. |
44 | template <typename T> |
45 | class Pool { |
46 | protected: |
47 | struct Item; |
48 | class Storage; |
49 | |
50 | public: |
51 | // A Loan is returned by the pool's borrow() function. |
52 | // Loans track the number of references to the loaned item, and return the |
53 | // item to the pool when the final Loan reference is dropped. |
54 | class Loan { |
55 | public: |
56 | MARL_NO_EXPORT inline Loan() = default; |
57 | MARL_NO_EXPORT inline Loan(Item*, const std::shared_ptr<Storage>&); |
58 | MARL_NO_EXPORT inline Loan(const Loan&); |
59 | MARL_NO_EXPORT inline Loan(Loan&&); |
60 | MARL_NO_EXPORT inline ~Loan(); |
61 | MARL_NO_EXPORT inline Loan& operator=(const Loan&); |
62 | MARL_NO_EXPORT inline Loan& operator=(Loan&&); |
63 | MARL_NO_EXPORT inline T& operator*(); |
64 | MARL_NO_EXPORT inline T* operator->() const; |
65 | MARL_NO_EXPORT inline T* get() const; |
66 | MARL_NO_EXPORT inline void reset(); |
67 | |
68 | private: |
69 | Item* item = nullptr; |
70 | std::shared_ptr<Storage> storage; |
71 | }; |
72 | |
73 | protected: |
74 | Pool() = default; |
75 | |
76 | // The shared storage between the pool and all loans. |
77 | class Storage { |
78 | public: |
79 | virtual ~Storage() = default; |
80 | virtual void return_(Item*) = 0; |
81 | }; |
82 | |
83 | // The backing data of a single item in the pool. |
84 | struct Item { |
85 | // get() returns a pointer to the item's data. |
86 | MARL_NO_EXPORT inline T* get(); |
87 | |
88 | // construct() calls the constructor on the item's data. |
89 | MARL_NO_EXPORT inline void construct(); |
90 | |
91 | // destruct() calls the destructor on the item's data. |
92 | MARL_NO_EXPORT inline void destruct(); |
93 | |
94 | using Data = typename aligned_storage<sizeof(T), alignof(T)>::type; |
95 | Data data; |
96 | std::atomic<int> refcount = {0}; |
97 | Item* next = nullptr; // pointer to the next free item in the pool. |
98 | }; |
99 | }; |
100 | |
101 | // Loan<T> is an alias to Pool<T>::Loan. |
102 | template <typename T> |
103 | using Loan = typename Pool<T>::Loan; |
104 | |
105 | //////////////////////////////////////////////////////////////////////////////// |
106 | // Pool<T>::Item |
107 | //////////////////////////////////////////////////////////////////////////////// |
108 | template <typename T> |
109 | T* Pool<T>::Item::get() { |
110 | return reinterpret_cast<T*>(&data); |
111 | } |
112 | |
113 | template <typename T> |
114 | void Pool<T>::Item::construct() { |
115 | new (&data) T; |
116 | } |
117 | |
118 | template <typename T> |
119 | void Pool<T>::Item::destruct() { |
120 | get()->~T(); |
121 | } |
122 | |
123 | //////////////////////////////////////////////////////////////////////////////// |
124 | // Pool<T>::Loan |
125 | //////////////////////////////////////////////////////////////////////////////// |
126 | template <typename T> |
127 | Pool<T>::Loan::Loan(Item* item, const std::shared_ptr<Storage>& storage) |
128 | : item(item), storage(storage) { |
129 | item->refcount++; |
130 | } |
131 | |
132 | template <typename T> |
133 | Pool<T>::Loan::Loan(const Loan& other) |
134 | : item(other.item), storage(other.storage) { |
135 | if (item != nullptr) { |
136 | item->refcount++; |
137 | } |
138 | } |
139 | |
140 | template <typename T> |
141 | Pool<T>::Loan::Loan(Loan&& other) : item(other.item), storage(other.storage) { |
142 | other.item = nullptr; |
143 | other.storage = nullptr; |
144 | } |
145 | |
146 | template <typename T> |
147 | Pool<T>::Loan::~Loan() { |
148 | reset(); |
149 | } |
150 | |
151 | template <typename T> |
152 | void Pool<T>::Loan::reset() { |
153 | if (item != nullptr) { |
154 | auto refs = --item->refcount; |
155 | MARL_ASSERT(refs >= 0, "reset() called on zero-ref pool item" ); |
156 | if (refs == 0) { |
157 | storage->return_(item); |
158 | } |
159 | item = nullptr; |
160 | storage = nullptr; |
161 | } |
162 | } |
163 | |
164 | template <typename T> |
165 | typename Pool<T>::Loan& Pool<T>::Loan::operator=(const Loan& rhs) { |
166 | reset(); |
167 | if (rhs.item != nullptr) { |
168 | item = rhs.item; |
169 | storage = rhs.storage; |
170 | rhs.item->refcount++; |
171 | } |
172 | return *this; |
173 | } |
174 | |
175 | template <typename T> |
176 | typename Pool<T>::Loan& Pool<T>::Loan::operator=(Loan&& rhs) { |
177 | reset(); |
178 | std::swap(item, rhs.item); |
179 | std::swap(storage, rhs.storage); |
180 | return *this; |
181 | } |
182 | |
183 | template <typename T> |
184 | T& Pool<T>::Loan::operator*() { |
185 | return *item->get(); |
186 | } |
187 | |
188 | template <typename T> |
189 | T* Pool<T>::Loan::operator->() const { |
190 | return item->get(); |
191 | } |
192 | |
193 | template <typename T> |
194 | T* Pool<T>::Loan::get() const { |
195 | return item ? item->get() : nullptr; |
196 | } |
197 | |
198 | //////////////////////////////////////////////////////////////////////////////// |
199 | // BoundedPool<T, N, POLICY> |
200 | //////////////////////////////////////////////////////////////////////////////// |
201 | |
202 | // BoundedPool<T, N, POLICY> is a pool of items of type T, with a maximum |
203 | // capacity of N items. |
204 | // BoundedPool<> is initially populated with N default-constructed items. |
205 | // POLICY controls whether pool items are constructed and destructed each |
206 | // time they are borrowed from and returned to the pool. |
207 | template <typename T, int N, PoolPolicy POLICY = PoolPolicy::Reconstruct> |
208 | class BoundedPool : public Pool<T> { |
209 | public: |
210 | using Item = typename Pool<T>::Item; |
211 | using Loan = typename Pool<T>::Loan; |
212 | |
213 | MARL_NO_EXPORT inline BoundedPool(Allocator* allocator = Allocator::Default); |
214 | |
215 | // borrow() borrows a single item from the pool, blocking until an item is |
216 | // returned if the pool is empty. |
217 | MARL_NO_EXPORT inline Loan borrow() const; |
218 | |
219 | // borrow() borrows count items from the pool, blocking until there are at |
220 | // least count items in the pool. The function f() is called with each |
221 | // borrowed item. |
222 | // F must be a function with the signature: void(T&&) |
223 | template <typename F> |
224 | MARL_NO_EXPORT inline void borrow(size_t count, const F& f) const; |
225 | |
226 | // tryBorrow() attempts to borrow a single item from the pool without |
227 | // blocking. |
228 | // The boolean of the returned pair is true on success, or false if the pool |
229 | // is empty. |
230 | MARL_NO_EXPORT inline std::pair<Loan, bool> tryBorrow() const; |
231 | |
232 | private: |
233 | class Storage : public Pool<T>::Storage { |
234 | public: |
235 | MARL_NO_EXPORT inline Storage(Allocator* allocator); |
236 | MARL_NO_EXPORT inline ~Storage(); |
237 | MARL_NO_EXPORT inline void return_(Item*) override; |
238 | |
239 | Item items[N]; |
240 | marl::mutex mutex; |
241 | ConditionVariable returned; |
242 | Item* free = nullptr; |
243 | }; |
244 | std::shared_ptr<Storage> storage; |
245 | }; |
246 | |
247 | template <typename T, int N, PoolPolicy POLICY> |
248 | BoundedPool<T, N, POLICY>::Storage::Storage(Allocator* allocator) |
249 | : returned(allocator) { |
250 | for (int i = 0; i < N; i++) { |
251 | if (POLICY == PoolPolicy::Preserve) { |
252 | items[i].construct(); |
253 | } |
254 | items[i].next = this->free; |
255 | this->free = &items[i]; |
256 | } |
257 | } |
258 | |
259 | template <typename T, int N, PoolPolicy POLICY> |
260 | BoundedPool<T, N, POLICY>::Storage::~Storage() { |
261 | if (POLICY == PoolPolicy::Preserve) { |
262 | for (int i = 0; i < N; i++) { |
263 | items[i].destruct(); |
264 | } |
265 | } |
266 | } |
267 | |
268 | template <typename T, int N, PoolPolicy POLICY> |
269 | BoundedPool<T, N, POLICY>::BoundedPool( |
270 | Allocator* allocator /* = Allocator::Default */) |
271 | : storage(allocator->make_shared<Storage>(allocator)) {} |
272 | |
273 | template <typename T, int N, PoolPolicy POLICY> |
274 | typename BoundedPool<T, N, POLICY>::Loan BoundedPool<T, N, POLICY>::borrow() |
275 | const { |
276 | Loan out; |
277 | borrow(1, [&](Loan&& loan) { out = std::move(loan); }); |
278 | return out; |
279 | } |
280 | |
281 | template <typename T, int N, PoolPolicy POLICY> |
282 | template <typename F> |
283 | void BoundedPool<T, N, POLICY>::borrow(size_t n, const F& f) const { |
284 | marl::lock lock(storage->mutex); |
285 | for (size_t i = 0; i < n; i++) { |
286 | storage->returned.wait(lock, [&] { return storage->free != nullptr; }); |
287 | auto item = storage->free; |
288 | storage->free = storage->free->next; |
289 | if (POLICY == PoolPolicy::Reconstruct) { |
290 | item->construct(); |
291 | } |
292 | f(std::move(Loan(item, storage))); |
293 | } |
294 | } |
295 | |
296 | template <typename T, int N, PoolPolicy POLICY> |
297 | std::pair<typename BoundedPool<T, N, POLICY>::Loan, bool> |
298 | BoundedPool<T, N, POLICY>::tryBorrow() const { |
299 | Item* item = nullptr; |
300 | { |
301 | marl::lock lock(storage->mutex); |
302 | if (storage->free == nullptr) { |
303 | return std::make_pair(Loan(), false); |
304 | } |
305 | item = storage->free; |
306 | storage->free = storage->free->next; |
307 | item->pool = this; |
308 | } |
309 | if (POLICY == PoolPolicy::Reconstruct) { |
310 | item->construct(); |
311 | } |
312 | return std::make_pair(Loan(item, storage), true); |
313 | } |
314 | |
315 | template <typename T, int N, PoolPolicy POLICY> |
316 | void BoundedPool<T, N, POLICY>::Storage::return_(Item* item) { |
317 | if (POLICY == PoolPolicy::Reconstruct) { |
318 | item->destruct(); |
319 | } |
320 | { |
321 | marl::lock lock(mutex); |
322 | item->next = free; |
323 | free = item; |
324 | } |
325 | returned.notify_one(); |
326 | } |
327 | |
328 | //////////////////////////////////////////////////////////////////////////////// |
329 | // UnboundedPool |
330 | //////////////////////////////////////////////////////////////////////////////// |
331 | |
332 | // UnboundedPool<T, POLICY> is a pool of items of type T. |
333 | // UnboundedPool<> will automatically allocate more items if the pool becomes |
334 | // empty. |
335 | // POLICY controls whether pool items are constructed and destructed each |
336 | // time they are borrowed from and returned to the pool. |
337 | template <typename T, PoolPolicy POLICY = PoolPolicy::Reconstruct> |
338 | class UnboundedPool : public Pool<T> { |
339 | public: |
340 | using Item = typename Pool<T>::Item; |
341 | using Loan = typename Pool<T>::Loan; |
342 | |
343 | MARL_NO_EXPORT inline UnboundedPool( |
344 | Allocator* allocator = Allocator::Default); |
345 | |
346 | // borrow() borrows a single item from the pool, automatically allocating |
347 | // more items if the pool is empty. |
348 | // This function does not block. |
349 | MARL_NO_EXPORT inline Loan borrow() const; |
350 | |
351 | // borrow() borrows count items from the pool, calling the function f() with |
352 | // each borrowed item. |
353 | // F must be a function with the signature: void(T&&) |
354 | // This function does not block. |
355 | template <typename F> |
356 | MARL_NO_EXPORT inline void borrow(size_t n, const F& f) const; |
357 | |
358 | private: |
359 | class Storage : public Pool<T>::Storage { |
360 | public: |
361 | MARL_NO_EXPORT inline Storage(Allocator* allocator); |
362 | MARL_NO_EXPORT inline ~Storage(); |
363 | MARL_NO_EXPORT inline void return_(Item*) override; |
364 | |
365 | Allocator* allocator; |
366 | marl::mutex mutex; |
367 | containers::vector<Item*, 4> items; |
368 | Item* free = nullptr; |
369 | }; |
370 | |
371 | Allocator* allocator; |
372 | std::shared_ptr<Storage> storage; |
373 | }; |
374 | |
375 | template <typename T, PoolPolicy POLICY> |
376 | UnboundedPool<T, POLICY>::Storage::Storage(Allocator* allocator) |
377 | : allocator(allocator), items(allocator) {} |
378 | |
379 | template <typename T, PoolPolicy POLICY> |
380 | UnboundedPool<T, POLICY>::Storage::~Storage() { |
381 | for (auto item : items) { |
382 | if (POLICY == PoolPolicy::Preserve) { |
383 | item->destruct(); |
384 | } |
385 | allocator->destroy(item); |
386 | } |
387 | } |
388 | |
389 | template <typename T, PoolPolicy POLICY> |
390 | UnboundedPool<T, POLICY>::UnboundedPool( |
391 | Allocator* allocator /* = Allocator::Default */) |
392 | : allocator(allocator), |
393 | storage(allocator->make_shared<Storage>(allocator)) {} |
394 | |
395 | template <typename T, PoolPolicy POLICY> |
396 | Loan<T> UnboundedPool<T, POLICY>::borrow() const { |
397 | Loan out; |
398 | borrow(1, [&](Loan&& loan) { out = std::move(loan); }); |
399 | return out; |
400 | } |
401 | |
402 | template <typename T, PoolPolicy POLICY> |
403 | template <typename F> |
404 | inline void UnboundedPool<T, POLICY>::borrow(size_t n, const F& f) const { |
405 | marl::lock lock(storage->mutex); |
406 | for (size_t i = 0; i < n; i++) { |
407 | if (storage->free == nullptr) { |
408 | auto count = std::max<size_t>(storage->items.size(), 32); |
409 | for (size_t j = 0; j < count; j++) { |
410 | auto item = allocator->create<Item>(); |
411 | if (POLICY == PoolPolicy::Preserve) { |
412 | item->construct(); |
413 | } |
414 | storage->items.push_back(item); |
415 | item->next = storage->free; |
416 | storage->free = item; |
417 | } |
418 | } |
419 | |
420 | auto item = storage->free; |
421 | storage->free = storage->free->next; |
422 | if (POLICY == PoolPolicy::Reconstruct) { |
423 | item->construct(); |
424 | } |
425 | f(std::move(Loan(item, storage))); |
426 | } |
427 | } |
428 | |
429 | template <typename T, PoolPolicy POLICY> |
430 | void UnboundedPool<T, POLICY>::Storage::return_(Item* item) { |
431 | if (POLICY == PoolPolicy::Reconstruct) { |
432 | item->destruct(); |
433 | } |
434 | marl::lock lock(mutex); |
435 | item->next = free; |
436 | free = item; |
437 | } |
438 | |
439 | } // namespace marl |
440 | |
441 | #endif // marl_pool_h |
442 | |