lockedpool.cpp raw

   1  // Copyright (c) 2016-2022 The Limenka 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  #include <support/lockedpool.h>
   6  #include <support/cleanse.h>
   7  
   8  #ifdef WIN32
   9  #include <windows.h>
  10  #else
  11  #include <sys/mman.h> // for mmap
  12  #include <sys/resource.h> // for getrlimit
  13  #include <limits.h> // for PAGESIZE
  14  #include <unistd.h> // for sysconf
  15  #endif
  16  
  17  #include <algorithm>
  18  #include <limits>
  19  #include <stdexcept>
  20  #include <utility>
  21  #ifdef ARENA_DEBUG
  22  #include <iomanip>
  23  #include <iostream>
  24  #endif
  25  
  26  LockedPoolManager* LockedPoolManager::_instance = nullptr;
  27  
  28  /*******************************************************************************/
  29  // Utilities
  30  //
  31  /** Align up to power of 2 */
  32  static inline size_t align_up(size_t x, size_t align)
  33  {
  34      return (x + align - 1) & ~(align - 1);
  35  }
  36  
  37  /*******************************************************************************/
  38  // Implementation: Arena
  39  
  40  Arena::Arena(void *base_in, size_t size_in, size_t alignment_in):
  41      base(base_in), end(static_cast<char*>(base_in) + size_in), alignment(alignment_in)
  42  {
  43      // Start with one free chunk that covers the entire arena
  44      auto it = size_to_free_chunk.emplace(size_in, base);
  45      chunks_free.emplace(base, it);
  46      chunks_free_end.emplace(static_cast<char*>(base) + size_in, it);
  47  }
  48  
  49  Arena::~Arena() = default;
  50  
  51  void* Arena::alloc(size_t size)
  52  {
  53      // Round to next multiple of alignment
  54      size = align_up(size, alignment);
  55  
  56      // Don't handle zero-sized chunks
  57      if (size == 0)
  58          return nullptr;
  59  
  60      // Pick a large enough free-chunk. Returns an iterator pointing to the first element that is not less than key.
  61      // This allocation strategy is best-fit. According to "Dynamic Storage Allocation: A Survey and Critical Review",
  62      // Wilson et. al. 1995, https://www.scs.stanford.edu/14wi-cs140/sched/readings/wilson.pdf, best-fit and first-fit
  63      // policies seem to work well in practice.
  64      auto size_ptr_it = size_to_free_chunk.lower_bound(size);
  65      if (size_ptr_it == size_to_free_chunk.end())
  66          return nullptr;
  67  
  68      // Create the used-chunk, taking its space from the end of the free-chunk
  69      const size_t size_remaining = size_ptr_it->first - size;
  70      char* const free_chunk = static_cast<char*>(size_ptr_it->second);
  71      auto allocated = chunks_used.emplace(free_chunk + size_remaining, size).first;
  72      chunks_free_end.erase(free_chunk + size_ptr_it->first);
  73      if (size_ptr_it->first == size) {
  74          // whole chunk is used up
  75          chunks_free.erase(size_ptr_it->second);
  76      } else {
  77          // still some memory left in the chunk
  78          auto it_remaining = size_to_free_chunk.emplace(size_remaining, size_ptr_it->second);
  79          chunks_free[size_ptr_it->second] = it_remaining;
  80          chunks_free_end.emplace(free_chunk + size_remaining, it_remaining);
  81      }
  82      size_to_free_chunk.erase(size_ptr_it);
  83  
  84      return allocated->first;
  85  }
  86  
  87  void Arena::free(void *ptr)
  88  {
  89      // Freeing the nullptr pointer is OK.
  90      if (ptr == nullptr) {
  91          return;
  92      }
  93  
  94      // Remove chunk from used map
  95      auto i = chunks_used.find(ptr);
  96      if (i == chunks_used.end()) {
  97          throw std::runtime_error("Arena: invalid or double free");
  98      }
  99      auto freed = std::make_pair(static_cast<char*>(i->first), i->second);
 100      chunks_used.erase(i);
 101  
 102      // coalesce freed with previous chunk
 103      auto prev = chunks_free_end.find(freed.first);
 104      if (prev != chunks_free_end.end()) {
 105          freed.first -= prev->second->first;
 106          freed.second += prev->second->first;
 107          size_to_free_chunk.erase(prev->second);
 108          chunks_free_end.erase(prev);
 109      }
 110  
 111      // coalesce freed with chunk after freed
 112      auto next = chunks_free.find(freed.first + freed.second);
 113      if (next != chunks_free.end()) {
 114          freed.second += next->second->first;
 115          size_to_free_chunk.erase(next->second);
 116          chunks_free.erase(next);
 117      }
 118  
 119      // Add/set space with coalesced free chunk
 120      auto it = size_to_free_chunk.emplace(freed.second, freed.first);
 121      chunks_free[freed.first] = it;
 122      chunks_free_end[freed.first + freed.second] = it;
 123  }
 124  
 125  Arena::Stats Arena::stats() const
 126  {
 127      Arena::Stats r{ 0, 0, 0, chunks_used.size(), chunks_free.size() };
 128      for (const auto& chunk: chunks_used)
 129          r.used += chunk.second;
 130      for (const auto& chunk: chunks_free)
 131          r.free += chunk.second->first;
 132      r.total = r.used + r.free;
 133      return r;
 134  }
 135  
 136  #ifdef ARENA_DEBUG
 137  static void printchunk(void* base, size_t sz, bool used) {
 138      std::cout <<
 139          "0x" << std::hex << std::setw(16) << std::setfill('0') << base <<
 140          " 0x" << std::hex << std::setw(16) << std::setfill('0') << sz <<
 141          " 0x" << used << std::endl;
 142  }
 143  void Arena::walk() const
 144  {
 145      for (const auto& chunk: chunks_used)
 146          printchunk(chunk.first, chunk.second, true);
 147      std::cout << std::endl;
 148      for (const auto& chunk: chunks_free)
 149          printchunk(chunk.first, chunk.second->first, false);
 150      std::cout << std::endl;
 151  }
 152  #endif
 153  
 154  /*******************************************************************************/
 155  // Implementation: Win32LockedPageAllocator
 156  
 157  #ifdef WIN32
 158  /** LockedPageAllocator specialized for Windows.
 159   */
 160  class Win32LockedPageAllocator: public LockedPageAllocator
 161  {
 162  public:
 163      Win32LockedPageAllocator();
 164      void* AllocateLocked(size_t len, bool *lockingSuccess) override;
 165      void FreeLocked(void* addr, size_t len) override;
 166      size_t GetLimit() override;
 167  private:
 168      size_t page_size;
 169  };
 170  
 171  Win32LockedPageAllocator::Win32LockedPageAllocator()
 172  {
 173      // Determine system page size in bytes
 174      SYSTEM_INFO sSysInfo;
 175      GetSystemInfo(&sSysInfo);
 176      page_size = sSysInfo.dwPageSize;
 177  }
 178  void *Win32LockedPageAllocator::AllocateLocked(size_t len, bool *lockingSuccess)
 179  {
 180      len = align_up(len, page_size);
 181      void *addr = VirtualAlloc(nullptr, len, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
 182      if (addr) {
 183          // VirtualLock is used to attempt to keep keying material out of swap. Note
 184          // that it does not provide this as a guarantee, but, in practice, memory
 185          // that has been VirtualLock'd almost never gets written to the pagefile
 186          // except in rare circumstances where memory is extremely low.
 187          *lockingSuccess = VirtualLock(const_cast<void*>(addr), len) != 0;
 188      }
 189      return addr;
 190  }
 191  void Win32LockedPageAllocator::FreeLocked(void* addr, size_t len)
 192  {
 193      len = align_up(len, page_size);
 194      memory_cleanse(addr, len);
 195      VirtualUnlock(const_cast<void*>(addr), len);
 196  }
 197  
 198  size_t Win32LockedPageAllocator::GetLimit()
 199  {
 200      size_t min, max;
 201      if(GetProcessWorkingSetSize(GetCurrentProcess(), &min, &max) != 0) {
 202          return min;
 203      }
 204      return std::numeric_limits<size_t>::max();
 205  }
 206  #endif
 207  
 208  /*******************************************************************************/
 209  // Implementation: PosixLockedPageAllocator
 210  
 211  #ifndef WIN32
 212  /** LockedPageAllocator specialized for OSes that don't try to be
 213   * special snowflakes.
 214   */
 215  class PosixLockedPageAllocator: public LockedPageAllocator
 216  {
 217  public:
 218      PosixLockedPageAllocator();
 219      void* AllocateLocked(size_t len, bool *lockingSuccess) override;
 220      void FreeLocked(void* addr, size_t len) override;
 221      size_t GetLimit() override;
 222  private:
 223      size_t page_size;
 224  };
 225  
 226  PosixLockedPageAllocator::PosixLockedPageAllocator()
 227  {
 228      // Determine system page size in bytes
 229  #if defined(PAGESIZE) // defined in limits.h
 230      page_size = PAGESIZE;
 231  #else                   // assume some POSIX OS
 232      page_size = sysconf(_SC_PAGESIZE);
 233  #endif
 234  }
 235  
 236  void *PosixLockedPageAllocator::AllocateLocked(size_t len, bool *lockingSuccess)
 237  {
 238      void *addr;
 239      len = align_up(len, page_size);
 240      addr = mmap(nullptr, len, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0);
 241      if (addr == MAP_FAILED) {
 242          return nullptr;
 243      }
 244      if (addr) {
 245          *lockingSuccess = mlock(addr, len) == 0;
 246  #if defined(MADV_DONTDUMP) // Linux
 247          madvise(addr, len, MADV_DONTDUMP);
 248  #elif defined(MADV_NOCORE) // FreeBSD
 249          madvise(addr, len, MADV_NOCORE);
 250  #endif
 251      }
 252      return addr;
 253  }
 254  void PosixLockedPageAllocator::FreeLocked(void* addr, size_t len)
 255  {
 256      len = align_up(len, page_size);
 257      memory_cleanse(addr, len);
 258      munlock(addr, len);
 259      munmap(addr, len);
 260  }
 261  size_t PosixLockedPageAllocator::GetLimit()
 262  {
 263  #ifdef RLIMIT_MEMLOCK
 264      struct rlimit rlim;
 265      if (getrlimit(RLIMIT_MEMLOCK, &rlim) == 0) {
 266          if (rlim.rlim_cur != RLIM_INFINITY &&
 267              std::cmp_less_equal(rlim.rlim_cur, static_cast<rlim_t>(std::numeric_limits<size_t>::max()))) {
 268              return rlim.rlim_cur;
 269          }
 270      }
 271  #endif
 272      return std::numeric_limits<size_t>::max();
 273  }
 274  #endif
 275  
 276  /*******************************************************************************/
 277  // Implementation: LockedPool
 278  
 279  LockedPool::LockedPool(std::unique_ptr<LockedPageAllocator> allocator_in, LockingFailed_Callback lf_cb_in)
 280      : allocator(std::move(allocator_in)), lf_cb(lf_cb_in)
 281  {
 282  }
 283  
 284  LockedPool::~LockedPool() = default;
 285  
 286  void* LockedPool::alloc(size_t size)
 287  {
 288      std::lock_guard<std::mutex> lock(mutex);
 289  
 290      // Don't handle impossible sizes
 291      if (size == 0 || size > ARENA_SIZE)
 292          return nullptr;
 293  
 294      // Try allocating from each current arena
 295      for (auto &arena: arenas) {
 296          void *addr = arena.alloc(size);
 297          if (addr) {
 298              return addr;
 299          }
 300      }
 301      // If that fails, create a new one
 302      if (new_arena(ARENA_SIZE, ARENA_ALIGN)) {
 303          return arenas.back().alloc(size);
 304      }
 305      return nullptr;
 306  }
 307  
 308  void LockedPool::free(void *ptr)
 309  {
 310      std::lock_guard<std::mutex> lock(mutex);
 311      // TODO we can do better than this linear search by keeping a map of arena
 312      // extents to arena, and looking up the address.
 313      for (auto &arena: arenas) {
 314          if (arena.addressInArena(ptr)) {
 315              arena.free(ptr);
 316              return;
 317          }
 318      }
 319      throw std::runtime_error("LockedPool: invalid address not pointing to any arena");
 320  }
 321  
 322  LockedPool::Stats LockedPool::stats() const
 323  {
 324      std::lock_guard<std::mutex> lock(mutex);
 325      LockedPool::Stats r{0, 0, 0, cumulative_bytes_locked, 0, 0};
 326      for (const auto &arena: arenas) {
 327          Arena::Stats i = arena.stats();
 328          r.used += i.used;
 329          r.free += i.free;
 330          r.total += i.total;
 331          r.chunks_used += i.chunks_used;
 332          r.chunks_free += i.chunks_free;
 333      }
 334      return r;
 335  }
 336  
 337  bool LockedPool::new_arena(size_t size, size_t align)
 338  {
 339      bool locked;
 340      // If this is the first arena, handle this specially: Cap the upper size
 341      // by the process limit. This makes sure that the first arena will at least
 342      // be locked. An exception to this is if the process limit is 0:
 343      // in this case no memory can be locked at all so we'll skip past this logic.
 344      if (arenas.empty()) {
 345          size_t limit = allocator->GetLimit();
 346          if (limit > 0) {
 347              size = std::min(size, limit);
 348          }
 349      }
 350      void *addr = allocator->AllocateLocked(size, &locked);
 351      if (!addr) {
 352          return false;
 353      }
 354      if (locked) {
 355          cumulative_bytes_locked += size;
 356      } else if (lf_cb) { // Call the locking-failed callback if locking failed
 357          if (!lf_cb()) { // If the callback returns false, free the memory and fail, otherwise consider the user warned and proceed.
 358              allocator->FreeLocked(addr, size);
 359              return false;
 360          }
 361      }
 362      arenas.emplace_back(allocator.get(), addr, size, align);
 363      return true;
 364  }
 365  
 366  LockedPool::LockedPageArena::LockedPageArena(LockedPageAllocator *allocator_in, void *base_in, size_t size_in, size_t align_in):
 367      Arena(base_in, size_in, align_in), base(base_in), size(size_in), allocator(allocator_in)
 368  {
 369  }
 370  LockedPool::LockedPageArena::~LockedPageArena()
 371  {
 372      allocator->FreeLocked(base, size);
 373  }
 374  
 375  /*******************************************************************************/
 376  // Implementation: LockedPoolManager
 377  //
 378  LockedPoolManager::LockedPoolManager(std::unique_ptr<LockedPageAllocator> allocator_in):
 379      LockedPool(std::move(allocator_in), &LockedPoolManager::LockingFailed)
 380  {
 381  }
 382  
 383  bool LockedPoolManager::LockingFailed()
 384  {
 385      return true;
 386  }
 387  
 388  void LockedPoolManager::CreateInstance()
 389  {
 390      // Using a local static instance guarantees that the object is initialized
 391      // when it's first needed and also deinitialized after all objects that use
 392      // it are done with it.  I can think of one unlikely scenario where we may
 393      // have a static deinitialization order/problem, but the check in
 394      // LockedPoolManagerBase's destructor helps us detect if that ever happens.
 395  #ifdef WIN32
 396      std::unique_ptr<LockedPageAllocator> allocator(new Win32LockedPageAllocator());
 397  #else
 398      std::unique_ptr<LockedPageAllocator> allocator(new PosixLockedPageAllocator());
 399  #endif
 400      static LockedPoolManager instance(std::move(allocator));
 401      LockedPoolManager::_instance = &instance;
 402  }
 403