Line data    Source code

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