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            0 :     pub fn get(tenants_dir: &Utf8Path, mocked: Option<&mock::Behavior>) -> nix::Result<Self> {
      16            0 :         if let Some(mocked) = mocked {
      17            0 :             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            0 :     }
      24              : 
      25              :     // NB: allow() because the block count type is u32 on macOS.
      26              :     #[allow(clippy::useless_conversion, clippy::unnecessary_fallible_conversions)]
      27            0 :     pub fn blocks(&self) -> u64 {
      28            0 :         match self {
      29            0 :             Statvfs::Real(stat) => u64::try_from(stat.blocks()).unwrap(),
      30            0 :             Statvfs::Mock(stat) => stat.blocks,
      31              :         }
      32            0 :     }
      33              : 
      34              :     // NB: allow() because the block count type is u32 on macOS.
      35              :     #[allow(clippy::useless_conversion, clippy::unnecessary_fallible_conversions)]
      36            0 :     pub fn blocks_available(&self) -> u64 {
      37            0 :         match self {
      38            0 :             Statvfs::Real(stat) => u64::try_from(stat.blocks_available()).unwrap(),
      39            0 :             Statvfs::Mock(stat) => stat.blocks_available,
      40              :         }
      41            0 :     }
      42              : 
      43            0 :     pub fn fragment_size(&self) -> u64 {
      44            0 :         match self {
      45            0 :             Statvfs::Real(stat) => stat.fragment_size(),
      46            0 :             Statvfs::Mock(stat) => stat.fragment_size,
      47              :         }
      48            0 :     }
      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 camino::Utf8Path;
      60              :     use regex::Regex;
      61              :     use tracing::log::info;
      62              : 
      63              :     pub use pageserver_api::config::statvfs::mock::Behavior;
      64              : 
      65            0 :     pub fn get(tenants_dir: &Utf8Path, behavior: &Behavior) -> nix::Result<Statvfs> {
      66            0 :         info!("running mocked statvfs");
      67              : 
      68            0 :         match behavior {
      69              :             Behavior::Success {
      70            0 :                 blocksize,
      71            0 :                 total_blocks,
      72            0 :                 ref name_filter,
      73            0 :             } => {
      74            0 :                 let used_bytes = walk_dir_disk_usage(tenants_dir, name_filter.as_deref()).unwrap();
      75            0 : 
      76            0 :                 // round it up to the nearest block multiple
      77            0 :                 let used_blocks = (used_bytes + (blocksize - 1)) / blocksize;
      78            0 : 
      79            0 :                 if used_blocks > *total_blocks {
      80            0 :                     panic!(
      81            0 :                         "mocking error: used_blocks > total_blocks: {used_blocks} > {total_blocks}"
      82            0 :                     );
      83            0 :                 }
      84            0 : 
      85            0 :                 let avail_blocks = total_blocks - used_blocks;
      86            0 : 
      87            0 :                 Ok(Statvfs {
      88            0 :                     blocks: *total_blocks,
      89            0 :                     blocks_available: avail_blocks,
      90            0 :                     fragment_size: *blocksize,
      91            0 :                     block_size: *blocksize,
      92            0 :                 })
      93              :             }
      94              :             #[cfg(feature = "testing")]
      95            0 :             Behavior::Failure { mocked_error } => Err((*mocked_error).into()),
      96              :         }
      97            0 :     }
      98              : 
      99            0 :     fn walk_dir_disk_usage(path: &Utf8Path, name_filter: Option<&Regex>) -> anyhow::Result<u64> {
     100            0 :         let mut total = 0;
     101            0 :         for entry in walkdir::WalkDir::new(path) {
     102            0 :             let entry = entry?;
     103            0 :             if !entry.file_type().is_file() {
     104            0 :                 continue;
     105            0 :             }
     106            0 :             if !name_filter
     107            0 :                 .as_ref()
     108            0 :                 .map(|filter| filter.is_match(entry.file_name().to_str().unwrap()))
     109            0 :                 .unwrap_or(true)
     110              :             {
     111            0 :                 continue;
     112            0 :             }
     113            0 :             let m = match entry.metadata() {
     114            0 :                 Ok(m) => m,
     115            0 :                 Err(e) if is_not_found(&e) => {
     116            0 :                     // some temp file which got removed right as we are walking
     117            0 :                     continue;
     118              :                 }
     119            0 :                 Err(e) => {
     120            0 :                     return Err(anyhow::Error::new(e)
     121            0 :                         .context(format!("get metadata of {:?}", entry.path())))
     122              :                 }
     123              :             };
     124            0 :             total += m.len();
     125              :         }
     126            0 :         Ok(total)
     127            0 :     }
     128              : 
     129            0 :     fn is_not_found(e: &walkdir::Error) -> bool {
     130            0 :         let Some(io_error) = e.io_error() else {
     131            0 :             return false;
     132              :         };
     133            0 :         let kind = io_error.kind();
     134            0 :         matches!(kind, std::io::ErrorKind::NotFound)
     135            0 :     }
     136              : 
     137              :     pub struct Statvfs {
     138              :         pub blocks: u64,
     139              :         pub blocks_available: u64,
     140              :         pub fragment_size: u64,
     141              :         pub block_size: u64,
     142              :     }
     143              : }