allocator.c

#include "allocator.h"
#include <pthread.h>
#include <stdatomic.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <sys/mman.h>
#include <unistd.h>

#define LARGE_CACHE_SIZE 32 // no. of tcache blocks
#define NUM_BINS 8
#define ALIGN(size) (((size) + 7) & ~7)
typedef struct block_header {
  size_t size;
  int isfree;
  int ismmapped;
  int in_tcache;
  struct block_header *next;
  struct block_header *prev;
  struct block_header *bin_next;
  struct block_header *bin_prev;
} block_header_t;

static block_header_t *bins[NUM_BINS];

static block_header_t *large_cache_blocks[LARGE_CACHE_SIZE];
static int large_cache_count = 0;
#define TCACHE_MAX 512
#define TCACHE_SLOTS 128

static inline int tcache_index(size_t size) {
  if (size < 8)
    return -1;
  return (size >> 3) - 1;
}

typedef struct {
  block_header_t *slots[TCACHE_SLOTS];
  int count[TCACHE_SLOTS];
} tcache_t;
__thread tcache_t tcache = {0};

#define LARGE_ALLOC 131072

alloc_stats_t stats = {0};
__thread size_t local_hits = 0;
__thread size_t local_misses = 0;
__thread size_t local_allocs = 0;
__thread size_t local_frees = 0;

int bin_index(size_t size) {
  if (size <= 8)
    return 0;
  int index = 31 - __builtin_clz(size - 1) - 3;
  if (index >= NUM_BINS)
    return NUM_BINS - 1;
  return index;
}

void bin_insert(block_header_t *block) {
  int index = bin_index(block->size);
  block->bin_prev = NULL;
  block->bin_next = bins[index];
  if (bins[index])
    bins[index]->bin_prev = block;
  bins[index] = block;
}

void bin_remove(block_header_t *block) {
  int index = bin_index(block->size);

  if (!block->bin_prev && bins[index] != block)
    return;
  if (block->bin_prev)
    block->bin_prev->bin_next = block->bin_next;
  else
    bins[index] = block->bin_next;
  if (block->bin_next) {
    block->bin_next->bin_prev = block->bin_prev;
  }
  block->bin_next = NULL;
  block->bin_prev = NULL;
}

static block_header_t *lastblock = NULL;
pthread_mutex_t heaplock = PTHREAD_MUTEX_INITIALIZER;
#define CHUNK_SIZE                                                             \
  (2 * 1024 * 1024) // chunk size of 2mb for fixed chunks in mmap
static block_header_t *current_chunk = NULL;
static size_t chunk_remaining = 0;

