syd 3.52.0

//
// Syd: rock-solid application kernel
// benches/sys/fork.rs: fork microbenchmarks
//
// Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
// Based in part upon gVisor's fork_benchmark.cc which is:
//   Copyright 2020 The gVisor Authors.
//   SPDX-License-Identifier: Apache-2.0
//
// SPDX-License-Identifier: GPL-3.0

// This replicates the gVisor "fork" (and related) micro-benchmarks, including:
//   1) BM_CPUBoundUniprocess
//   2) BM_CPUBoundAsymmetric
//   3) BM_CPUBoundSymmetric
//   4) BM_ProcessSwitch
//   5) BM_ThreadSwitch
//   6) BM_ThreadStart
//   7) BM_ProcessLifecycle

use std::{
    hint::black_box,
    sync::{Arc, Barrier},
    thread,
    time::Duration,
};

use brunch::{benches, Bench};
use libc::{_exit, c_int, close, fork, pipe, read, waitpid, write, WEXITSTATUS, WIFEXITED};
use nix::errno::Errno;

/// A little CPU-bound "busy" function, mimicking gVisor's prime-like loops.
fn busy(max: i32) -> i32 {
    // Prevent the compiler from optimizing this out:
    let mut count = 0;
    for i in 1..max {
        for j in 2..(i / 2) {
            if i % j == 0 {
                count += 1;
            }
        }
    }
    // Use black_box to ensure the result isn't optimized away.
    black_box(count)
}

/// 1) CPU-bound uniprocess: Just run busy() in the same process.
fn bm_cpubound_uniprocess() {
    busy(250);
}

/// 2) CPU-bound Asymmetric: One fork child does all the busy() calls, while
///    the parent calls KeepRunningBatch, then waits for the child to exit.
fn bm_cpubound_asymmetric(iterations: usize) {
    unsafe {
        let child = fork();
        if child == 0 {
            // Child: do all the busy-loops, then _exit.
            for _ in 0..iterations {
                busy(250);
            }
            _exit(0);
        } else if child < 0 {
            panic!("fork() failed");
        } else {
            // Parent: keep "running" until child's loops are done, then wait.
            // In the gVisor code, they do KeepRunningBatch(max). We'll emulate it
            // by just letting the child do the heavy lifting. The parent just
            // waits below.
            let mut status: c_int = 0;
            let w = waitpid(child, &mut status as *mut c_int, 0);
            if w < 0 {
                panic!("waitpid() failed: {:?}", Errno::last());
            }
            if WIFEXITED(status) && WEXITSTATUS(status) == 0 {
                // Ok
            } else {
                panic!("Child did not exit(0).");
            }
        }
    }
}

/// 3) CPU-bound Symmetric: We fork N processes, dividing total iterations
///    among them. Each child does `cur` busy-loops and exits. The parent
///    calls KeepRunningBatch(cur) for each child that actually runs.
fn bm_cpubound_symmetric(procs: usize, max_iters: usize) {
    let mut children = Vec::new();
    let mut total_done = 0;

    // Distribute the total iterations among `procs`.
    for _ in 0..procs {
        // The next child will handle up to "remaining / #children_left".
        let remaining = max_iters - total_done;
        if remaining == 0 {
            break;
        }

        // Round up if needed:
        let cur = remaining / (procs - children.len());
        let cur = if cur == 0 { remaining } else { cur };
        total_done += cur;

        unsafe {
            let child = fork();
            if child == 0 {
                // Child
                for _ in 0..cur {
                    busy(250);
                }
                _exit(0);
            } else if child < 0 {
                panic!("fork() failed in symmetric");
            } else {
                // Parent
                if cur > 0 {
                    // Emulate KeepRunningBatch(cur). We'll just pretend we used
                    // up those iterations in the parent's benchmark loop.
                }
                children.push(child);
            }
        }
    }

    // Wait for them all.
    unsafe {
        for &ch in &children {
            let mut status: c_int = 0;
            let w = waitpid(ch, &mut status, 0);
            if w < 0 {
                panic!("waitpid() failed");
            }
            if WIFEXITED(status) && WEXITSTATUS(status) == 0 {
                // Ok
            } else {
                panic!("Child did not exit(0).");
            }
        }
    }
}

