Skip to content

Instantly share code, notes, and snippets.

@nadavrot

nadavrot/mymalloc.cc

Created Nov 29, 2020
Embed
What would you like to do?
A working malloc implementation in 250 LOC
#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <pthread.h>
#include <unistd.h>
namespace {
/// A Block is a header to a managed memory buffer. Blocks are arranged as
/// sorted linked list.
struct Block {
/// A pointer to the 'next free' buffer.
Block *next_;
/// The size of the buffer.
uint64_t size_;
Block() { init(0, nullptr); }
// Initialize the struct.
void init(uint64_t size, Block *next) {
size_ = size;
next_ = next;
}
/// \Returns a block pointer from a raw pointer.
static Block *fromRaw(void *ptr) {
return (Block *)(uint64_t(ptr) - sizeof(Block));
}
/// \Returns a raw pointer to the payload area.
void *toRaw() const { return begin(); }
/// Check that the block structure is correct. The block free list must be
/// sorted in an increasing address order, and blocks must not overlap.
void verify() {
#ifndef NDEBUG
Block *ptr = this;
while (ptr) {
if (ptr->next_) {
assert(ptr->next_ >= ptr->end() && "Next pointer inside the payload");
}
assert(uint64_t(ptr) % sizeof(uint64_t) == 0 && "Block is not aligned");
if (ptr->next_) {
assert(ptr < ptr->next_ && "Block list non-consecutive.");
}
ptr = ptr->next_;
}
#endif
}
/// \returns the start of the payload.
void *begin() const { return (void *)(uint64_t(this) + sizeof(Block)); }
/// \returns the end of the payload.
void *end() const { return (void *)(uint64_t(this) + sizeof(Block) + size_); }
/// Try to merge the current block to the next in the chain. \returns True if
/// the block was merged.
bool mergeToNext() {
verify();
if (end() == next_) {
size_ += next_->size_ + sizeof(Block);
next_ = next_->next_;
mergeToNext();
return true;
}
return false;
}
/// Add a block \p other to the linked list of blocks (free list). Notice that
/// this method is called on the first block in the chain. \returns The new
/// entry block.
Block *addToBlockList(Block *other) {
assert(other->next_ == nullptr);
verify();
// If we are returning a block with a lower address, insert the current
// block into the other block and return the new head.
if (other < this) {
other->next_ = this;
other->verify();
return other;
}
Block *ptr = this;
assert(ptr < other && "first block needs to have a lower address");
while (ptr) {
// Keep looking for the insertion point.
if (ptr->next_ && ptr->next_ < other) {
ptr = ptr->next_;
continue;
}
assert(ptr->end() <= other && "Inserting block into a block");
// Just insert the block.
other->next_ = ptr->next_;
ptr->next_ = other;
verify();
// Try to coalese the block forward and backwards.
other->mergeToNext();
mergeToNext();
return this;
}
assert(ptr->next_ == nullptr && "Unexpected list structure");
ptr->next_ = other;
return this;
}
/// Allocate a new buffer from the free list, with the payload size \p size.
/// \return the pointer to the buffer or nullptr if the allocation failed.
Block *getFirstAvail(uint64_t size) {
// Increase the size to make sure that all blocks are aligned.
uint64_t align = size % sizeof(Block);
if (align) {
size += sizeof(Block) - align;
}
// The size of the payload and the header.
uint64_t fullSize = size + sizeof(Block);
verify();
Block *ptr = this;
while (ptr) {
// Does the requested allocation fit in this block?
if (ptr->size_ > fullSize) {
// Add the new block at the end of the current payload.
auto e = uint64_t(ptr->end());
Block *newBlock = (Block *)(e - fullSize);
newBlock->size_ = size;
newBlock->next_ = nullptr;
ptr->size_ -= fullSize;
verify();
return newBlock;
}
ptr = ptr->next_;
}
return nullptr;
}
};
/// This class represents a working memory allocator that handles free/malloc
/// requests, and can support multiple concurrent requests.
class Allocator {
/// The main block, which is the entry to the free list. This block must have
/// the lowest buffer address.
Block *master_{nullptr};
/// A lock that protects all operations on the free-list.
pthread_mutex_t bmLock_;
/// New managed chunks use this value.
static constexpr int regionSize = 1 << 24;
/// Add a memory region to the free-list.
void addMemoryRegion(void *region, uint64_t size) {
Block *r = (Block *)region;
r->init(size - sizeof(Block), nullptr);
if (master_) {
master_ = master_->addToBlockList(r);
} else {
master_ = r;
}
}
/// Ask the kernel to provide a new buffer on the heap.
static void *newBufferOnHeap(size_t size) {
void *ptr = sbrk(size);
assert(ptr && "Invalid pointer");
return ptr;
}
public:
Allocator() { bmLock_ = PTHREAD_MUTEX_INITIALIZER; }
/// \returns the size of the underlying buffer. \ptr must be a valid pointer
/// that is backed by a block.
size_t getBufferSize(void *ptr) {
Block *b = Block::fromRaw(ptr);
return b->size_;
}
void *malloc(size_t size) {
pthread_mutex_lock(&bmLock_);
start:
if (Block *b = master_->getFirstAvail(size)) {
pthread_mutex_unlock(&bmLock_);
return b->toRaw();
}
// Not enough memory. Allocate a new heap region.
auto newChunkSize = std::max<size_t>(size, regionSize);
auto *ptr = newBufferOnHeap(newChunkSize);
addMemoryRegion(ptr, newChunkSize);
goto start;
}
void free(void *ptr) {
// We are allowed to free zeros.
if (!ptr) {
return;
}
pthread_mutex_lock(&bmLock_);
Block *b = Block::fromRaw(ptr);
assert(b);
master_ = master_->addToBlockList(b);
pthread_mutex_unlock(&bmLock_);
}
};
Allocator AA;
} // namespace
extern "C" void *malloc(size_t size) { return AA.malloc(size); }
extern "C" void free(void *ptr) { AA.free(ptr); }
extern "C" void *calloc(size_t nmemb, size_t size) {
auto sz = nmemb * size;
if (!sz) {
return nullptr;
}
auto *ptr = AA.malloc(sz);
bzero(ptr, sz);
return ptr;
}
extern "C" void *realloc(void *ptr, size_t size) {
if (!size) {
free(ptr);
return nullptr;
}
if (!ptr) {
return malloc(size);
}
auto origSz = AA.getBufferSize(ptr);
if (size <= origSz) {
return ptr;
}
if (void *newPtr = malloc(size)) {
memcpy(newPtr, ptr, origSz);
free(ptr);
return newPtr;
}
return nullptr;
}
extern "C" void *reallocarray(void *ptr, size_t nmemb, size_t size) {
return realloc(ptr, nmemb * size);
}
@nadavrot

This comment has been minimized.

Copy link
Owner Author

@nadavrot nadavrot commented Nov 29, 2020

Makefile:

libmymalloc.so: mymalloc.cc
gcc -shared -o libmymalloc.so mymalloc.cc -fPIC -pthread

run: libmymalloc.so
LD_PRELOAD=./libmymalloc.so python

clean:
rm -f libmymalloc.so

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment