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 }