// Copyright 2022 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
//! This module provides unified APIs for accessing thread/process related
//! information. Only Linux platform is implemented correctly, for other
//! platform, it only guarantees successful compilation.

use std::{io, io::Result, sync::Mutex, thread};

use collections::HashMap;
use tikv_alloc::{
    add_thread_memory_accessor, remove_thread_memory_accessor, thread_allocate_exclusive_arena,
    thread_allocate_shared_exclusive_arena,
};

/// A cross-platform CPU statistics data structure.
#[derive(Debug, Copy, Clone, Default, PartialEq)]
pub struct ThreadStat {
    // libc::clock_t is not used here because the definition of
    // clock_t is different on linux and bsd.
    pub s_time: i64,
    pub u_time: i64,
}

impl ThreadStat {
    /// Gets the cpu time in seconds.
    #[inline]
    pub fn total_cpu_time(&self) -> f64 {
        cpu_total(self.s_time, self.u_time)
    }
}

#[inline]
fn cpu_total(sys_time: i64, user_time: i64) -> f64 {
    (sys_time + user_time) as f64 / ticks_per_second() as f64
}

#[cfg(target_os = "linux")]
pub mod linux {
    #[inline]
    pub fn cpu_total(stat: &super::FullStat) -> f64 {
        super::cpu_total(stat.stime, stat.utime)
    }
}

#[cfg(target_os = "linux")]
mod imp {
    use std::{
        fs,
        io::{self, Error},
        iter::FromIterator,
    };

    use libc::c_int;
    pub use libc::pid_t as Pid;
    pub use procinfo::pid::{self, Stat as FullStat};

    lazy_static::lazy_static! {
        // getconf CLK_TCK
        static ref CLOCK_TICK: i64 = {
            unsafe {
                libc::sysconf(libc::_SC_CLK_TCK)
            }
        };

        static ref PROCESS_ID: Pid = unsafe { libc::getpid() };
    }

    #[inline]
    pub fn ticks_per_second() -> i64 {
        *CLOCK_TICK
    }

    /// Gets the ID of the current process.
    #[inline]
    pub fn process_id() -> Pid {
        *PROCESS_ID
    }

    /// Gets the ID of the current thread.
    #[inline]
    pub fn thread_id() -> Pid {
        thread_local! {
            static TID: Pid = unsafe { libc::syscall(libc::SYS_gettid) as Pid };
        }
        TID.with(|t| *t)
    }

    /// Gets thread ids of the given process id.
    /// WARN: Don't call this function frequently. Otherwise there will be a lot
    /// of memory fragments.
    pub fn thread_ids<C: FromIterator<Pid>>(pid: Pid) -> io::Result<C> {
        let dir = fs::read_dir(format!("/proc/{}/task", pid))?;
        Ok(dir
            .filter_map(|task| {
                let file_name = match task {
                    Ok(t) => t.file_name(),
                    Err(e) => {
                        error!("read task failed"; "pid" => pid, "err" => ?e);
                        return None;
                    }
                };

                match file_name.to_str() {
                    Some(tid) => match tid.parse() {
                        Ok(tid) => Some(tid),
                        Err(e) => {
                            error!("read task failed"; "pid" => pid, "err" => ?e);
                            None
                        }
                    },
                    None => {
                        error!("read task failed"; "pid" => pid);
                        None
                    }
                }
            })
            .collect())
    }

    pub fn full_thread_stat(pid: Pid, tid: Pid) -> io::Result<FullStat> {
        pid::stat_task(pid, tid)
    }

    pub fn set_priority(pri: i32) -> io::Result<()> {
        // Unsafe due to FFI.
        unsafe {
            let tid = libc::syscall(libc::SYS_gettid);
            if libc::setpriority(libc::PRIO_PROCESS, tid as u32, pri) != 0 {
                let e = Error::last_os_error();
                return Err(e);
            }
            Ok(())
        }
    }

