[Pyroscope] Rust 메모리 사용량 수집 구현

Pyroscope Rust 에이전트는 현재 CPU 사용량에 대한 수집만을 제공한다.

당연하다. Rust 같은 시스템 수준 언어에서는 메모리에 대한 관리자가 따로 없어서 메모리 사용량을 수집할 수 있는 뭔가의 중간 레이어가 없기 때문이다.

그래서 Rust에서 메모리 프로파일링을 하려면 jemalloc 같은 서드파티 메모리 allocator를 사용해야 한다.
여기서는 jemalloc을 써서 메모리 pprof를 뽑고, pyroscope에 집어넣어서 실시간 시각화를 구성하는 방법까지를 다뤄본다.

사전 조건과 종속성

jemalloc은 사실상 Linux에서만 사용 가능한 환경이다. 그래서 Linux만 된다고 보면 된다.

tokio = { version = "1.0", features = ["full"] }
tikv-jemallocator = { version = "0.6.0", features = ["profiling", "unprefixed_malloc_on_supported_platforms"] }
jemalloc_pprof = { version = "0.8.1", features = ["symbolize"] }
reqwest = "0.12.23"

그리고 jemalloc 관련된 모듈 몇개를 추가해준다.

jemalloc 활성화

전역변수 기반으로 제어할 수 있게 되어있다. 적당히 이런 식으로 쳐준다.

#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

#[allow(non_upper_case_globals)]
#[unsafe(export_name = "malloc_conf")]
pub static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:20\0";

샘플링 간격은 20으로 했는데, 필요에 따라 조정해도 된다.

pyroscope 수집 구현

자, 그럼 pyroscope에 데이터를 부어주는 백그라운드 스레드를 만들어보자.
10초를 간격으로 해서 쌓인 메모리 덤프를 pprof로 뽑고, 그걸 pyroscope에 부어주는 형태가 된다.

pyroscope에 붓는건 ingest API를 쓰면 되는데, 복잡하진 않다.

async fn send_to_pyroscope(
    pprof_data: Vec<u8>,
    from: u64,
    until: u64,
) -> Result<(), Box<dyn std::error::Error>> {
    let client = reqwest::Client::new();
    let pyroscope_url = "http://localhost:4040/ingest";

    let response = client
        .post(pyroscope_url)
        .header("Content-Type", "application/octet-stream")
        .query(&[
            ("name", "rust-app{}"),
            ("from", &from.to_string()),
            ("until", &until.to_string()),
            ("format", "pprof"),
        ])
        .body(pprof_data)
        .send()
        .await?;

    if response.status().is_success() {
        println!("Successfully sent pprof data to Pyroscope");
    } else {
        println!("Failed to send data: {}", response.status());
    }

    Ok(())
}

from - until로 초 단위 시간을 넣고, pprof를 바이트열로 request body에 때려넣기만 하면 된다.

다음이 저걸 활용한 setup 함수다.

fn setup_pyroscope() {
    tokio::spawn(async move {
        let mut prof_ctl = jemalloc_pprof::PROF_CTL.as_ref().unwrap().lock().await;

        if prof_ctl.activated() {
            println!("Profiling is activated.");
        } else {
            panic!("Profiling is not activated.");
        }

        if let Err(e) = prof_ctl.activate() {
            eprintln!("Failed to activate profiler: {}", e);
        }

        loop {
            let start = SystemTime::now();
            tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
            let end = SystemTime::now();

            let pprof_data = prof_ctl.dump_pprof().expect("Failed to dump pprof data");

            println!("Dumped pprof data of size: {}", pprof_data.len());

            let start = start.duration_since(UNIX_EPOCH).unwrap().as_secs();
            let end = end.duration_since(UNIX_EPOCH).unwrap().as_secs();

            if let Err(e) = send_to_pyroscope(pprof_data, start, end).await {
                eprintln!("Error sending pprof data: {}", e);
            }
        }
    });
}

계속 루프 뺑뺑이 돌면서 10초 간격으로 부어주도록 구성했다.

그럼 저게 돌아가는지도 한번 보자.
메모리를 점유하는 waste_memory라는 공통 함수를 두고, 메모리를 다르게 사용하는 a와 b라는 함수를 루프마다 돌리게 했다.

fn waste_memory(bytes: usize) {
    let mut vec = Vec::new();
    for _ in 0..bytes {
        vec.push(0u8);
    }

    println!("Wasted {} bytes of memory", bytes);
}

fn a() {
    waste_memory(1024 * 1024 * 30); // Waste 30 MB of memory
}

fn b() {
    waste_memory(1024 * 1024 * 500); // Waste 500 MB of memory
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    setup_pyroscope();

    loop {
        a();
        b();

        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
    }
}

그리고 실행해보면

이런 식으로 계속 쌓일 것이다.

그럼 이렇게 flamegraph로 시각화해서 볼 수도 있고, 기간 필터를 걸어서 특정 부분만 볼 수도 있다.

전체 테스트 코드

use std::time::{SystemTime, UNIX_EPOCH};

#[cfg(not(target_env = "msvc"))]
#[global_allocator]
static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

#[allow(non_upper_case_globals)]
#[unsafe(export_name = "malloc_conf")]
pub static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:20\0";

fn setup_pyroscope() {
    tokio::spawn(async move {
        let mut prof_ctl = jemalloc_pprof::PROF_CTL.as_ref().unwrap().lock().await;

        if prof_ctl.activated() {
            println!("Profiling is activated.");
        } else {
            panic!("Profiling is not activated.");
        }

        if let Err(e) = prof_ctl.activate() {
            eprintln!("Failed to activate profiler: {}", e);
        }

        loop {
            let start = SystemTime::now();
            tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
            let end = SystemTime::now();

            let pprof_data = prof_ctl.dump_pprof().expect("Failed to dump pprof data");

            println!("Dumped pprof data of size: {}", pprof_data.len());

            let start = start.duration_since(UNIX_EPOCH).unwrap().as_secs();
            let end = end.duration_since(UNIX_EPOCH).unwrap().as_secs();

            if let Err(e) = send_to_pyroscope(pprof_data, start, end).await {
                eprintln!("Error sending pprof data: {}", e);
            }
        }
    });
}

async fn send_to_pyroscope(
    pprof_data: Vec<u8>,
    from: u64,
    until: u64,
) -> Result<(), Box<dyn std::error::Error>> {
    let client = reqwest::Client::new();
    let pyroscope_url = "http://localhost:4040/ingest";

    let response = client
        .post(pyroscope_url)
        .header("Content-Type", "application/octet-stream")
        .query(&[
            ("name", "rust-app{}"),
            ("from", &from.to_string()),
            ("until", &until.to_string()),
            ("format", "pprof"),
        ])
        .body(pprof_data)
        .send()
        .await?;

    if response.status().is_success() {
        println!("Successfully sent pprof data to Pyroscope");
    } else {
        println!("Failed to send data: {}", response.status());
    }

    Ok(())
}

fn waste_memory(bytes: usize) {
    let mut vec = Vec::new();
    for _ in 0..bytes {
        vec.push(0u8);
    }

    println!("Wasted {} bytes of memory", bytes);
}

fn a() {
    waste_memory(1024 * 1024 * 30); // Waste 100 MB of memory
}

fn b() {
    waste_memory(1024 * 1024 * 500); // Waste 200 MB of memory
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    setup_pyroscope();

    loop {
        a();
        b();

        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
    }
}

참조
https://linux.die.net/man/3/jemalloc
https://crates.io/crates/jemalloc_pprof