block_header_t *reqestspace(block_header_t *last, size_t size) {
  // block_header_t *block = sbrk(0);
  // if (block == (void *)-1) {
  //   return NULL;
  // }
  size_t needed = ALIGN(sizeof(block_header_t) + size);
  if (chunk_remaining < needed) {
    if (chunk_remaining >= sizeof(block_header_t) + 8) {
      block_header_t *leftover = (block_header_t *)current_chunk;
      size_t usable = chunk_remaining - sizeof(block_header_t);
      usable = usable & ~7;
      if (usable >= 8) {
        leftover->size = usable;
        leftover->isfree = 1;
        leftover->ismmapped = 0;
        leftover->prev = NULL;
        leftover->next = NULL;
        leftover->bin_next = NULL;
        leftover->bin_prev = NULL;
        bin_insert(leftover);
      }
    }
    current_chunk = mmap(NULL, CHUNK_SIZE, PROT_READ | PROT_WRITE,
                         MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (current_chunk == MAP_FAILED)
      return NULL;
    chunk_remaining = CHUNK_SIZE;
  }
  block_header_t *block = (block_header_t *)current_chunk;
  current_chunk = (block_header_t *)((char *)current_chunk + needed);
  chunk_remaining -= needed;
  block->size = ALIGN(size);
  block->isfree = 0;
  block->ismmapped = 0;
  block->prev = last;
  block->next = NULL;
  block->bin_next = NULL;
  block->bin_prev = NULL;
  if (last) {
    last->next = block;
  }
  lastblock = block;
  return block;
};

void splitblocks(block_header_t *block, size_t size) {
  block_header_t *leftover = (block_header_t *)((char *)(block + 1) + size);
  leftover->size = block->size - size - sizeof(block_header_t);
  leftover->isfree = 1;
  leftover->next = block->next;
  leftover->prev = block;
  leftover->ismmapped = 0;
  leftover->bin_prev = NULL;
  leftover->bin_next = NULL;
  if (block->next) {
    block->next->prev = leftover;
  }
  block->next = leftover;
  block->size = size;
  bin_insert(leftover);
  if (block == lastblock)
    lastblock = leftover;
}

void coalesce(block_header_t *block) {

  // prev is free
  if (block->prev && block->isfree && block->prev->isfree &&
      (char *)(block->prev + 1) + block->prev->size == (char *)block) {

    bin_remove(block->prev);
    bin_remove(block);
    if (block == lastblock) {
      lastblock = block->prev;
    }
    block->prev->size =
        block->prev->size + sizeof(block_header_t) + block->size;
    block->prev->next = block->next;
    if (block->next) {
      block->next->prev = block->prev;
    }
    bin_insert(block->prev);
    block = block->prev;
  }
  // next is free
  if (block->next && block->isfree && block->next->isfree &&
      ((char *)(block + 1) + block->size == (char *)block->next)) {
    bin_remove(block->next);
    bin_remove(block);
    if (block->next == lastblock)
      lastblock = block;
    block->size = block->size + sizeof(block_header_t) + block->next->size;
    block->next = block->next->next;
    if (block->next)
      block->next->prev = block;
    bin_insert(block);
  }
}

block_header_t *find_free_block(size_t size) {
  int index = bin_index(size);
  while (index < NUM_BINS) {
    block_header_t *block = bins[index];
    while (block) {
      if (block->size >= size) {
        return block;
      }
      block = block->bin_next;
    }
    index++;
  }
  return NULL;
}

static pthread_key_t stats_key;
static pthread_once_t stats_key_once = PTHREAD_ONCE_INIT;
static void flush_thread_stats(void *arg) {
  (void)arg;
  atomic_fetch_add(&stats.alloc_count, local_allocs);
  atomic_fetch_add(&stats.free_count, local_frees);
  atomic_fetch_add(&stats.tcache_hits, local_hits);
  atomic_fetch_add(&stats.tcache_misses, local_misses);
  local_hits = local_allocs = local_frees = local_misses = 0;
}

static void init_stats_key(void) {
  pthread_key_create(&stats_key, flush_thread_stats);
}
void flush_tcache(int idx) {
  atomic_fetch_add(&stats.alloc_count, local_allocs);
  atomic_fetch_add(&stats.free_count, local_frees);
  atomic_fetch_add(&stats.tcache_hits, local_hits);
  atomic_fetch_add(&stats.tcache_misses, local_misses);
  local_frees = local_misses = local_hits = local_allocs = 0;
  pthread_mutex_lock(&heaplock);
  block_header_t *block = tcache.slots[idx];
  while (block) {
    block_header_t *next = block->bin_next;
    block->in_tcache = 0;
    block->bin_next = NULL;
    block->bin_prev = NULL;
    bin_insert(block);
    block = next;
  }
  tcache.slots[idx] = NULL;
  tcache.count[idx] = 0;
  pthread_mutex_unlock(&heaplock);
}

__thread int stats_registered = 0;

void *mymalloc(size_t size) {
  if (!stats_registered) {
    pthread_once(&stats_key_once, init_stats_key);
    pthread_setspecific(stats_key, (void *)1);
    stats_registered = 1;
  }
  if (size == 0)
    return NULL;
  size = ALIGN(size);

  if (size < LARGE_ALLOC) {
    int slot = tcache_index(size);
    if (slot >= 0 && slot < TCACHE_SLOTS && tcache.slots[slot]) {
      block_header_t *block = tcache.slots[slot];
      tcache.slots[slot] = block->bin_next;
      tcache.count[slot]--;
      block->isfree = 0;
      block->in_tcache = 0;
      block->bin_next = NULL;
      // atomic_fetch_add(&stats.alloc_count, 1);
      // atomic_fetch_add(&stats.tcache_hits, 1);
      local_hits++;
      local_allocs++;
      return (void *)(block + 1);
    }
  }

  if (size >= LARGE_ALLOC) {
    pthread_mutex_lock(&heaplock);
    block_header_t *block;
    for (int i = 0; i < large_cache_count; i++) {
      if (large_cache_blocks[i]->size >= size &&
          large_cache_blocks[i]->size <= size * 2) {
        block = large_cache_blocks[i];
        large_cache_blocks[i] = large_cache_blocks[--large_cache_count];
        pthread_mutex_unlock(&heaplock);
        size_t wasted = block->size - size;
        atomic_fetch_add(&stats.large_cache_wasted_bytes, wasted);
        atomic_fetch_add(&stats.large_cache_hits, 1);
        // atomic_fetch_add(&stats.tcache_hits, 1);
        // atomic_fetch_add(&stats.alloc_count, 1);
        local_hits++;
        local_allocs++;
        return (void *)(block + 1);
      }
    }
    pthread_mutex_unlock(&heaplock);
    block = mmap(NULL, sizeof(block_header_t) + size, PROT_READ | PROT_WRITE,
                 MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (block == MAP_FAILED) {
      return NULL;
    }
    block->size = size;
    block->ismmapped = 1;
    block->isfree = 0;
    block->next = NULL;
    block->prev = NULL;
    block->bin_next = NULL;
    block->bin_prev = NULL;
    atomic_fetch_add(&stats.mmap_calls, 1);
    // atomic_fetch_add(&stats.alloc_count, 1);
    local_allocs++;
    return (void *)(block + 1);
  }

  // atomic_fetch_add(&stats.tcache_misses, 1);
  local_misses++;
  pthread_mutex_lock(&heaplock);

  block_header_t *block = find_free_block(size);

  if (block) {
    bin_remove(block); // remove block before splitting to avoid bin confusions
    if (block->size >= size + sizeof(block_header_t) + 8) {

      splitblocks(block, size);
    }
    block->isfree = 0;
  } else {

    block = reqestspace(lastblock, size);
    if (!block) {
      pthread_mutex_unlock(&heaplock);
      return NULL;
    }
  }
  // atomic_fetch_add(&stats.alloc_count, 1);
  local_allocs++;
  pthread_mutex_unlock(&heaplock);
  return (void *)(block + 1);
}

void myfree(void *ptr) {
  if (!ptr) {
    return;
  }
  block_header_t *block = (block_header_t *)ptr - 1;
  if (!block->ismmapped) {
    int slot = tcache_index(block->size);
    if (slot >= 0 && slot < TCACHE_SLOTS) {
      if (tcache.count[slot] >= TCACHE_MAX) {
        flush_tcache(slot);
      }
      block->in_tcache = 1;
      block->bin_next = tcache.slots[slot];
      block->bin_prev = NULL;
      tcache.slots[slot] = block;
      tcache.count[slot]++;
      // atomic_fetch_add(&stats.free_count, 1);
      local_frees++;
      return;
    }
  }
  pthread_mutex_lock(&heaplock);
  if (block->ismmapped) {
    if (large_cache_count < LARGE_CACHE_SIZE) {
      large_cache_blocks[large_cache_count++] = block;
      // atomic_fetch_add(&stats.free_count, 1);
      local_frees++;
    } else {
      munmap(block, sizeof(block_header_t) + block->size);
      atomic_fetch_add(&stats.munmap_calls, 1);
      // atomic_fetch_add(&stats.free_count, 1);
      local_frees++;
    }
    pthread_mutex_unlock(&heaplock);
    return;
  }
  block->isfree = 1;
  bin_insert(block);
  coalesce(block);
  // atomic_fetch_add(&stats.free_count, 1);
  local_frees++;
  pthread_mutex_unlock(&heaplock);
}

void *mycalloc(size_t n, size_t size) {
  if (n != 0 && size > SIZE_MAX / n)
    return NULL;
  if (n == 0 || size == 0)
    return NULL;
  void *block = mymalloc(n * size);
  if (!block)
    return NULL;
  memset(block, 0, n * size);
  return block;
}

void *myrealloc(void *blk, size_t size) {
  if (size == 0) {
    myfree(blk);
    return NULL;
  }
  block_header_t *block = (block_header_t *)blk - 1;
  /*
   * in_tache is a way to differentiate between tcached blocks and global list
   * blocks
   */

  size_t oldsize = block->size;
  pthread_mutex_lock(&heaplock);
  if (block->next && block->next->isfree && !block->next->in_tcache &&
      (block->next->bin_prev != NULL ||
       bins[bin_index(block->next->size)] == block->next) &&
      (char *)(block + 1) + block->size == (char *)block->next &&
      block->size + sizeof(block_header_t) + block->next->size >= size) {
    if ((char *)(block + 1) + block->size == (char *)block->next &&
        block->next->isfree &&
        block->size + sizeof(block_header_t) + block->next->size >= size) {
      bin_remove(block->next);
      if (block->next == lastblock)
        lastblock = block;
      block->size = block->size + sizeof(block_header_t) + block->next->size;
      block->next = block->next->next;
      if (block->next)
        block->next->prev = block;
      if (block->size > size + sizeof(block_header_t) + 8) {
        splitblocks(block, size);
      }
      pthread_mutex_unlock(&heaplock);
      return (void *)(block + 1);
    }
  }
  pthread_mutex_unlock(&heaplock);
  void *newblock = mymalloc(size);
  if (!newblock)
    return NULL;
  memcpy(newblock, blk, oldsize < size ? oldsize : size);
  myfree(blk);
  return (void *)newblock;
}

void allocator_print_stats(void) {
  atomic_fetch_add(&stats.alloc_count, local_allocs);
  atomic_fetch_add(&stats.free_count, local_frees);
  atomic_fetch_add(&stats.tcache_hits, local_hits);
  atomic_fetch_add(&stats.tcache_misses, local_misses);
  local_frees = local_misses = local_hits = local_allocs = 0;

  size_t hits = atomic_load(&stats.tcache_hits);
  size_t misses = atomic_load(&stats.tcache_misses);
  size_t total = hits + misses;
  size_t allocs = atomic_load(&stats.alloc_count);
  size_t frees = atomic_load(&stats.free_count);
  printf("\n ===allocator stats===\n");
  printf("tcache hits: %zu\n", hits);
  printf("tcache misses: %zu\n", misses);
  printf("tcache hit rate : %.1f%%\n", total ? 100.0 * hits / total : 0.0);
  printf("total allocs:      %zu\n", allocs);
  printf("total frees:       %zu\n", frees);
  printf("live allocations:  %zu\n", allocs - frees);
  printf("mmap calls: %zu\n", atomic_load(&stats.mmap_calls));
  printf("munmap calls: %zu\n", atomic_load(&stats.munmap_calls));
  printf("large cache hits: %zu\n", atomic_load(&stats.large_cache_hits));
  printf("large cache wasted : %zu bytes (%.1f MB)\n",
         atomic_load(&stats.large_cache_wasted_bytes),
         atomic_load(&stats.large_cache_wasted_bytes) / 1e6);
  // live bytes is bugged causing size_t overflow
  //  printf("live bytes: %zu\n", atomic_load(&stats.live_bytes));
}