/// A helper that just runs the read->write loop in a child or thread, until
/// we can't read anymore.
fn switch_child_loop(read_fd: c_int, write_fd: c_int) {
    let mut buf = [0u8; 1];
    loop {
        let n = unsafe { read(read_fd, buf.as_mut_ptr() as *mut _, 1) };
        if n == 0 {
            // EOF
            break;
        } else if n < 0 {
            // read error
            let e = Errno::last();
            panic!("Child read() error: {:?}", e);
        }
        // Now write the same byte out.
        let w = unsafe { write(write_fd, buf.as_ptr() as *const _, 1) };
        if w < 0 {
            // If EPIPE, the chain is done
            let e = Errno::last();
            if e == Errno::EPIPE {
                break;
            }
            panic!("Child write() error: {:?}", e);
        }
        if w == 0 {
            break;
        }
    }
}

/// 4) BM_ProcessSwitch: We form a ring of processes and pipes, passing a
///    single byte around among them to measure context-switch overhead.
fn bm_process_switch(num_processes: usize, iterations: usize) {
    if num_processes < 2 {
        return; // must have >=2
    }
    // Create pipes (read_fds[i], write_fds[i]) for i in [0..num_processes].
    let mut read_fds = Vec::with_capacity(num_processes);
    let mut write_fds = Vec::with_capacity(num_processes);

    unsafe {
        // First pipe belongs to this process (index 0).
        for _ in 0..num_processes {
            let mut fds = [0; 2];
            if pipe(fds.as_mut_ptr()) < 0 {
                panic!("pipe() failed");
            }
            read_fds.push(fds[0]);
            write_fds.push(fds[1]);
        }

        let mut children = Vec::new();
        // We already "are" process index 0. We'll fork the other processes.
        for i in 1..num_processes {
            let read_index = i;
            let write_index = (i + 1) % num_processes;
            let child = fork();
            if child == 0 {
                // Child
                // Close all other fds except read_index, write_index
                for j in 0..num_processes {
                    if j != read_index {
                        close(read_fds[j]);
                    }
                    if j != write_index {
                        close(write_fds[j]);
                    }
                }
                switch_child_loop(read_fds[read_index], write_fds[write_index]);
                _exit(0);
            } else if child < 0 {
                panic!("fork() failed in BM_ProcessSwitch");
            } else {
                children.push(child);
            }
        }

        // Now in the parent (index 0):
        // We'll read from read_fds[0], write to write_fds[1].
        let read_idx = 0;
        let write_idx = 1;

        // Kickstart: write one byte to write_idx
        let mut c = [b'a'];
        if write(write_fds[write_idx], c.as_ptr() as *const _, 1) != 1 {
            panic!("initial write failed");
        }

        // Do the loop for "iterations".
        for _ in 0..iterations {
            if read(read_fds[read_idx], c.as_mut_ptr() as *mut _, 1) != 1 {
                panic!("read in parent failed");
            }
            if write(write_fds[write_idx], c.as_ptr() as *const _, 1) != 1 {
                panic!("write in parent failed");
            }
        }

        // Close everything so children exit.
        for i in 0..num_processes {
            close(read_fds[i]);
            close(write_fds[i]);
        }

        // Wait for children
        for &ch in &children {
            let mut status: c_int = 0;
            if waitpid(ch, &mut status, 0) < 0 {
                panic!("waitpid failed in BM_ProcessSwitch");
            }
            if !WIFEXITED(status) || WEXITSTATUS(status) != 0 {
                panic!("child exit code not 0");
            }
        }
    }
}