    pub fn get_priority() -> io::Result<i32> {
        // Unsafe due to FFI.
        unsafe {
            let tid = libc::syscall(libc::SYS_gettid);
            clear_errno();
            let ret = libc::getpriority(libc::PRIO_PROCESS, tid as u32);
            if ret == -1 {
                let e = Error::last_os_error();
                if let Some(errno) = e.raw_os_error() {
                    if errno != 0 {
                        return Err(e);
                    }
                }
            }
            Ok(ret)
        }
    }

    // Sadly the std lib does not have any support for setting `errno`, so we
    // have to implement this ourselves.
    extern "C" {
        #[link_name = "__errno_location"]
        fn errno_location() -> *mut c_int;
    }

    fn clear_errno() {
        // Unsafe due to FFI.
        unsafe {
            *errno_location() = 0;
        }
    }

    #[cfg(test)]
    mod tests {
        use std::io::ErrorKind;

        use super::*;
        use crate::sys::HIGH_PRI;

        #[test]
        fn test_set_priority() {
            // priority is a value in range -20 to 19, the default priority
            // is 0, lower priorities cause more favorable scheduling.
            assert_eq!(get_priority().unwrap(), 0);
            set_priority(10).unwrap();
            assert_eq!(get_priority().unwrap(), 10);

            // only users who have `SYS_NICE_CAP` capability can increase priority.
            let ret = set_priority(HIGH_PRI);
            if let Err(e) = ret {
                assert_eq!(e.kind(), ErrorKind::PermissionDenied);
            } else {
                assert_eq!(get_priority().unwrap(), HIGH_PRI);
            }
        }
    }
}

#[cfg(target_os = "macos")]
// `deprecated`: suppress warnings for libc::mach_* functions (libc recommends mach2 crate).
#[allow(bad_style, deprecated)]
mod imp {
    use std::{io, iter::FromIterator, mem::size_of, ptr::null_mut, slice};

    use libc::*;

    pub type Pid = mach_port_t;

    type task_inspect_t = mach_port_t;
    type thread_act_t = mach_port_t;
    type thread_act_array_t = *mut thread_act_t;

    extern "C" {
        fn task_threads(
            target_task: task_inspect_t,
            act_list: *mut thread_act_array_t,
            act_listCnt: *mut mach_msg_type_number_t,
        ) -> kern_return_t;
    }

    const MICRO_SEC_PER_SEC: i64 = 1_000_000;

    #[derive(Default)]
    pub struct FullStat {
        pub stime: i64,
        pub utime: i64,
        pub command: String,
    }

    /// Unlike Linux, the unit of `stime` and `utime` is microseconds instead of
    /// ticks. See [`full_thread_stat()`]
    #[inline]
    pub fn ticks_per_second() -> i64 {
        MICRO_SEC_PER_SEC
    }

    /// Gets the ID of the current process.
    #[inline]
    pub fn process_id() -> Pid {
        #[allow(deprecated)]
        unsafe {
            mach_task_self_
        }
    }

    /// Gets the ID of the current thread.
    #[inline]
    pub fn thread_id() -> Pid {
        #[allow(deprecated)]
        unsafe {
            mach_thread_self()
        }
    }

    pub fn thread_ids<C: FromIterator<Pid>>(pid: Pid) -> io::Result<C> {
        // https://www.gnu.org/software/hurd/gnumach-doc/Task-Information.html

        unsafe {
            let mut act_list: thread_act_array_t = null_mut();
            let mut act_count: mach_msg_type_number_t = 0;
            let ret = task_threads(pid, &mut act_list, &mut act_count);
            if ret != KERN_SUCCESS {
                return Err(io::Error::from_raw_os_error(ret));
            }

            let pids = slice::from_raw_parts_mut(act_list, act_count as _)
                .iter()
                .copied()
                .collect();

            vm_deallocate(
                pid,
                act_list as vm_address_t,
                size_of::<thread_act_t>() * act_count as usize,
            );

            Ok(pids)
        }
    }

