memalloc.c

// memalloc - a general purpose memory allocator

#include "memalloc.h"
#include <math.h>
#include <pthread.h>
#include <stdalign.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#define PAGE_SIZE 4096 // 4K

#define MIN_CHUNK_SIZE 16
#define CHUNK_PAD 8 // 8 bytes padding in chunk for alignment to 16 bytes boundary

/*The fastbin block size should be equivalent to fit upto 10 chunks.*/
#define INITIAL_CHUNK_COUNT 10

/* Header (16 bytes) and data block layout:
 *
 * --------------------->chunk
 * size_t size (8 bytes)
 * ---------------------
 * Padding (8 bytes)
 * ---------------------
 * FASTBIN_DATA_BLOCK
 *  */
#define CHUNK_HEADER_SIZE 16
#define FASTBIN_DATA_BLOCK(chunk) ((char *)chunk + sizeof(size_t) + CHUNK_PAD)
#define GET_HEADER_POINTER(chunk) ((char *)chunk - (sizeof(size_t) + CHUNK_PAD))

/*Fast path for allocation request of size <=1024 bytes*/
#define FASTBIN_MAX_LIMIT 1024

#define GET_FASTBIN_OFFSET(size) ((size - 1) / MIN_CHUNK_SIZE)

typedef struct _el_fastbin {
  /*Points at the starting address of the fastbin block of a specified size.*/
  void *block;
  /*Freelist of the fastbin of the specified size.*/
  void *freelist;
  /*Top of the used block*/
  void *top;
} el_fastbin_t;

typedef struct _th_cache {
  /*fast_bin has pointers to chunks from 16 bytes to 1024 bytes spaced at 16 bytes ie 16, 32, 48 ... 1024.*/
  el_fastbin_t fast_bin[64];
  struct _th_cache *next;
  struct _th_cache *prev;
} th_cache_t; // thread local storage

typedef struct _memalloc_ctx {
  pthread_key_t th_key;
  pthread_once_t once;
  /*Beginning of the heap memory region*/
  void *heap;
  /*Pointer to the top of the used heap*/
  void *top;
  /*Pointer to the recently created thread cache block*/
  th_cache_t *recent_th_cache;
  pthread_mutex_t mtx_memalloc_ctx_t;
} memalloc_ctx_t;

static memalloc_ctx_t memalloc_ctx = {.once = PTHREAD_ONCE_INIT};

/*Increments the program break then updates required members in `memalloc_ctx`*/
static void incr_pgbrk(size_t size) {
  void *block = sbrk(size);
  if (block == (void *)-1) {
    perror("sbrk");
  }
  if (!memalloc_ctx.heap && !memalloc_ctx.top) {
    memalloc_ctx.heap = block;
    memalloc_ctx.top = block;
  }
}

/*Returns a block of `size` bytes from the address space-
 * This function updates the current `top` from memalloc_ctx then returns address
 * of the previous `top`.*/
static void *get_block(size_t size) {
  pthread_mutex_lock(&memalloc_ctx.mtx_memalloc_ctx_t);
  /*Get a huge heap (4 times PAGE_SIZE) if heap is being incremented the very first time.*/
  if (!memalloc_ctx.heap && !memalloc_ctx.top) {
    incr_pgbrk(PAGE_SIZE * 4);
  }
  /* Increment heap by PAGE_SIZE if the currently unused heap region is not enough for the block `size` request
   * i.e. `sbrk(0)-memalloc_ctx.top <= size`
   * Note: The `incr_pgbr` policy might need to be updated or rethink later on!*/
  if ((char *)sbrk(0) - (char *)memalloc_ctx.top <= size) {
    incr_pgbrk(PAGE_SIZE);
  }
  void *cur_top = memalloc_ctx.top;
  memalloc_ctx.top = cur_top + size;
  pthread_mutex_unlock(&memalloc_ctx.mtx_memalloc_ctx_t);
  return cur_top;
}

/*Fastpath allocation strategy-
 * Allocate in fastbin.*/