/// 5) BM_ThreadSwitch: same ring approach, but with threads instead of processes.
fn bm_thread_switch(num_threads: usize, iterations: usize) {
    if num_threads < 2 {
        return;
    }

    // We create `num_threads` pipes, then spawn threads 1..num_threads. The main
    // thread is index 0.
    let mut read_fds = Vec::new();
    let mut write_fds = Vec::new();

    // Each pipe is used by exactly one "slot".
    unsafe {
        for _ in 0..num_threads {
            let mut fds = [0; 2];
            if pipe(fds.as_mut_ptr()) < 0 {
                panic!("pipe() failed for thread_switch");
            }
            read_fds.push(fds[0]);
            write_fds.push(fds[1]);
        }
    }

    let mut handles = Vec::with_capacity(num_threads - 1);

    // For thread i from 1..num_threads:
    for i in 1..num_threads {
        // read from read_idx = i, write to write_idx = (i + 1) % num_threads
        let read_idx = i;
        let write_idx = (i + 1) % num_threads;
        let rfd = read_fds[read_idx];
        let wfd = write_fds[write_idx];

        // Move fd ownership into the thread
        let handle = thread::spawn(move || {
            switch_child_loop(rfd, wfd);
            // Close at the end to ensure no leaks
            unsafe {
                close(rfd);
                close(wfd);
            }
        });
        handles.push(handle);
    }

    // The main thread is index 0:
    let read_idx = 0;
    let write_idx = 1;
    // Kickstart:
    let c = [b'a'];
    unsafe {
        if write(write_fds[write_idx], c.as_ptr() as *const _, 1) != 1 {
            panic!("thread main initial write failed");
        }
    }

    // Loop for "iterations".
    let mut c = [0u8; 1];
    for _ in 0..iterations {
        unsafe {
            if read(read_fds[read_idx], c.as_mut_ptr() as *mut _, 1) != 1 {
                panic!("thread main read failed");
            }
            if write(write_fds[write_idx], c.as_ptr() as *const _, 1) != 1 {
                panic!("thread main write failed");
            }
        }
    }

    // Close main's fds to kill the ring.
    unsafe {
        close(read_fds[read_idx]);
        close(write_fds[write_idx]);
    }

    // Join all threads.
    for h in handles {
        let _ = h.join();
    }
}

/// 6) BM_ThreadStart: repeatedly create N threads that do basically nothing
///    except wait on a barrier, then the main thread rejoins them.
fn bm_thread_start(num_threads: usize, iterations: usize) {
    for _ in 0..iterations {
        // We'll barrier with (num_threads + 1) total.
        let barrier = Arc::new(Barrier::new(num_threads + 1));

        // Spawn N threads:
        let mut threads = Vec::with_capacity(num_threads);
        for _ in 0..num_threads {
            let b = barrier.clone();
            threads.push(thread::spawn(move || {
                // Wait on the barrier; after the last arrives, barrier is destroyed
                b.wait();
            }));
        }

        // Main thread also waits:
        barrier.wait();

        // Join them all:
        for t in threads {
            let _ = t.join();
        }
    }
}

/// 7) BM_ProcessLifecycle: fork + exit + wait, repeated for `num_procs` procs each iteration.
fn bm_process_lifecycle(num_procs: usize, iterations: usize) {
    unsafe {
        let mut pids = Vec::with_capacity(num_procs);
        for _ in 0..iterations {
            pids.clear();
            for _i in 0..num_procs {
                let pid = fork();
                if pid == 0 {
                    _exit(0);
                } else if pid < 0 {
                    panic!("fork() failed in process_lifecycle");
                } else {
                    pids.push(pid);
                }
            }
            // Wait for them
            for &p in &pids {
                let mut status = 0;
                let w = waitpid(p, &mut status, 0);
                if w < 0 {
                    panic!("waitpid() failed in process_lifecycle");
                }
                if !WIFEXITED(status) || WEXITSTATUS(status) != 0 {
                    panic!("child exit code not 0 in process_lifecycle");
                }
            }
        }
    }
}

