Skip to main content

mz_ore/
cgroup.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License in the LICENSE file at the
6// root of this repository, or online at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Linux cgroup detection utilities.
17use std::fs::File;
18use std::io::{BufRead, BufReader};
19use std::path::{Path, PathBuf};
20
21/// An entry in /proc/self/cgroup
22#[derive(Debug, PartialEq)]
23pub struct CgroupEntry {
24    subsystems: Vec<String>,
25    root: PathBuf,
26}
27
28impl CgroupEntry {
29    fn from_line(line: String) -> Option<CgroupEntry> {
30        let mut fields = line.split(':');
31        let subsystems = fields
32            .nth(1)?
33            .split(',')
34            .filter(|s| !s.is_empty())
35            .map(|s| s.to_owned())
36            .collect();
37        let root = PathBuf::from(fields.next()?);
38        Some(CgroupEntry { subsystems, root })
39    }
40}
41
42/// Parses /proc/self/cgroup into a `Vec<CgroupEntry>`, if the file exists.
43pub fn parse_proc_self_cgroup() -> Option<Vec<CgroupEntry>> {
44    let file = File::open("/proc/self/cgroup").ok()?;
45    let file = BufReader::new(file);
46    Some(
47        file.lines()
48            .map_while(Result::ok)
49            .filter_map(CgroupEntry::from_line)
50            .collect(),
51    )
52}
53
54/// An entry in /proc/self/mountinfo
55#[derive(Debug, PartialEq)]
56pub struct Mountinfo {
57    root: PathBuf,
58    mount_point: PathBuf,
59    fs_type: String,
60    super_opts: Vec<String>,
61}
62
63impl Mountinfo {
64    fn from_line(line: String) -> Option<Mountinfo> {
65        // https://www.kernel.org/doc/Documentation/filesystems/proc.txt
66        let mut split = line.split(" - ");
67
68        let mut mount_fields = split.next()?.split(' ');
69        let root = PathBuf::from(mount_fields.nth(3)?);
70        let mount_point = PathBuf::from(mount_fields.next()?);
71
72        let mut fs_fields = split.next()?.split(' ');
73
74        let fs_type = fs_fields.next()?.split('.').next()?.to_owned();
75        let super_opts: Vec<String> = fs_fields
76            .nth(1)?
77            .split(',')
78            .filter(|s| !s.is_empty())
79            .map(|s| s.to_owned())
80            .collect();
81
82        Some(Mountinfo {
83            root,
84            mount_point,
85            fs_type,
86            super_opts,
87        })
88    }
89}
90
91/// Parses /proc/self/mountinfo into vectors of Mountinfo objects,
92/// returning (v2_mounts, v1_mounts).
93pub fn parse_proc_self_mountinfo() -> Option<(Vec<Mountinfo>, Vec<Mountinfo>)> {
94    let file = File::open("/proc/self/mountinfo").ok()?;
95    let file = BufReader::new(file);
96    Some(
97        file.lines()
98            .map_while(Result::ok)
99            .filter_map(Mountinfo::from_line)
100            .filter(|mi| mi.fs_type == "cgroup" || mi.fs_type == "cgroup2")
101            .partition(|mi| mi.fs_type == "cgroup2"),
102    )
103}
104
105/// Represents a cgroup memory limit, with both ram and swap maximums if they exist.
106/// Fields should be None if a limit does not exist or when running on a platform without cgroup
107/// support (ie: non-Linux platforms).
108#[derive(Debug)]
109pub struct MemoryLimit {
110    /// Maximum RAM limit, in bytes, if a limit exists.
111    pub max: Option<usize>,
112    /// Maximum swap limit, in bytes, if a limit exists.
113    pub swap_max: Option<usize>,
114}
115
116fn read_file_to_usize<P: AsRef<Path>>(path: P) -> Option<usize> {
117    match std::fs::read_to_string(path.as_ref()) {
118        Ok(s) => s.trim().parse::<usize>().ok(),
119        Err(_) => None,
120    }
121}
122
123/// Finds the mountpoint corresponding to the provided cgroup v2,
124/// and reads the memory limits within.
125fn read_v2_memory_limit(cgroups: &[CgroupEntry], mounts: &[Mountinfo]) -> Option<MemoryLimit> {
126    // cgroups v2 only supports a single cgroup per process
127    let mount = mounts.first()?;
128    if mount.root != cgroups.first()?.root {
129        // We don't support mixed v2/v1
130        return None;
131    }
132    let mount_point = &mount.mount_point;
133    let controllers = std::fs::read_to_string(mount_point.join("cgroup.controllers")).ok()?;
134    let mut controllers = controllers.trim().split(' ');
135    if controllers.any(|c| c == "memory") {
136        let max = read_file_to_usize(mount_point.join("memory.max"));
137        // Unlike v1, this is only the swap, not swap + memory.
138        let swap_max = read_file_to_usize(mount_point.join("memory.swap.max"));
139        return Some(MemoryLimit { max, swap_max });
140    }
141    None
142}
143
144/// Finds the cgroup v1 and mountpoint combination containing the memory controller,
145/// and reads the memory limits within.
146fn read_v1_memory_limit(cgroups: &[CgroupEntry], mounts: &[Mountinfo]) -> Option<MemoryLimit> {
147    // https://www.kernel.org/doc/Documentation/cgroup-v1/memory.txt
148    let memory_cgroup = cgroups
149        .into_iter()
150        .find(|cgroup| cgroup.subsystems.iter().any(|s| s == "memory"))?;
151    let memory_mount = mounts
152        .iter()
153        .find(|mi| mi.root == memory_cgroup.root && mi.super_opts.iter().any(|o| o == "memory"))?;
154    let mount_point = &memory_mount.mount_point;
155    let max = read_file_to_usize(mount_point.join("memory.limit_in_bytes"));
156    // This is memory + swap, not just swap.
157    let memsw_max = read_file_to_usize(mount_point.join("memory.memsw.limit_in_bytes"));
158    let swap_max = match (max, memsw_max) {
159        (Some(max), Some(memsw_max)) => Some(memsw_max - max),
160        _ => None,
161    };
162    Some(MemoryLimit { max, swap_max })
163}
164/// Returns the cgroup (v1 or v2) memory limit if it exists.
165pub fn detect_memory_limit() -> Option<MemoryLimit> {
166    let (v2_mounts, v1_mounts) = parse_proc_self_mountinfo()?;
167    let cgroups = parse_proc_self_cgroup()?;
168
169    if !v2_mounts.is_empty() {
170        return read_v2_memory_limit(&cgroups, &v2_mounts);
171    }
172    read_v1_memory_limit(&cgroups, &v1_mounts)
173}
174
175#[cfg(test)]
176mod tests {
177    use std::path::PathBuf;
178
179    use super::{CgroupEntry, Mountinfo};
180
181    #[crate::test]
182    fn test_cgroup_from_line() {
183        // cgroups v2
184        assert_eq!(
185            CgroupEntry::from_line("0::/".to_owned()),
186            Some(CgroupEntry {
187                subsystems: vec![],
188                root: PathBuf::from("/"),
189            })
190        );
191
192        // cgroups v1
193        assert_eq!(
194            CgroupEntry::from_line("6:cpu,cpuacct:/kubepods/pod5b977639-f878-469b-94ee-47a4aa7e597a/dd55abbabd99bcb4d2ce17ffa77d6f811c90e09202f537c273962a8259cac8a0".to_owned()),
195            Some(CgroupEntry {
196                subsystems: vec!["cpu".to_owned(), "cpuacct".to_owned()],
197                root: PathBuf::from("/kubepods/pod5b977639-f878-469b-94ee-47a4aa7e597a/dd55abbabd99bcb4d2ce17ffa77d6f811c90e09202f537c273962a8259cac8a0"),
198            })
199        );
200        assert_eq!(
201            CgroupEntry::from_line("5:memory:/kubepods/pod5b977639-f878-469b-94ee-47a4aa7e597a/dd55abbabd99bcb4d2ce17ffa77d6f811c90e09202f537c273962a8259cac8a0".to_owned()),
202            Some(CgroupEntry {
203                subsystems: vec!["memory".to_owned()],
204                root: PathBuf::from("/kubepods/pod5b977639-f878-469b-94ee-47a4aa7e597a/dd55abbabd99bcb4d2ce17ffa77d6f811c90e09202f537c273962a8259cac8a0"),
205            })
206        );
207    }
208
209    #[crate::test]
210    fn test_mountinfo_from_line() {
211        // Mount with optional field (master:305)
212        assert_eq!(Mountinfo::from_line("863 758 0:63 / / rw,relatime master:305 - overlay overlay rw,seclabel,lowerdir=/var/lib/docker/overlay2/l/SUKWDHL7W7YZCJ6YI66I7Z5PR2:/var/lib/docker/overlay2/l/ORL2I23UNUGM7FYF4BSL5JUCAB:/var/lib/docker/overlay2/l/LLKK3J2EHGPF5IGGDSAQGRFHLV:/var/lib/docker/overlay2/l/JEQIUQIQTVNRBAGCU7SLV4KK4K:/var/lib/docker/overlay2/l/5DS7KSJCA7BHWAYWII7BI5DBC5:/var/lib/docker/overlay2/l/ZAGXZ62GNFPZFLNUDZ3JOZIMYR:/var/lib/docker/overlay2/l/6WVXMD372IA24ZXRWGGTIPEQPA,upperdir=/var/lib/docker/overlay2/5c7734eb769484f3469b234181365466eb30bcd7f31c912f4250c8d701637ee4/diff,workdir=/var/lib/docker/overlay2/5c7734eb769484f3469b234181365466eb30bcd7f31c912f4250c8d701637ee4/work".to_owned()),
213        Some(Mountinfo{
214            root: PathBuf::from("/"),
215            mount_point: PathBuf::from("/"),
216            fs_type: "overlay".to_owned(),
217            super_opts: vec![
218                "rw".to_owned(),
219                "seclabel".to_owned(),
220                "lowerdir=/var/lib/docker/overlay2/l/SUKWDHL7W7YZCJ6YI66I7Z5PR2:/var/lib/docker/overlay2/l/ORL2I23UNUGM7FYF4BSL5JUCAB:/var/lib/docker/overlay2/l/LLKK3J2EHGPF5IGGDSAQGRFHLV:/var/lib/docker/overlay2/l/JEQIUQIQTVNRBAGCU7SLV4KK4K:/var/lib/docker/overlay2/l/5DS7KSJCA7BHWAYWII7BI5DBC5:/var/lib/docker/overlay2/l/ZAGXZ62GNFPZFLNUDZ3JOZIMYR:/var/lib/docker/overlay2/l/6WVXMD372IA24ZXRWGGTIPEQPA".to_owned(),
221                "upperdir=/var/lib/docker/overlay2/5c7734eb769484f3469b234181365466eb30bcd7f31c912f4250c8d701637ee4/diff".to_owned(),
222                "workdir=/var/lib/docker/overlay2/5c7734eb769484f3469b234181365466eb30bcd7f31c912f4250c8d701637ee4/work".to_owned(),
223            ],
224        })
225        );
226
227        // cgroups v2
228        assert_eq!(Mountinfo::from_line("868 867 0:27 / /sys/fs/cgroup ro,nosuid,nodev,noexec,relatime - cgroup2 cgroup rw,seclabel,nsdelegate,memory_recursiveprot".to_owned()), Some(Mountinfo{
229
230            root: PathBuf::from("/"),
231            mount_point: PathBuf::from("/sys/fs/cgroup"),
232            fs_type: "cgroup2".to_owned(),
233            super_opts: vec![
234                "rw".to_owned(),
235                "seclabel".to_owned(),
236                "nsdelegate".to_owned(),
237                "memory_recursiveprot".to_owned(),
238            ],
239        }));
240
241        // cgroups v1
242        assert_eq!(Mountinfo::from_line("702 697 0:30 /kubepods/pod5b977639-f878-469b-94ee-47a4aa7e597a/dd55abbabd99bcb4d2ce17ffa77d6f811c90e09202f537c273962a8259cac8a0 /sys/fs/cgroup/memory ro,nosuid,nodev,noexec,relatime master:13 - cgroup cgroup rw,memory".to_owned()), Some(Mountinfo{
243
244            root: PathBuf::from("/kubepods/pod5b977639-f878-469b-94ee-47a4aa7e597a/dd55abbabd99bcb4d2ce17ffa77d6f811c90e09202f537c273962a8259cac8a0"),
245            mount_point: PathBuf::from("/sys/fs/cgroup/memory"),
246            fs_type: "cgroup".to_owned(),
247            super_opts: vec![
248                "rw".to_owned(),
249                "memory".to_owned(),
250            ],
251        }));
252    }
253}