Line data Source code
1 : //! Wrapper around nix::sys::statvfs::Statvfs that allows for mocking.
2 :
3 : use camino::Utf8Path;
4 :
5 : pub enum Statvfs {
6 : Real(nix::sys::statvfs::Statvfs),
7 : Mock(mock::Statvfs),
8 : }
9 :
10 : // NB: on macOS, the block count type of struct statvfs is u32.
11 : // The workaround seems to be to use the non-standard statfs64 call.
12 : // Sincce it should only be a problem on > 2TiB disks, let's ignore
13 : // the problem for now and upcast to u64.
14 : impl Statvfs {
15 9 : pub fn get(tenants_dir: &Utf8Path, mocked: Option<&mock::Behavior>) -> nix::Result<Self> {
16 9 : if let Some(mocked) = mocked {
17 9 : Ok(Statvfs::Mock(mock::get(tenants_dir, mocked)?))
18 : } else {
19 0 : Ok(Statvfs::Real(nix::sys::statvfs::statvfs(
20 0 : tenants_dir.as_std_path(),
21 0 : )?))
22 : }
23 9 : }
24 :
25 : // NB: allow() because the block count type is u32 on macOS.
26 : #[allow(clippy::useless_conversion, clippy::unnecessary_fallible_conversions)]
27 8 : pub fn blocks(&self) -> u64 {
28 8 : match self {
29 0 : Statvfs::Real(stat) => u64::try_from(stat.blocks()).unwrap(),
30 8 : Statvfs::Mock(stat) => stat.blocks,
31 : }
32 8 : }
33 :
34 : // NB: allow() because the block count type is u32 on macOS.
35 : #[allow(clippy::useless_conversion, clippy::unnecessary_fallible_conversions)]
36 8 : pub fn blocks_available(&self) -> u64 {
37 8 : match self {
38 0 : Statvfs::Real(stat) => u64::try_from(stat.blocks_available()).unwrap(),
39 8 : Statvfs::Mock(stat) => stat.blocks_available,
40 : }
41 8 : }
42 :
43 16 : pub fn fragment_size(&self) -> u64 {
44 16 : match self {
45 0 : Statvfs::Real(stat) => stat.fragment_size(),
46 16 : Statvfs::Mock(stat) => stat.fragment_size,
47 : }
48 16 : }
49 :
50 0 : pub fn block_size(&self) -> u64 {
51 0 : match self {
52 0 : Statvfs::Real(stat) => stat.block_size(),
53 0 : Statvfs::Mock(stat) => stat.block_size,
54 : }
55 0 : }
56 : }
57 :
58 : pub mod mock {
59 : use anyhow::Context;
60 : use camino::Utf8Path;
61 : use regex::Regex;
62 : use tracing::log::info;
63 :
64 32 : #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
65 : #[serde(tag = "type")]
66 : pub enum Behavior {
67 : Success {
68 : blocksize: u64,
69 : total_blocks: u64,
70 : name_filter: Option<utils::serde_regex::Regex>,
71 : },
72 : Failure {
73 : mocked_error: MockedError,
74 : },
75 : }
76 :
77 2 : #[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
78 : #[allow(clippy::upper_case_acronyms)]
79 : pub enum MockedError {
80 : EIO,
81 : }
82 :
83 : impl From<MockedError> for nix::Error {
84 1 : fn from(e: MockedError) -> Self {
85 1 : match e {
86 1 : MockedError::EIO => nix::Error::EIO,
87 1 : }
88 1 : }
89 : }
90 :
91 9 : pub fn get(tenants_dir: &Utf8Path, behavior: &Behavior) -> nix::Result<Statvfs> {
92 9 : info!("running mocked statvfs");
93 :
94 9 : match behavior {
95 : Behavior::Success {
96 8 : blocksize,
97 8 : total_blocks,
98 8 : ref name_filter,
99 8 : } => {
100 8 : let used_bytes = walk_dir_disk_usage(tenants_dir, name_filter.as_deref()).unwrap();
101 8 :
102 8 : // round it up to the nearest block multiple
103 8 : let used_blocks = (used_bytes + (blocksize - 1)) / blocksize;
104 8 :
105 8 : if used_blocks > *total_blocks {
106 0 : panic!(
107 0 : "mocking error: used_blocks > total_blocks: {used_blocks} > {total_blocks}"
108 0 : );
109 8 : }
110 8 :
111 8 : let avail_blocks = total_blocks - used_blocks;
112 8 :
113 8 : Ok(Statvfs {
114 8 : blocks: *total_blocks,
115 8 : blocks_available: avail_blocks,
116 8 : fragment_size: *blocksize,
117 8 : block_size: *blocksize,
118 8 : })
119 : }
120 1 : Behavior::Failure { mocked_error } => Err((*mocked_error).into()),
121 : }
122 9 : }
123 :
124 8 : fn walk_dir_disk_usage(path: &Utf8Path, name_filter: Option<&Regex>) -> anyhow::Result<u64> {
125 8 : let mut total = 0;
126 331 : for entry in walkdir::WalkDir::new(path) {
127 331 : let entry = entry?;
128 331 : if !entry.file_type().is_file() {
129 56 : continue;
130 275 : }
131 275 : if !name_filter
132 275 : .as_ref()
133 275 : .map(|filter| filter.is_match(entry.file_name().to_str().unwrap()))
134 275 : .unwrap_or(true)
135 : {
136 48 : continue;
137 227 : }
138 227 : total += entry
139 227 : .metadata()
140 227 : .with_context(|| format!("get metadata of {:?}", entry.path()))?
141 227 : .len();
142 : }
143 8 : Ok(total)
144 8 : }
145 :
146 : pub struct Statvfs {
147 : pub blocks: u64,
148 : pub blocks_available: u64,
149 : pub fragment_size: u64,
150 : pub block_size: u64,
151 : }
152 : }
|
