malloc_hook.cpp

#include "plot_actions.hpp"
#include <cerrno>
#include <chrono>
#include <cstdlib>
#include <deque>
#include <mutex>
#include <new>
#if __unix__
# include <unistd.h>
#elif _WIN32
# include <windows.h>
#endif
#include "alloc_action.hpp"

namespace {

uint32_t get_thread_id() {
#if __unix__
    return gettid();
#elif _WIN32
    return GetCurrentThreadId();
#else
    return 0;
#endif
}

struct alignas(64) PerThreadData {
    std::recursive_mutex lock;
    std::deque<AllocAction> actions;
    bool enable = false;
};

struct GlobalData {
    std::mutex lock;

    static inline size_t const kPerThreadsCount = 32;
    PerThreadData per_threads[kPerThreadsCount];

    GlobalData() {
        for (size_t i = 0; i < kPerThreadsCount; ++i) {
            per_threads[i].enable = true;
        }
    }

    ~GlobalData() {
        for (size_t i = 0; i < kPerThreadsCount; ++i) {
            per_threads[i].enable = false;
        }
        std::deque<AllocAction> actions;
        for (size_t i = 0; i < kPerThreadsCount; ++i) {
            auto &their_actions = per_threads[i].actions;
            actions.insert(actions.end(), their_actions.begin(),
                           their_actions.end());
        }
        plot_alloc_actions(actions);
    }
} global;

struct EnableGuard {
    uint32_t tid;
    bool was_enable;
    PerThreadData &per_thread;

    EnableGuard()
        : tid(get_thread_id()),
          per_thread(global.per_threads[((size_t)tid * 17) %
                                        global.kPerThreadsCount]) {
        per_thread.lock.lock();
        was_enable = per_thread.enable;
        per_thread.enable = false;
    }

    explicit operator bool() const {
        return was_enable;
    }

    void on(AllocOp op, void *ptr, size_t size, size_t align,
            void *caller) const {
        if (ptr) {
            auto now = std::chrono::high_resolution_clock::now();
            int64_t time =
                std::chrono::duration_cast<std::chrono::nanoseconds>(
                    now.time_since_epoch())
                    .count();
            per_thread.actions.push_back(
                AllocAction{op, tid, ptr, size, align, caller, time});
        }
    }

    ~EnableGuard() {
        per_thread.enable = was_enable;
        per_thread.lock.unlock();
    }
};

} // namespace

#if __GNUC__
extern "C" void *__libc_malloc(size_t size) noexcept;
extern "C" void __libc_free(void *ptr) noexcept;
extern "C" void *__libc_calloc(size_t nmemb, size_t size) noexcept;
extern "C" void *__libc_realloc(void *ptr, size_t size) noexcept;
extern "C" void *__libc_reallocarray(void *ptr, size_t nmemb,
                                     size_t size) noexcept;
extern "C" void *__libc_valloc(size_t size) noexcept;
extern "C" void *__libc_memalign(size_t align, size_t size) noexcept;
# define REAL_LIBC(name)      __libc_##name
# define MAY_OVERRIDE_MALLOC  1
# define MAY_SUPPORT_MEMALIGN 1
# define RETURN_ADDRESS       __builtin_return_address(0)
#elif _MSC_VER
static void *msvc_malloc(size_t size) noexcept {
    return HeapAlloc(GetProcessHeap(), 0, size);
}

static void *msvc_calloc(size_t nmemb, size_t size) noexcept {
    return HeapAlloc(GetProcessHeap(), HEAP_ZERO_MEMORY, nmemb * size);
}

static void msvc_free(void *ptr) noexcept {
    HeapFree(GetProcessHeap(), 0, ptr);
}

static void *msvc_realloc(void *ptr, size_t size) noexcept {
    return HeapReAlloc(GetProcessHeap(), 0, ptr, size);
}

static void *msvc_reallocarray(void *ptr, size_t nmemb, size_t size) noexcept {
    return msvc_realloc(ptr, nmemb * size);
}

