use std::{
fs::{self, File, OpenOptions},
io::{Error, ErrorKind, Read, Seek, SeekFrom, Write},
ops::{Deref, DerefMut},
sync::LazyLock,
thread,
time::Duration,
};
use arraystring::error;
use core_affinity::CoreId;
use fs2::FileExt;
use crate::fast_thread_pool::FILE_CORE_AFFINITY;
pub static CORES: LazyLock<Vec<CoreId>> = LazyLock::new(|| {
let core_ids = core_affinity::get_core_ids().unwrap_or_else(|| {
warn!("get core ids from core_affinity failed, use default empty vector");
vec![]
});
debug!(
"use core_affinity core_ids: {:?}",
core_ids.iter().map(|x| x.id).collect::<Vec<_>>()
);
core_ids
});
struct LockFile {
file: File,
locked: bool,
}
impl LockFile {
fn open(path: &str) -> Result<Self, Error> {
let file = OpenOptions::new().read(true).write(true).open(path)?;
match file.try_lock_exclusive() {
Ok(_) => Ok(LockFile { file, locked: true }),
Err(e) => Err(e),
}
}
}
impl Deref for LockFile {
type Target = File;
fn deref(&self) -> &Self::Target {
&self.file
}
}
impl DerefMut for LockFile {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.file
}
}
impl Drop for LockFile {
fn drop(&mut self) {
if self.locked {
let _ = fs2::FileExt::unlock(&self.file);
}
}
}
pub fn use_last_core(use_name: &str) -> usize {
use_last_core2(use_name, 1)[0]
}
pub fn use_last_core2(use_name: &str, count: usize) -> Vec<usize> {
let core = CORES.clone();
_ = fs::File::create_new(FILE_CORE_AFFINITY);
const MAX_RETRY: usize = 10;
let mut lock_file = {
let mut retry_count = 0;
loop {
if retry_count >= MAX_RETRY {
error!(
"open core_affinity file overflow max retry {MAX_RETRY}, conot continue execute, return default core: 0"
);
let use_core = core.last().map(|x| x.clone()).unwrap_or_else(|| {
warn!("get cpu core number failed; get use default core: 0");
core_affinity::CoreId { id: 0 }
});
return vec![use_core.id];
}
match LockFile::open(FILE_CORE_AFFINITY) {
Ok(file) => break file,
Err(e) => {
warn!(
".core_affinity is locked, waiting to release: {}, retry times: {}/{}",
e,
retry_count + 1,
MAX_RETRY
);
retry_count += 1;
thread::sleep(Duration::from_millis(150));
}
}
}
};
let mut file_content = String::new();
let use_core_ids = {
if let Err(e) = lock_file.read_to_string(&mut file_content) {
error!("读取core_affinity文件内容失败: {},无法继续执行", e);
let use_core = core.last().map(|x| x.clone()).unwrap_or_else(|| {
warn!("获取cpu核心数失败");
core_affinity::CoreId { id: 0 }
});
return vec![use_core.id];
}
let mut used_cores = Vec::new();
for line in file_content.lines() {
if line.trim().is_empty() || line.trim().starts_with("realtime_system") {
continue;
}
let parts: Vec<&str> = line.split(',').collect();
for i in 1..parts.len() {
if let Ok(core_id) = parts[i].trim().parse::<usize>() {
used_cores.push(core_id);
}
}
}
let mut all_cores: Vec<usize> = core.iter().map(|c| c.id).collect();
all_cores.sort_by(|a, b| b.cmp(a));
if all_cores.is_empty() {
let default_cores = (0..count).collect::<Vec<_>>();
warn!("没有可用的核心,使用默认核心: {:?}", default_cores);
return default_cores;
}
let available_cores: Vec<usize> = all_cores
.iter()
.filter(|&id| !used_cores.contains(id))
.cloned()
.collect();
let mut selected_cores = Vec::with_capacity(count);
if available_cores.is_empty() {
let last_core = used_cores.last().cloned().unwrap_or(all_cores[0]);
let mut last_index = 0;
for (i, &core_id) in all_cores.iter().enumerate() {
if core_id == last_core {
last_index = i;
break;
}
}
let mut start_index = (last_index + 1) % all_cores.len();
for i in 0..count {
let current_index = (start_index + i) % all_cores.len();
selected_cores.push(all_cores[current_index]);
}
} else {
let mut remaining = count;
for &core_id in &available_cores {
if remaining == 0 {
break;
}
selected_cores.push(core_id);
remaining -= 1;
}
if remaining > 0 {
let last_core = if selected_cores.is_empty() {
used_cores.last().cloned().unwrap_or(all_cores[0])
} else {
selected_cores.last().cloned().unwrap()
};
let mut last_index = 0;
for (i, &core_id) in all_cores.iter().enumerate() {
if core_id == last_core {
last_index = i;
break;
}
}
let mut start_index = (last_index + 1) % all_cores.len();
for i in 0..remaining {
let current_index = (start_index + i) % all_cores.len();
selected_cores.push(all_cores[current_index]);
}
}
}
if let Err(e) = lock_file.seek(SeekFrom::End(0)) {
error!("移动文件指针到文件末尾失败: {},无法写入新数据", e);
return selected_cores;
}
let mut write_content = format!("{use_name}");
for core_id in &selected_cores {
write_content.push_str(&format!(",{}", core_id));
}
write_content.push_str("\n");
if let Err(e) = lock_file.write_all(write_content.as_bytes()) {
error!("写入数据到core_affinity文件失败: {}", e);
return selected_cores;
}
if let Err(e) = lock_file.flush() {
error!("刷新文件缓冲区失败: {}", e);
}
selected_cores
};
debug!("{use_name} use_cores: {:?}", use_core_ids);
use_core_ids
}
fn read_linux_system_cpu_cores() -> Option<Vec<CoreId>> {
std::fs::read_to_string("/sys/devices/system/cpu/present")
.ok()
.and_then(|content| {
parse_cpu_range(&content)
.map(|cores| cores.into_iter().map(|id| CoreId { id }).collect())
})
}
fn parse_cpu_range(content: &str) -> Option<Vec<usize>> {
let mut cpus = Vec::new();
let content = content.trim();
if content.is_empty() {
return None;
}
for part in content.split(',') {
if let Some((start, end)) = part.split_once('-') {
let start = start.parse::<usize>().ok()?;
let end = end.parse::<usize>().ok()?;
cpus.extend(start..=end);
} else {
let cpu = part.parse::<usize>().ok()?;
cpus.push(cpu);
}
}
Some(cpus)
}
#[cfg(feature = "deal_physical_cpu")]
pub fn get_core_skip() -> usize {
let core_ids = num_cpus::get();
let core_physical = num_cpus::get_physical();
if core_ids / core_physical == 2 {
warn!("core_ids: {core_ids}, core_physical: {core_physical}; skip 2");
2
} else {
1
}
}
#[cfg(not(feature = "deal_physical_cpu"))]
pub fn get_core_skip() -> usize {
1
}