resource_pool.h

Go to the documentation of this file.

/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */
 
#pragma once
 
#include <dispenso/platform.h>
#include <dispenso/tsan_annotations.h>
#include <moodycamel/blockingconcurrentqueue.h>
 
namespace dispenso {
 
template <typename T>
class ResourcePool;
 
template <typename T>
class Resource {
 public:
  Resource(Resource&& other) : resource_(other.resource_), pool_(other.pool_) {
    other.resource_ = nullptr;
  }
 
  Resource& operator=(Resource&& other) {
    if (&other != this) {
      recycle();
      resource_ = other.resource_;
      pool_ = other.pool_;
      other.resource_ = nullptr;
    }
    return *this;
  }
 
  T& get() {
    return *resource_;
  }
 
  ~Resource() {
    recycle();
  }
 
 private:
  Resource(T* res, ResourcePool<T>* pool) : resource_(res), pool_(pool) {}
 
  void recycle();
 
  T* resource_;
  ResourcePool<T>* pool_;
 
  friend class ResourcePool<T>;
};
 
template <typename T>
class ResourcePool {
 public:
  template <typename F>
  ResourcePool(size_t size, const F& init)
      : pool_(size),
        backingResources_(reinterpret_cast<char*>(
            detail::alignedMallocOrThrow(size * detail::alignToCacheLine(sizeof(T))))),
        size_(size) {
    char* buf = backingResources_;
 
    // There are three reasons we create our own buffer and use placement new:
    // 1. We want to be able to handle non-movable non-copyable objects
    //   * Note that we could do this with std::deque
    // 2. We want to minimize memory allocations, since that can be a common point of contention in
    //    multithreaded programs.
    // 3. We can easily ensure that the objects are cache aligned to help avoid false sharing.
 
    // If init() or a T constructor throws, the destructor will not run, so release what has been
    // built so far here.
    struct Cleanup {
      char* backing;
      size_t constructed;
      ~Cleanup() {
        if (!backing) {
          return;
        }
        for (size_t i = 0; i < constructed; ++i) {
          reinterpret_cast<T*>(backing + i * detail::alignToCacheLine(sizeof(T)))->~T();
        }
        detail::alignedFree(backing);
      }
    } cleanup{backingResources_, 0};
 
    while (cleanup.constructed < size) {
      T* t = new (buf) T(init());
      ++cleanup.constructed;
      pool_.enqueue(t);
      buf += detail::alignToCacheLine(sizeof(T));
    }
    cleanup.backing = nullptr;
  }
 
  Resource<T> acquire() {
    T* t;
    DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_BEGIN();
    pool_.wait_dequeue(t);
    DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_END();
    return Resource<T>(t, this);
  }
 
  ~ResourcePool() {
    assert(pool_.size_approx() == size_);
    for (size_t i = 0; i < size_; ++i) {
      T* t;
      DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_BEGIN();
      pool_.wait_dequeue(t);
      DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_END();
      t->~T();
    }
    detail::alignedFree(backingResources_);
  }
 
 private:
  void recycle(T* t) {
    DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_BEGIN();
    pool_.enqueue(t);
    DISPENSO_TSAN_ANNOTATE_IGNORE_WRITES_END();
  }
 
  moodycamel::BlockingConcurrentQueue<T*> pool_;
  char* backingResources_;
  size_t size_;
 
  friend class Resource<T>;
};
 
template <typename T>
void Resource<T>::recycle() {
  if (resource_) {
    pool_->recycle(resource_);
  }
}
 
} // namespace dispenso