hyperhive/hivectl/src/quota.rs

112 lines
3.9 KiB
Rust

//! `hivectl quota` — per-agent disk accounting + optional quotas via
//! btrfs qgroups. The daemon holds the privileged helper that reads /
//! sets qgroups; hivectl relays the request and formats the reply.
use std::path::Path;
use anyhow::{Context as _, Result, bail};
use crate::util::daemon_request;
/// `quota enable` — turn on btrfs qgroup accounting (operator opt-in).
pub(crate) async fn quota_enable(socket: &Path) -> Result<()> {
daemon_request(socket, hive_host_sock::HostRequest::QuotaEnable, "quota").await
}
/// `quota show [name]` — report per-agent disk usage from btrfs qgroups.
/// The daemon resolves the agent set + reads each subvolume's usage (it
/// holds the privileged helper); the client just formats the returned rows.
pub(crate) async fn quota_show(socket: &Path, name: Option<&str>) -> Result<()> {
let resp = crate::client::request(
socket,
hive_host_sock::HostRequest::QuotaShow {
name: name.map(str::to_owned),
},
)
.await
.with_context(|| format!("connect to daemon socket {}", socket.display()))?;
if !resp.ok {
// Carries the "btrfs quota not enabled — run `hivectl quota enable`
// first" hint when qgroups are off.
bail!("{}", resp.error.as_deref().unwrap_or("quota show failed"));
}
let rows = resp.quota.unwrap_or_default();
if rows.is_empty() {
println!("no agents with a state dir found");
return Ok(());
}
for row in &rows {
match (row.referenced, row.exclusive) {
(Some(rfer), Some(excl)) => println!(
"{:<12} referenced {:>10} exclusive {:>10}",
row.agent,
human_bytes(rfer),
human_bytes(excl)
),
// Plain-dir agent (no qgroup) — the daemon set an explanatory note.
_ => println!(
"{:<12} ({})",
row.agent,
row.note.as_deref().unwrap_or("no qgroup data")
),
}
}
Ok(())
}
/// `quota limit <name> <size>` — set or clear an agent's disk quota.
pub(crate) async fn quota_limit(socket: &Path, name: &str, size: &str) -> Result<()> {
let limit = parse_quota_size(size)?;
daemon_request(
socket,
hive_host_sock::HostRequest::QuotaLimit {
name: name.to_owned(),
limit,
},
"quota",
)
.await?;
match limit {
Some(n) => println!("set {name} quota to {} ({n} bytes)", human_bytes(n)),
None => println!("cleared {name} quota"),
}
Ok(())
}
/// Parse a quota size: a byte count, a `K`/`M`/`G`/`T`-suffixed value
/// (powers of 1024), or `none` to clear. Rejects overflow + bad input.
fn parse_quota_size(s: &str) -> Result<Option<u64>> {
let t = s.trim();
if t.eq_ignore_ascii_case("none") {
return Ok(None);
}
let (num, mult) = match t.chars().last() {
Some('K' | 'k') => (&t[..t.len() - 1], 1024u64),
Some('M' | 'm') => (&t[..t.len() - 1], 1024 * 1024),
Some('G' | 'g') => (&t[..t.len() - 1], 1024 * 1024 * 1024),
Some('T' | 't') => (&t[..t.len() - 1], 1024u64 * 1024 * 1024 * 1024),
_ => (t, 1),
};
let n: u64 = num.trim().parse().with_context(|| {
format!("invalid size {s:?} — use a byte count, a K/M/G/T suffix (e.g. 5G), or `none`")
})?;
n.checked_mul(mult)
.map(Some)
.with_context(|| format!("size {s:?} overflows u64"))
}
/// Format a byte count as a human-friendly KiB/MiB/GiB/TiB string.
#[allow(clippy::cast_precision_loss)] // display-only; precision loss is cosmetic
fn human_bytes(n: u64) -> String {
const UNITS: [&str; 5] = ["B", "KiB", "MiB", "GiB", "TiB"];
if n < 1024 {
return format!("{n} B");
}
let mut v = n as f64;
let mut i = 0;
while v >= 1024.0 && i < UNITS.len() - 1 {
v /= 1024.0;
i += 1;
}
format!("{v:.1} {}", UNITS[i])
}