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