[pbs-devel] [PATCH proxmox 1/2] http: factor out PBS shared rate limiter implementation

[Christian Ebner <c.ebner@proxmox.com>]

Moves the current shared rate limiter implementation from the Proxmox
Backup Server into proxmox-http for it to be reusable, e.g. for s3
client rate limiting.

Signed-off-by: Christian Ebner <c.ebner@proxmox.com>
---
Note: Unsure if also the magic number for the mmapped files should be
adapted, leaving it as is for full backwards compat for now.

 proxmox-http/Cargo.toml                 |   7 ++
 proxmox-http/debian/control             |  18 ++++
 proxmox-http/src/lib.rs                 |   5 +
 proxmox-http/src/shared_rate_limiter.rs | 129 ++++++++++++++++++++++++
 4 files changed, 159 insertions(+)
 create mode 100644 proxmox-http/src/shared_rate_limiter.rs

diff --git a/proxmox-http/Cargo.toml b/proxmox-http/Cargo.toml
index 1717673c..e70b1fc3 100644
--- a/proxmox-http/Cargo.toml
+++ b/proxmox-http/Cargo.toml
@@ -21,6 +21,7 @@ http-body-util = { workspace = true, optional = true }
 hyper = { workspace = true, optional = true }
 hyper-util = { workspace = true, optional = true, features = ["http2"] }
 native-tls = { workspace = true, optional = true }
+nix = { workspace = true, optional = true }
 openssl =  { version = "0.10", optional = true }
 serde_json = { workspace = true, optional = true }
 sync_wrapper = { workspace = true, optional = true }
@@ -32,6 +33,7 @@ url = { workspace = true, optional = true }
 
 proxmox-async = { workspace = true, optional = true }
 proxmox-base64 = { workspace = true, optional = true }
+proxmox-shared-memory = { workspace = true, optional = true }
 proxmox-sys = { workspace = true, optional = true }
 proxmox-io = { workspace = true, optional = true }
 proxmox-lang = { workspace = true, optional = true }
@@ -54,6 +56,11 @@ body = [
     "sync_wrapper?/futures",
 ]
 rate-limiter = ["dep:hyper"]
+shared-rate-limiter = [
+    "dep:nix",
+    "dep:proxmox-shared-memory",
+    "dep:proxmox-sys",
+]
 rate-limited-stream = [
     "dep:tokio",
     "dep:hyper-util",
diff --git a/proxmox-http/debian/control b/proxmox-http/debian/control
index f50f8ea7..e3996f6c 100644
--- a/proxmox-http/debian/control
+++ b/proxmox-http/debian/control
@@ -30,6 +30,7 @@ Suggests:
  librust-proxmox-http+proxmox-async-dev (= ${binary:Version}),
  librust-proxmox-http+rate-limited-stream-dev (= ${binary:Version}),
  librust-proxmox-http+rate-limiter-dev (= ${binary:Version}),
+ librust-proxmox-http+shared-rate-limiter-dev (= ${binary:Version}),
  librust-proxmox-http+websocket-dev (= ${binary:Version})
 Provides:
  librust-proxmox-http+default-dev (= ${binary:Version}),
@@ -203,6 +204,23 @@ Description: Proxmox HTTP library - feature "rate-limiter"
  This metapackage enables feature "rate-limiter" for the Rust proxmox-http
  crate, by pulling in any additional dependencies needed by that feature.
 
+Package: librust-proxmox-http+shared-rate-limiter-dev
+Architecture: any
+Multi-Arch: same
+Depends:
+ ${misc:Depends},
+ librust-proxmox-http-dev (= ${binary:Version}),
+ librust-nix-0.29+default-dev,
+ librust-proxmox-shared-memory-1+default-dev,
+ librust-proxmox-sys-1+default-dev
+Provides:
+ librust-proxmox-http-1+shared-rate-limiter-dev (= ${binary:Version}),
+ librust-proxmox-http-1.0+shared-rate-limiter-dev (= ${binary:Version}),
+ librust-proxmox-http-1.0.2+shared-rate-limiter-dev (= ${binary:Version})
+Description: Proxmox HTTP library - feature "shared-rate-limiter"
+ This metapackage enables feature "shared-rate-limiter" for the Rust proxmox-
+ http crate, by pulling in any additional dependencies needed by that feature.
+
 Package: librust-proxmox-http+websocket-dev
 Architecture: any
 Multi-Arch: same
diff --git a/proxmox-http/src/lib.rs b/proxmox-http/src/lib.rs
index 8b6953b0..ce6a77a1 100644
--- a/proxmox-http/src/lib.rs
+++ b/proxmox-http/src/lib.rs
@@ -31,6 +31,11 @@ mod rate_limiter;
 #[cfg(feature = "rate-limiter")]
 pub use rate_limiter::{RateLimit, RateLimiter, RateLimiterVec, ShareableRateLimit};
 
+#[cfg(feature = "shared-rate-limiter")]
+mod shared_rate_limiter;
+#[cfg(feature = "shared-rate-limiter")]
+pub use shared_rate_limiter::SharedRateLimiter;
+
 #[cfg(feature = "rate-limited-stream")]
 mod rate_limited_stream;
 #[cfg(feature = "rate-limited-stream")]
diff --git a/proxmox-http/src/shared_rate_limiter.rs b/proxmox-http/src/shared_rate_limiter.rs
new file mode 100644
index 00000000..7be321a5
--- /dev/null
+++ b/proxmox-http/src/shared_rate_limiter.rs
@@ -0,0 +1,129 @@
+//! Rate limiter designed for shared memory
+
+use std::mem::MaybeUninit;
+use std::path::Path;
+use std::time::{Duration, Instant};
+
+use anyhow::{bail, Error};
+use nix::sys::stat::Mode;
+use nix::unistd::User;
+
+use proxmox_sys::fs::{create_path, CreateOptions};
+use proxmox_shared_memory::{check_subtype, initialize_subtype};
+use proxmox_shared_memory::{Init, SharedMemory, SharedMutex};
+
+use crate::{RateLimit, RateLimiter, ShareableRateLimit};
+
+/// Magic number for shared rate limiter exposed file mappings
+///
+/// Generated by `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];
+
+// 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(())
+        }
+    }
+}
+
+/// Rate limiter designed for shared memory ([SharedMemory])
+///
+/// The actual [RateLimiter] is protected by a [SharedMutex] and
+/// implements [Init]. This way we can share the limiter between
+/// different processes.
+pub struct SharedRateLimiter {
+    shmem: SharedMemory<SharedRateLimiterData>,
+}
+
+impl SharedRateLimiter {
+    /// Creates a new mmap'ed instance.
+    ///
+    /// Data is mapped in `<base_path>/<name>` using
+    /// `TMPFS`.
+    pub fn mmap_shmem<P: AsRef<Path>>(
+        name: &str,
+        rate: u64,
+        burst: u64,
+        user: User,
+        base_path: P,
+    ) -> Result<Self, Error> {
+        let mut path = base_path.as_ref().to_path_buf();
+
+        let dir_opts = CreateOptions::new()
+            .perm(Mode::from_bits_truncate(0o770))
+            .owner(user.uid)
+            .group(user.gid);
+
+        create_path(&path, Some(dir_opts), 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)
+    }
+}
-- 
2.47.2



_______________________________________________
pbs-devel mailing list
pbs-devel@lists.proxmox.com
https://lists.proxmox.com/cgi-bin/mailman/listinfo/pbs-devel