src/shared_rate_limiter.rs

   1 use std::path::PathBuf;
   2 use std::mem::MaybeUninit;
   3 use std::time::{Instant, Duration};
   4
   5 use anyhow::{bail, Error};
   6 use nix::sys::stat::Mode;
   7
   8 use proxmox::tools::fs::{create_path, CreateOptions};
   9
  10 use proxmox_http::client::{RateLimit, RateLimiter, ShareableRateLimit};
  11 use proxmox_shared_memory::{Init, SharedMemory, SharedMutex};
  12 use proxmox_shared_memory::{check_subtype, initialize_subtype};
  13
  14 // openssl::sha::sha256(b"Proxmox Backup SharedRateLimiter v1.0")[0..8];
  15 pub const PROXMOX_BACKUP_SHARED_RATE_LIMITER_MAGIC_1_0: [u8; 8] = [6, 58, 213, 96, 161, 122, 130, 117];
  16
  17 const BASE_PATH: &str = pbs_buildcfg::rundir!("/shmem/tbf");
  18
  19 // Wrap RateLimiter, so that we can provide an Init impl
  20 #[repr(C)]
  21 struct WrapLimiter(RateLimiter);
  22
  23 impl Init for WrapLimiter {
  24     fn initialize(this: &mut MaybeUninit<Self>) {
  25         // default does not matter here, because we override later
  26         this.write(WrapLimiter(RateLimiter::new(1_000_000, 1_000_000)));
  27     }
  28 }
  29
  30 #[repr(C)]
  31 struct SharedRateLimiterData {
  32     magic: [u8; 8],
  33     tbf: SharedMutex<WrapLimiter>,
  34     padding: [u8; 4096 - 104],
  35 }
  36
  37 impl Init for SharedRateLimiterData {
  38     fn initialize(this: &mut MaybeUninit<Self>) {
  39         unsafe {
  40             let me = &mut *this.as_mut_ptr();
  41             me.magic = PROXMOX_BACKUP_SHARED_RATE_LIMITER_MAGIC_1_0;
  42             initialize_subtype(&mut me.tbf);
  43         }
  44     }
  45
  46     fn check_type_magic(this: &MaybeUninit<Self>) -> Result<(), Error> {
  47         unsafe {
  48             let me = &*this.as_ptr();
  49             if me.magic != PROXMOX_BACKUP_SHARED_RATE_LIMITER_MAGIC_1_0 {
  50                 bail!("SharedRateLimiterData: wrong magic number");
  51             }
  52             check_subtype(&me.tbf)?;
  53             Ok(())
  54         }
  55     }
  56 }
  57
  58 pub struct SharedRateLimiter {
  59     shmem: SharedMemory<SharedRateLimiterData>
  60 }
  61
  62 impl SharedRateLimiter {
  63
  64     pub fn mmap_shmem(name: &str, rate: u64, burst: u64) -> Result<Self, Error> {
  65         let mut path = PathBuf::from(BASE_PATH);
  66
  67         let user = pbs_config::backup_user()?;
  68
  69         let dir_opts = CreateOptions::new()
  70             .perm(Mode::from_bits_truncate(0o770))
  71             .owner(user.uid)
  72             .group(user.gid);
  73
  74         create_path(
  75             &path,
  76             Some(dir_opts.clone()),
  77             Some(dir_opts))?;
  78
  79         path.push(name);
  80
  81         let file_opts = CreateOptions::new()
  82             .perm(Mode::from_bits_truncate(0o660))
  83             .owner(user.uid)
  84             .group(user.gid);
  85
  86         let shmem: SharedMemory<SharedRateLimiterData> =
  87             SharedMemory::open(&path, file_opts)?;
  88
  89         shmem.data().tbf.lock().0.update_rate(rate, burst);
  90
  91         Ok(Self { shmem })
  92     }
  93 }
  94
  95 impl ShareableRateLimit for SharedRateLimiter {
  96     fn update_rate(&self, rate: u64, bucket_size: u64) {
  97         self.shmem.data().tbf.lock().0
  98             .update_rate(rate, bucket_size);
  99     }
 100
 101     fn traffic(&self) -> u64 {
 102         self.shmem.data().tbf.lock().0
 103             .traffic()
 104     }
 105
 106     fn register_traffic(&self, current_time: Instant, data_len: u64) -> Duration {
 107         self.shmem.data().tbf.lock().0
 108             .register_traffic(current_time, data_len)
 109     }
 110 }