// | |
// Simple allocator for small memory blocks (from 1 to 255 bytes). | |
// It keeps freelists of small blocks, built on top of 8K pages. | |
// The allocator only grows - empty pages are never freed back | |
// to the system. | |
// | |
// This allocator could easily be made lockless with atomic freelists. | |
// | |
#include <cstdint> | |
#include <cstdlib> | |
#include <cstring> | |
#include <cstdio> | |
// ---------------------------------------------------------------------------- | |
#if defined(__clang__) && defined(__has_feature) | |
#if __has_feature(address_sanitizer) | |
#include <sanitizer/asan_interface.h> | |
#define TINY_ALLOC_ASAN 1 | |
#endif // address_sanitizer | |
#endif // __clang__ && __has_feature | |
#if !defined(TINY_ALLOC_ASAN) | |
#define ASAN_POISON_MEMORY_REGION(addr, size) /* nothing */ | |
#define ASAN_UNPOISON_MEMORY_REGION(addr, size) /* nothing */ | |
#define ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE /* nothing */ | |
#else // TINY_ALLOC_ASAN | |
// ASAN_POISON_MEMORY_REGION() already defined in asan_interface.h | |
// ASAN_UNPOISON_MEMORY_REGION() already defined in asan_interface.h | |
#define ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE __attribute__((no_sanitize("address"))) | |
#endif // TINY_ALLOC_ASAN | |
// ---------------------------------------------------------------------------- | |
class TinyAllocator final | |
{ | |
private: | |
enum Constants | |
{ | |
kHeaderSize = 2, | |
kAlign = 8, | |
kTinyAllocTag = 0xAA, | |
kInvalidAllocTag = 0xDD, | |
kPageSize = 8192, | |
}; | |
struct Page | |
{ | |
Page * Next; | |
std::uint8_t Memory[kPageSize - sizeof(Page *)]; | |
}; | |
static_assert(sizeof(Page) == kPageSize, "Wrong size for TinyAllocator::Page"); | |
private: | |
static constexpr std::size_t AlignSize(std::size_t bytes) | |
{ | |
return ((bytes + kAlign - 1) & ~(kAlign - 1)); | |
} | |
static constexpr std::size_t TinyAlign(std::size_t bytes) | |
{ | |
return AlignSize(bytes + kHeaderSize) - kHeaderSize; | |
} | |
static Page * AllocPage() | |
{ | |
Page * p = static_cast<Page *>(std::malloc(sizeof(Page))); | |
if (p == nullptr) | |
{ | |
return nullptr; | |
} | |
p->Next = nullptr; | |
std::memset(p->Memory, kInvalidAllocTag, sizeof(Page::Memory)); | |
ASAN_POISON_MEMORY_REGION(p, sizeof(Page)); | |
return p; | |
} | |
static void FreePage(Page * p) | |
{ | |
ASAN_UNPOISON_MEMORY_REGION(p, sizeof(Page)); | |
std::free(p); | |
} | |
Page * m_FirstUsedPage; | |
Page * m_CurrentPage; | |
std::size_t m_CurrentPageOffset; | |
void * m_FirstFree[(256 / kAlign) + 1]; | |
public: | |
TinyAllocator() | |
: m_FirstUsedPage{ nullptr } | |
, m_CurrentPage{ AllocPage() } | |
, m_CurrentPageOffset{ TinyAlign(0) } | |
{ | |
std::memset(m_FirstFree, 0, sizeof(m_FirstFree)); | |
} | |
ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE ~TinyAllocator() | |
{ | |
for (Page * p = m_FirstUsedPage; p != nullptr;) | |
{ | |
Page * next = p->Next; | |
FreePage(p); | |
p = next; | |
} | |
} | |
ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE void * Allocate(std::size_t bytes) | |
{ | |
if (bytes > 255) | |
{ | |
std::fprintf(stderr, "TinyAllocator::Allocate - Max allocation size is 255. Requested %zu bytes", bytes); | |
std::abort(); | |
} | |
// We need the at least the size of a pointer for the free list. | |
if (bytes < sizeof(void *)) | |
{ | |
bytes = sizeof(void *); | |
} | |
// Increase the number of bytes if necessary to make sure the next tiny allocation is aligned. | |
bytes = TinyAlign(bytes); | |
// Append to the free list: | |
auto * tinyBlock = static_cast<std::uint8_t *>(m_FirstFree[bytes / kAlign]); | |
if (tinyBlock != nullptr) | |
{ | |
std::size_t * link = reinterpret_cast<std::size_t *>(tinyBlock + kHeaderSize); | |
tinyBlock[1] = kTinyAllocTag; // Allocation identifier | |
m_FirstFree[bytes / kAlign] = reinterpret_cast<void *>(*link); | |
void * userPtr = reinterpret_cast<void *>(link); | |
ASAN_UNPOISON_MEMORY_REGION(userPtr, bytes); | |
return userPtr; | |
} | |
// Check if we need to allocate a new page: | |
const std::size_t bytesLeftInCurrPage = sizeof(Page::Memory) - m_CurrentPageOffset; | |
if (bytes >= bytesLeftInCurrPage) | |
{ | |
Page * newPage = AllocPage(); | |
if (newPage == nullptr) | |
{ | |
return nullptr; | |
} | |
m_CurrentPage->Next = m_FirstUsedPage; | |
m_FirstUsedPage = m_CurrentPage; | |
m_CurrentPage = newPage; | |
// Make sure the first allocation is aligned. | |
m_CurrentPageOffset = TinyAlign(0); | |
} | |
// Allocate from page: | |
tinyBlock = m_CurrentPage->Memory + m_CurrentPageOffset; | |
tinyBlock[0] = std::uint8_t(bytes / kAlign); // Write # of bytes/kAlign | |
tinyBlock[1] = kTinyAllocTag; // Allocation identifier | |
m_CurrentPageOffset += bytes + kHeaderSize; // Increase the offset on the current page | |
void * userPtr = (tinyBlock + kHeaderSize); // Skip the first two bytes header | |
ASAN_UNPOISON_MEMORY_REGION(userPtr, bytes); | |
return userPtr; | |
} | |
ASAN_NO_SANITIZE_ADDRESS_ATTRIBUTE void Free(void * ptr) | |
{ | |
auto * tinyBlock = static_cast<std::uint8_t *>(ptr) - kHeaderSize; | |
// Check allocation tag: | |
if (tinyBlock[1] != kTinyAllocTag) | |
{ | |
std::fprintf(stderr, "TinyAllocator::Free - Invalid header tag at %p", ptr); | |
std::abort(); | |
} | |
// Clear tag: | |
tinyBlock[1] = kInvalidAllocTag; | |
// Index into the table with free tiny memory blocks: | |
const std::size_t index = tinyBlock[0]; | |
// Check if the index is valid: | |
if (index > (256 / kAlign)) | |
{ | |
std::fprintf(stderr, "TinyAllocator::Free - Invalid list index at %p", ptr); | |
std::abort(); | |
} | |
std::size_t * link = reinterpret_cast<std::size_t *>(ptr); | |
*link = reinterpret_cast<std::size_t>(m_FirstFree[index]); | |
m_FirstFree[index] = tinyBlock; | |
ASAN_POISON_MEMORY_REGION(ptr, index * kAlign); | |
} | |
// Disallow copy. | |
TinyAllocator(const TinyAllocator & other) = delete; | |
TinyAllocator& operator = (const TinyAllocator & other) = delete; | |
}; | |
// ---------------------------------------------------------------------------- | |
int main() | |
{ | |
TinyAllocator tAlloc; | |
void* ptrs[256]; | |
std::printf("Allocating...\n"); | |
for (int i = 0; i < 256; ++i) | |
{ | |
ptrs[i] = tAlloc.Allocate(i); | |
std::memset(ptrs[i], 0xCD, i); | |
} | |
std::printf("Freeing...\n"); | |
for (int i = 0; i < 256; ++i) | |
{ | |
tAlloc.Free(ptrs[i]); | |
} | |
char * p64 = static_cast<char *>(tAlloc.Allocate(64)); | |
for (int i = 0; i < 64; ++i) | |
{ | |
p64[i] = 'X'; | |
} | |
tAlloc.Free(p64); | |
//p64[0] = 'Y'; ASan ERROR: use after free | |
std::printf("Done!\n"); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment