lemoncake

rbtree based memory allocator
git clone git@git.2f30.org/lemoncake.git
Log | Files | Refs | README | LICENSE

lemoncake.c (9697B)


      1 /* See LICENSE file for copyright and license details. */
      2 #include <sys/mman.h>
      3 #include <sys/stat.h>
      4 #include <fcntl.h>
      5 #include <unistd.h>
      6 #include <stdio.h>
      7 #include <stdlib.h>
      8 #include <string.h>
      9 #include <stdint.h>
     10 #include <stdbool.h>
     11 #include <stdarg.h>
     12 #include <errno.h>
     13 #include "tree.h"
     14 #include "spinlock.h"
     15 
     16 /* Memory alignment for the allocation functions */
     17 enum { ALIGN = 4 * sizeof(size_t) };
     18 
     19 static void writelog(int fd, const char *fmt, ...);
     20 static void dumpstats(void);
     21 
     22 struct lemon_stats {
     23 	/* # of sbrk calls */
     24 	unsigned long nr_sbrk;
     25 	/* # of mmap calls */
     26 	unsigned long nr_mmap;
     27 	/* # of malloc calls */
     28 	unsigned long nr_malloc;
     29 	/* # of realloc calls */
     30 	unsigned long nr_realloc;
     31 	/* # of free calls */
     32 	unsigned long nr_free;
     33 	/* # of nodes in the allocation tree */
     34 	unsigned long nr_at_nodes;
     35 	/* # of nodes in the free tree */
     36 	unsigned long nr_ft_nodes;
     37 	/* # of new allocations */
     38 	unsigned long nr_alloc_new;
     39 	/* # of allocations satisfied from the free tree */
     40 	unsigned long nr_alloc_free;
     41 	/* # of shrinked reallocs */
     42 	unsigned long nr_shrink_realloc;
     43 	/* # of invalid frees */
     44 	unsigned long nr_invalid_free;
     45 };
     46 
     47 static bool init;
     48 static struct lemon_stats st;
     49 static bool strict = false;
     50 
     51 struct node {
     52 	void *buf;
     53 	size_t sz;
     54 	RB_ENTRY(node) entry;
     55 };
     56 
     57 RB_HEAD(free_tree, node) ft = RB_INITIALIZER(&ft);
     58 RB_HEAD(alloc_tree, node) at = RB_INITIALIZER(&at);
     59 static spinlock_t rblock;
     60 
     61 static int ft_cmp(struct node *a, struct node *b);
     62 RB_PROTOTYPE(free_tree, node, entry, ft_cmp)
     63 
     64 static int at_cmp(struct node *a, struct node *b);
     65 RB_PROTOTYPE(alloc_tree, node, entry, at_cmp)
     66 
     67 RB_GENERATE(free_tree, node, entry, ft_cmp)
     68 /* These are ordered by `sz' */
     69 static int
     70 ft_cmp(struct node *a, struct node *b)
     71 {
     72 	if (a->sz < b->sz)
     73 		return -1;
     74 	if (a->sz > b->sz)
     75 		return 1;
     76 	return 0;
     77 }
     78 
     79 RB_GENERATE(alloc_tree, node, entry, at_cmp)
     80 /* These are ordered by address */
     81 static int
     82 at_cmp(struct node *a, struct node *b)
     83 {
     84 	if (a->buf < b->buf)
     85 		return -1;
     86 	if (a->buf > b->buf)
     87 		return 1;
     88 	return 0;
     89 }
     90 
     91 static inline void *
     92 align_pointer(size_t align, void *p)
     93 {
     94 	return (void *)(((uintptr_t)p + align) & ~(align - 1));
     95 }
     96 
     97 static inline void *
     98 alloc_object(size_t sz)
     99 {
    100 	void *base, *p;
    101 
    102 	base = sbrk(sz + ALIGN);
    103 	if (base == (void *)-1)
    104 		return NULL;
    105 	p = base;
    106 	st.nr_sbrk++;
    107 	return align_pointer(ALIGN, p);
    108 }
    109 
    110 static inline void *
    111 mmap_aligned(size_t align, size_t sz)
    112 {
    113 	void *p;
    114 
    115 	/* align should be a power of two */
    116 	if ((align - 1) & align)
    117 		return NULL;
    118 	p = mmap(0, sz + align, PROT_READ | PROT_WRITE,
    119 		 MAP_PRIVATE | MAP_ANON, -1, 0);
    120 	if (p == MAP_FAILED)
    121 		return NULL;
    122 	st.nr_mmap++;
    123 	return align_pointer(align, p);
    124 }
    125 
    126 static int
    127 lemon_init(void)
    128 {
    129 	char *p;
    130 
    131 	if (init)
    132 		return 0;
    133 	p = getenv("LEMONCAKE_DEBUG");
    134 	if (p)
    135 		atexit(dumpstats);
    136 	init = true;
    137 	return 0;
    138 }
    139 
    140 void *
    141 malloc(size_t sz)
    142 {
    143 	struct node n, *an, *res;
    144 	void *p;
    145 
    146 	if (lemon_init())
    147 		return NULL;
    148 	if (!sz)
    149 		return NULL;
    150 	if (sz > SIZE_MAX / 2)
    151 		return NULL;
    152 	lock(&rblock);
    153 	/* Lookup in the free tree for a block greater
    154 	 * than or equal to `sz' bytes */
    155 	n.sz = sz;
    156 	res = RB_NFIND(free_tree, &ft, &n);
    157 	if (!res) {
    158 		/* No available free block, create a new block
    159 		 * and add it to the alloc tree */
    160 		an = alloc_object(sizeof(*an));
    161 		if (!an) {
    162 			unlock(&rblock);
    163 			return NULL;
    164 		}
    165 		p = mmap_aligned(ALIGN, sz);
    166 		if (!p) {
    167 			unlock(&rblock);
    168 			return NULL;
    169 		}
    170 		an->buf = p;
    171 		an->sz = sz;
    172 		RB_INSERT(alloc_tree, &at, an);
    173 		unlock(&rblock);
    174 		st.nr_malloc++;
    175 		st.nr_alloc_new++;
    176 		return an->buf;
    177 	}
    178 	an = RB_REMOVE(free_tree, &ft, res);
    179 	RB_INSERT(alloc_tree, &at, an);
    180 	unlock(&rblock);
    181 	st.nr_malloc++;
    182 	st.nr_alloc_free++;
    183 	return an->buf;
    184 }
    185 
    186 static inline void
    187 align_node(struct node *n)
    188 {
    189 	void *p = n->buf;
    190 
    191 	if (n->sz < ALIGN)
    192 		return;
    193 	n->buf = align_pointer(ALIGN, n->buf);
    194 	n->sz = n->sz - ((uintptr_t)n->buf - (uintptr_t)p);
    195 }
    196 
    197 void *
    198 realloc(void *oldp, size_t sz)
    199 {
    200 	struct node n, *res;
    201 	struct node *oldan, *newan;
    202 	struct node *fn, *newfn;
    203 
    204 	if (lemon_init())
    205 		return NULL;
    206 	if (!oldp)
    207 		return malloc(sz);
    208 	if (!sz) {
    209 		if (oldp)
    210 			free(oldp);
    211 		return NULL;
    212 	}
    213 	if (sz > SIZE_MAX / 2)
    214 		return NULL;
    215 	lock(&rblock);
    216 	n.buf = oldp;
    217 	res = RB_FIND(alloc_tree, &at, &n);
    218 	if (res) {
    219 		if (res->sz == sz) {
    220 			unlock(&rblock);
    221 			st.nr_realloc++;
    222 			return res->buf;
    223 		}
    224 		if (res->sz > sz) {
    225 			size_t diff = res->sz - sz;
    226 			/* We cannot ensure the alignment of the free
    227 			 * block if its size is less than ALIGN - so just
    228 			 * don't split the node in that case */
    229 			if (diff < ALIGN) {
    230 				unlock(&rblock);
    231 				st.nr_realloc++;
    232 				return res->buf;
    233 			}
    234 			newfn = alloc_object(sizeof(*newfn));
    235 			if (!newfn) {
    236 				unlock(&rblock);
    237 				return NULL;
    238 			}
    239 			res->sz = sz;
    240 			newfn->buf = res->buf + sz;
    241 			newfn->sz = diff;
    242 			align_node(newfn);
    243 			RB_INSERT(free_tree, &ft, newfn);
    244 			unlock(&rblock);
    245 			st.nr_realloc++;
    246 			st.nr_shrink_realloc++;
    247 			return res->buf;
    248 		}
    249 		oldan = res;
    250 		/* Lookup in the free tree for a block greater
    251 		 * than or equal to `sz' bytes */
    252 		n.sz = sz;
    253 		res = RB_NFIND(free_tree, &ft, &n);
    254 		if (!res) {
    255 			/* No available free block, create a new block
    256 			 * and add it to the alloc tree */
    257 			newan = alloc_object(sizeof(*newan));
    258 			if (!newan) {
    259 				unlock(&rblock);
    260 				return NULL;
    261 			}
    262 			newan->buf = mmap_aligned(ALIGN, sz);
    263 			if (!newan->buf) {
    264 				unlock(&rblock);
    265 				return NULL;
    266 			}
    267 			newan->sz = sz;
    268 			RB_INSERT(alloc_tree, &at, newan);
    269 			st.nr_alloc_new++;
    270 		} else {
    271 			/* Grab the block from the free tree instead */
    272 			newan = RB_REMOVE(free_tree, &ft, res);
    273 			RB_INSERT(alloc_tree, &at, newan);
    274 			st.nr_alloc_free++;
    275 		}
    276 		/* Copy over the contents from `oldp' to the
    277 		 * new memory block */
    278 		memcpy(newan->buf, oldan->buf,
    279 		       sz < oldan->sz ? sz : oldan->sz);
    280 		/* Return `oldp' to the free tree */
    281 		fn = RB_REMOVE(alloc_tree, &at, oldan);
    282 		RB_INSERT(free_tree, &ft, fn);
    283 		unlock(&rblock);
    284 		st.nr_realloc++;
    285 		return newan->buf;
    286 	}
    287 	unlock(&rblock);
    288 	return NULL;
    289 }
    290 
    291 void *
    292 calloc(size_t nmemb, size_t sz)
    293 {
    294 	void *p;
    295 
    296 	if (lemon_init())
    297 		return NULL;
    298 	p = malloc(nmemb * sz);
    299 	if (!p)
    300 		return NULL;
    301 	memset(p, 0, nmemb * sz);
    302 	return p;
    303 }
    304 
    305 void
    306 free(void *p)
    307 {
    308 	struct node n, *fn, *res;
    309 
    310 	if (lemon_init())
    311 		return;
    312 	if (!p)
    313 		return;
    314 	lock(&rblock);
    315 	n.buf = p;
    316 	res = RB_FIND(alloc_tree, &at, &n);
    317 	if (res) {
    318 		fn = RB_REMOVE(alloc_tree, &at, res);
    319 		RB_INSERT(free_tree, &ft, fn);
    320 		st.nr_free++;
    321 	} else {
    322 		st.nr_invalid_free++;
    323 		if (strict) {
    324 			writelog(STDERR_FILENO, "*** Freeing invalid pointer: %p, aborting ***\n", p);
    325 			_Exit(1);
    326 		}
    327 	}
    328 	unlock(&rblock);
    329 }
    330 
    331 void
    332 cfree(void *p)
    333 {
    334 	return free(p);
    335 }
    336 
    337 void *
    338 memalign(size_t align, size_t sz)
    339 {
    340 	struct node *an;
    341 	void *p;
    342 
    343 	if (lemon_init())
    344 		return NULL;
    345 	if (((align - 1) & align))
    346 		return NULL;
    347 	if (align < sizeof(void *))
    348 		return NULL;
    349 	if (align > SIZE_MAX - align)
    350 		return NULL;
    351 	/* Just allocate a new block, we don't care to look
    352 	 * for a block in the free tree as it might not be properly
    353 	 * aligned.  The previous implementation could cope with
    354 	 * that but it was sort of hackish.  There are few calls to
    355 	 * posix_memalign() in most cases, so the overhead should
    356 	 * not really matter. */
    357 	an = alloc_object(sizeof(*an));
    358 	if (!an)
    359 		return NULL;
    360 	p = mmap_aligned(align, sz);
    361 	if (!p)
    362 		return NULL;
    363 	an->buf = p;
    364 	an->sz = sz;
    365 	lock(&rblock);
    366 	RB_INSERT(alloc_tree, &at, an);
    367 	unlock(&rblock);
    368 	st.nr_alloc_new++;
    369 	return p;
    370 }
    371 
    372 void *
    373 aligned_alloc(size_t align, size_t sz)
    374 {
    375 	if (sz % align)
    376 		return NULL;
    377 	return memalign(align, sz);
    378 }
    379 
    380 void *
    381 valloc(size_t sz)
    382 {
    383 	return memalign(sysconf(_SC_PAGESIZE), sz);
    384 }
    385 
    386 void *
    387 pvalloc(size_t sz)
    388 {
    389 	long pagesize = sysconf(_SC_PAGESIZE);
    390 	sz = pagesize * ((sz + pagesize - 1) / pagesize);
    391 	return valloc(sz);
    392 }
    393 
    394 int
    395 posix_memalign(void **memptr, size_t align, size_t sz)
    396 {
    397 	*memptr = NULL;
    398 	if (((align - 1) & align))
    399 		return EINVAL;
    400 	if (align < sizeof(void *))
    401 		return EINVAL;
    402 	if (sz > SIZE_MAX - align)
    403 		return ENOMEM;
    404 	*memptr = memalign(align, sz);
    405 	if (!*memptr)
    406 		return ENOMEM;
    407 	return 0;
    408 }
    409 
    410 size_t
    411 malloc_usable_size(void *p)
    412 {
    413 	struct node n, *res;
    414 
    415 	if (lemon_init())
    416 		return 0;
    417 	if (!p)
    418 		return 0;
    419 	lock(&rblock);
    420 	n.buf = p;
    421 	res = RB_FIND(alloc_tree, &at, &n);
    422 	if (res) {
    423 		unlock(&rblock);
    424 		return res->sz;
    425 	}
    426 	unlock(&rblock);
    427 	return 0;
    428 }
    429 
    430 size_t
    431 malloc_size(void *p)
    432 {
    433 	return malloc_usable_size(p);
    434 }
    435 
    436 static void
    437 writelog(int fd, const char *fmt, ...)
    438 {
    439 	va_list ap;
    440 	char buf[BUFSIZ];
    441 
    442 	va_start(ap, fmt);
    443 	vsnprintf(buf, sizeof(buf), fmt, ap);
    444 	va_end(ap);
    445 	write(fd, buf, strlen(buf));
    446 }
    447 
    448 static void
    449 dumpstats(void)
    450 {
    451 	int fd;
    452 	struct node *n;
    453 
    454 	fd = open("lemoncake.out", O_WRONLY | O_CREAT | O_APPEND, 0644);
    455 	if (fd < 0)
    456 		return;
    457 	writelog(fd, "*** Lemoncake stats [process: %s pid: %lu] ***\n",
    458 		 getenv("_") ? getenv("_") : "<unknown>", (unsigned long)getpid());
    459 	writelog(fd, "sbrk calls: %lu\n", st.nr_sbrk);
    460 	writelog(fd, "mmap calls: %lu\n", st.nr_mmap);
    461 	writelog(fd, "malloc calls: %lu\n", st.nr_malloc);
    462 	writelog(fd, "realloc calls: %lu\n", st.nr_realloc);
    463 	writelog(fd, "shrinked realloc calls (no memcpy() required): %lu\n",
    464 		 st.nr_shrink_realloc);
    465 	writelog(fd, "free calls: %lu\n", st.nr_free);
    466 	writelog(fd, "new allocations: %lu\n", st.nr_alloc_new);
    467 	writelog(fd, "allocations from the free tree: %lu\n",
    468 		 st.nr_alloc_free);
    469 	writelog(fd, "invalid free calls: %lu\n", st.nr_invalid_free);
    470 	RB_FOREACH(n, alloc_tree, &at)
    471 		st.nr_at_nodes++;
    472 	RB_FOREACH(n, free_tree, &ft)
    473 		st.nr_ft_nodes++;
    474 	writelog(fd, "nodes in the allocation tree: %lu\n",
    475 		 st.nr_at_nodes);
    476 	writelog(fd, "nodes in the free tree: %lu\n",
    477 		 st.nr_ft_nodes);
    478 	close(fd);
    479 }