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
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>
Storage(Allocator * allocator)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>
~Storage()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>
BoundedPool(Allocator * allocator)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>
borrow()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>
borrow(size_t n,const F & 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>
tryBorrow()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>
return_(Item * item)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>
Storage(Allocator * allocator)376 UnboundedPool<T, POLICY>::Storage::Storage(Allocator* allocator)
377 : allocator(allocator), items(allocator) {}
378
379 template <typename T, PoolPolicy POLICY>
~Storage()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>
UnboundedPool(Allocator * allocator)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>
borrow()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>
borrow(size_t n,const F & 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>
return_(Item * item)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