effective_limits/
lib.rs

1#![warn(clippy::all)]
2// for error_chain!
3#![recursion_limit = "1024"]
4
5use std::cmp::min;
6
7use cfg_if::cfg_if;
8
9cfg_if! {
10    if #[cfg(any(windows,
11                 target_os="macos",
12                 target_os="linux",
13                 target_os="freebsd",
14                 target_os="illumos",
15                 target_os="solaris",
16                 target_os="netbsd",
17                ))] {
18        #[derive(thiserror::Error, Debug)]
19        pub enum Error {
20            #[error("sysinfo failure")]
21            SysInfo(#[from] ::sys_info::Error),
22            #[error("io error")]
23            IoError(#[from] std::io::Error),
24            #[error("io error {1} ({0:?})")]
25            IoExplainedError(#[source] std::io::Error, String),
26            #[error("utf8 error")]
27            Utf8Error(#[from] std::str::Utf8Error),
28            #[error("u64 parse error on '{1}' ({0:?})")]
29            U64Error(#[source] std::num::ParseIntError, String),
30        }
31    } else {
32        #[derive(thiserror::Error, Debug)]
33        pub enum Error {
34            #[error("no system information on this platform")]
35            SysInfo,
36            #[error("io error")]
37            IoError(#[from] std::io::Error),
38            #[error("io error {1} ({0:?})")]
39            IoExplainedError(#[source] std::io::Error, String),
40            #[error("utf8 error")]
41            Utf8Error(#[from] std::str::Utf8Error),
42            #[error("u64 parse error on '{1}' ({0:?})")]
43            U64Error(#[source] std::num::ParseIntError, String),
44        }
45    }
46}
47
48pub type Result<R> = std::result::Result<R, Error>;
49
50#[allow(dead_code)]
51fn min_opt(left: u64, right: Option<u64>) -> u64 {
52    match right {
53        None => left,
54        Some(right) => min(left, right),
55    }
56}
57
58#[allow(dead_code)]
59#[cfg(unix)]
60fn ulimited_memory() -> Result<Option<u64>> {
61    let mut out = libc::rlimit {
62        rlim_cur: 0,
63        rlim_max: 0,
64    };
65    // https://github.com/rust-lang/libc/pull/1919
66    cfg_if!(
67    if #[cfg( target_os="netbsd")] {
68    // https://github.com/NetBSD/src/blob/f869ef2144970023b53d335d9a23ecf100d4b973/sys/sys/resource.h#L98
69    let rlimit_as = 10;
70        }
71    else {
72    let rlimit_as = libc::RLIMIT_AS;
73    }
74    );
75    match unsafe { libc::getrlimit(rlimit_as, &mut out as *mut libc::rlimit) } {
76        0 => Ok(()),
77        _ => Err(std::io::Error::last_os_error()),
78    }?;
79    let address_limit = match out.rlim_cur {
80        libc::RLIM_INFINITY => None,
81        _ => Some(out.rlim_cur as u64),
82    };
83    let mut out = libc::rlimit {
84        rlim_cur: 0,
85        rlim_max: 0,
86    };
87    match unsafe { libc::getrlimit(libc::RLIMIT_DATA, &mut out as *mut libc::rlimit) } {
88        0 => Ok(()),
89        _ => Err(std::io::Error::last_os_error()),
90    }?;
91    let data_limit = match out.rlim_cur {
92        libc::RLIM_INFINITY => address_limit,
93        _ => Some(out.rlim_cur as u64),
94    };
95    Ok(address_limit
96        .or(data_limit)
97        .map(|left| min_opt(left, data_limit)))
98}
99
100#[cfg(not(unix))]
101fn win_err<T>(fn_name: &str) -> Result<T> {
102    Err(Error::IoExplainedError(
103        std::io::Error::last_os_error(),
104        fn_name.into(),
105    ))
106}
107
108#[cfg(not(unix))]
109fn ulimited_memory() -> Result<Option<u64>> {
110    use std::mem::size_of;
111
112    use winapi::shared::minwindef::{FALSE, LPVOID};
113    use winapi::shared::ntdef::NULL;
114    use winapi::um::jobapi::IsProcessInJob;
115    use winapi::um::jobapi2::QueryInformationJobObject;
116    use winapi::um::processthreadsapi::GetCurrentProcess;
117    use winapi::um::winnt::{
118        JobObjectExtendedLimitInformation, JOBOBJECT_EXTENDED_LIMIT_INFORMATION,
119        JOB_OBJECT_LIMIT_PROCESS_MEMORY,
120    };
121
122    let mut in_job = 0;
123    match unsafe { IsProcessInJob(GetCurrentProcess(), NULL, &mut in_job) } {
124        FALSE => win_err("IsProcessInJob"),
125        _ => Ok(()),
126    }?;
127    if in_job == FALSE {
128        return Ok(None);
129    }
130    let mut job_info = winapi::um::winnt::JOBOBJECT_EXTENDED_LIMIT_INFORMATION {
131        ..Default::default()
132    };
133    let mut written: u32 = 0;
134    match unsafe {
135        QueryInformationJobObject(
136            NULL,
137            JobObjectExtendedLimitInformation,
138            &mut job_info as *mut JOBOBJECT_EXTENDED_LIMIT_INFORMATION as LPVOID,
139            size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32,
140            &mut written,
141        )
142    } {
143        FALSE => win_err("QueryInformationJobObject"),
144        _ => Ok(()),
145    }?;
146    if job_info.BasicLimitInformation.LimitFlags & JOB_OBJECT_LIMIT_PROCESS_MEMORY
147        == JOB_OBJECT_LIMIT_PROCESS_MEMORY
148    {
149        Ok(Some(job_info.ProcessMemoryLimit as u64))
150    } else {
151        Ok(None)
152    }
153}
154
155/// How much memory is effectively available for this process to use,
156/// considering the physical machine and ulimits, but not the impact of noisy
157/// neighbours, swappiness and so on. The goal is to have a good chance of
158/// avoiding failed allocations without requiring either developer or user
159/// a-priori selection of memory limits.
160pub fn memory_limit() -> Result<u64> {
161    cfg_if! {
162        if #[cfg(any(windows,
163                     target_os="macos",
164                     target_os="linux",
165                     target_os="freebsd",
166                     target_os="illumos",
167                     target_os="solaris",
168                     target_os="netbsd",
169                    ))] {
170            let info = sys_info::mem_info()?;
171            let total_ram = info.total * 1024;
172            let ulimit_mem = ulimited_memory()?;
173            Ok(min_opt(total_ram, ulimit_mem))
174        } else {
175            // https://github.com/FillZpp/sys-info-rs/issues/72
176            Err(Error::SysInfo)
177        }
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use std::env;
184    #[cfg(unix)]
185    use std::os::unix::process::CommandExt;
186    #[cfg(windows)]
187    use std::os::windows::process::CommandExt;
188    use std::path::PathBuf;
189    use std::process::Command;
190    use std::str;
191
192    #[cfg(windows)]
193    use winapi::shared::minwindef::{DWORD, FALSE, LPVOID};
194    #[cfg(windows)]
195    use winapi::shared::ntdef::NULL;
196
197    use super::*;
198
199    #[cfg(any(
200        windows,
201        target_os = "macos",
202        target_os = "linux",
203        target_os = "freebsd",
204        target_os = "illumos",
205        target_os = "solaris",
206        target_os = "netbsd",
207    ))]
208    #[test]
209    fn it_works() -> Result<()> {
210        assert_ne!(0, memory_limit()?);
211        Ok(())
212    }
213
214    #[test]
215    fn test_min_opt() {
216        assert_eq!(0, min_opt(0, None));
217        assert_eq!(0, min_opt(0, Some(1)));
218        assert_eq!(1, min_opt(2, Some(1)));
219    }
220
221    fn test_process_path() -> Option<PathBuf> {
222        env::current_exe().ok().and_then(|p| {
223            p.parent().map(|p| {
224                p.with_file_name("test-limited")
225                    .with_extension(env::consts::EXE_EXTENSION)
226            })
227        })
228    }
229
230    fn read_test_process(ulimit: Option<u64>) -> Result<u64> {
231        // Spawn the test helper and read it's result.
232        let path = test_process_path().unwrap();
233        let mut cmd = Command::new(&path);
234        let output = match ulimit {
235            Some(ulimit) => {
236                #[cfg(windows)]
237                {
238                    use std::mem::size_of;
239                    use std::process::Stdio;
240
241                    cmd.creation_flags(winapi::um::winbase::CREATE_SUSPENDED);
242                    let job = match unsafe {
243                        winapi::um::winbase::CreateJobObjectA(
244                            NULL as *mut winapi::um::minwinbase::SECURITY_ATTRIBUTES,
245                            NULL as *const i8,
246                        )
247                    } {
248                        NULL => win_err("CreateJobObjectA"),
249                        handle => Ok(handle),
250                    }?;
251                    let mut job_info = winapi::um::winnt::JOBOBJECT_EXTENDED_LIMIT_INFORMATION {
252                        BasicLimitInformation:
253                            winapi::um::winnt::JOBOBJECT_BASIC_LIMIT_INFORMATION {
254                                LimitFlags: winapi::um::winnt::JOB_OBJECT_LIMIT_PROCESS_MEMORY,
255                                ..Default::default()
256                            },
257                        ProcessMemoryLimit: ulimit as usize,
258                        ..Default::default()
259                    };
260                    match unsafe {
261                        winapi::um::jobapi2::SetInformationJobObject(
262                            job,
263                            winapi::um::winnt::JobObjectExtendedLimitInformation,
264                            &mut job_info
265                                as *mut winapi::um::winnt::JOBOBJECT_EXTENDED_LIMIT_INFORMATION
266                                as LPVOID,
267                            size_of::<winapi::um::winnt::JOBOBJECT_EXTENDED_LIMIT_INFORMATION>()
268                                as u32,
269                        )
270                    } {
271                        FALSE => win_err("SetInformationJobObject"),
272                        _ => Ok(()),
273                    }?;
274                    let child = cmd
275                        .stdin(Stdio::null())
276                        .stdout(Stdio::piped())
277                        .stderr(Stdio::piped())
278                        .spawn()
279                        .map_err(|e| {
280                            crate::Error::IoExplainedError(e, "error spawning helper".into())
281                        })?;
282                    let childhandle = match unsafe {
283                        winapi::um::processthreadsapi::OpenProcess(
284                            winapi::um::winnt::JOB_OBJECT_ASSIGN_PROCESS
285                        // The docs say only JOB_OBJECT_ASSIGN_PROCESS is
286                        // needed, but access denied is returned unless more
287                        // permissions are requested, and the actual set needed
288                        // is not documented.
289                            | winapi::um::winnt::PROCESS_ALL_ACCESS,
290                            FALSE,
291                            child.id(),
292                        )
293                    } {
294                        NULL => win_err("OpenProcess"),
295                        handle => Ok(handle),
296                    }?;
297                    println!("assigning job {} pid {}", job as u32, childhandle as u32);
298                    let res =
299                        unsafe { winapi::um::jobapi2::AssignProcessToJobObject(job, childhandle) };
300                    match res {
301                        FALSE => win_err("AssignProcessToJobObject"),
302                        _ => Ok(()),
303                    }?;
304                    let mut tid: DWORD = 0;
305                    let tool = match unsafe {
306                        winapi::um::tlhelp32::CreateToolhelp32Snapshot(
307                            winapi::um::tlhelp32::TH32CS_SNAPTHREAD,
308                            0,
309                        )
310                    } {
311                        winapi::um::handleapi::INVALID_HANDLE_VALUE => {
312                            win_err("CreateToolhelp32Snapshot")
313                        }
314                        handle => Ok(handle),
315                    }?;
316                    let mut te = winapi::um::tlhelp32::THREADENTRY32 {
317                        dwSize: size_of::<winapi::um::tlhelp32::THREADENTRY32>() as u32,
318                        ..Default::default()
319                    };
320                    match unsafe { winapi::um::tlhelp32::Thread32First(tool, &mut te) } {
321                        FALSE => win_err("Thread32First"),
322                        _ => Ok(()),
323                    }?;
324                    while {
325                        if te.dwSize >= 16 /* owner proc id field offset */ &&te.th32OwnerProcessID == child.id()
326                        {
327                            tid = te.th32ThreadID;
328                            // a break here would be nice.
329                        };
330                        te.dwSize = size_of::<winapi::um::tlhelp32::THREADENTRY32>() as u32;
331                        match unsafe { winapi::um::tlhelp32::Thread32Next(tool, &mut te) } {
332                            FALSE => {
333                                let err = unsafe { winapi::um::errhandlingapi::GetLastError() };
334                                match err {
335                                    winapi::shared::winerror::ERROR_NO_MORE_FILES => Ok(false),
336                                    _ => win_err("Thread32Next"),
337                                }
338                            }
339                            _ => Ok(true),
340                        }?
341                    } {}
342                    match unsafe { winapi::um::handleapi::CloseHandle(tool) } {
343                        FALSE => win_err("CloseHandle"),
344                        _ => Ok(()),
345                    }?;
346                    let thread = match unsafe {
347                        winapi::um::processthreadsapi::OpenThread(
348                            winapi::um::winnt::THREAD_SUSPEND_RESUME,
349                            FALSE,
350                            tid,
351                        )
352                    } {
353                        NULL => win_err("OpenThread"),
354                        handle => Ok(handle),
355                    }?;
356
357                    match unsafe { winapi::um::processthreadsapi::ResumeThread(thread) } {
358                        std::u32::MAX => win_err("ResumeThread"),
359                        _ => Ok(()),
360                    }?;
361                    child.wait_with_output().map_err(|e| {
362                        crate::Error::IoExplainedError(e, "error waiting for child".into())
363                    })?
364                }
365                #[cfg(unix)]
366                {
367                    use std::io::Error;
368                    // https://github.com/rust-lang/libc/pull/1919
369                    cfg_if!(
370                    if #[cfg( target_os="netbsd")] {
371                    let rlimit_as = 10;
372                        }
373                    else {
374                    let rlimit_as = libc::RLIMIT_AS;
375                    }
376                    );
377                    unsafe {
378                        cmd.pre_exec(move || {
379                            let lim = libc::rlimit {
380                                rlim_cur: ulimit,
381                                rlim_max: libc::RLIM_INFINITY,
382                            };
383                            match libc::setrlimit(rlimit_as, &lim as *const libc::rlimit) {
384                                0 => Ok(()),
385                                _ => Err(Error::last_os_error()),
386                            }
387                        });
388                    }
389                    cmd.output().map_err(|e| {
390                        crate::Error::IoExplainedError(e, "error running helper".into())
391                    })?
392                }
393            }
394            None => cmd
395                .output()
396                .map_err(|e| crate::Error::IoExplainedError(e, "error running helper".into()))?,
397        };
398        assert_eq!(true, output.status.success());
399        eprintln!("stderr {}", str::from_utf8(&output.stderr).unwrap());
400        let limit_bytes = output.stdout;
401        let limit: u64 = str::from_utf8(&limit_bytes)?
402            .trim()
403            .parse()
404            .map_err(|e| Error::U64Error(e, str::from_utf8(&limit_bytes).unwrap().into()))?;
405
406        Ok(limit)
407    }
408
409    #[cfg(any(
410        windows,
411        target_os = "macos",
412        target_os = "linux",
413        target_os = "freebsd",
414        target_os = "illumos",
415        target_os = "solaris",
416    ))]
417    #[test]
418    fn test_no_ulimit() -> Result<()> {
419        // This test depends on the dev environment being run uncontained.
420        let info = sys_info::mem_info()?;
421        let total_ram = info.total * 1024;
422        let limit = read_test_process(None)?;
423        assert_eq!(total_ram, limit);
424        Ok(())
425    }
426
427    #[test]
428    fn test_ulimit() -> Result<()> {
429        // Page size rounding
430        let limit = read_test_process(Some(99_999_744))?;
431        assert_eq!(99_999_744, limit);
432        Ok(())
433    }
434}
effective_limits/lib.rs

effective_limits/
lib.rs
