Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
//
// 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