Search code examples
rustcpu-usage

Limit the cpurate in a Rust program


I want to limit the cpurate of a rust program I am running. I found ways to do it in the command, but was wondering if there was a way to do it inside the program using a built in rust module.


Solution

  • Typically, this is not the concern of an individual program. The simplest and most flexible approach is usually to set resource limits when running a program (e.g. via Docker settings) and allow the program to just use the resources it's been given.

    With that out of the way, let's assume that you really need to be able to do this. On Linux, you can use cgroups to set resource limits.

    There are a few cgroups crates for Rust, of varying ages and quality. Some are native and some are wrappers for C libraries but none of them have been recently updated. Even so, I just picked one that had examples in its documentation (controlgroup) and was able to make this work:

    use controlgroup::v1::Builder; // controlgroup 0.3.0
    use signal_hook:: register; // signal-hook 0.1.16
    
    fn main() {
        set_cpu_percent(10.0);
        
        // Do your work here
    }
    
    fn set_cpu_percent(percent: f64) {
        let mut cgroups = Builder::new("my-cgroups".into())
            .cpu()
            .cfs_quota_us((percent * 1000.0) as i64)
            .cfs_period_us(100_000)
            .done()
            .build()
            .unwrap();
    
        // Add the current task to crgoups
        let pid = std::process::id().into();
        cgroups.add_task(pid).unwrap();
    
        // Do some clean up, e.g. if you terminate the process with ctrl+C
        unsafe {
            register(signal_hook::SIGINT, {
                // register handler is Fn, so cannot mutate its environment
                let cgroups = std::sync::Mutex::new(cgroups);
                move || {
                    let mut cgroups = cgroups.lock().unwrap();
                    cgroups.remove_task(pid).unwrap();
                    cgroups.delete().unwrap();
                    std::process::exit(0);
                }
            })
            .unwrap()
        };
    }
    

    Note that I didn't do any error handling, and you should see this as only a proof-of-concept, not production quality code. The reason I added a SIGINT cleanup handler was because I otherwise got errors the second time I ran this, due to not cleaning up the previous run. I'm not an expert in this area and it's likely that there are a lot of other edge cases you'll need to consider in order to make this robust.