当前位置：首页 > article >正文

jemalloc替换标准库 malloc等函数的三种方式

article 2024/10/28 18:52:02

1、LD_PRELOAD

使用LD_PRELOAD可以做到无侵入式替换，只需要在运行程序前设置env ，export LD_PRELOAD=/path/to/jemalloc
注：编译jemalloc时不设置--with-jemalloc-prefix

#include <cstdlib>
#include <stdio.h>
#include <string.h>
#include <iostream>

// 当使用 export LD_PRELOAD= 预加载jemalloc库时
// malloc 将会被 jemalloc的符号替换, 做到无需修改源代码即可使用jemalloc

/**
 * @brief 原理
 * 
 * 加载顺序
 * 1、可执行程序
 * 2、LD_PRELOAD 库
 * 3、其他库，根据其在链接时的顺序
 * 
 * 当多个库中有相同符号时, 动态链接器会选择第一个加载的符号. 故使用 LD_PRELOAD 可以替换标准库的 malloc
 */

int main(int argc, char **argv)
{
    size_t size = 1024;
    void* ptr = malloc(size);
    if (!ptr) {
        std::cerr << "Memory allocation failed" << std::endl;
        return 1;
    }

    memset(ptr, 0, size);

    free(ptr);
    return 0;
}

2、代码中调用je_malloc

这种做法是侵入式的，如果没有jemalloc库将无法启动程序，另外编译代码时需要设置 --with-jemalloc-prefix=je_

#include <cstdlib>
#include <stdio.h>
#include <string.h>
#include <iostream>

#include <jemalloc/jemalloc.h>

// 主动链接jemalloc时需要在编译时带上 --with-jemalloc-prefix=je_
// 不带时编译出来的符号就是 malloc, 防止冲突, 带上后符号是 je_malloc

int main()
{
    size_t size = 1024;
    void* ptr = je_malloc(size);
    if (!ptr) {
        std::cerr << "Memory allocation failed" << std::endl;
        return 1;
    }

    memset(ptr, 0, size);

    ::atexit([](){
        uint64_t epoch = 1;
        size_t sz = sizeof(epoch);
        je_mallctl("epoch", &epoch, &sz, &epoch, sz);

        size_t allocated, active, mapped;
        sz = sizeof(size_t);
        je_mallctl("stats.allocated", &allocated, &sz, NULL, 0);
        je_mallctl("stats.active", &active, &sz, NULL, 0);
        je_mallctl("stats.mapped", &mapped, &sz, NULL, 0);

        printf("allocated/active/mapped: %zu/%zu/%zu\n", allocated, active, mapped);
    });

    // 释放内存
    je_free(ptr);
    return 0;
}

3、使用hook方式

侵入式的，并且存在一定的不安全性

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <string>

#include <signal.h>

#include <unistd.h>
#include <dlfcn.h>

#include <jemalloc/jemalloc.h>

// 这种情况需要将 malloc/calloc 等全都 hook, 假设没有 hook realloc, 第三方库使用此函数申请的内存使用free会出现段错误
// c++ 重载 operator new/operator delete 一般是成对的, 默认的 new/delete 底层使用 malloc 无需考虑

struct ReplaceMallocFree
{
    using SelfMalloc = void *(*)(size_t);
    using SelfFree = void (*)(void *);

    ReplaceMallocFree()
    {
        m_malloc = (SelfMalloc)dlsym(RTLD_NEXT, "malloc");
        m_free = (SelfFree)dlsym(RTLD_NEXT, "free");

        assert(m_malloc && m_free);

        printf("malloc = %p, free = %p\n", m_malloc, m_free);
    }

    SelfMalloc  m_malloc;
    SelfFree    m_free;
};

ReplaceMallocFree g_replaceMallocFree;

#ifdef __cplusplus
extern "C" {
#endif

void* malloc(size_t size)
{
    // printf 有调用malloc行为, 会导致无限递归出现段错误
    write(STDOUT_FILENO, "---> malloc\n", 12);
    return je_malloc(size);
}

void free(void *ptr)
{
    write(STDOUT_FILENO, "---> free\n", 12);
    je_free(ptr);
}

#ifdef __cplusplus
}
#endif


void catch_sig(int32_t sig)
{
    if (sig == SIGABRT) {
        printf("\nSUCCESS\n");
    }

    exit(0);
}

int main()
{
    // 标准库的 free 检测到异常会调用 abort
    signal(SIGABRT, catch_sig);

    size_t size = 1024;
    void* ptr = malloc(size);
    if (!ptr) {
        printf("Memory allocation failed\n");
        return 1;
    }

    memset(ptr, 0, size);

    uint32_t *pp = new uint32_t;
    free(pp);

    printf("-------------\n");

    g_replaceMallocFree.m_free(ptr);
    return 0;
}