Skip to main content

rustpython_host_env/
resource.rs

1use std::io;
2
3use crate::os::CheckLibcResult;
4
5pub use libc::{
6    RLIM_INFINITY, RLIMIT_AS, RLIMIT_CORE, RLIMIT_CPU, RLIMIT_DATA, RLIMIT_FSIZE, RLIMIT_MEMLOCK,
7    RLIMIT_NOFILE, RLIMIT_NPROC, RLIMIT_RSS, RLIMIT_STACK, c_long, rlim_t, rlimit, timeval,
8};
9
10#[cfg(any(target_os = "linux", target_os = "android", target_os = "emscripten"))]
11pub use libc::{RLIMIT_MSGQUEUE, RLIMIT_NICE, RLIMIT_RTPRIO, RLIMIT_SIGPENDING};
12
13#[cfg(target_os = "linux")]
14pub use libc::RLIMIT_RTTIME;
15
16#[cfg(any(
17    target_os = "freebsd",
18    target_os = "netbsd",
19    target_os = "solaris",
20    target_os = "illumos"
21))]
22pub use libc::RLIMIT_SBSIZE;
23
24#[cfg(any(target_os = "freebsd", target_os = "solaris", target_os = "illumos"))]
25pub use libc::{RLIMIT_NPTS, RLIMIT_SWAP};
26
27#[cfg(any(target_os = "solaris", target_os = "illumos"))]
28pub use libc::RLIMIT_VMEM;
29
30#[cfg(any(target_os = "linux", target_os = "emscripten", target_os = "freebsd"))]
31pub use libc::RUSAGE_THREAD;
32
33#[cfg(not(any(target_os = "windows", target_os = "redox")))]
34pub use libc::{RUSAGE_CHILDREN, RUSAGE_SELF};
35
36#[cfg(target_os = "android")]
37pub const RLIM_NLIMITS: libc::c_int = 16;
38
39#[derive(Debug, Clone, Copy)]
40pub struct RUsage {
41    pub ru_utime: libc::timeval,
42    pub ru_stime: libc::timeval,
43    pub ru_maxrss: libc::c_long,
44    pub ru_ixrss: libc::c_long,
45    pub ru_idrss: libc::c_long,
46    pub ru_isrss: libc::c_long,
47    pub ru_minflt: libc::c_long,
48    pub ru_majflt: libc::c_long,
49    pub ru_nswap: libc::c_long,
50    pub ru_inblock: libc::c_long,
51    pub ru_oublock: libc::c_long,
52    pub ru_msgsnd: libc::c_long,
53    pub ru_msgrcv: libc::c_long,
54    pub ru_nsignals: libc::c_long,
55    pub ru_nvcsw: libc::c_long,
56    pub ru_nivcsw: libc::c_long,
57}
58
59impl RUsage {
60    pub fn utime_secs(&self) -> f64 {
61        timeval_to_secs(self.ru_utime)
62    }
63
64    pub fn stime_secs(&self) -> f64 {
65        timeval_to_secs(self.ru_stime)
66    }
67
68    pub fn total_cpu_duration(&self) -> Option<core::time::Duration> {
69        let utime = timeval_to_nanos(self.ru_utime)?;
70        let stime = timeval_to_nanos(self.ru_stime)?;
71        let total = utime.checked_add(stime)?;
72        Some(core::time::Duration::from_nanos(total as u64))
73    }
74}
75
76fn timeval_to_secs(tv: libc::timeval) -> f64 {
77    tv.tv_sec as f64 + (tv.tv_usec as f64 / 1_000_000.0)
78}
79
80fn timeval_to_nanos(tv: libc::timeval) -> Option<i64> {
81    let secs = widen_to_i64(tv.tv_sec)?.checked_mul(1_000_000_000)?;
82    let usecs = widen_to_i64(tv.tv_usec)?.checked_mul(1_000)?;
83    secs.checked_add(usecs)
84}
85
86fn widen_to_i64<T: Copy + TryInto<i64>>(v: T) -> Option<i64> {
87    v.try_into().ok()
88}
89
90impl From<libc::rusage> for RUsage {
91    fn from(rusage: libc::rusage) -> Self {
92        Self {
93            ru_utime: rusage.ru_utime,
94            ru_stime: rusage.ru_stime,
95            ru_maxrss: rusage.ru_maxrss,
96            ru_ixrss: rusage.ru_ixrss,
97            ru_idrss: rusage.ru_idrss,
98            ru_isrss: rusage.ru_isrss,
99            ru_minflt: rusage.ru_minflt,
100            ru_majflt: rusage.ru_majflt,
101            ru_nswap: rusage.ru_nswap,
102            ru_inblock: rusage.ru_inblock,
103            ru_oublock: rusage.ru_oublock,
104            ru_msgsnd: rusage.ru_msgsnd,
105            ru_msgrcv: rusage.ru_msgrcv,
106            ru_nsignals: rusage.ru_nsignals,
107            ru_nvcsw: rusage.ru_nvcsw,
108            ru_nivcsw: rusage.ru_nivcsw,
109        }
110    }
111}
112
113pub fn getrusage(who: i32) -> io::Result<RUsage> {
114    let mut rusage = core::mem::MaybeUninit::<libc::rusage>::uninit();
115    unsafe { libc::getrusage(who, rusage.as_mut_ptr()) }.check_libc_neg()?;
116    Ok(unsafe { rusage.assume_init() }.into())
117}
118
119pub fn getrlimit(resource: libc::rlim_t) -> io::Result<libc::rlimit> {
120    let mut rlimit = core::mem::MaybeUninit::<libc::rlimit>::uninit();
121    unsafe { libc::getrlimit(resource as _, rlimit.as_mut_ptr()) }.check_libc_neg()?;
122    Ok(unsafe { rlimit.assume_init() })
123}
124
125pub fn setrlimit(resource: libc::rlim_t, limits: libc::rlimit) -> io::Result<()> {
126    unsafe { libc::setrlimit(resource as _, &limits) }.check_libc_neg()?;
127    Ok(())
128}
129
130#[cfg(not(any(target_os = "redox", target_os = "wasi")))]
131pub fn disable_core_dumps() {
132    let rl = libc::rlimit {
133        rlim_cur: 0,
134        rlim_max: 0,
135    };
136    unsafe {
137        let _ = libc::setrlimit(libc::RLIMIT_CORE, &rl);
138    }
139}