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>
get()109 T* Pool<T>::Item::get() {
110 return reinterpret_cast<T*>(&data);
111 }
112
113 template <typename T>
construct()114 void Pool<T>::Item::construct() {
115 new (&data) T;
116 }
117
118 template <typename T>
destruct()119 void Pool<T>::Item::destruct() {
120 get()->~T();
121 }
122
123 ////////////////////////////////////////////////////////////////////////////////
124 // Pool<T>::Loan
125 ////////////////////////////////////////////////////////////////////////////////
126 template <typename T>
Loan(Item * item,const std::shared_ptr<Storage> & storage)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>
Loan(const Loan & other)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>
Loan(Loan && other)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>
~Loan()147 Pool<T>::Loan::~Loan() {
148 reset();
149 }
150
151 template <typename T>
reset()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>
get()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 // We cannot copy this as the Item pointers would be shared and
239 // deleted at a wrong point. We cannot move this because we return
240 // pointers into items[N].
241 MARL_NO_EXPORT inline Storage(const Storage&) = delete;
242 MARL_NO_EXPORT inline Storage& operator=(const Storage&) = delete;
243
244 Item items[N];
245 marl::mutex mutex;
246 ConditionVariable returned;
247 Item* free = nullptr;
248 };
249 std::shared_ptr<Storage> storage;
250 };
251
252 template <typename T, int N, PoolPolicy POLICY>
Storage(Allocator * allocator)253 BoundedPool<T, N, POLICY>::Storage::Storage(Allocator* allocator)
254 : returned(allocator) {
255 for (int i = 0; i < N; i++) {
256 if (POLICY == PoolPolicy::Preserve) {
257 items[i].construct();
258 }
259 items[i].next = this->free;
260 this->free = &items[i];
261 }
262 }
263
264 template <typename T, int N, PoolPolicy POLICY>
~Storage()265 BoundedPool<T, N, POLICY>::Storage::~Storage() {
266 if (POLICY == PoolPolicy::Preserve) {
267 for (int i = 0; i < N; i++) {
268 items[i].destruct();
269 }
270 }
271 }
272
273 template <typename T, int N, PoolPolicy POLICY>
BoundedPool(Allocator * allocator)274 BoundedPool<T, N, POLICY>::BoundedPool(
275 Allocator* allocator /* = Allocator::Default */)
276 : storage(allocator->make_shared<Storage>(allocator)) {}
277
278 template <typename T, int N, PoolPolicy POLICY>
borrow()279 typename BoundedPool<T, N, POLICY>::Loan BoundedPool<T, N, POLICY>::borrow()
280 const {
281 Loan out;
282 borrow(1, [&](Loan&& loan) { out = std::move(loan); });
283 return out;
284 }
285
286 template <typename T, int N, PoolPolicy POLICY>
287 template <typename F>
borrow(size_t n,const F & f)288 void BoundedPool<T, N, POLICY>::borrow(size_t n, const F& f) const {
289 marl::lock lock(storage->mutex);
290 for (size_t i = 0; i < n; i++) {
291 storage->returned.wait(lock, [&] { return storage->free != nullptr; });
292 auto item = storage->free;
293 storage->free = storage->free->next;
294 if (POLICY == PoolPolicy::Reconstruct) {
295 item->construct();
296 }
297 f(std::move(Loan(item, storage)));
298 }
299 }
300
301 template <typename T, int N, PoolPolicy POLICY>
302 std::pair<typename BoundedPool<T, N, POLICY>::Loan, bool>
tryBorrow()303 BoundedPool<T, N, POLICY>::tryBorrow() const {
304 Item* item = nullptr;
305 {
306 marl::lock lock(storage->mutex);
307 if (storage->free == nullptr) {
308 return std::make_pair(Loan(), false);
309 }
310 item = storage->free;
311 storage->free = storage->free->next;
312 item->pool = this;
313 }
314 if (POLICY == PoolPolicy::Reconstruct) {
315 item->construct();
316 }
317 return std::make_pair(Loan(item, storage), true);
318 }
319
320 template <typename T, int N, PoolPolicy POLICY>
return_(Item * item)321 void BoundedPool<T, N, POLICY>::Storage::return_(Item* item) {
322 if (POLICY == PoolPolicy::Reconstruct) {
323 item->destruct();
324 }
325 {
326 marl::lock lock(mutex);
327 item->next = free;
328 free = item;
329 }
330 returned.notify_one();
331 }
332
333 ////////////////////////////////////////////////////////////////////////////////
334 // UnboundedPool
335 ////////////////////////////////////////////////////////////////////////////////
336
337 // UnboundedPool<T, POLICY> is a pool of items of type T.
338 // UnboundedPool<> will automatically allocate more items if the pool becomes
339 // empty.
340 // POLICY controls whether pool items are constructed and destructed each
341 // time they are borrowed from and returned to the pool.
342 template <typename T, PoolPolicy POLICY = PoolPolicy::Reconstruct>
343 class UnboundedPool : public Pool<T> {
344 public:
345 using Item = typename Pool<T>::Item;
346 using Loan = typename Pool<T>::Loan;
347
348 MARL_NO_EXPORT inline UnboundedPool(
349 Allocator* allocator = Allocator::Default);
350
351 // borrow() borrows a single item from the pool, automatically allocating
352 // more items if the pool is empty.
353 // This function does not block.
354 MARL_NO_EXPORT inline Loan borrow() const;
355
356 // borrow() borrows count items from the pool, calling the function f() with
357 // each borrowed item.
358 // F must be a function with the signature: void(T&&)
359 // This function does not block.
360 template <typename F>
361 MARL_NO_EXPORT inline void borrow(size_t n, const F& f) const;
362
363 private:
364 class Storage : public Pool<T>::Storage {
365 public:
366 MARL_NO_EXPORT inline Storage(Allocator* allocator);
367 MARL_NO_EXPORT inline ~Storage();
368 MARL_NO_EXPORT inline void return_(Item*) override;
369 // We cannot copy this as the Item pointers would be shared and
370 // deleted at a wrong point. We could move this but would have to take
371 // extra care no Item pointers are left in the moved-out object.
372 MARL_NO_EXPORT inline Storage(const Storage&) = delete;
373 MARL_NO_EXPORT inline Storage& operator=(const Storage&) = delete;
374
375 Allocator* allocator;
376 marl::mutex mutex;
377 containers::vector<Item*, 4> items;
378 Item* free = nullptr;
379 };
380
381 Allocator* allocator;
382 std::shared_ptr<Storage> storage;
383 };
384
385 template <typename T, PoolPolicy POLICY>
Storage(Allocator * allocator)386 UnboundedPool<T, POLICY>::Storage::Storage(Allocator* allocator)
387 : allocator(allocator), items(allocator) {}
388
389 template <typename T, PoolPolicy POLICY>
~Storage()390 UnboundedPool<T, POLICY>::Storage::~Storage() {
391 for (auto item : items) {
392 if (POLICY == PoolPolicy::Preserve) {
393 item->destruct();
394 }
395 allocator->destroy(item);
396 }
397 }
398
399 template <typename T, PoolPolicy POLICY>
UnboundedPool(Allocator * allocator)400 UnboundedPool<T, POLICY>::UnboundedPool(
401 Allocator* allocator /* = Allocator::Default */)
402 : allocator(allocator),
403 storage(allocator->make_shared<Storage>(allocator)) {}
404
405 template <typename T, PoolPolicy POLICY>
borrow()406 Loan<T> UnboundedPool<T, POLICY>::borrow() const {
407 Loan out;
408 borrow(1, [&](Loan&& loan) { out = std::move(loan); });
409 return out;
410 }
411
412 template <typename T, PoolPolicy POLICY>
413 template <typename F>
borrow(size_t n,const F & f)414 inline void UnboundedPool<T, POLICY>::borrow(size_t n, const F& f) const {
415 marl::lock lock(storage->mutex);
416 for (size_t i = 0; i < n; i++) {
417 if (storage->free == nullptr) {
418 auto count = std::max<size_t>(storage->items.size(), 32);
419 for (size_t j = 0; j < count; j++) {
420 auto item = allocator->create<Item>();
421 if (POLICY == PoolPolicy::Preserve) {
422 item->construct();
423 }
424 storage->items.push_back(item);
425 item->next = storage->free;
426 storage->free = item;
427 }
428 }
429
430 auto item = storage->free;
431 storage->free = storage->free->next;
432 if (POLICY == PoolPolicy::Reconstruct) {
433 item->construct();
434 }
435 f(std::move(Loan(item, storage)));
436 }
437 }
438
439 template <typename T, PoolPolicy POLICY>
return_(Item * item)440 void UnboundedPool<T, POLICY>::Storage::return_(Item* item) {
441 if (POLICY == PoolPolicy::Reconstruct) {
442 item->destruct();
443 }
444 marl::lock lock(mutex);
445 item->next = free;
446 free = item;
447 }
448
449 } // namespace marl
450
451 #endif // marl_pool_h
452