free_cpus/
lib.rs

1//! This crate exports a utility function that returns the set of available core IDs on the machine.
2//!
3//! The set is determined by iterating over the `/proc` directory's process statuses.
4//! The logic is heavily inspired by [AFL++'s code](https://github.com/AFLplusplus/AFLplusplus/blob/85c5b5218c6a7b2289f309fbd1625a5d0a602a00/src/afl-fuzz-init.c#L109-L452).
5use anyhow::{Context, Result};
6use std::{
7    collections::HashSet,
8    fs::{read_dir, File},
9    io::{BufRead, BufReader},
10};
11
12/// Returns the set of available core IDs in your Linux machine.
13///
14/// Logic inspired by [AFL++'s code](https://github.com/AFLplusplus/AFLplusplus/blob/85c5b5218c6a7b2289f309fbd1625a5d0a602a00/src/afl-fuzz-init.c#L109-L452)
15#[cfg(target_os = "linux")]
16pub fn get() -> Result<HashSet<usize>> {
17    let mut cpu_used = HashSet::new();
18    let proc = read_dir("/proc").context("Could not read /proc, we cannot safely assign cores")?;
19    // We iterate over every directory's status file in /proc
20    for status in proc
21        .flatten()
22        .filter(|p| p.metadata().map(|p| p.is_dir()).unwrap_or(false))
23        .filter_map(|p| p.file_name().to_str().map(String::from))
24        .flat_map(|pid| File::open(format!("/proc/{pid}/status")))
25        .map(BufReader::new)
26    {
27        let mut has_vmsize = false;
28        for line in status.lines().flatten() {
29            if line.contains("VmSize:\t") {
30                has_vmsize = true;
31            }
32            if has_vmsize
33                && line.contains("Cpus_allowed_list:\t")
34                && !line.contains('-')
35                && !line.contains(',')
36            {
37                if let Some(id_str) = line.strip_prefix("Cpus_allowed_list:\t") {
38                    if let Ok(id) = id_str.parse::<usize>() {
39                        cpu_used.insert(id);
40                    }
41                }
42            }
43        }
44    }
45    let cores = HashSet::from_iter(0..num_cpus::get());
46    Ok(HashSet::from_iter(cores.difference(&cpu_used).copied()))
47}
48
49#[cfg(not(target_os = "linux"))]
50pub fn get() -> HashSet<usize> {
51    unimplemented!("free-cpus is only implemented for Linux");
52}
free_cpus/lib.rs

free_cpus/
lib.rs
