Line data    Source code

       1             : // Copyright (c) 2014-2023 The Dash Core developers
       2             : // Distributed under the MIT software license, see the accompanying
       3             : // file COPYING or http://www.opensource.org/licenses/mit-license.php.
       4             : 
       5             : #ifndef BITCOIN_SUPPORT_ALLOCATORS_MT_POOLED_SECURE_H
       6             : #define BITCOIN_SUPPORT_ALLOCATORS_MT_POOLED_SECURE_H
       7             : 
       8             : #include <support/allocators/pooled_secure.h>
       9             : 
      10             : #include <thread>
      11             : #include <mutex>
      12             : 
      13             : //
      14             : // Manages a pool of pools to balance allocation between those when multiple threads are involved
      15             : // This allocator is fully thread safe
      16             : //
      17             : template <typename T>
      18             : struct mt_pooled_secure_allocator : public std::allocator<T> {
      19             :     using base = std::allocator<T>;
      20             :     using traits = std::allocator_traits<base>;
      21             :     using size_type = typename traits::size_type;
      22             :     using difference_type = typename traits::difference_type;
      23             :     using pointer = typename traits::pointer;
      24             :     using const_pointer = typename traits::const_pointer;
      25             :     using value_type = typename traits::value_type;
      26          20 :     mt_pooled_secure_allocator(size_type nrequested_size = 32,
      27             :                                size_type nnext_size = 32,
      28             :                                size_type nmax_size = 0,
      29             :                                size_type pools_count = std::thread::hardware_concurrency()) noexcept
      30          10 :     {
      31             :         // we add enough bytes to the requested size so that we can store the bucket as well
      32          10 :         nrequested_size += sizeof(size_t);
      33             : 
      34          10 :         if (pools_count == 0) {
      35           1 :             pools_count = 1;
      36           1 :         }
      37          10 :         pools.resize(pools_count);
      38         155 :         for (size_t i = 0; i < pools_count; i++) {
      39         145 :             pools[i] = std::make_unique<internal_pool>(nrequested_size, nnext_size, nmax_size);
      40         145 :         }
      41          20 :     }
      42          20 :     ~mt_pooled_secure_allocator() noexcept {}
      43             : 
      44         634 :     T* allocate(std::size_t n, const void* hint = nullptr)
      45             :     {
      46         634 :         size_t bucket = get_bucket();
      47         634 :         std::lock_guard<std::mutex> lock(pools[bucket]->mutex);
      48         634 :         uint8_t* ptr = pools[bucket]->allocate(n * sizeof(T) + sizeof(size_t));
      49         634 :         *(size_t*)ptr = bucket;
      50         634 :         return static_cast<T*>(ptr + sizeof(size_t));
      51         634 :     }
      52             : 
      53         634 :     void deallocate(T* p, std::size_t n)
      54             :     {
      55         634 :         if (!p) {
      56           0 :             return;
      57             :         }
      58         634 :         uint8_t* ptr = (uint8_t*)p - sizeof(size_t);
      59         634 :         size_t bucket = *(size_t*)ptr;
      60         634 :         std::lock_guard<std::mutex> lock(pools[bucket]->mutex);
      61         634 :         pools[bucket]->deallocate(ptr, n * sizeof(T));
      62         634 :     }
      63             : 
      64             : private:
      65         634 :     size_t get_bucket()
      66             :     {
      67         634 :         size_t x = std::hash<std::thread::id>{}(std::this_thread::get_id());
      68         634 :         return x % pools.size();
      69             :     }
      70             : 
      71             :     struct internal_pool : pooled_secure_allocator<uint8_t> {
      72         435 :         internal_pool(size_type nrequested_size,
      73             :                       size_type nnext_size,
      74             :                       size_type nmax_size) :
      75         145 :                       pooled_secure_allocator(nrequested_size, nnext_size, nmax_size)
      76         145 :         {
      77         290 :         }
      78             :         std::mutex mutex;
      79             :     };
      80             : 
      81             : private:
      82             :     std::vector<std::unique_ptr<internal_pool>> pools;
      83             : };
      84             : 
      85             : #endif // BITCOIN_SUPPORT_ALLOCATORS_MT_POOLED_SECURE_H