111 lines
3.1 KiB
Rust
111 lines
3.1 KiB
Rust
use std::path::PathBuf;
|
|
use std::mem::MaybeUninit;
|
|
use std::time::{Instant, Duration};
|
|
|
|
use anyhow::{bail, Error};
|
|
use nix::sys::stat::Mode;
|
|
|
|
use proxmox::tools::fs::{create_path, CreateOptions};
|
|
|
|
use proxmox_http::client::{RateLimit, RateLimiter, ShareableRateLimit};
|
|
use proxmox_shared_memory::{Init, SharedMemory, SharedMutex};
|
|
use proxmox_shared_memory::{check_subtype, initialize_subtype};
|
|
|
|
// openssl::sha::sha256(b"Proxmox Backup SharedRateLimiter v1.0")[0..8];
|
|
pub const PROXMOX_BACKUP_SHARED_RATE_LIMITER_MAGIC_1_0: [u8; 8] = [6, 58, 213, 96, 161, 122, 130, 117];
|
|
|
|
const BASE_PATH: &str = pbs_buildcfg::rundir!("/shmem/tbf");
|
|
|
|
// Wrap RateLimiter, so that we can provide an Init impl
|
|
#[repr(C)]
|
|
struct WrapLimiter(RateLimiter);
|
|
|
|
impl Init for WrapLimiter {
|
|
fn initialize(this: &mut MaybeUninit<Self>) {
|
|
// default does not matter here, because we override later
|
|
this.write(WrapLimiter(RateLimiter::new(1_000_000, 1_000_000)));
|
|
}
|
|
}
|
|
|
|
#[repr(C)]
|
|
struct SharedRateLimiterData {
|
|
magic: [u8; 8],
|
|
tbf: SharedMutex<WrapLimiter>,
|
|
padding: [u8; 4096 - 104],
|
|
}
|
|
|
|
impl Init for SharedRateLimiterData {
|
|
fn initialize(this: &mut MaybeUninit<Self>) {
|
|
unsafe {
|
|
let me = &mut *this.as_mut_ptr();
|
|
me.magic = PROXMOX_BACKUP_SHARED_RATE_LIMITER_MAGIC_1_0;
|
|
initialize_subtype(&mut me.tbf);
|
|
}
|
|
}
|
|
|
|
fn check_type_magic(this: &MaybeUninit<Self>) -> Result<(), Error> {
|
|
unsafe {
|
|
let me = &*this.as_ptr();
|
|
if me.magic != PROXMOX_BACKUP_SHARED_RATE_LIMITER_MAGIC_1_0 {
|
|
bail!("SharedRateLimiterData: wrong magic number");
|
|
}
|
|
check_subtype(&me.tbf)?;
|
|
Ok(())
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct SharedRateLimiter {
|
|
shmem: SharedMemory<SharedRateLimiterData>
|
|
}
|
|
|
|
impl SharedRateLimiter {
|
|
|
|
pub fn mmap_shmem(name: &str, rate: u64, burst: u64) -> Result<Self, Error> {
|
|
let mut path = PathBuf::from(BASE_PATH);
|
|
|
|
let user = pbs_config::backup_user()?;
|
|
|
|
let dir_opts = CreateOptions::new()
|
|
.perm(Mode::from_bits_truncate(0o770))
|
|
.owner(user.uid)
|
|
.group(user.gid);
|
|
|
|
create_path(
|
|
&path,
|
|
Some(dir_opts.clone()),
|
|
Some(dir_opts))?;
|
|
|
|
path.push(name);
|
|
|
|
let file_opts = CreateOptions::new()
|
|
.perm(Mode::from_bits_truncate(0o660))
|
|
.owner(user.uid)
|
|
.group(user.gid);
|
|
|
|
let shmem: SharedMemory<SharedRateLimiterData> =
|
|
SharedMemory::open(&path, file_opts)?;
|
|
|
|
shmem.data().tbf.lock().0.update_rate(rate, burst);
|
|
|
|
Ok(Self { shmem })
|
|
}
|
|
}
|
|
|
|
impl ShareableRateLimit for SharedRateLimiter {
|
|
fn update_rate(&self, rate: u64, bucket_size: u64) {
|
|
self.shmem.data().tbf.lock().0
|
|
.update_rate(rate, bucket_size);
|
|
}
|
|
|
|
fn traffic(&self) -> u64 {
|
|
self.shmem.data().tbf.lock().0
|
|
.traffic()
|
|
}
|
|
|
|
fn register_traffic(&self, current_time: Instant, data_len: u64) -> Duration {
|
|
self.shmem.data().tbf.lock().0
|
|
.register_traffic(current_time, data_len)
|
|
}
|
|
}
|