A working malloc implementation in 250 LOC

mymalloc.cc

#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);
}

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