Allokáció-optimalizálás Rustban: SmallVec, Bumpalo és a heap elkerülése

Miért drága a heap allokáció a hot path-okon

Amikor Vec::new()-t hívunk és utána push-olunk, a legtöbb esetben ez teljesen rendben van. De ha ezt egy ciklusban, másodpercenként milliószor tesszük – például egy parserben, ahol minden tokenhez létrehozol egy kis vektort –, akkor a malloc/free hívások kezdenek dominálni a profilban.

A heap allokátor (alapból a rendszer allokátora, pl. glibc malloc vagy jemalloc) minden híváskor:

• zárolást végezhet (thread-safe allokátoroknál),
• keres egy megfelelő méretű szabad blokkot,
• könyvelést vezet a metaadatokról,
• cache-missek sorozatát okozhatja, mert a heap memória szétszórtan van.

Ezzel szemben a stack-en történő foglalás praktikusan ingyenes: egy pointer mozgatása. Ha tudjuk, hogy egy adatstruktúra a legtöbb esetben kicsi marad, érdemes elkerülni a heap-et, és csak akkor "spillelni" oda, ha valóban muszáj.

A SmallVec crate: stack-alapú tárolás kis méretekhez

A smallvec crate egy Vec-hez nagyon hasonló API-t ad, de van egy trükkje: van egy beépített, inline tárolt kapacitása. Amíg az elemek száma nem lépi túl ezt a méretet, semmi nem kerül a heap-re.

Adjuk hozzá a Cargo.toml-hoz:

[dependencies]
smallvec = "1.11"

És egy egyszerű használat:

use smallvec::{SmallVec, smallvec};

fn collect_tokens() -> SmallVec<[u32; 8]> {
    let mut buf: SmallVec<[u32; 8]> = smallvec![];
    for i in 0..5 {
        buf.push(i);
    }
    buf
}

fn main() {
    let tokens = collect_tokens();
    println!("{:?}", tokens);
    // 5 elem, mind a stack-en, nincs heap allokáció
}

A SmallVec<[u32; 8]> azt jelenti, hogy maximum 8 darab u32 fér el inline, stack-en tárolva. Ha 9. elemet is push-olunk, a struktúra automatikusan átvált heap-alapú tárolásra – transzparensen, a hívó kódnak semmit nem kell tudnia erről.

use smallvec::SmallVec;

fn main() {
    let mut v: SmallVec<[i32; 4]> = SmallVec::new();
    for i in 0..10 {
        v.push(i);
    }
    // 4 elem után heap-re spillel, de az API ugyanaz maradt
    assert!(v.spilled());
    println!("{:?}", v);
}

Arena allokáció bemutatása a Bumpalo crate-tel

Míg a SmallVec egy-egy vektor allokációját spórolja meg, a bumpalo egy egészen más filozófiát képvisel: az arena allokációt. Az ötlet az, hogy egy nagy memóriablokkot (arénát) kérünk le egyszer a rendszertől, aztán minden további "allokáció" csak egy pointer előreléptetése az arénán belül ("bump" allokáció) – rendkívül gyors, mert nincs se lock, se bookkeeping, se free-lista keresés.

A cserébe: az arénán belüli objektumokat egyenként nem lehet felszabadítani. Amikor végzel, az egész arénát egyszerre dobod el (drop), vagy reset()-eled, ha újra akarod használni.

[dependencies]
bumpalo = "3.14"

use bumpalo::Bump;

struct Node<'a> {
    value: i32,
    children: Vec<&'a Node<'a>>,
}

fn build_tree<'a>(bump: &'a Bump) -> &'a Node<'a> {
    let leaf1 = bump.alloc(Node { value: 1, children: vec![] });
    let leaf2 = bump.alloc(Node { value: 2, children: vec![] });
    let root = bump.alloc(Node {
        value: 0,
        children: vec![leaf1, leaf2],
    });
    root
}

fn main() {
    let bump = Bump::new();
    let tree = build_tree(&bump);
    println!("root value: {}", tree.value);
    // amikor bump kikerül a scope-ból, minden node egyszerre szabadul fel
}

Az alloc hívás egy &'a T-t ad vissza, ahol a lifetime a Bump-hoz van kötve. Ez elegáns megoldás: a borrow checker garantálja, hogy nem tudsz az arénán túl élő referenciát tartani.

Gyakorlati példa: egy parser allokációinak csökkentése

Nézzünk egy egyszerűsített tokenizert, amely egy szöveget bont fel szavakra és számokra. A "naiv" verzió minden tokenhez heap-allokált Vec<Token>-et épít, illetve minden token belsejében is allokálhat.

#[derive(Debug)]
enum Token {
    Word(String),
    Number(i64),
}

