A src/util/mem.c

@@ -0,0 +1,331 @@ 
+#define _GNU_SOURCE
+#include <dlfcn.h>
+#include <endian.h>
+#include <fcntl.h>
+#include <pthread.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/stat.h>
+#include <sys/time.h>
+#include <sys/types.h>
+#include <time.h>
+#include <unistd.h>
+#include <sys/mman.h>
+#include <limits.h>
+#include <signal.h>
+
+#include "bt.h"
+#include "bool.h"
+
+#define UNUSED(x) ((void)x)
+
+void ERR(const char *s)
+{
+    if (!s) {
+        ERR("(null)");
+        return;
+    }
+    ssize_t ret = write(STDERR_FILENO, s, strlen(s));
+    UNUSED(ret);
+    ret = fsync(STDERR_FILENO);
+    UNUSED(ret);
+}
+
+void utoa(unsigned long n, char *s)
+{
+    if (n == 0) {
+        *s++ = '0';
+        *s = 0;
+        return;
+    }
+    char buffer[128] = {0};
+    char *digit;
+    for (digit = buffer; n; digit++, n/=10) {
+        *digit = '0' + (n % 10);
+    }
+    digit--;
+    while (digit >= buffer) {
+        *s++ = *digit--;
+    }
+    *s = 0;
+}
+
+unsigned long uabs(long n)
+{
+    if (n == LONG_MIN)
+        return ((unsigned long)LONG_MAX) + 1;
+    return (unsigned long)labs(n);
+}
+
+void itoa(long n, char *s)
+{
+    if (n < 0) {
+        *s++ = '-';
+    }
+    utoa(uabs(n), s);
+}
+
+__attribute__((noreturn))
+static void crash(char *file, int line, char *msg)
+{
+    ERR(file);
+    ERR(":");
+    char buf[256];
+    itoa(line, buf);
+    ERR(buf);
+    ERR(" ");
+    ERR(msg);
+    _exit(1);
+}
+
+#define ASSERT_OK(call) if (0 != (call)) crash(__FILE__, __LINE__, #call)
+#define ASSERT_FD(fd) if ((fd) == -1) crash(__FILE__, __LINE__, "bad file descriptor")
+#define ASSERT_PID(pid) if ((pid) == -1) crash(__FILE__, __LINE__, "bad PID")
+#define ASSERT(b) if (!(b)) crash(__FILE__, __LINE__, "assert failed")
+
+void *return_null(void *x, void *y)
+{
+    return 0;
+}
+
+void load_func_or_crash(void **result, const char *name)
+{
+    static int inside_loader = 0;
+    if (inside_loader) {
+        *result = (void*)return_null;
+        return;
+    } else if (*result == (void*)return_null) {
+        *result = 0;
+    }
+    if (*result)
+        return;
+    dlerror();
+    inside_loader++;
+    *result = dlsym(RTLD_NEXT, name);
+    inside_loader--;
+    char *err = dlerror();
+    if (err) {
+        ERR("Failed to load ");
+        ERR(name);
+        ERR(" error: ");
+        ERR(err);
+        ERR("\n");
+        exit(1);
+    }
+}
+
+static int fd = -1;
+static u_int64_t tstart = 0;
+static bool stop_alloc_log = false;
+
+static void write_char(char c)
+{
+    static char buffer[4096] = {0};
+    static size_t count = 0;
+    if (c) {
+        buffer[count++] = c;
+    }
+    if (!count)
+        return;
+    if (c == '\n' || c == 0 || count >= sizeof(buffer)) {
+        ssize_t ret = write(fd, buffer, count);
+        ASSERT(ret > 0);
+        count = 0;
+    }
+}
+
+static void write_string(const char *s)
+{
+    for (; *s; s++)
+        write_char(*s);
+}
+
+static void write_word(const char *s)
+{
+    write_string(s);
+    write_string(" ");
+}
+
+static char hex(u_int8_t value)
+{
+    if (value < 10)
+        return '0' + (char)value;
+    return 'a' + (char)(value - 10);
+}
+
+static void write_hex(u_int64_t v)
+{
+    write_string("0x");
+    if (!v) {
+        write_string("0 ");
+        return;
+    }
+    v = htobe64(v);
+    int leading_zero = 1;
+    while (v) {
+        u_int8_t byte = v & 0xff;
+        if (byte)
+            leading_zero = 0;
+        if (!leading_zero) {
+            write_char(hex(byte >> 4));
+            write_char(hex(byte & 0xf));
+        }
+        v = v >> 8;
+    }
+    write_string(" ");
+}
+
+static void write_backtrace(int skip)
+{
+    struct backtrace bt;
+    get_backtrace(&bt, skip + 1);
+    for (size_t i = 0; i < bt.frame_count; i++) {
+        write_word(bt.frames[i].name);
+    }
+}
+
+static void log_alloc_finish(void);
+
+static u_int64_t current_time_ms()
+{
+    struct timespec t = {0, 0};
+    clock_gettime(CLOCK_MONOTONIC, &t);
+    u_int64_t result = t.tv_sec * 1000 + t.tv_nsec / 1000 / 1000;
+    return result;
+}
+
+static void log_alloc_init(u_int64_t t)
+{
+    if (stop_alloc_log)
+        return;
+    if (fd == -1) {
+        int pfd[2] = {-1, -1};
+        ASSERT_OK(pipe(pfd));
+        pid_t child = fork();
+        ASSERT_PID(child);
+        if (child == 0) {
+            // child
+            close(pfd[1]);
+            ASSERT_FD(dup2(pfd[0], STDIN_FILENO));
+            ASSERT_OK(close(pfd[0]));
+            int out = open("mem.log.gz", O_APPEND | O_CLOEXEC | O_CREAT | O_WRONLY | O_TRUNC, 0600);
+            ASSERT_FD(out);
+            ASSERT_FD(dup2(out, STDOUT_FILENO));
+            sigset_t mask;
+            sigfillset(&mask);
+            sigprocmask(SIG_SETMASK, &mask, NULL);
+            ASSERT_OK(execlp("gzip", "gzip", "-c", NULL));
+            _exit(1);
+        } else {
+            // parent
+            close(pfd[0]);
+            fd = pfd[1];
+        }
+        atexit(log_alloc_finish);
+        tstart = t;
+        write_string("# function time_ms result ptr size count backtrace\n");
+    }
+}
+
+static void log_alloc_locked(const char *func_name, void *result, void *ptr, size_t size, size_t count)
+{
+    if (stop_alloc_log)
+        return;
+    u_int64_t t = current_time_ms();
+    if (fd == -1)
+        log_alloc_init(t);
+    if (fd == -1)
+        return;
+    if (func_name) {
+        write_word(func_name);
+        write_hex((u_int64_t)(t - tstart));
+        write_hex((u_int64_t)result);
+        write_hex((u_int64_t)ptr);
+        write_hex((u_int64_t)size);
+        write_hex((u_int64_t)count);
+        write_backtrace(2);
+        write_string("\n");
+    } else {
+        write_string("# done\n");
+        close(fd);
+        fd = -1;
+        stop_alloc_log = true;
+    }
+}
+
+static int pid = -1;
+static void log_alloc(const char *func_name, void *result, void *ptr, size_t size, size_t count)
+{
+    static pthread_mutex_t mutex_global = PTHREAD_MUTEX_INITIALIZER;
+    static pthread_mutex_t mutex_recursive;
+    static pthread_mutex_t mutex_nonrecursive = PTHREAD_MUTEX_INITIALIZER;
+    static bool mutexes_initialized = false;
+
+    pthread_mutex_lock(&mutex_global);
+    {
+        if (!mutexes_initialized) {
+            pthread_mutexattr_t attr;
+            pthread_mutexattr_init(&attr);
+            pthread_mutexattr_settype(&attr, PTHREAD_MUTEX_RECURSIVE);
+            pthread_mutex_init(&mutex_recursive, &attr);
+            mutexes_initialized = true;
+            pid = getpid();
+        }
+    }
+    pthread_mutex_unlock(&mutex_global);
+
+    // Do not log from forked processes.
+    if (pid != getpid())
+        return;
+
+    pthread_mutex_lock(&mutex_recursive);
+    int ret = pthread_mutex_trylock(&mutex_nonrecursive);
+    if (ret == 0) {
+        log_alloc_locked(func_name, result, ptr, size, count);
+        pthread_mutex_unlock(&mutex_nonrecursive);
+    }
+    pthread_mutex_unlock(&mutex_recursive);
+}
+
+static void log_alloc_finish()
+{
+    log_alloc(0, 0, 0, 0, 0);
+}
+
+void *malloc(size_t size)
+{
+    static void *(*original)(size_t size) = 0;
+    load_func_or_crash((void *)&original, __FUNCTION__);
+    void *result = original(size);
+    log_alloc(__FUNCTION__, result, 0, size, 0);
+    return result;
+}
+
+void *realloc(void *ptr, size_t size)
+{
+    static void *(*original)(void *p, size_t size) = 0;
+    load_func_or_crash((void *)&original, __FUNCTION__);
+    void *result = original(ptr, size);
+    log_alloc(__FUNCTION__, result, ptr, size, 0);
+    return result;
+}
+
+void *calloc(size_t nmemb, size_t size)
+{
+    static void *(*original)(size_t nmemb, size_t size) = 0;
+    load_func_or_crash((void *)&original, __FUNCTION__);
+    void *result = original(nmemb, size);
+    log_alloc(__FUNCTION__, result, 0, size, nmemb);
+    return result;
+}
+
+void free(void *ptr)
+{
+    static void *(*original)(void *p) = 0;
+    load_func_or_crash((void *)&original, __FUNCTION__);
+    if (!original) {
+        return;
+    }
+    original(ptr);
+    log_alloc(__FUNCTION__, 0, ptr, 0, 0);
+}