Simulating worst case for linux kernel RSS memory optimisation update.

main.rs

#!/usr/bin/env rust-script
//! ```cargo
//! [dependencies]
//! clap = { version = "4.4.9", features = ["std", "derive", "help"] }
//! memmap = "0.7.0"
//! procfs = { version = "0.16.0", default-features = false }
//! thousands = "0.2.0"
//! ```
///  
/// This app can be executed with `rust-script` tool like this:
/// rust-script script.rs -t 512 -p 256 -i 63
///  
use std::fs::read_to_string;
use std::sync::{Arc, Barrier};
use std::thread;
use std::time::Duration;

use clap::Parser;
use memmap::MmapOptions;
use thousands::Separable;

const PAGE_SIZE: usize = 4096;

#[derive(Parser)]
#[command()]
struct Conf {
    #[arg(short = 't', default_value_t = 32, help = "")]
    number_of_threads: usize,
    #[arg(
        short,
        default_value_t = 128,
        help = "number of allocated pages per thread, and this number must be >= iter_limit"
    )]
    pages_chunk_per_thread: usize,
    #[arg(short, default_value_t = 1, help = "number of pages to touch")]
    iter_limit: usize,

    #[arg(
        short,
        default_value_t = 0,
        help = "sleep time before each thread continue execution"
    )]
    sleep_time: u64,

    #[arg(skip)]
    chunk_size: usize,
}

fn read_statm_rss() -> i64 {
    let me = procfs::process::Process::myself().unwrap();
    let stat = me.stat().unwrap();

    stat.rss as i64 * PAGE_SIZE as i64
}

fn read_smap_rollups() -> i64 {
    let s = read_to_string("/proc/self/smaps_rollup").unwrap();
    let mut i = s.lines();

    i.next().unwrap();
    let rss_line = i.next().unwrap();
    let rss = rss_line
        .split_whitespace()
        .nth(1)
        .unwrap()
        .parse::<i64>()
        .unwrap()
        * 1024;
    rss
}

fn worker(thread_id: usize, barrier: Arc<Barrier>, mem_chunk: &mut [u8], cfg: &Conf) {
    thread::sleep(Duration::from_millis(cfg.sleep_time) * thread_id as u32);
    let f = || {
        let rss_1 = read_statm_rss();
        let rss_2 = read_smap_rollups();
        let diff = rss_2 - rss_1;
        println!(
            "/proc/self/stat RSS: {}\n/proc/self/smaps_rollup RSS: {}\ndiff: {}",
            rss_1.separate_with_underscores(),
            rss_2.separate_with_underscores(),
            diff.separate_with_underscores(),
        );
    };

    let available_pages = mem_chunk.len() / PAGE_SIZE;

    for i in 0..available_pages {
        let pos = i * PAGE_SIZE;
        if i >= cfg.iter_limit {
            break;
        }

        mem_chunk[pos] = 1;
    }

    barrier.wait();

    if thread_id == 0 {
        f();
    }

    barrier.wait();
}

fn main() {
    let mut cli_m = Conf::parse();
    cli_m.chunk_size = cli_m.pages_chunk_per_thread * PAGE_SIZE;

    let cli = cli_m;
    assert!(cli.pages_chunk_per_thread >= cli.iter_limit);

    let barrier: Arc<Barrier> = Arc::new(Barrier::new(cli.number_of_threads));

    let mut map = MmapOptions::new()
        .len(cli.chunk_size * cli.number_of_threads)
        .map_anon()
        .unwrap();

    thread::scope(|s| {
        let mut chunk = map.chunks_mut(cli.chunk_size);
        for i in 0..cli.number_of_threads {
            let b = Arc::clone(&barrier);

            let y: &mut [u8] = chunk.next().unwrap();
            let cfg = &cli;
            s.spawn(move || worker(i, b, y, cfg));
        }
    });
}