1*01826a49SYabin Cui /* 2*01826a49SYabin Cui * Copyright (c) Meta Platforms, Inc. and affiliates. 3*01826a49SYabin Cui * All rights reserved. 4*01826a49SYabin Cui * 5*01826a49SYabin Cui * This source code is licensed under both the BSD-style license (found in the 6*01826a49SYabin Cui * LICENSE file in the root directory of this source tree) and the GPLv2 (found 7*01826a49SYabin Cui * in the COPYING file in the root directory of this source tree). 8*01826a49SYabin Cui */ 9*01826a49SYabin Cui #pragma once 10*01826a49SYabin Cui 11*01826a49SYabin Cui #include <cassert> 12*01826a49SYabin Cui #include <functional> 13*01826a49SYabin Cui #include <memory> 14*01826a49SYabin Cui #include <mutex> 15*01826a49SYabin Cui #include <vector> 16*01826a49SYabin Cui 17*01826a49SYabin Cui namespace pzstd { 18*01826a49SYabin Cui 19*01826a49SYabin Cui /** 20*01826a49SYabin Cui * An unbounded pool of resources. 21*01826a49SYabin Cui * A `ResourcePool<T>` requires a factory function that takes allocates `T*` and 22*01826a49SYabin Cui * a free function that frees a `T*`. 23*01826a49SYabin Cui * Calling `ResourcePool::get()` will give you a new `ResourcePool::UniquePtr` 24*01826a49SYabin Cui * to a `T`, and when it goes out of scope the resource will be returned to the 25*01826a49SYabin Cui * pool. 26*01826a49SYabin Cui * The `ResourcePool<T>` *must* survive longer than any resources it hands out. 27*01826a49SYabin Cui * Remember that `ResourcePool<T>` hands out mutable `T`s, so make sure to clean 28*01826a49SYabin Cui * up the resource before or after every use. 29*01826a49SYabin Cui */ 30*01826a49SYabin Cui template <typename T> 31*01826a49SYabin Cui class ResourcePool { 32*01826a49SYabin Cui public: 33*01826a49SYabin Cui class Deleter; 34*01826a49SYabin Cui using Factory = std::function<T*()>; 35*01826a49SYabin Cui using Free = std::function<void(T*)>; 36*01826a49SYabin Cui using UniquePtr = std::unique_ptr<T, Deleter>; 37*01826a49SYabin Cui 38*01826a49SYabin Cui private: 39*01826a49SYabin Cui std::mutex mutex_; 40*01826a49SYabin Cui Factory factory_; 41*01826a49SYabin Cui Free free_; 42*01826a49SYabin Cui std::vector<T*> resources_; 43*01826a49SYabin Cui unsigned inUse_; 44*01826a49SYabin Cui 45*01826a49SYabin Cui public: 46*01826a49SYabin Cui /** 47*01826a49SYabin Cui * Creates a `ResourcePool`. 48*01826a49SYabin Cui * 49*01826a49SYabin Cui * @param factory The function to use to create new resources. 50*01826a49SYabin Cui * @param free The function to use to free resources created by `factory`. 51*01826a49SYabin Cui */ ResourcePool(Factory factory,Free free)52*01826a49SYabin Cui ResourcePool(Factory factory, Free free) 53*01826a49SYabin Cui : factory_(std::move(factory)), free_(std::move(free)), inUse_(0) {} 54*01826a49SYabin Cui 55*01826a49SYabin Cui /** 56*01826a49SYabin Cui * @returns A unique pointer to a resource. The resource is null iff 57*01826a49SYabin Cui * there are no available resources and `factory()` returns null. 58*01826a49SYabin Cui */ get()59*01826a49SYabin Cui UniquePtr get() { 60*01826a49SYabin Cui std::lock_guard<std::mutex> lock(mutex_); 61*01826a49SYabin Cui if (!resources_.empty()) { 62*01826a49SYabin Cui UniquePtr resource{resources_.back(), Deleter{*this}}; 63*01826a49SYabin Cui resources_.pop_back(); 64*01826a49SYabin Cui ++inUse_; 65*01826a49SYabin Cui return resource; 66*01826a49SYabin Cui } 67*01826a49SYabin Cui UniquePtr resource{factory_(), Deleter{*this}}; 68*01826a49SYabin Cui ++inUse_; 69*01826a49SYabin Cui return resource; 70*01826a49SYabin Cui } 71*01826a49SYabin Cui ~ResourcePool()72*01826a49SYabin Cui ~ResourcePool() noexcept { 73*01826a49SYabin Cui assert(inUse_ == 0); 74*01826a49SYabin Cui for (const auto resource : resources_) { 75*01826a49SYabin Cui free_(resource); 76*01826a49SYabin Cui } 77*01826a49SYabin Cui } 78*01826a49SYabin Cui 79*01826a49SYabin Cui class Deleter { 80*01826a49SYabin Cui ResourcePool *pool_; 81*01826a49SYabin Cui public: Deleter(ResourcePool & pool)82*01826a49SYabin Cui explicit Deleter(ResourcePool &pool) : pool_(&pool) {} 83*01826a49SYabin Cui operator()84*01826a49SYabin Cui void operator() (T *resource) { 85*01826a49SYabin Cui std::lock_guard<std::mutex> lock(pool_->mutex_); 86*01826a49SYabin Cui // Make sure we don't put null resources into the pool 87*01826a49SYabin Cui if (resource) { 88*01826a49SYabin Cui pool_->resources_.push_back(resource); 89*01826a49SYabin Cui } 90*01826a49SYabin Cui assert(pool_->inUse_ > 0); 91*01826a49SYabin Cui --pool_->inUse_; 92*01826a49SYabin Cui } 93*01826a49SYabin Cui }; 94*01826a49SYabin Cui }; 95*01826a49SYabin Cui 96*01826a49SYabin Cui } 97