# define REAL_LIBC(name)      msvc_##name
# define MAY_OVERRIDE_MALLOC  1
# define MAY_SUPPORT_MEMALIGN 0

# include <intrin.h>

# pragma intrinsic(_ReturnAddress)
# define RETURN_ADDRESS _ReturnAddress()

#else
# define REAL_LIBC(name)      name
# define MAY_OVERRIDE_MALLOC  0
# define MAY_SUPPORT_MEMALIGN 0
# define RETURN_ADDRESS       ((void *)1)
#endif

#if MAY_OVERRIDE_MALLOC
extern "C" void *malloc(size_t size) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(malloc)(size);
    if (ena) {
        ena.on(AllocOp::Malloc, ptr, size, kNone, RETURN_ADDRESS);
    }
    return ptr;
}

extern "C" void free(void *ptr) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::Free, ptr, kNone, kNone, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

extern "C" void *calloc(size_t nmemb, size_t size) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(calloc)(nmemb, size);
    if (ena) {
        ena.on(AllocOp::Malloc, ptr, nmemb * size, kNone, RETURN_ADDRESS);
    }
    return ptr;
}

extern "C" void *realloc(void *ptr, size_t size) noexcept {
    EnableGuard ena;
    void *new_ptr = REAL_LIBC(realloc)(ptr, size);
    if (ena) {
        ena.on(AllocOp::Malloc, new_ptr, size, kNone, RETURN_ADDRESS);
        if (new_ptr) {
            ena.on(AllocOp::Free, ptr, kNone, kNone, RETURN_ADDRESS);
        }
    }
    return new_ptr;
}

extern "C" void *reallocarray(void *ptr, size_t nmemb, size_t size) noexcept {
    EnableGuard ena;
    void *new_ptr = REAL_LIBC(reallocarray)(ptr, nmemb, size);
    if (ena) {
        ena.on(AllocOp::Malloc, new_ptr, nmemb * size, kNone, RETURN_ADDRESS);
        if (new_ptr) {
            ena.on(AllocOp::Free, ptr, kNone, kNone, RETURN_ADDRESS);
        }
    }
    return new_ptr;
}

# if MAY_SUPPORT_MEMALIGN
extern "C" void *valloc(size_t size) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(valloc)(size);
    if (ena) {
#  if __unix__
        size_t pagesize = sysconf(_SC_PAGESIZE);
#  elif _WIN32
        SYSTEM_INFO info;
        info.dwPageSize = kNone;
        GetSystemInfo(&info);
        size_t pagesize = info.dwPageSize;
#  else
        size_t pagesize = 0;
#  endif
        ena.on(AllocOp::Malloc, ptr, size, pagesize, RETURN_ADDRESS);
    }
    return ptr;
}

extern "C" void *memalign(size_t align, size_t size) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(memalign)(align, size);
    if (ena) {
        ena.on(AllocOp::Malloc, ptr, size, align, RETURN_ADDRESS);
    }
    return ptr;
}

extern "C" void *aligned_alloc(size_t align, size_t size) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(memalign)(align, size);
    if (ena) {
        ena.on(AllocOp::Malloc, ptr, size, align, RETURN_ADDRESS);
    }
    return ptr;
}

extern "C" int posix_memalign(void **memptr, size_t align,
                              size_t size) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(memalign)(align, size);
    if (ena) {
        ena.on(AllocOp::Malloc, *memptr, size, align, RETURN_ADDRESS);
    }
    int ret = 0;
    if (!ptr) {
        ret = errno;
    } else {
        *memptr = ptr;
    }
    return ret;
}
# endif
#endif

void operator delete(void *ptr) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::Delete, ptr, kNone, kNone, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void operator delete[](void *ptr) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::DeleteArray, ptr, kNone, kNone, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void operator delete(void *ptr, std::nothrow_t const &) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::Delete, ptr, kNone, kNone, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void operator delete[](void *ptr, std::nothrow_t const &) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::DeleteArray, ptr, kNone, kNone, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void *operator new(size_t size) noexcept(false) {
    EnableGuard ena;
    void *ptr = REAL_LIBC(malloc)(size);
    if (ena) {
        ena.on(AllocOp::New, ptr, size, kNone, RETURN_ADDRESS);
    }
    if (ptr == nullptr) {
        throw std::bad_alloc();
    }
    return ptr;
}