    pub fn full_thread_stat(_pid: Pid, tid: Pid) -> io::Result<FullStat> {
        // https://www.gnu.org/software/hurd/gnumach-doc/Thread-Information.html

        unsafe {
            let flavor = THREAD_BASIC_INFO;
            let mut info = std::mem::zeroed::<thread_basic_info>();
            let mut thread_info_cnt = THREAD_BASIC_INFO_COUNT;

            let ret = thread_info(
                tid,
                flavor as task_flavor_t,
                (&mut info as *mut _) as thread_info_t,
                &mut thread_info_cnt,
            );
            if ret != KERN_SUCCESS {
                return Err(io::Error::from_raw_os_error(ret));
            }

            Ok(FullStat {
                stime: info.system_time.seconds as i64 * 1_000_000
                    + info.system_time.microseconds as i64,
                utime: info.user_time.seconds as i64 * 1_000_000
                    + info.user_time.microseconds as i64,
                ..Default::default()
            })
        }
    }

    pub fn set_priority(_: i32) -> io::Result<()> {
        Ok(())
    }

    pub fn get_priority() -> io::Result<i32> {
        Ok(0)
    }
}

#[cfg(not(any(target_os = "linux", target_os = "macos")))]
mod imp {
    use std::{io, iter::FromIterator};

    pub type Pid = u32;

    #[derive(Default)]
    pub struct FullStat {
        pub stime: i64,
        pub utime: i64,
        pub command: String,
    }

    lazy_static::lazy_static! {
        static ref PROCESS_ID: Pid = std::process::id();
    }

    #[inline]
    pub fn ticks_per_second() -> i64 {
        1
    }

    /// Gets the ID of the current process.
    #[inline]
    pub fn process_id() -> Pid {
        *PROCESS_ID
    }

    /// Gets the ID of the current thread.
    #[inline]
    pub fn thread_id() -> Pid {
        u64::from(std::thread::current().id().as_u64()) as Pid
    }

    pub fn thread_ids<C: FromIterator<Pid>>(_pid: Pid) -> io::Result<C> {
        Ok(std::iter::once(thread_id()).collect())
    }

    pub fn full_thread_stat(_pid: Pid, _tid: Pid) -> io::Result<FullStat> {
        Ok(FullStat::default())
    }

    pub fn set_priority(_: i32) -> io::Result<()> {
        Ok(())
    }

    pub fn get_priority() -> io::Result<i32> {
        Ok(0)
    }
}

pub use self::imp::*;

pub fn thread_stat(pid: Pid, tid: Pid) -> io::Result<ThreadStat> {
    let full_stat = full_thread_stat(pid, tid)?;
    Ok(ThreadStat {
        s_time: full_stat.stime as _,
        u_time: full_stat.utime as _,
    })
}

pub fn current_thread_stat() -> io::Result<ThreadStat> {
    thread_stat(process_id(), thread_id())
}

pub trait StdThreadBuildWrapper {
    /// Extends thread spawning functionality with additional control over arena
    /// sharing.
    ///
    /// # Arguments
    ///
    /// * `share_arena` - Boolean flag indicating whether to share the memory
    ///   arena with the new thread
    /// * `f` - The closure to be executed in the new thread
    fn spawn_wrapper_ext<F, T>(self, share_arena: bool, f: F) -> io::Result<thread::JoinHandle<T>>
    where
        F: FnOnce() -> T,
        F: Send + 'static,
        T: Send + 'static;

    fn spawn_wrapper<F, T>(self, f: F) -> io::Result<thread::JoinHandle<T>>
    where
        F: FnOnce() -> T,
        F: Send + 'static,
        T: Send + 'static;
}

pub trait ThreadBuildWrapper {
    fn after_start_wrapper<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static;

    fn before_stop_wrapper<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static;

    /// Register all system hooks along with a custom hook pair.
    fn with_sys_and_custom_hooks<F1, F2>(&mut self, after_start: F1, before_end: F2) -> &mut Self
    where
        F1: Fn() + Send + Sync + 'static,
        F2: Fn() + Send + Sync + 'static;

    /// Register some generic hooks like memory tracing or thread lifetime
    /// tracing.
    fn with_sys_hooks(&mut self) -> &mut Self {
        self.with_sys_and_custom_hooks(|| {}, || {})
    }
}

lazy_static::lazy_static! {
    pub static ref THREAD_NAME_HASHMAP: Mutex<HashMap<Pid, String>> = Mutex::new(HashMap::default());
    pub static ref THREAD_START_HOOKS: Mutex<Vec<Box<dyn Fn() + Sync + Send>>> = Mutex::new(Vec::new());
}

