// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

#[cfg(target_os = "linux")]
mod cgroup;
pub mod cpu_time;
pub mod disk;
pub mod inspector;
pub mod ioload;
#[cfg(any(test, feature = "testexport"))]
pub mod process;
pub mod thread;

// re-export some traits for ease of use
#[cfg(target_os = "linux")]
use std::{fs, os::unix::fs::MetadataExt};
use std::{
    path::Path,
    sync::atomic::{AtomicU64, Ordering},
};

use fail::fail_point;
#[cfg(target_os = "linux")]
use lazy_static::lazy_static;
use sysinfo::RefreshKind;
pub use sysinfo::{CpuExt, DiskExt, NetworkExt, ProcessExt, SystemExt};

use crate::config::ReadableSize;

pub const HIGH_PRI: i32 = -1;
const CPU_CORES_QUOTA_ENV_VAR_KEY: &str = "TIKV_CPU_CORES_QUOTA";

static GLOBAL_MEMORY_USAGE: AtomicU64 = AtomicU64::new(0);
static MEMORY_USAGE_HIGH_WATER: AtomicU64 = AtomicU64::new(u64::MAX);
static MEMORY_USAGE_THROTTLING_LEVEL: AtomicU64 = AtomicU64::new(u64::MAX);

#[cfg(target_os = "linux")]
lazy_static! {
    static ref SELF_CGROUP: cgroup::CGroupSys = cgroup::CGroupSys::new().unwrap_or_default();
}

pub struct SysQuota;
impl SysQuota {
    #[cfg(target_os = "linux")]
    pub fn cpu_cores_quota() -> f64 {
        let mut cpu_num = num_cpus::get() as f64;
        let cpuset_cores = SELF_CGROUP.cpuset_cores().len() as f64;
        let cpu_quota = SELF_CGROUP.cpu_quota().unwrap_or(0.);

        if cpuset_cores != 0. {
            cpu_num = cpu_num.min(cpuset_cores);
        }

        if cpu_quota != 0. {
            cpu_num = cpu_num.min(cpu_quota);
        }

        limit_cpu_cores_quota_by_env_var(cpu_num)
    }

    #[cfg(target_os = "linux")]
    pub fn cpu_cores_quota_current() -> f64 {
        let cgroup = cgroup::CGroupSys::new().unwrap_or_default();
        let mut cpu_num = num_cpus::get() as f64;
        let cpuset_cores = cgroup.cpuset_cores().len() as f64;
        let cpu_quota = cgroup.cpu_quota().unwrap_or(0.);

        if cpuset_cores != 0. {
            cpu_num = cpu_num.min(cpuset_cores);
        }

        if cpu_quota != 0. {
            cpu_num = cpu_num.min(cpu_quota);
        }

        limit_cpu_cores_quota_by_env_var(cpu_num)
    }

    #[cfg(not(target_os = "linux"))]
    pub fn cpu_cores_quota() -> f64 {
        let cpu_num = num_cpus::get() as f64;
        limit_cpu_cores_quota_by_env_var(cpu_num)
    }

    #[cfg(not(target_os = "linux"))]
    pub fn cpu_cores_quota_current() -> f64 {
        let cpu_num = num_cpus::get() as f64;
        limit_cpu_cores_quota_by_env_var(cpu_num)
    }

    #[cfg(target_os = "linux")]
    pub fn memory_limit_in_bytes() -> u64 {
        let total_mem = Self::sysinfo_memory_limit_in_bytes();
        if let Some(cgroup_memory_limit) = SELF_CGROUP.memory_limit_in_bytes() {
            std::cmp::min(total_mem, cgroup_memory_limit)
        } else {
            total_mem
        }
    }

    #[cfg(target_os = "linux")]
    pub fn memory_limit_in_bytes_current() -> u64 {
        let cgroup = cgroup::CGroupSys::new().unwrap_or_default();
        let total_mem = Self::sysinfo_memory_limit_in_bytes();
        if let Some(cgroup_memory_limit) = cgroup.memory_limit_in_bytes() {
            std::cmp::min(total_mem, cgroup_memory_limit)
        } else {
            total_mem
        }
    }

    #[cfg(not(target_os = "linux"))]
    pub fn memory_limit_in_bytes() -> u64 {
        Self::sysinfo_memory_limit_in_bytes()
    }

    #[cfg(not(target_os = "linux"))]
    pub fn memory_limit_in_bytes_current() -> u64 {
        Self::sysinfo_memory_limit_in_bytes()
    }

    pub fn log_quota() {
        #[cfg(target_os = "linux")]
        info!(
            "cgroup quota: memory={:?}, cpu={:?}, cores={:?}",
            SELF_CGROUP.memory_limit_in_bytes(),
            SELF_CGROUP.cpu_quota(),
            SELF_CGROUP.cpuset_cores(),
        );

        info!(
            "memory limit in bytes: {}, cpu cores quota: {}",
            Self::memory_limit_in_bytes(),
            Self::cpu_cores_quota()
        );
    }

    fn sysinfo_memory_limit_in_bytes() -> u64 {
        let system = sysinfo::System::new_with_specifics(RefreshKind::new().with_memory());
        system.total_memory()
    }
}

/// Get the current global memory usage in bytes. Users need to call
/// `record_global_memory_usage` to refresh it periodically.
pub fn get_global_memory_usage() -> u64 {
    GLOBAL_MEMORY_USAGE.load(Ordering::Acquire)
}