void *operator new[](size_t size) noexcept(false) {
    EnableGuard ena;
    void *ptr = REAL_LIBC(malloc)(size);
    if (ena) {
        ena.on(AllocOp::NewArray, ptr, size, kNone, RETURN_ADDRESS);
    }
    if (ptr == nullptr) {
        throw std::bad_alloc();
    }
    return ptr;
}

void *operator new(size_t size, std::nothrow_t const &) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(malloc)(size);
    if (ena) {
        ena.on(AllocOp::New, ptr, size, kNone, RETURN_ADDRESS);
    }
    return ptr;
}

void *operator new[](size_t size, std::nothrow_t const &) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(malloc)(size);
    if (ena) {
        ena.on(AllocOp::NewArray, ptr, size, kNone, RETURN_ADDRESS);
    }
    return ptr;
}

#if (__cplusplus >= 201402L || _MSC_VER >= 1916)
void operator delete(void *ptr, size_t size) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::Delete, ptr, size, kNone, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void operator delete[](void *ptr, size_t size) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::DeleteArray, ptr, size, kNone, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}
#endif

#if (__cplusplus > 201402L || defined(__cpp_aligned_new))
# if MAY_SUPPORT_MEMALIGN
void operator delete(void *ptr, std::align_val_t align) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::Delete, ptr, kNone, (size_t)align, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void operator delete[](void *ptr, std::align_val_t align) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::DeleteArray, ptr, kNone, (size_t)align, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void operator delete(void *ptr, size_t size, std::align_val_t align) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::Delete, ptr, size, (size_t)align, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void operator delete[](void *ptr, size_t size,
                       std::align_val_t align) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::DeleteArray, ptr, size, (size_t)align, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void operator delete(void *ptr, std::align_val_t align,
                     std::nothrow_t const &) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::Delete, ptr, kNone, (size_t)align, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void operator delete[](void *ptr, std::align_val_t align,
                       std::nothrow_t const &) noexcept {
    EnableGuard ena;
    if (ena) {
        ena.on(AllocOp::DeleteArray, ptr, kNone, (size_t)align, RETURN_ADDRESS);
    }
    REAL_LIBC(free)(ptr);
}

void *operator new(size_t size, std::align_val_t align) noexcept(false) {
    EnableGuard ena;
    void *ptr = REAL_LIBC(memalign)((size_t)align, size);
    if (ena) {
        ena.on(AllocOp::New, ptr, size, (size_t)align, RETURN_ADDRESS);
    }
    if (ptr == nullptr) {
        throw std::bad_alloc();
    }
    return ptr;
}

void *operator new[](size_t size, std::align_val_t align) noexcept(false) {
    EnableGuard ena;
    void *ptr = REAL_LIBC(memalign)((size_t)align, size);
    if (ena) {
        ena.on(AllocOp::NewArray, ptr, size, (size_t)align, RETURN_ADDRESS);
    }
    if (ptr == nullptr) {
        throw std::bad_alloc();
    }
    return ptr;
}

void *operator new(size_t size, std::align_val_t align,
                   std::nothrow_t const &) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(memalign)((size_t)align, size);
    if (ena) {
        ena.on(AllocOp::New, ptr, size, (size_t)align, RETURN_ADDRESS);
    }
    return ptr;
}

void *operator new[](size_t size, std::align_val_t align,
                     std::nothrow_t const &) noexcept {
    EnableGuard ena;
    void *ptr = REAL_LIBC(memalign)((size_t)align, size);
    if (ena) {
        ena.on(AllocOp::NewArray, ptr, size, (size_t)align, RETURN_ADDRESS);
    }
    return ptr;
}
# endif
#endif