pub fn hook_thread_start(f: Box<dyn Fn() + Sync + Send>) {
    THREAD_START_HOOKS.lock().unwrap().push(f);
}

pub(crate) fn call_thread_start_hooks() {
    for f in THREAD_START_HOOKS.lock().unwrap().iter() {
        f();
    }
}

pub fn add_thread_name_to_map() {
    if let Some(name) = std::thread::current().name() {
        let tid = thread_id();
        THREAD_NAME_HASHMAP
            .lock()
            .unwrap()
            .insert(tid, name.to_string());
        debug!("tid {} thread name is {}", tid, name);
    }
}

pub(crate) fn remove_thread_name_from_map() {
    let tid = thread_id();
    THREAD_NAME_HASHMAP.lock().unwrap().remove(&tid);
}

/// Thread prefix for gRPC server worker threads.
pub const GRPC_SERVER_THREAD: &str = "grpc-server";

/// Thread prefix for unified read pool worker threads.
pub const UNIFIED_READ_POOL_THREAD: &str = "unified-read-pool";

const LINUX_THREAD_NAME_MAX_LEN: usize = 15;

/// Returns whether `thread_name` belongs to a thread family whose prefix is
/// `prefix`.
///
/// Linux `/proc` thread names are truncated to 15 visible bytes, so callers can
/// pass either a stable family prefix or the full thread prefix and still match
/// truncated names such as `unified-read-po`.
#[inline]
pub fn matches_thread_name_prefix(thread_name: &str, prefix: &str) -> bool {
    if thread_name.starts_with(prefix) {
        return true;
    }
    if prefix.len() <= LINUX_THREAD_NAME_MAX_LEN {
        return false;
    }
    let truncated_prefix = match prefix.get(..LINUX_THREAD_NAME_MAX_LEN) {
        Some(prefix) => prefix,
        None => prefix
            .char_indices()
            .map(|(index, _)| index)
            .take_while(|index| *index <= LINUX_THREAD_NAME_MAX_LEN)
            .last()
            .map(|index| &prefix[..index])
            .unwrap_or(""),
    };
    thread_name.starts_with(truncated_prefix)
}

impl StdThreadBuildWrapper for std::thread::Builder {
    fn spawn_wrapper_ext<F, T>(self, share_arena: bool, f: F) -> io::Result<thread::JoinHandle<T>>
    where
        F: FnOnce() -> T,
        F: Send + 'static,
        T: Send + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.spawn(move || {
            call_thread_start_hooks();
            // SAFETY: we will call `remove_thread_memory_accessor` at defer.
            unsafe { add_thread_memory_accessor() };
            if share_arena {
                // Binds this thread to a shared arena.
                thread_allocate_shared_exclusive_arena().unwrap();
            } else {
                thread_allocate_exclusive_arena().unwrap();
            }
            add_thread_name_to_map();
            defer! {{
                remove_thread_name_from_map();
                remove_thread_memory_accessor();
            }};
            f()
        })
    }

    fn spawn_wrapper<F, T>(self, f: F) -> Result<std::thread::JoinHandle<T>>
    where
        F: FnOnce() -> T,
        F: Send + 'static,
        T: Send + 'static,
    {
        self.spawn_wrapper_ext(false, f)
    }
}

/// Creates and spawns a new thread with an automatically generated name based
/// on the calling crate name. This is a convenience wrapper around
/// std::thread::Builder that:
/// - Sets the thread name to the calling crate name
/// - Handles the thread spawning details
#[macro_export]
macro_rules! spawn_anonymous_thread_with {
    ($body:expr) => {{
        const THREAD_PREFIX: &'static str = "tmp";

        let thread_name = format!(
            "{}-{}",
            option_env!("CARGO_PKG_NAME").unwrap_or("unknown"),
            THREAD_PREFIX
        );
        std::thread::Builder::new()
            .name(thread_name.into())
            .spawn_wrapper_ext(true, $body)
            .unwrap()
    }};
}

