process_limiter 0.1.0

A lib which can directly limit the specified process occupancy, and support multithreading and dynamically changing occupancy
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use crate::{core, misc, LimitInfo, Task, TaskStatus};
use std::error::Error;
use std::sync::{Arc, Mutex};
use std::thread;

use sysinfo::{Pid, PidExt, ProcessExt, ProcessRefreshKind, System, SystemExt};

// What the Process needs to be refreshed
macro_rules! PROCESS_REFRESH {
    () => {
        ProcessRefreshKind::new().with_cpu().with_user()
    };
}

pub trait TaskExt {
    fn new(system: Arc<Mutex<System>>) -> Self;
    fn start(&mut self, pid: u32);
    fn pause(&mut self) -> Result<(), Box<dyn Error>>;
    fn remuse(&mut self) -> Result<(), Box<dyn Error>>;
    fn set_target(&mut self, new_target: f32) -> Result<(), Box<dyn Error>>;
    fn status(&mut self) -> TaskStatus;
}

impl TaskExt for Task {
    fn new(system: Arc<Mutex<System>>) -> Self {
        Self {
            system,
            thread: None,
            target: Arc::new(Mutex::new(1.0)),
        }
    }
    fn start(&mut self, pid: u32) {
        let target = self.target.clone();
        let system = self.system.clone();
        self.thread = Some(thread::spawn(move || {
            let mut info = LimitInfo::default();
            loop {
                let mut system = if let Ok(o) = system.lock() {
                    o
                } else {
                    continue;
                };
                if !system.refresh_process_specifics(Pid::from_u32(pid), PROCESS_REFRESH!()) {
                    break;
                };
                let process = if let Some(o) = system.process(Pid::from_u32(pid)) {
                    o
                } else {
                    break;
                };

                info.update_current_usage(process.cpu_usage());
                if let Ok(o) = target.lock() {
                    info.update_taregt_usage(*o);
                } else {
                    continue;
                }
                let work_slice = core::process::limit_process(process, &mut info);
                info.update_work_slice(work_slice);
            }
        }));
    }
    fn pause(&mut self) -> Result<(), Box<dyn Error>> {
        let pid = self.get_thread_pid()?;
        let mut system = if let Ok(o) = self.system.lock() {
            o
        } else {
            return Err("Failed to get system".into());
        };
        misc::stop_by_pid(&mut system, pid)
    }
    fn remuse(&mut self) -> Result<(), Box<dyn Error>> {
        let pid = self.get_thread_pid()?;
        let mut system = if let Ok(o) = self.system.lock() {
            o
        } else {
            return Err("Failed to get system".into());
        };
        misc::cont_by_pid(&mut system, pid)
    }
    fn set_target(&mut self, new_target: f32) -> Result<(), Box<dyn Error>> {
        let mut target = match self.target.lock() {
            Ok(o) => o,
            Err(_e) => {
                return Err("Failed to get the lock".into());
            }
        };
        *target = new_target;
        Ok(())
    }
    fn status(&mut self) -> TaskStatus {
        let pid = if let Ok(o) = self.get_thread_pid() {
            o
        } else {
            return TaskStatus::NeedInit;
        };
        let mut system = if let Ok(o) = self.system.lock() {
            o
        } else {
            return TaskStatus::NeedInit;
        };
        match misc::status_by_pid(&mut system, pid) {
            Ok(o) => o,
            Err(_) => TaskStatus::NeedInit,
        }
    }
}

impl Task {
    fn get_thread_pid(&mut self) -> Result<u32, Box<dyn Error>> {
        if self.thread.is_none() {
            return Err("Limiter process does not exist".into());
        }
        if let Some(o) = &self.thread {
            Ok(misc::get_thread_pid(o))
        } else {
            Err("Limiter process does not exist".into())
        }
    }
}