[Rust] allocation 수준 최적화

프로그램의 규모가 커지다보면, 생각보다 많은 부하를 잡아먹을 수 있는게 바로 메모리의 할당 문제다.

그래서 C++ 같은 언어에서도 placement new와 같은 구문을 사용해서 메모리 풀을 미리 만들고, 그 풀 내에서 빠르게 데이터를 넣고 꺼내쓰는 패턴을 사용하곤 한다.
https://blog.naver.com/sssang97/222724831217

Rust에서도 이와 같은 기능을 어느정도 제공한다.

아래는 비효율적인 할당을 반복하는 예시 코드다.
일반적이지는 않다.

fn main() {
    let mut total: i128 = 0;

    let start = std::time::Instant::now();

    for _ in 0..1000000 {
        let vec = Vec::<i32>::with_capacity(30000000); // 몹시 무거운 할당
        total += vec.capacity() as i128;
    }

    let end = std::time::Instant::now();
    let elapsed = end - start;
    println!("Elapsed: {:?}", elapsed);

    println!("Total: {}", total);
}

루프 내에서 길이 30000000짜리의 벡터를 계속해서 생성을 하고 있는데, 한두번이면 몰라도 저런 대규모의 할당이 중첩되면 상당한 오버헤드를 불러일으키기 마련이다.

그래서 저걸 실행하면 거의

몇초에 달하는 매우 느린 속도를 보여줄 것이다.
이럴때는 일종의 버퍼 메모리를 미리 할당하고, 그걸 메모리풀로 활용해서 할당을 하는 것이 효율적일 수 있다.

Global Allocator 정의

Global Allocator는 전역적인 할당에 대한 설정을 덮어씌우게 해주는 기능이다.
사실 프로그램 전체에 대해서 할당 기능을 overwrite해버리는 기능이기 때문에 이러한 예제에는 적합하지는 않다.
보통 jemalloc처럼 조금 더 효율적인 할당을 할 수 있는 전역 allocator를 갖다쓰는 형태로 사용을 할 것이다.

우선은 최대한 단순한 형태로 정리를 해보겠다. 제대로 구현하려면 규모가 너무 커진다.

우선 메모리 풀 사이즈를 적절히 지정해준다.
벡터 사이즈가 30000000에, element 사이즈가 4바이트이니 사이즈는 이렇게 잡아주면 될 것이다.

그리고 Allocator 객체를 정의한다.
UnsafeCell을 통해 저장공간을 만들어주는 것이다.

그리고 GlobalAlloc를 구현한다.

이걸 통해서 할당과 할당해제가 일어났을때의 동작을 정의하는 것이다.
이 경우에는 딱 저 사이즈의 공간만 제공해주면 되기 때문에, alloc에서도 포인터만 주고, dealloc에서도 뭘 따로 하지는 않았다.

하지만 제대로 allocator를 구현할 때는 alloc에서 미사용중인 unsafecell 메모리의 특정 지점을 찾아서 포인터를 반환해주고,
dealloc을 통해서는 사용중으로 표시되었던 unsafecell 메모리 영역을 미사용으로 표시해주는 기능 등을 구현해야 할 것이다.

그리고 특수한 매크로 태그를 통해 전역변수를 초기화하면, 기본작업은 끝이다.

프로그램의 시작과 동시에 해당 메모리가 적재될 것이다.

그리고 이대로 다시 돌려보면

매우 빨라진 것을 볼 수 있다.
불필요한 할당과 할당해제가 전부 사라지고, 사실상 같은 메모리를 그대로 쓰기만 했기 때문이다.

전체 코드다.

use std::alloc::{GlobalAlloc, Layout};
use std::cell::UnsafeCell;

const MEMORY_SIZE: usize = 120000000;

#[repr(C, align(4096))]
struct SimpleAllocator {
    arena: UnsafeCell<[u8; MEMORY_SIZE]>,
}

#[global_allocator]
static ALLOCATOR: SimpleAllocator = SimpleAllocator {
    arena: UnsafeCell::new([0x55; MEMORY_SIZE]),
};

unsafe impl Sync for SimpleAllocator {}

unsafe impl GlobalAlloc for SimpleAllocator {
    unsafe fn alloc(&self, _layout: Layout) -> *mut u8 {
        self.arena.get().cast::<u8>()
    }

    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
}

fn main() {
    let mut total: i128 = 0;

    let start = std::time::Instant::now();

    for _ in 0..1000000 {
        let vec = Vec::<i32>::with_capacity(30000000);
        total += vec.capacity() as i128;
    }

    let end = std::time::Instant::now();
    let elapsed = end - start;
    println!("Elapsed: {:?}", elapsed);

    println!("Total: {}", total);
}

Global Allocator 말고도 local allocator가 nightly에서 개발중이긴 한데, 아직 완성도가 낮은 것 같다.

참조
https://doc.rust-lang.org/1.9.0/book/custom-allocators.html
https://www.reddit.com/r/rust/comments/pwef20/placement_new/
https://os.phil-opp.com/heap-allocation/
https://doc.rust-lang.org/std/alloc/trait.GlobalAlloc.html