impl ThreadBuildWrapper for tokio::runtime::Builder {
    fn with_sys_and_custom_hooks<F1, F2>(&mut self, start: F1, end: F2) -> &mut Self
    where
        F1: Fn() + Send + Sync + 'static,
        F2: Fn() + Send + Sync + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.on_thread_start(move || {
            call_thread_start_hooks();
            // SAFETY: we will call `remove_thread_memory_accessor` at
            // `before-stop_wrapper`.
            // FIXME: What if the user only calls `after_start_wrapper`?
            unsafe { add_thread_memory_accessor() };
            thread_allocate_exclusive_arena().unwrap();
            add_thread_name_to_map();
            start();
        })
        .on_thread_stop(move || {
            end();
            remove_thread_name_from_map();
            remove_thread_memory_accessor();
        })
    }

    fn after_start_wrapper<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.on_thread_start(move || {
            call_thread_start_hooks();
            // SAFETY: we will call `remove_thread_memory_accessor` at
            // `before-stop_wrapper`.
            // FIXME: What if the user only calls `after_start_wrapper`?
            unsafe { add_thread_memory_accessor() };
            thread_allocate_exclusive_arena().unwrap();
            add_thread_name_to_map();
            f();
        })
    }

    fn before_stop_wrapper<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.on_thread_stop(move || {
            f();
            remove_thread_name_from_map();
            remove_thread_memory_accessor();
        })
    }
}

impl ThreadBuildWrapper for futures::executor::ThreadPoolBuilder {
    fn with_sys_and_custom_hooks<F1, F2>(&mut self, start: F1, end: F2) -> &mut Self
    where
        F1: Fn() + Send + Sync + 'static,
        F2: Fn() + Send + Sync + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.after_start(move |_| {
            call_thread_start_hooks();
            // SAFETY: we will call `remove_thread_memory_accessor` at
            // `before-stop_wrapper`.
            // FIXME: What if the user only calls `after_start_wrapper`?
            unsafe { add_thread_memory_accessor() };
            thread_allocate_exclusive_arena().unwrap();
            add_thread_name_to_map();
            start();
        })
        .before_stop(move |_| {
            end();
            remove_thread_name_from_map();
            remove_thread_memory_accessor();
        })
    }

    fn after_start_wrapper<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.after_start(move |_| {
            call_thread_start_hooks();
            // SAFETY: we will call `remove_thread_memory_accessor` at
            // `before-stop_wrapper`.
            // FIXME: What if the user only calls `after_start_wrapper`?
            unsafe { add_thread_memory_accessor() };
            thread_allocate_exclusive_arena().unwrap();
            add_thread_name_to_map();
            f();
        })
    }

    fn before_stop_wrapper<F>(&mut self, f: F) -> &mut Self
    where
        F: Fn() + Send + Sync + 'static,
    {
        #[allow(clippy::disallowed_methods)]
        self.before_stop(move |_| {
            f();
            remove_thread_name_from_map();
            remove_thread_memory_accessor();
        })
    }
}

#[cfg(test)]
mod tests {
    use std::{
        collections::HashSet,
        sync,
        sync::{Arc, Condvar, Mutex},
    };

    use futures::executor::block_on;

    use super::*;
    use crate::yatp_pool::{DefaultTicker, YatpPoolBuilder};

    #[test]
    fn test_thread_id() {
        let id = thread_id();
        assert_ne!(id, 0);
        std::thread::spawn(move || {
            // Two threads should have different ids.
            assert_ne!(thread_id(), id);
        })
        .join()
        .unwrap();
    }

    #[test]
    fn test_matches_thread_name_prefix() {
        assert!(matches_thread_name_prefix(
            "unified-read-pool-1",
            "unified-read"
        ));
        assert!(matches_thread_name_prefix(
            "unified-read-po",
            "unified-read-pool"
        ));
        assert!(matches_thread_name_prefix("sched-worker-pool-1", "sched"));
        assert!(matches_thread_name_prefix("sched-pool-1", "sched"));
        assert!(!matches_thread_name_prefix(
            "foo-unified-read",
            "unified-read"
        ));
        assert!(matches_thread_name_prefix(
            "abcdefghijklmn",
            "abcdefghijklmné"
        ));
    }