fn tokenize_naive(input: &str) -> Vec<Token> {
    let mut tokens = Vec::new();
    for word in input.split_whitespace() {
        if let Ok(n) = word.parse::<i64>() {
            tokens.push(Token::Number(n));
        } else {
            tokens.push(Token::Word(word.to_string()));
        }
    }
    tokens
}

Ebben két helyen allokálunk feleslegesen sokat:

1. A Vec::new() a legtöbb rövid inputnál csak néhány tokent tartalmaz, mégis heap-re megy azonnal az első push-tól.
2. A word.to_string() minden szóhoz külön heap-allokációt csinál.

Cseréljük ki a token-listát SmallVec-re, és a szó-tárolást intézzük a Bumpalo arénával, hogy a Token::Word egy &str-t tartalmazzon, ami az arénába mutat:

use bumpalo::Bump;
use smallvec::SmallVec;

#[derive(Debug)]
enum Token<'a> {
    Word(&'a str),
    Number(i64),
}

fn tokenize_fast<'a>(input: &'a str, bump: &'a Bump) -> SmallVec<[Token<'a>; 16]> {
    let mut tokens: SmallVec<[Token<'a>; 16]> = SmallVec::new();
    for word in input.split_whitespace() {
        if let Ok(n) = word.parse::<i64>() {
            tokens.push(Token::Number(n));
        } else {
            // az arénából allokálunk egy másolatot, nem a heap-ről
            let stored: &'a str = bump.alloc_str(word);
            tokens.push(Token::Word(stored));
        }
    }
    tokens
}

fn main() {
    let bump = Bump::new();
    let input = "alma 42 korte 7 dinnye";
    let tokens = tokenize_fast(input, &bump);
    println!("{:?}", tokens);
}

Itt két optimalizálás történt egyszerre:

• A tokenlista maga max 16 elemig a stacken él (SmallVec<[Token; 16]>), csak nagyon hosszú bemenetnél spillel a heap-re.
• Az egyes Word variánsok szövege nem String (ami saját heap-allokációt csinál), hanem bump.alloc_str – ami az aréna egyetlen nagy blokkjából oszt le helyet, nem külön allokációkkal.

Mérés előtte-utána Criterion benchmarkkal

Sose higgy a "nyilván gyorsabb" érvnek – mérj! A Criterion a de facto standard benchmark keretrendszer Rustban.

[dev-dependencies]
criterion = "0.5"

[[bench]]
name = "tokenize_bench"
harness = false

use criterion::{black_box, criterion_group, criterion_main, Criterion};
use bumpalo::Bump;

fn bench_naive(c: &mut Criterion) {
    let input = "alma 42 korte 7 dinnye barack 100 szilva 3 cseresznye";
    c.bench_function("tokenize_naive", |b| {
        b.iter(|| tokenize_naive(black_box(input)))
    });
}

fn bench_fast(c: &mut Criterion) {
    let input = "alma 42 korte 7 dinnye barack 100 szilva 3 cseresznye";
    c.bench_function("tokenize_fast", |b| {
        b.iter(|| {
            let bump = Bump::new();
            tokenize_fast(black_box(input), &bump)
        })
    });
}

criterion_group!(benches, bench_naive, bench_fast);
criterion_main!(benches);

Egy tipikus, kisebb bemenetekre futtatott mérésnél (a saját gépemen, egy hétköznapi laptopon mérve, tehát nem szentírás, csak illusztráció) valahogy így nézhet ki a kép:

• tokenize_naive: ~380 ns / iteráció, allokációk száma: token-lista + minden szóra egy String
• tokenize_fast: ~95 ns / iteráció, allokáció: csak az aréna első blokkja (amortizálva a további hívások között, ha újrahasznosítod a Bump-ot)

A különbség forrása nem varázslat: egyszerűen kevesebb malloc/free hívás történik, és a memória-hozzáférés is lokalizáltabb, ami jobban kedvez a CPU cache-nek.

Összefoglalás

A heap allokáció elkerülése nem premature optimization, ha a profilod ezt indokolja. A SmallVec ideális, amikor tudod, hogy a kollekcióid túlnyomó része kicsi marad, és nem akarsz minden apró vektorért a heap-hez fordulni. A Bumpalo arénák pedig ott jönnek jól, ahol sok rövid életű, egyszerre felszabadítható objektumot hozol létre – tipikusan parserekben, AST-építésnél, vagy batch-jellegű feldolgozásoknál.

A lényeg: ne találgass, hanem mérj Criterionnal, nézd meg a perf stat-ban az allokációk számát, és csak akkor nyúlj ezekhez az eszközökhöz, ha a számok ezt indokolják. A borrow checker és a lifetime-rendszer itt is a barátod – garantálja, hogy az arénás optimalizálás ne vezessen use-after-free hibákhoz, amit egy C++ programozónak külön kellene kézzel garantálnia.