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

// A micro-benchmark that approximates the gVisor getdents
// micro-benchmarks. We create a directory containing N files, then
// measure how quickly we can read all directory entries either by
// reusing the same FD each time (GetdentsSameFD) or by opening a new FD
// per iteration (GetdentsNewFD).

use std::{
    env,
    fs::{self, File},
    path::{Path, PathBuf},
    time::SystemTime,
};

use brunch::{benches, Bench};
use libc::{c_void, SYS_getdents64};
use nix::{
    fcntl::{open, OFlag},
    sys::stat::Mode,
    unistd::{lseek, unlinkat, Whence},
};

/// We use a fixed buffer size of 65536 bytes, as in the original benchmark.
const BUFFER_SIZE: usize = 65536;

/// Create a directory (under `env::temp_dir()`) and fill it with `count` files.
/// Returns the path to the directory and the list of file names (not absolute).
fn create_directory(count: usize) -> (PathBuf, Vec<String>) {
    // Make a unique top-level directory in /tmp.
    let mut dir = env::temp_dir();
    let unique = format!(
        "syd_getdents_bench_count_{}_{}",
        count,
        SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    );
    dir.push(unique);

    fs::create_dir_all(&dir).unwrap_or_else(|_| panic!("Failed to create directory: {:?}", &dir));

    // Open the directory FD with O_DIRECTORY so we can do unlinkat later.
    let dfd = open(&dir, OFlag::O_RDONLY | OFlag::O_DIRECTORY, Mode::empty())
        .expect("Failed to open directory FD");

    // Create N files in this directory.
    let mut files = Vec::with_capacity(count);
    for i in 0..count {
        let filename = format!("file_{}", i);
        // Use mknod (via nix) to create a regular file. Alternatively, just do `File::create`.
        // But we'll follow the style of the original test as closely as we can.
        let file_path = dir.join(&filename);
        File::create(&file_path)
            .unwrap_or_else(|_| panic!("Failed to create file: {:?}", &file_path));
        files.push(filename);
    }

    // Close the FD so it doesn't remain open for the entire run—some benchmarks
    // want to open/close in different ways.
    let _ = nix::unistd::close(dfd);
    (dir, files)
}

/// Remove all `files` from `dir`, then remove `dir` itself.
fn cleanup_directory(dir: &Path, files: &[String]) {
    // We re-open the directory with O_DIRECTORY for unlinkat.
    let dfd = open(dir, OFlag::O_RDONLY | OFlag::O_DIRECTORY, Mode::empty());
    if let Ok(dfd) = dfd {
        for fname in files {
            let _ = unlinkat(
                Some(dfd),
                Path::new(fname),
                nix::unistd::UnlinkatFlags::NoRemoveDir,
            );
        }
        // For the directory itself:
        let _ = nix::unistd::close(dfd);
    }
    // Finally remove the top-level directory.
    let _ = fs::remove_dir_all(dir);
}

/// For "GetdentsSameFD", we open the directory once, lseek to 0 before each read,
/// and read until no more entries are returned.
fn getdents_same_fd(dir: &Path) {
    // Open once (like the original code). Then each iteration does lseek + read.
    let fd = open(dir, OFlag::O_RDONLY | OFlag::O_DIRECTORY, Mode::empty())
        .expect("Failed to open directory (same-FD)");
    let mut buffer = vec![0_u8; BUFFER_SIZE];

    // Seek back to the start of the directory.
    lseek(fd, 0, Whence::SeekSet).expect("Failed to lseek to start");

    // Repeatedly call getdents until ret == 0.
    loop {
        let ret = unsafe {
            libc::syscall(
                SYS_getdents64,
                fd,
                buffer.as_mut_ptr() as *mut c_void,
                BUFFER_SIZE,
            )
        };
        if ret < 0 {
            // If we get an error, bail out. We mimic SyscallSucceeds() -> expect OK.
            panic!("getdents_same_fd: syscall error, returned {}", ret);
        } else if ret == 0 {
            // No more entries to read.
            break;
        }
        // Keep reading until empty.
    }

    // Clean up.
    nix::unistd::close(fd).expect("Failed to close directory FD");
}

/// For "GetdentsNewFD", we open the directory fresh on each iteration, read all
/// entries, and then close it.
fn getdents_new_fd(dir: &Path) {
    // Open fresh.
    let fd = open(dir, OFlag::O_RDONLY | OFlag::O_DIRECTORY, Mode::empty())
        .expect("Failed to open directory (new-FD)");
    let mut buffer = vec![0_u8; BUFFER_SIZE];

    loop {
        let ret = unsafe {
            libc::syscall(
                SYS_getdents64,
                fd,
                buffer.as_mut_ptr() as *mut c_void,
                BUFFER_SIZE,
            )
        };
        if ret < 0 {
            panic!("getdents_new_fd: syscall error, returned {}", ret);
        } else if ret == 0 {
            break;
        }
    }

    // Close.
    nix::unistd::close(fd).expect("Failed to close directory FD");
}

fn main() {
    // We'll test several "counts" (number of files in the directory):
    let counts = [1_usize, 16, 256, 4096];

    // Prepare setups for "SameFD" approach:
    let mut same_fd_setups = Vec::new();
    for &count in &counts {
        let (dir, files) = create_directory(count);
        same_fd_setups.push((count, dir, files));
    }

    // Prepare setups for "NewFD" approach:
    let mut new_fd_setups = Vec::new();
    for &count in &counts {
        let (dir, files) = create_directory(count);
        new_fd_setups.push((count, dir, files));
    }

    // We'll define separate benches for each count and each approach.
    benches!(
        inline:

        // -- SAME FD BENCHMARKS --
        Bench::new("GetdentsSameFD(count=1)").run(|| {
            getdents_same_fd(&same_fd_setups[0].1);
        }),
        Bench::new("GetdentsSameFD(count=16)").run(|| {
            getdents_same_fd(&same_fd_setups[1].1);
        }),
        Bench::new("GetdentsSameFD(count=256)").run(|| {
            getdents_same_fd(&same_fd_setups[2].1);
        }),
        Bench::new("GetdentsSameFD(count=4096)").run(|| {
            getdents_same_fd(&same_fd_setups[3].1);
        }),

        // -- NEW FD BENCHMARKS --
        Bench::new("GetdentsNewFD(count=1)").run(|| {
            getdents_new_fd(&new_fd_setups[0].1);
        }),
        Bench::new("GetdentsNewFD(count=16)").run(|| {
            getdents_new_fd(&new_fd_setups[1].1);
        }),
        Bench::new("GetdentsNewFD(count=256)").run(|| {
            getdents_new_fd(&new_fd_setups[2].1);
        }),
        Bench::new("GetdentsNewFD(count=4096)").run(|| {
            getdents_new_fd(&new_fd_setups[3].1);
        }),
    );

    // Cleanup after benchmarks.
    for (_, dir, files) in same_fd_setups {
        cleanup_directory(&dir, &files);
    }
    for (_, dir, files) in new_fd_setups {
        cleanup_directory(&dir, &files);
    }
}
