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
24namespace 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.
29enum 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<>.
44template <typename T>
45class 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.
102template <typename T>
103using Loan = typename Pool<T>::Loan;
104
105////////////////////////////////////////////////////////////////////////////////
106// Pool<T>::Item
107////////////////////////////////////////////////////////////////////////////////
108template <typename T>
109T* Pool<T>::Item::get() {
110 return reinterpret_cast<T*>(&data);
111}
112
113template <typename T>
114void Pool<T>::Item::construct() {
115 new (&data) T;
116}
117
118template <typename T>
119void Pool<T>::Item::destruct() {
120 get()->~T();
121}
122
123////////////////////////////////////////////////////////////////////////////////
124// Pool<T>::Loan
125////////////////////////////////////////////////////////////////////////////////
126template <typename T>
127Pool<T>::Loan::Loan(Item* item, const std::shared_ptr<Storage>& storage)
128 : item(item), storage(storage) {
129 item->refcount++;
130}
131
132template <typename T>
133Pool<T>::Loan::Loan(const Loan& other)
134 : item(other.item), storage(other.storage) {
135 if (item != nullptr) {
136 item->refcount++;
137 }
138}
139
140template <typename T>
141Pool<T>::Loan::Loan(Loan&& other) : item(other.item), storage(other.storage) {
142 other.item = nullptr;
143 other.storage = nullptr;
144}
145
146template <typename T>
147Pool<T>::Loan::~Loan() {
148 reset();
149}
150
151template <typename T>
152void 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
164template <typename T>
165typename 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
175template <typename T>
176typename 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
183template <typename T>
184T& Pool<T>::Loan::operator*() {
185 return *item->get();
186}
187
188template <typename T>
189T* Pool<T>::Loan::operator->() const {
190 return item->get();
191}
192
193template <typename T>
194T* 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.
207template <typename T, int N, PoolPolicy POLICY = PoolPolicy::Reconstruct>
208class 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
247template <typename T, int N, PoolPolicy POLICY>
248BoundedPool<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
259template <typename T, int N, PoolPolicy POLICY>
260BoundedPool<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
268template <typename T, int N, PoolPolicy POLICY>
269BoundedPool<T, N, POLICY>::BoundedPool(
270 Allocator* allocator /* = Allocator::Default */)
271 : storage(allocator->make_shared<Storage>(allocator)) {}
272
273template <typename T, int N, PoolPolicy POLICY>
274typename 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
281template <typename T, int N, PoolPolicy POLICY>
282template <typename F>
283void 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
296template <typename T, int N, PoolPolicy POLICY>
297std::pair<typename BoundedPool<T, N, POLICY>::Loan, bool>
298BoundedPool<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
315template <typename T, int N, PoolPolicy POLICY>
316void 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.
337template <typename T, PoolPolicy POLICY = PoolPolicy::Reconstruct>
338class 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
375template <typename T, PoolPolicy POLICY>
376UnboundedPool<T, POLICY>::Storage::Storage(Allocator* allocator)
377 : allocator(allocator), items(allocator) {}
378
379template <typename T, PoolPolicy POLICY>
380UnboundedPool<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
389template <typename T, PoolPolicy POLICY>
390UnboundedPool<T, POLICY>::UnboundedPool(
391 Allocator* allocator /* = Allocator::Default */)
392 : allocator(allocator),
393 storage(allocator->make_shared<Storage>(allocator)) {}
394
395template <typename T, PoolPolicy POLICY>
396Loan<T> UnboundedPool<T, POLICY>::borrow() const {
397 Loan out;
398 borrow(1, [&](Loan&& loan) { out = std::move(loan); });
399 return out;
400}
401
402template <typename T, PoolPolicy POLICY>
403template <typename F>
404inline 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
429template <typename T, PoolPolicy POLICY>
430void 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