// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Wrappers for CPU affinity functions.

use std::iter::FromIterator;
use std::mem;

use libc::{cpu_set_t, sched_setaffinity, CPU_SET, CPU_SETSIZE, CPU_ZERO, EINVAL};

use crate::{errno_result, Error, Result};

// This is needed because otherwise the compiler will complain that the
// impl doesn't reference any types from inside this crate.
struct CpuSet(cpu_set_t);

impl FromIterator<usize> for CpuSet {
    fn from_iter<I: IntoIterator<Item = usize>>(cpus: I) -> Self {
        // cpu_set_t is a C struct and can be safely initialized with zeroed memory.
        let mut cpuset: cpu_set_t = unsafe { mem::zeroed() };
        // Safe because we pass a valid cpuset pointer.
        unsafe { CPU_ZERO(&mut cpuset) };
        for cpu in cpus {
            // Safe because we pass a valid cpuset pointer and cpu index.
            unsafe { CPU_SET(cpu, &mut cpuset) };
        }
        CpuSet(cpuset)
    }
}

/// Set the CPU affinity of the current thread to a given set of CPUs.
///
/// # Examples
///
/// Set the calling thread's CPU affinity so it will run on only CPUs
/// 0, 1, 5, and 6.
///
/// ```
/// # use sys_util::set_cpu_affinity;
///   set_cpu_affinity(vec![0, 1, 5, 6]).unwrap();
/// ```
pub fn set_cpu_affinity<I: IntoIterator<Item = usize>>(cpus: I) -> Result<()> {
    let CpuSet(cpuset) = cpus
        .into_iter()
        .map(|cpu| {
            if cpu < CPU_SETSIZE as usize {
                Ok(cpu)
            } else {
                Err(Error::new(EINVAL))
            }
        })
        .collect::<Result<CpuSet>>()?;

    // Safe because we pass 0 for the current thread, and cpuset is a valid pointer and only
    // used for the duration of this call.
    let res = unsafe { sched_setaffinity(0, mem::size_of_val(&cpuset), &cpuset) };

    if res != 0 {
        errno_result()
    } else {
        Ok(())
    }
}

/// Enable experimental core scheduling for the current thread.
///
/// If succesful, the kernel should not schedule this thread with any other thread within the same
/// SMT core.
#[cfg(feature = "chromeos")]
pub fn enable_core_scheduling() -> Result<()> {
    use libc::prctl;
    const PR_SET_CORE_SCHED: i32 = 0x200;
    let ret = unsafe { prctl(PR_SET_CORE_SCHED, 1) };
    if ret == -1 {
        errno_result()
    } else {
        Ok(())
    }
}