static void *fastpath_allocation(th_cache_t *tcache, int size) {
  int offset = GET_FASTBIN_OFFSET(size);
  int block_size = MIN_CHUNK_SIZE * (offset + 1);
  int chunk_size = block_size + CHUNK_HEADER_SIZE;
  int fastbin_size = chunk_size * INITIAL_CHUNK_COUNT;
  el_fastbin_t *fastbin_slot = tcache->fast_bin + offset;

  if (fastbin_slot->freelist) {
    void *cur = fastbin_slot->freelist;
    fastbin_slot->freelist = *(void **)((char *)cur + CHUNK_HEADER_SIZE);
    return FASTBIN_DATA_BLOCK(cur);
  }

  if (!fastbin_slot->block) {
    fastbin_slot->block = get_block(fastbin_size);
    fastbin_slot->freelist = NULL;
    fastbin_slot->top = fastbin_slot->block;
  }

  if ((char *)fastbin_slot->block + fastbin_size == fastbin_slot->top) {
    /*This condition is true if allocation request of a `size` bytes from a thread has been more than
     * INITIAL_CHUNK_COUNT times. So, there are already INITIAL_CHUNK_COUNT number of chunks of a specific `size`
     * allocated in a fastbin of a thread. In this case, the allocator first gets the new block using `get_block`.
     * Then it updates (links) the padding bytes of the first chunk of the new block to the `fastbin_slot->block`.
     * Then updates `fastbin_slot->block` to point at the new block.
     *
     * Making use of the padding bytes as a next pointer will save space in the `el_fastbin_t` struct. The reason
     * for using the first chunk for this is to follow spatial locality principle.
     *
     * Todo [improvement]: May be later on we can steal the block from other fastbin_slot of the same tcache or even
     * fastbin_slot of another tcache probably. For ex: If 16 bytes and 32 bytes are allocated in the fast path and then
     * already upto 10 allocation requests are made for 16 bytes size but only a few 32 bytes allocations are made then
     * we can steal chunks from 32 bytes block of the same or different tcache.
     * i.e. first try, steal_block() if not possible then back to get_block()
     */

    void *new_block = get_block(fastbin_size);
    /*Padding byte of the new block's first chunk*/
    void **padding_byte = (void *)((char *)new_block + sizeof(size_t));
    *padding_byte = fastbin_slot->block;

    fastbin_slot->block = new_block;
    fastbin_slot->top = new_block;
  }
  void *chunk = fastbin_slot->top;
  size_t *p = chunk;
  *p = chunk_size;
  fastbin_slot->top += chunk_size;
  return FASTBIN_DATA_BLOCK(chunk);
}

/*Function parameter to pass in pthread_once.*/
static void init_once(void) {
  int c = pthread_mutex_init(&memalloc_ctx.mtx_memalloc_ctx_t, NULL);
  if (c != 0) {
    perror("pthread_mutex_init");
  }

  c = pthread_key_create(&memalloc_ctx.th_key, NULL);
  if (c != 0) {
    perror("pthread_key_create");
  }
}

/*Initialize thread specific cache block and required meta data*/
void *init_tcache() {
  int c = pthread_once(&memalloc_ctx.once, init_once);
  if (c != 0) {
    perror("pthread_once");
  }
  th_cache_t *tcache;
  tcache = pthread_getspecific(memalloc_ctx.th_key);
  if (!tcache) {
    tcache = get_block(sizeof(th_cache_t));
    pthread_mutex_lock(&memalloc_ctx.mtx_memalloc_ctx_t);
    if (!memalloc_ctx.recent_th_cache) {
      tcache->next = NULL;
      tcache->prev = NULL;
    } else {
      memalloc_ctx.recent_th_cache->next = tcache;
      tcache->prev = memalloc_ctx.recent_th_cache;
      tcache->next = NULL;
    }
    memalloc_ctx.recent_th_cache = tcache;
    pthread_mutex_unlock(&memalloc_ctx.mtx_memalloc_ctx_t);
    int c = pthread_setspecific(memalloc_ctx.th_key, tcache);
    if (c != 0) {
      perror("pthread_setspecific");
    }
  }
  return tcache;
}

void fastbin_free(void *header, size_t chunk_size) {
  th_cache_t *tcache = init_tcache();
  int offset = GET_FASTBIN_OFFSET(chunk_size - CHUNK_HEADER_SIZE);
  el_fastbin_t *fastbin_t = tcache->fast_bin + offset;
  void *cur = fastbin_t->freelist;
  fastbin_t->freelist = header;
  *(void **)((char *)header + CHUNK_HEADER_SIZE) = cur;
}

void memfree(void *chunk) {
  void *header = GET_HEADER_POINTER(chunk);
  size_t chunk_size = *(size_t *)header;
  if (chunk_size <= FASTBIN_MAX_LIMIT) {
    fastbin_free(header, chunk_size);
  }
}

void *memalloc(size_t size) {
  th_cache_t *tcache = init_tcache();

  /*Follow fastpath for allocation request of size <=FASTPATH_MAX_LIMIT*/
  if (size <= FASTBIN_MAX_LIMIT) {
    return fastpath_allocation(tcache, size);
  }

  return NULL;
}

void *th1(void *arg) {
  int *p = memalloc(sizeof(int));
  *p = 31;
  printf("p: %d\n", *p);
  return NULL;
}

void *th2(void *arg) {
  int *q = memalloc(sizeof(int));
  *q = 59;
  printf("q: %d\n", *q);
  return NULL;
}

int main(void) {
  pthread_t t1, t2;
  pthread_create(&t1, NULL, th1, NULL);
  pthread_create(&t2, NULL, th2, NULL);
  pthread_join(t1, NULL);
  pthread_join(t2, NULL);

  return 0;
}