    #[test]
    #[cfg(any(target_os = "linux", target_os = "macos"))]
    fn test_thread_ids() {
        const THREAD_COUNT: usize = 10;
        let (tx, rx) = crossbeam::channel::bounded(THREAD_COUNT);

        #[allow(clippy::mutex_atomic)]
        let stop_threads_cvar = Arc::new((Mutex::new(false), Condvar::new()));

        let threads = (0..THREAD_COUNT)
            .map(|_| {
                let tx = tx.clone();
                let stop_threads_cvar = stop_threads_cvar.clone();
                std::thread::spawn(move || {
                    tx.send(thread_id()).unwrap();

                    let (lock, cvar) = &*stop_threads_cvar;
                    let _guard = cvar
                        .wait_while(lock.lock().unwrap(), |stop| !*stop)
                        .unwrap();
                })
            })
            .collect::<Vec<_>>();

        let actual_tids = rx.iter().take(THREAD_COUNT).collect::<Vec<_>>();
        let ids = thread_ids::<HashSet<_>>(process_id()).unwrap();
        for tid in &actual_tids {
            assert!(ids.contains(tid));
        }

        {
            let (lock, cvar) = &*stop_threads_cvar;
            *lock.lock().unwrap() = true;
            cvar.notify_all();
        }
        for thread in threads {
            thread.join().unwrap();
        }

        let stopped_tids = actual_tids;
        let ids = thread_ids::<HashSet<_>>(process_id()).unwrap();
        for tid in &stopped_tids {
            assert!(!ids.contains(tid));
        }
    }

    #[test]
    fn test_thread_name_wrapper() {
        let thread_name = "thread_for_test";

        let (tx, rx) = sync::mpsc::sync_channel(10);

        let get_name = move || {
            let tid = thread_id();
            if let Some(name) = THREAD_NAME_HASHMAP.lock().unwrap().get(&tid) {
                tx.clone().send(name.to_string()).unwrap();
            } else {
                panic!("thread not found");
            }
        };

        // test std thread builder
        std::thread::Builder::new()
            .name(thread_name.to_string())
            .spawn_wrapper(get_name.clone())
            .unwrap()
            .join()
            .unwrap();

        let name = rx.recv().unwrap();
        assert_eq!(name, thread_name);

        // test Yatp
        let get_name_fn = get_name.clone();
        block_on(
            YatpPoolBuilder::new(DefaultTicker)
                .name_prefix(thread_name)
                .after_start(|| {})
                .before_stop(|| {})
                .build_future_pool()
                .spawn_handle(async move { get_name_fn() })
                .unwrap(),
        )
        .unwrap();

        let name = rx.recv().unwrap();
        assert!(name.contains(thread_name));

        // test tokio thread builder
        let get_name_fn = get_name.clone();
        block_on(
            tokio::runtime::Builder::new_multi_thread()
                .thread_name(thread_name)
                .after_start_wrapper(|| {})
                .before_stop_wrapper(|| {})
                .build()
                .unwrap()
                .spawn(async move { get_name_fn() }),
        )
        .unwrap();

        let name = rx.recv().unwrap();
        assert_eq!(name, thread_name);

        // test tokio thread builder on sys hooks
        let get_name_fn = get_name.clone();
        block_on(
            tokio::runtime::Builder::new_multi_thread()
                .thread_name(thread_name)
                .with_sys_hooks()
                .build()
                .unwrap()
                .spawn(async move { get_name_fn() }),
        )
        .unwrap();

        let name = rx.recv().unwrap();
        assert_eq!(name, thread_name);

        // test anonymous std::thread using shared arena
        tikv_alloc::set_thread_exclusive_arena(true);
        for _ in 0..3 {
            let get_name = get_name.clone();
            let closure_func = move || {
                spawn_anonymous_thread_with!(get_name);
            };
            let tmp_name = format!("{}-tmp", std::env::var("CARGO_PKG_NAME").unwrap());
            closure_func();
            let name = rx.recv().unwrap();
            assert_eq!(name, tmp_name);
        }
    }
}