fn main() {
    benches!(
        inline:

        // 1) BM_CPUBoundUniprocess
        Bench::new("BM_CPUBoundUniprocess").run(|| {
            bm_cpubound_uniprocess();
        }),

        // 2) BM_CPUBoundAsymmetric
        // We'll pick an arbitrary iteration count, e.g. 100, for demonstration.
        Bench::new("BM_CPUBoundAsymmetric").run(|| {
            bm_cpubound_asymmetric(100);
        }),

        // 3) BM_CPUBoundSymmetric: We'll do 2..16 processes in separate benches.
        Bench::new("BM_CPUBoundSymmetric(2 procs)").run(|| {
            bm_cpubound_symmetric(2, 100);
        }),
        Bench::new("BM_CPUBoundSymmetric(4 procs)").run(|| {
            bm_cpubound_symmetric(4, 100);
        }),
        Bench::new("BM_CPUBoundSymmetric(8 procs)").run(|| {
            bm_cpubound_symmetric(8, 100);
        }),
        Bench::new("BM_CPUBoundSymmetric(16 procs)").run(|| {
            bm_cpubound_symmetric(16, 100);
        }),

        // 4) BM_ProcessSwitch: We'll do 2..16 processes with some iteration count, e.g. 1000.
        Bench::new("BM_ProcessSwitch(2 procs)").run(|| {
            bm_process_switch(2, 1000);
        }),
        Bench::new("BM_ProcessSwitch(4 procs)").run(|| {
            bm_process_switch(4, 1000);
        }),
        Bench::new("BM_ProcessSwitch(8 procs)").run(|| {
            bm_process_switch(8, 1000);
        }),
        Bench::new("BM_ProcessSwitch(16 procs)").run(|| {
            bm_process_switch(16, 1000);
        }),

        // 5) BM_ThreadSwitch: We'll do 2..16 threads, 1000 iterations.
        Bench::new("BM_ThreadSwitch(2 threads)").run(|| {
            bm_thread_switch(2, 1000);
        }),
        Bench::new("BM_ThreadSwitch(4 threads)").run(|| {
            bm_thread_switch(4, 1000);
        }),
        Bench::new("BM_ThreadSwitch(8 threads)").run(|| {
            bm_thread_switch(8, 1000);
        }),
        Bench::new("BM_ThreadSwitch(16 threads)").run(|| {
            bm_thread_switch(16, 1000);
        }),

        // 6) BM_ThreadStart: Range(1..2048)? We'll pick a few points.
        Bench::new("BM_ThreadStart(1)").run(|| {
            bm_thread_start(1, 10);
        }),
        Bench::new("BM_ThreadStart(64)")
            .with_timeout(Duration::from_secs(30))
            .run(|| {
                bm_thread_start(64, 10);
            }),
        Bench::new("BM_ThreadStart(128)")
            .with_timeout(Duration::from_secs(30))
            .run(|| {
                bm_thread_start(128, 10);
            }),
        Bench::new("BM_ThreadStart(1024)")
            .with_timeout(Duration::from_secs(30))
            .run(|| {
                bm_thread_start(1024, 10);
            }),

        // 7) BM_ProcessLifecycle: Range(1..512)? We'll pick a few points.
        Bench::new("BM_ProcessLifecycle(1 proc)").run(|| {
            bm_process_lifecycle(1, 10);
        }),
        Bench::new("BM_ProcessLifecycle(64 procs)")
            .with_timeout(Duration::from_secs(30))
            .run(|| {
                bm_process_lifecycle(64, 10);
            }),
        Bench::new("BM_ProcessLifecycle(128 procs)")
            .with_timeout(Duration::from_secs(60))
            .run(|| {
                bm_process_lifecycle(128, 10);
            }),
        Bench::new("BM_ProcessLifecycle(512 procs)")
            .with_timeout(Duration::from_secs(150))
            .run(|| {
                bm_process_lifecycle(512, 10);
            }),
    );
}