hipthread 0.1.3

#![no_std]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]

use core::mem::{self, MaybeUninit};
use core::ffi::c_char;
use hierr::{Error, Result};

mod pthread;
mod stdlib;

mod thrd;
pub use thrd::*;

mod local;
pub use local::*;

mod mutex;
pub use mutex::*;

mod sem;
pub use sem::*;

mod once;
pub use once::*;

/// 获取当前线程的唯一标识.
pub fn thrd_self() -> u64 {
    unsafe { pthread::pthread_self() as u64 }
}

/// 获取当前进程可用的cpu核数
pub fn sched_cpu_count() -> Result<usize> {
    let mut cpuset = MaybeUninit::<pthread::cpu_set_t>::zeroed();
    let size = mem::size_of_val(&cpuset);
    let ret =
        unsafe { pthread::sched_getaffinity(pthread::getpid(), size, cpuset.assume_init_mut()) };
    if ret == 0 {
        Ok((unsafe { pthread::CPU_COUNT(cpuset.assume_init_ref()) }) as usize)
    } else {
        Err(Error::last())
    }
}

/// 获取当前进程同cpu核的亲和性.
pub fn sched_getaffinity(cpus: &mut [usize]) -> Result<&[usize]> {
    let mut cpuset = MaybeUninit::<pthread::cpu_set_t>::zeroed();
    let size = mem::size_of_val(&cpuset);
    let ret =
        unsafe { pthread::sched_getaffinity(pthread::getpid(), size, cpuset.assume_init_mut()) };
    if ret != 0 {
        return Err(Error::last());
    }
    let cnt = cpus
        .len()
        .min(unsafe { pthread::CPU_COUNT(cpuset.assume_init_ref()) } as usize);
    if cnt == 0 {
        return Ok(&[]);
    }

    let mut idx = 0;
    for n in 0.. {
        if unsafe { pthread::CPU_ISSET(n, cpuset.assume_init_ref()) } {
            cpus[idx] = n;
            idx += 1;
            if idx == cnt {
                return Ok(&cpus[..cnt]);
            }
        }
    }
    unreachable!();
}

/// 设置当前进程同cpu核的亲和性.
pub fn sched_setaffinity(cpus: &[usize]) -> Result<()> {
    let mut cpuset = MaybeUninit::<pthread::cpu_set_t>::zeroed();
    for cpu in cpus {
        unsafe {
            pthread::CPU_SET(*cpu, cpuset.assume_init_mut());
        }
    }
    let size = mem::size_of_val(&cpuset);
    let ret =
        unsafe { pthread::sched_setaffinity(pthread::getpid(), size, cpuset.assume_init_ref()) };
    if ret == 0 {
        Ok(())
    } else {
        Err(Error::last())
    }
}

/// 设置当前线程核cpu核的亲和性. 输入值范围为[0, sched_cpu_count()).
pub fn thrd_setaffinity(cpu: usize) -> Result<()> {
    let mut cpuset = MaybeUninit::<pthread::cpu_set_t>::zeroed();
    let size = mem::size_of_val(&cpuset);
    let ret =
        unsafe { pthread::sched_getaffinity(pthread::getpid(), size, cpuset.assume_init_mut()) };
    if ret != 0 {
        return Err(Error::last());
    }
    let cnt = unsafe { pthread::CPU_COUNT(cpuset.assume_init_ref()) };
    if cpu >= cnt as usize {
        return Err(Error::inval());
    }

    let mut idx = 0;
    for n in 0.. {
        if unsafe { pthread::CPU_ISSET(n, cpuset.assume_init_ref()) } {
            if idx == cpu {
                unsafe { pthread::CPU_ZERO(cpuset.assume_init_mut()) };
                unsafe { pthread::CPU_SET(n, cpuset.assume_init_mut()) };
                // 0表示为当前线程.
                let ret = unsafe { pthread::sched_setaffinity(0, size, cpuset.assume_init_ref()) };
                if ret == 0 {
                    return Ok(());
                } else {
                    return Err(Error::last());
                }
            }
            idx += 1;
        }
    }
    unreachable!();
}

pub fn thrd_setname(name: &str) {
    let mut n = [0_u8; 16];
    let len = name.len().min(n.len());
    (&mut n[..len]).copy_from_slice(&name.as_bytes()[..len]);
    n[15] = 0;
    unsafe { pthread::pthread_setname_np(pthread::pthread_self(), n.as_ptr().cast::<c_char>()); }
}

pub fn thrd_getname(name: &mut [u8]) -> &str {
    if name.is_empty() {
        return "";
    }
    let mut n = [0_u8; 16];
    unsafe { pthread::pthread_getname_np(pthread::pthread_self(), n.as_mut_ptr().cast::<c_char>(), n.len()) };
    let len = unsafe { strlen(n.as_ptr()) as usize };
    let len = len.min(name.len());
    (&mut name[..len]).copy_from_slice(&n[..len]);
    ::core::str::from_utf8(&name[..len]).unwrap_or("<unknown>")
}

extern "C" {
    fn strlen(s: *const u8) -> usize;
}

#[cfg(test)]
mod test {
    use crate::*;
    extern crate std;
    use std::println;

    #[test]
    fn test_name() {
        let name = "hello world";
        thrd_setname(name);
        let mut buf = [0_u8; 16];
        let n = thrd_getname(&mut buf);
        assert_eq!(n, name);

        let n = thrd_getname(&mut []);
        assert_eq!(n, "");

        let n = thrd_getname(&mut buf[..1]);
        assert_eq!(n, "h");

        let n = thrd_getname(&mut buf[..5]);
        assert_eq!(n, "hello");
    }

    #[test]
    fn test_sched() {
        test_thrd_affinity();
        test_process_affinity();
    }

    fn test_thrd_affinity() {
        let cnt = sched_cpu_count().unwrap();
        assert!(cnt > 0);
        for n in 0..cnt {
            let ret = thrd_setaffinity(n);
            assert!(ret.is_ok());
            let new_cnt = sched_cpu_count().unwrap();
            assert_eq!(cnt, new_cnt);
        }
        let ret = thrd_setaffinity(cnt);
        assert!(ret.is_err());
        let new_cnt = sched_cpu_count().unwrap();
        assert_eq!(cnt, new_cnt);
        println!("test_thrd_setaffinity - sched_cpu_count: {cnt}");
    }

    fn test_process_affinity() {
        let cnt = sched_cpu_count().unwrap();
        assert!(cnt > 0);

        let mut cpus = [0_usize; 24];
        let ret = sched_getaffinity(&mut cpus);
        assert!(ret.is_ok());
        assert_eq!(ret.unwrap().len(), cnt);

        println!("sched_getaffinity: {:?}", ret.unwrap());

        let ret = sched_setaffinity(&mut[]);
        assert!(ret.is_err());

        if cnt == 1 {
            return;
        }

        let ret = sched_setaffinity(&mut cpus[..cnt - 1]);
        assert!(ret.is_ok());

        let ret = sched_getaffinity(&mut cpus);
        assert!(ret.is_ok());
        assert_eq!(ret.unwrap().len(), cnt - 1);

        let ret = sched_getaffinity(&mut cpus[..1]);
        assert!(ret.is_ok());
        assert_eq!(ret.unwrap().len(), 1);
    }

}