pool_allocator.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 <atomic>
#include <functional>
#include <vector>
 
#include <dispenso/platform.h>
 
namespace dispenso {
 
template <bool kThreadSafe>
class PoolAllocatorT {
 public:
  DISPENSO_DLL_ACCESS PoolAllocatorT(
      size_t chunkSize,
      size_t allocSize,
      std::function<void*(size_t)> allocFunc,
      std::function<void(void*)> deallocFunc);
 
  DISPENSO_DLL_ACCESS char* alloc();
 
  DISPENSO_DLL_ACCESS void dealloc(char* ptr);
 
  DISPENSO_DLL_ACCESS void clear();
 
  size_t totalChunkCapacity() const {
    return (backingAllocs2_.size() + backingAllocs_.size()) * chunksPerAlloc_;
  }
  DISPENSO_DLL_ACCESS ~PoolAllocatorT();
 
 private:
  const size_t chunkSize_;
  const size_t allocSize_;
  const size_t chunksPerAlloc_;
 
  std::function<void*(size_t)> allocFunc_;
  std::function<void(void*)> deallocFunc_;
 
  // Use of a spin lock was found to be faster than std::mutex in benchmarks.
  alignas(kCacheLineSize) std::atomic<uint32_t> backingAllocLock_{0};
  std::vector<char*> backingAllocs_;
  std::vector<char*> backingAllocs2_;
 
  std::vector<char*> chunks_;
};
 
using PoolAllocator = PoolAllocatorT<true>;
using NoLockPoolAllocator = PoolAllocatorT<false>;
 
} // namespace dispenso