/// Record the current global memory usage of the process.
#[cfg(target_os = "linux")]
pub fn record_global_memory_usage() {
    let s = procinfo::pid::statm_self().unwrap();
    let usage = s.resident * page_size::get();
    GLOBAL_MEMORY_USAGE.store(usage as u64, Ordering::Release);
}

#[cfg(not(target_os = "linux"))]
pub fn record_global_memory_usage() {
    GLOBAL_MEMORY_USAGE.store(0, Ordering::Release);
}

/// Register the high water mark so that `memory_usage_reaches_high_water` is
/// available.
pub fn register_memory_usage_high_water(mark: u64) {
    MEMORY_USAGE_HIGH_WATER.store(mark, Ordering::Release);
}

pub fn memory_usage_reaches_high_water(usage: &mut u64) -> bool {
    fail_point!("memory_usage_reaches_high_water", |_| true);
    *usage = get_global_memory_usage();
    *usage >= MEMORY_USAGE_HIGH_WATER.load(Ordering::Acquire)
}

/// Register the throttling level so that
/// `memory_usage_over_throttling_level` is available.
pub fn register_memory_usage_throttling_level(mark: u64) {
    MEMORY_USAGE_THROTTLING_LEVEL.store(mark, Ordering::Release);
}

/// Returns the currently registered throttling level.
pub fn memory_usage_throttling_level() -> u64 {
    MEMORY_USAGE_THROTTLING_LEVEL.load(Ordering::Acquire)
}

/// Returns whether the latest process memory usage plus `additional_bytes`
/// would reach the registered throttling level.
pub fn memory_usage_would_reach_throttling_level(additional_bytes: u64) -> bool {
    get_global_memory_usage().saturating_add(additional_bytes) >= memory_usage_throttling_level()
}

/// Returns the memory usage over the throttling level and the throttling level.
pub fn memory_usage_over_throttling_level()
-> Option<(u64 /* overshoot */, u64 /* throttling level */)> {
    let usage = get_global_memory_usage();
    let throttling_level = memory_usage_throttling_level();
    if usage <= throttling_level {
        None
    } else {
        Some((usage - throttling_level, throttling_level))
    }
}

fn limit_cpu_cores_quota_by_env_var(quota: f64) -> f64 {
    match std::env::var(CPU_CORES_QUOTA_ENV_VAR_KEY)
        .ok()
        .and_then(|value| value.parse().ok())
    {
        Some(env_var_quota) if quota.is_sign_positive() => f64::min(quota, env_var_quota),
        _ => quota,
    }
}

fn read_size_in_cache(level: usize, field: &str) -> Option<u64> {
    std::fs::read_to_string(format!(
        "/sys/devices/system/cpu/cpu0/cache/index{}/{}",
        level, field
    ))
    .ok()
    .and_then(|s| s.parse::<ReadableSize>().ok())
    .map(|s| s.0)
}

/// Gets the size of given level cache.
///
/// It will only return `Some` on Linux.
pub fn cache_size(level: usize) -> Option<u64> {
    read_size_in_cache(level, "size")
}

/// Gets the size of given level cache line.
///
/// It will only return `Some` on Linux.
pub fn cache_line_size(level: usize) -> Option<u64> {
    read_size_in_cache(level, "coherency_line_size")
}

#[cfg(target_os = "linux")]
pub fn path_in_diff_mount_point(path1: impl AsRef<Path>, path2: impl AsRef<Path>) -> bool {
    let (path1, path2) = (path1.as_ref(), path2.as_ref());
    let device1 = match fs::metadata(path1) {
        Ok(metadata) => metadata.dev(),
        Err(error) => {
            warn!(
                "Get device ID error for path {}, {}",
                path1.display(),
                error
            );
            return false;
        }
    };
    let device2 = match fs::metadata(path2) {
        Ok(metadata) => metadata.dev(),
        Err(error) => {
            warn!(
                "Get device ID error for path {}, {}",
                path2.display(),
                error
            );
            return false;
        }
    };
    device1 != device2
}

#[cfg(not(target_os = "linux"))]
pub fn path_in_diff_mount_point(_path1: impl AsRef<Path>, _path2: impl AsRef<Path>) -> bool {
    false
}

pub fn set_memory_usage_for_test(usage: u64) {
    GLOBAL_MEMORY_USAGE.store(usage, Ordering::Release);
}

#[cfg(all(test, target_os = "linux"))]
mod tests {
    use super::*;

    #[test]
    fn test_path_in_diff_mount_point() {
        let (empty_path1, path2) = ("", "/");
        let result = path_in_diff_mount_point(empty_path1, path2);
        assert!(!result);

        let (no_mount_point_path, path2) = ("no_mount_point_path_w943nn", "/");
        let result = path_in_diff_mount_point(no_mount_point_path, path2);
        assert!(!result);

        let (not_existed_path, path2) = ("/non_existed_path_eu2yndh", "/");
        let result = path_in_diff_mount_point(not_existed_path, path2);
        assert!(!result);

        let result = path_in_diff_mount_point(path2, not_existed_path);
        assert!(!result);

        let (normal_path1, normal_path2) = ("/", "/");
        let result = path_in_diff_mount_point(normal_path1, normal_path2);
        assert!(!result);

        let root_device = fs::metadata(normal_path1).unwrap().dev();
        let different_device_path = ["/dev", "/proc", "/sys"]
            .iter()
            .find(|path| fs::metadata(path).is_ok_and(|metadata| metadata.dev() != root_device))
            .unwrap();
        let result = path_in_diff_mount_point(normal_path1, different_device_path);
        assert!(result);
    }
}