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