below_btrfs/
lib.rs

1// Copyright (c) Facebook, Inc. and its affiliates.
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 at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#![deny(clippy::all)]
16
17use std::collections::HashMap;
18use std::fs::File;
19use std::os::unix::io::AsRawFd;
20use std::path::Path;
21use std::path::PathBuf;
22use std::rc::Rc;
23
24use rand_distr::Distribution;
25use rand_distr::Uniform;
26use slog::error;
27use slog::warn;
28
29pub mod btrfs_api;
30
31#[cfg(not(fbcode_build))]
32pub use btrfs_api::open_source::btrfs_sys::*;
33#[cfg(fbcode_build)]
34pub use btrfs_sys::*;
35
36mod types;
37pub use types::*;
38
39#[cfg(test)]
40mod test;
41use thiserror::Error;
42
43#[derive(Error, Debug)]
44pub enum Error {
45    #[error("Invalid file format: {0:?}")]
46    InvalidFileFormat(PathBuf),
47    #[error("{1:?}: {0:?}")]
48    IoError(PathBuf, #[source] std::io::Error),
49    #[error("Failed call to btrfs")]
50    SysError(btrfs_api::Error),
51}
52
53impl From<btrfs_api::Error> for Error {
54    fn from(item: btrfs_api::Error) -> Self {
55        Error::SysError(item)
56    }
57}
58
59pub type Result<T> = std::result::Result<T, Error>;
60
61pub const DEFAULT_ROOT: &str = "/";
62pub const DEFAULT_SAMPLES: u64 = 100;
63pub const DEFAULT_MIN_PCT: f64 = 0.0;
64
65// The SampleTree structure stores a hierarchical structure
66// of path names that we have some size estimations for. This
67// is supposed to follow the structure of files in the subvolume
68struct SampleTree {
69    // total number of samples under this tree.
70    total: usize,
71    children: HashMap<String, SampleTree>,
72}
73
74impl Default for SampleTree {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80impl SampleTree {
81    fn new() -> Self {
82        Self {
83            total: 0,
84            children: HashMap::new(),
85        }
86    }
87
88    // path implements an iterator trait. This method is recursive and completes because path.next()
89    // consumes one path instance from the iterator at a time.
90    fn add<'a>(&mut self, mut path: impl Iterator<Item = &'a str>) {
91        if let Some(p) = path.next() {
92            self.total += 1;
93            self.children.entry(p.to_string()).or_default().add(path);
94        }
95    }
96
97    // This method parses the sample tree and outputs the paths corresponding to subvolumes that occupy
98    // more than min_disk_fraction of the disk.
99    fn convert(
100        &self,
101        total_samples: usize,
102        total_length: u64,
103        min_disk_fraction: Option<f64>,
104    ) -> Result<BtrfsMap> {
105        let mut btrfs_map: BtrfsMap = Default::default();
106        self.convert_internal(
107            total_samples,
108            total_length,
109            min_disk_fraction,
110            "".to_string(),
111            &mut btrfs_map,
112        )?;
113
114        Ok(btrfs_map)
115    }
116
117    fn convert_internal(
118        &self,
119        total_samples: usize,
120        total_length: u64,
121        min_disk_fraction: Option<f64>,
122        base_path: String,
123        btrfs_map: &mut BtrfsMap,
124    ) -> Result<()> {
125        for (p, child_tree) in &self.children {
126            let dfraction = (child_tree.total as f64) / (total_samples as f64);
127            let dbytes = (total_length as f64 * dfraction) as u64;
128
129            match min_disk_fraction {
130                Some(min_disk_fraction) if dfraction < min_disk_fraction => continue,
131                _ => {}
132            }
133
134            let path = format!("{}/{}", base_path, p);
135
136            let btrfs_stat = BtrfsStat {
137                name: Some(path.clone()),
138                disk_fraction: Some(dfraction * 100.0),
139                disk_bytes: Some(dbytes),
140            };
141
142            btrfs_map.insert(path.clone(), btrfs_stat);
143
144            child_tree.convert_internal(
145                total_samples,
146                total_length,
147                min_disk_fraction,
148                path,
149                btrfs_map,
150            )?;
151        }
152
153        Ok(())
154    }
155}
156
157// This structure contains for each btrfs instance a hashmap of the subvolume ids and
158// their respective paths.
159struct Roots {
160    fd: i32,
161    // hashmap key is subvolume id and value is vector with the path of that subvolume
162    m: HashMap<u64, Rc<Vec<String>>>,
163}
164
165impl Roots {
166    fn new(fd: i32) -> Self {
167        Self {
168            fd,
169            m: HashMap::from([(BTRFS_FS_TREE_OBJECTID as u64, Rc::new(Vec::new()))]),
170        }
171    }
172
173    fn get_root(&mut self, root_id: u64) -> Result<Rc<Vec<String>>> {
174        match self.m.get(&root_id) {
175            Some(path) => Ok(Rc::clone(path)),
176            None => {
177                let root_backref = btrfs_api::find_root_backref(self.fd, root_id)?;
178                match root_backref {
179                    Some((name, parent_id)) => {
180                        let rec_root = self.get_root(parent_id)?;
181                        let mut path = Vec::clone(&rec_root);
182                        path.push(name);
183                        let path_rc = Rc::new(path);
184                        self.m.insert(root_id, path_rc.clone());
185                        Ok(path_rc)
186                    }
187                    None => Err(Error::SysError(btrfs_api::Error::SysError(
188                        nix::errno::Errno::ENOENT,
189                    ))),
190                }
191            }
192        }
193    }
194}
195
196pub struct BtrfsReader {
197    samples: u64,
198    min_pct: f64,
199    path: PathBuf,
200    logger: slog::Logger,
201}
202
203impl BtrfsReader {
204    pub fn new(samples: u64, min_pct: f64, logger: slog::Logger) -> BtrfsReader {
205        BtrfsReader::new_with_path(DEFAULT_ROOT.to_string(), samples, min_pct, logger)
206    }
207
208    pub fn new_with_path(
209        p: String,
210        samples: u64,
211        min_pct: f64,
212        logger: slog::Logger,
213    ) -> BtrfsReader {
214        BtrfsReader {
215            samples,
216            min_pct,
217            path: p.into(),
218            logger,
219        }
220    }
221
222    pub fn sample(&self) -> Result<BtrfsMap> {
223        let f = File::open(&self.path).map_err(|e| self.io_error(&self.path, e))?;
224
225        let fd = f.as_raw_fd();
226
227        #[derive(Debug)]
228        struct ChunkInfo {
229            pos: u64,
230            chunk_offset: u64,
231            chunk_length: u64,
232            chunk_type: u64,
233        }
234
235        let samples = self.samples;
236        let mut chunks = Vec::<ChunkInfo>::new();
237        let mut total_chunk_length = 0;
238        let mut chunks_size = 0;
239        btrfs_api::tree_search_cb(
240            fd,
241            btrfs_api::BTRFS_CHUNK_TREE_OBJECTID as u64,
242            btrfs_api::SearchKey::ALL,
243            |sh, data| {
244                if sh.type_ == btrfs_api::BTRFS_CHUNK_ITEM_KEY {
245                    let chunk = unsafe { &*(data.as_ptr() as *const btrfs_api::btrfs_chunk) };
246                    chunks.push(ChunkInfo {
247                        pos: total_chunk_length,
248                        chunk_offset: sh.offset,
249                        chunk_length: chunk.length,
250                        chunk_type: chunk.type_,
251                    });
252                    chunks_size += 1;
253                    total_chunk_length += chunk.length;
254                };
255            },
256        )
257        .map_err(Error::SysError)?;
258
259        let mut roots = Roots::new(fd);
260        let uniform = Uniform::new(0, total_chunk_length);
261        let mut rng = rand::thread_rng();
262
263        let mut sample_tree = SampleTree::new();
264        let mut total_samples = 0;
265
266        let mut random_positions = Vec::new();
267        for _ in 0..samples {
268            random_positions.push(uniform.sample(&mut rng));
269        }
270        random_positions.sort_unstable();
271
272        let mut chunk_idx = 0;
273        for random_position in &random_positions {
274            while random_position > &(chunks[chunk_idx].pos + chunks[chunk_idx].chunk_length) {
275                chunk_idx += 1;
276            }
277
278            let random_chunk = &chunks[chunk_idx];
279            total_samples += 1;
280            match (random_chunk.chunk_type as u32) & btrfs_api::BTRFS_BLOCK_GROUP_TYPE_MASK {
281                btrfs_api::BTRFS_BLOCK_GROUP_DATA => {
282                    let random_offset =
283                        random_chunk.chunk_offset + (random_position - random_chunk.pos);
284                    let mut err = Ok(());
285                    btrfs_api::logical_ino(fd, random_offset, false, |res| match res {
286                        Ok(inodes) => {
287                            for inode in inodes {
288                                btrfs_api::ino_lookup(fd, inode.root, inode.inum, |res| match res {
289                                    Ok(path) => match roots.get_root(inode.root) {
290                                        Ok(root_path) => {
291                                            let root_path_it = root_path.iter().map(|s| s.as_str());
292                                            let inode_path = path
293                                                .to_str()
294                                                .expect("Could not convert path to string")
295                                                .split('/')
296                                                .filter(|s| !s.is_empty());
297                                            sample_tree.add(root_path_it.chain(inode_path));
298                                        }
299                                        Err(e) => {
300                                            err = Err(e);
301                                        }
302                                    },
303                                    Err(btrfs_api::Error::SysError(nix::errno::Errno::ENOENT)) => {}
304                                    Err(e) => {
305                                        warn!(
306                                            self.logger,
307                                            "INO_LOOKUP Returned error {} for inode.root {} and inode.inum {}",
308                                            e,
309                                            inode.root,
310                                            inode.inum
311                                        );
312                                    }
313                                })
314                            }
315                        }
316                        Err(btrfs_api::Error::SysError(nix::errno::Errno::ENOENT)) => {}
317                        Err(e) => {
318                            warn!(
319                                self.logger,
320                                "LOGICAL_INO returned error {} for random offset {} ",
321                                e,
322                                random_offset
323                            );
324                        }
325                    });
326                    err?;
327                }
328                btrfs_api::BTRFS_BLOCK_GROUP_METADATA => {}
329                btrfs_api::BTRFS_BLOCK_GROUP_SYSTEM => {}
330                _ => {}
331            };
332        }
333
334        sample_tree.convert(
335            total_samples,
336            total_chunk_length,
337            Some(self.min_pct / 100.0),
338        )
339    }
340
341    fn io_error<P: AsRef<Path>>(&self, file_name: P, e: std::io::Error) -> Error {
342        let mut p = self.path.clone();
343        p.push(file_name);
344        Error::IoError(p, e)
345    }
346}
below_btrfs/lib.rs

below_btrfs/
lib.rs
