use std::num::NonZeroU64;
#[derive(
Clone, Copy, Debug, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize,
)]
#[serde(rename_all = "snake_case")]
pub enum Filesystem {
#[default]
Raw,
Btrfs,
}
impl Filesystem {
pub(crate) fn cache_tag(self) -> &'static str {
match self {
Filesystem::Raw => "raw",
Filesystem::Btrfs => "btrfs",
}
}
pub(crate) fn mkfs_binary_name(self) -> Option<&'static str> {
match self {
Filesystem::Raw => None,
Filesystem::Btrfs => Some("mkfs.btrfs"),
}
}
}
#[derive(
Clone, Copy, Debug, Default, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize,
)]
pub struct DiskThrottle {
pub iops: Option<NonZeroU64>,
pub bytes_per_sec: Option<NonZeroU64>,
pub iops_burst_capacity: Option<NonZeroU64>,
pub bytes_burst_capacity: Option<NonZeroU64>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum ThrottleDimension {
Iops,
Bytes,
}
impl ThrottleDimension {
pub fn burst_field(self) -> &'static str {
match self {
ThrottleDimension::Iops => "iops_burst_capacity",
ThrottleDimension::Bytes => "bytes_burst_capacity",
}
}
pub fn rate_field(self) -> &'static str {
match self {
ThrottleDimension::Iops => "iops",
ThrottleDimension::Bytes => "bytes_per_sec",
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
pub enum DiskThrottleValidationError {
#[error(
"{burst_field} ({burst}) must be >= {rate_field} ({rate}), \
or pass 0 to clear the burst override",
burst_field = dimension.burst_field(),
rate_field = dimension.rate_field(),
)]
BurstBelowRate {
dimension: ThrottleDimension,
burst: u64,
rate: u64,
},
#[error(
"{burst_field} set without {rate_field} refill rate, \
or pass 0 to clear the burst override",
burst_field = dimension.burst_field(),
rate_field = dimension.rate_field(),
)]
BurstWithoutRate {
dimension: ThrottleDimension,
},
}
impl DiskThrottleValidationError {
pub fn dimension(&self) -> ThrottleDimension {
match self {
DiskThrottleValidationError::BurstBelowRate { dimension, .. } => *dimension,
DiskThrottleValidationError::BurstWithoutRate { dimension } => *dimension,
}
}
}
impl DiskThrottle {
pub fn validate(&self) -> Result<(), DiskThrottleValidationError> {
if let Some(burst) = self.iops_burst_capacity {
match self.iops {
Some(rate) if burst < rate => {
return Err(DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Iops,
burst: burst.get(),
rate: rate.get(),
});
}
None => {
return Err(DiskThrottleValidationError::BurstWithoutRate {
dimension: ThrottleDimension::Iops,
});
}
_ => {}
}
}
if let Some(burst) = self.bytes_burst_capacity {
match self.bytes_per_sec {
Some(rate) if burst < rate => {
return Err(DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Bytes,
burst: burst.get(),
rate: rate.get(),
});
}
None => {
return Err(DiskThrottleValidationError::BurstWithoutRate {
dimension: ThrottleDimension::Bytes,
});
}
_ => {}
}
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
pub struct DiskConfig {
pub capacity_mb: u32,
pub filesystem: Filesystem,
pub throttle: DiskThrottle,
pub read_only: bool,
pub name: Option<String>,
pub no_auto_mount: bool,
}
impl Default for DiskConfig {
fn default() -> Self {
DiskConfig {
capacity_mb: 256,
filesystem: Filesystem::Raw,
throttle: DiskThrottle::default(),
read_only: false,
name: None,
no_auto_mount: false,
}
}
}
impl DiskConfig {
#[must_use = "builder methods consume self; bind the result"]
pub fn capacity_mb(mut self, mb: u32) -> Self {
self.capacity_mb = mb;
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn filesystem(mut self, fs: Filesystem) -> Self {
self.filesystem = fs;
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn iops(mut self, iops: u64) -> Self {
self.throttle.iops = NonZeroU64::new(iops);
if self.throttle.iops.is_none() {
self.throttle.iops_burst_capacity = None;
}
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn bytes_per_sec(mut self, bytes_per_sec: u64) -> Self {
self.throttle.bytes_per_sec = NonZeroU64::new(bytes_per_sec);
if self.throttle.bytes_per_sec.is_none() {
self.throttle.bytes_burst_capacity = None;
}
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn iops_burst_capacity(mut self, capacity: u64) -> Self {
self.throttle.iops_burst_capacity = NonZeroU64::new(capacity);
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn bytes_burst_capacity(mut self, capacity: u64) -> Self {
self.throttle.bytes_burst_capacity = NonZeroU64::new(capacity);
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn read_only(mut self) -> Self {
self.read_only = true;
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
#[must_use = "builder methods consume self; bind the result"]
pub fn no_auto_mount(mut self) -> Self {
self.no_auto_mount = true;
self
}
#[allow(dead_code)]
pub(crate) fn auto_mount_path(&self) -> String {
match self.name.as_deref() {
Some(n) => format!("/mnt/{n}"),
None => "/mnt/disk0".to_string(),
}
}
pub(crate) fn capacity_bytes(&self) -> u64 {
(self.capacity_mb as u64) << 20
}
#[allow(dead_code)]
pub(crate) fn capacity_sectors(&self) -> u64 {
self.capacity_bytes() / 512
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn default_is_256mb_raw() {
let d = DiskConfig::default();
assert_eq!(d.capacity_mb, 256);
assert_eq!(d.filesystem, Filesystem::Raw);
assert_eq!(d.throttle, DiskThrottle::default());
assert!(!d.read_only);
assert!(d.name.is_none());
}
#[test]
fn capacity_helpers() {
let d = DiskConfig::default();
assert_eq!(d.capacity_bytes(), 256 * 1024 * 1024);
assert_eq!(d.capacity_sectors(), 524_288);
let d = DiskConfig::default().capacity_mb(512);
assert_eq!(d.capacity_bytes(), 512 * 1024 * 1024);
assert_eq!(d.capacity_sectors(), 1_048_576);
}
#[test]
fn filesystem_builder_sets_variant() {
let d = DiskConfig::default().filesystem(Filesystem::Btrfs);
assert_eq!(d.filesystem, Filesystem::Btrfs);
let d = d.filesystem(Filesystem::Raw);
assert_eq!(d.filesystem, Filesystem::Raw);
}
#[test]
fn builder_chain() {
let d = DiskConfig::default()
.capacity_mb(128)
.iops(1000)
.bytes_per_sec(10 * 1024 * 1024)
.read_only();
assert_eq!(d.capacity_mb, 128);
assert_eq!(d.filesystem, Filesystem::Raw);
assert_eq!(d.throttle.iops, NonZeroU64::new(1000));
assert_eq!(d.throttle.bytes_per_sec, NonZeroU64::new(10 * 1024 * 1024));
assert!(d.read_only);
}
#[test]
fn iops_zero_becomes_none() {
let d = DiskConfig::default().iops(0);
assert!(d.throttle.iops.is_none());
let d = DiskConfig::default().bytes_per_sec(0);
assert!(d.throttle.bytes_per_sec.is_none());
}
#[test]
fn filesystem_default_is_raw() {
assert_eq!(Filesystem::default(), Filesystem::Raw);
}
#[test]
fn filesystem_serde_snake_case() {
assert_eq!(serde_json::to_string(&Filesystem::Raw).unwrap(), r#""raw""#);
assert_eq!(
serde_json::to_string(&Filesystem::Btrfs).unwrap(),
r#""btrfs""#
);
let parsed: Filesystem = serde_json::from_str(r#""raw""#).unwrap();
assert_eq!(parsed, Filesystem::Raw);
let parsed: Filesystem = serde_json::from_str(r#""btrfs""#).unwrap();
assert_eq!(parsed, Filesystem::Btrfs);
}
#[test]
fn filesystem_cache_tag_round_trips_serde_name() {
for fs in [Filesystem::Raw, Filesystem::Btrfs] {
let json = serde_json::to_string(&fs).unwrap();
let stripped = json.trim_matches('"');
assert_eq!(fs.cache_tag(), stripped, "cache_tag drift for {fs:?}");
}
}
#[test]
fn throttle_default_is_unthrottled() {
let t = DiskThrottle::default();
assert!(t.iops.is_none());
assert!(t.bytes_per_sec.is_none());
assert!(t.iops_burst_capacity.is_none());
assert!(t.bytes_burst_capacity.is_none());
}
#[test]
fn iops_zero_serde_roundtrip() {
let original = DiskConfig::default().iops(0).bytes_per_sec(0);
let json = serde_json::to_string(&original).expect("serialize");
let parsed: DiskConfig = serde_json::from_str(&json).expect("deserialize");
assert!(parsed.throttle.iops.is_none());
assert!(parsed.throttle.bytes_per_sec.is_none());
assert_eq!(parsed, original);
}
#[test]
fn disk_config_full_serde_roundtrip() {
let original = DiskConfig {
capacity_mb: 256,
filesystem: Filesystem::Raw,
throttle: DiskThrottle {
iops: NonZeroU64::new(2_500),
bytes_per_sec: NonZeroU64::new(50 * 1024 * 1024),
iops_burst_capacity: NonZeroU64::new(10_000),
bytes_burst_capacity: NonZeroU64::new(200 * 1024 * 1024),
},
read_only: true,
name: Some("data-disk".to_string()),
no_auto_mount: false,
};
let json = serde_json::to_string(&original).expect("serialize DiskConfig");
let parsed: DiskConfig = serde_json::from_str(&json).expect("deserialize DiskConfig");
assert_eq!(parsed, original);
assert_eq!(parsed.capacity_mb, original.capacity_mb);
assert_eq!(parsed.filesystem, original.filesystem);
assert_eq!(parsed.throttle.iops, original.throttle.iops);
assert_eq!(
parsed.throttle.bytes_per_sec,
original.throttle.bytes_per_sec
);
assert_eq!(
parsed.throttle.iops_burst_capacity,
original.throttle.iops_burst_capacity
);
assert_eq!(
parsed.throttle.bytes_burst_capacity,
original.throttle.bytes_burst_capacity
);
assert_eq!(parsed.read_only, original.read_only);
assert_eq!(parsed.name, original.name);
assert_eq!(parsed.name.as_deref(), Some("data-disk"));
}
#[test]
fn disk_config_default_unthrottled_serde_roundtrip() {
let original = DiskConfig::default();
assert!(original.throttle.iops.is_none());
assert!(original.throttle.bytes_per_sec.is_none());
assert!(original.name.is_none());
let json = serde_json::to_string(&original).expect("serialize default DiskConfig");
let parsed: DiskConfig =
serde_json::from_str(&json).expect("deserialize default DiskConfig");
assert_eq!(parsed, original);
assert_eq!(parsed.capacity_mb, original.capacity_mb);
assert_eq!(parsed.filesystem, original.filesystem);
assert!(parsed.throttle.iops.is_none());
assert!(parsed.throttle.bytes_per_sec.is_none());
assert!(parsed.throttle.iops_burst_capacity.is_none());
assert!(parsed.throttle.bytes_burst_capacity.is_none());
assert_eq!(parsed.read_only, original.read_only);
assert!(parsed.name.is_none());
}
#[test]
fn name_builder_sets_label() {
let d = DiskConfig::default().name("data-disk");
assert_eq!(d.name.as_deref(), Some("data-disk"));
let d = DiskConfig::default().name(String::from("log-disk"));
assert_eq!(d.name.as_deref(), Some("log-disk"));
let d = DiskConfig::default().name("first").name("second");
assert_eq!(d.name.as_deref(), Some("second"));
}
#[test]
fn burst_capacity_builders_set_fields() {
let d = DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(5_000)
.bytes_per_sec(10 * 1024 * 1024)
.bytes_burst_capacity(50 * 1024 * 1024);
assert_eq!(d.throttle.iops, NonZeroU64::new(1_000));
assert_eq!(d.throttle.iops_burst_capacity, NonZeroU64::new(5_000));
assert_eq!(d.throttle.bytes_per_sec, NonZeroU64::new(10 * 1024 * 1024));
assert_eq!(
d.throttle.bytes_burst_capacity,
NonZeroU64::new(50 * 1024 * 1024)
);
}
#[test]
fn burst_capacity_zero_becomes_none() {
let d = DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(5_000)
.iops_burst_capacity(0);
assert!(d.throttle.iops_burst_capacity.is_none());
let d = DiskConfig::default()
.bytes_per_sec(1_000)
.bytes_burst_capacity(5_000)
.bytes_burst_capacity(0);
assert!(d.throttle.bytes_burst_capacity.is_none());
}
#[test]
fn burst_capacity_default_is_none() {
let d = DiskConfig::default();
assert!(d.throttle.iops_burst_capacity.is_none());
assert!(d.throttle.bytes_burst_capacity.is_none());
}
#[test]
fn clearing_iops_clears_iops_burst() {
let d = DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(5_000)
.iops(0);
assert!(d.throttle.iops.is_none());
assert!(
d.throttle.iops_burst_capacity.is_none(),
"clearing iops must also clear iops_burst_capacity \
so validate() doesn't fail with a stale-burst error",
);
let d = DiskConfig::default()
.bytes_per_sec(2_000)
.bytes_burst_capacity(8_000)
.iops(0);
assert!(d.throttle.bytes_per_sec.is_some());
assert!(d.throttle.bytes_burst_capacity.is_some());
}
#[test]
fn clearing_bytes_per_sec_clears_bytes_burst() {
let d = DiskConfig::default()
.bytes_per_sec(2_000)
.bytes_burst_capacity(8_000)
.bytes_per_sec(0);
assert!(d.throttle.bytes_per_sec.is_none());
assert!(
d.throttle.bytes_burst_capacity.is_none(),
"clearing bytes_per_sec must also clear \
bytes_burst_capacity",
);
let d = DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(5_000)
.bytes_per_sec(0);
assert!(d.throttle.iops.is_some());
assert!(d.throttle.iops_burst_capacity.is_some());
}
#[test]
fn clearing_rate_leaves_throttle_validate_clean() {
let throttle = DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(5_000)
.bytes_per_sec(2_000)
.bytes_burst_capacity(8_000)
.iops(0)
.bytes_per_sec(0)
.throttle;
assert!(throttle.iops.is_none());
assert!(throttle.bytes_per_sec.is_none());
assert!(throttle.iops_burst_capacity.is_none());
assert!(throttle.bytes_burst_capacity.is_none());
throttle
.validate()
.expect("post-clear throttle must validate clean");
}
#[test]
fn validate_accepts_burst_at_or_above_rate() {
DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(1_000)
.throttle
.validate()
.expect("burst == iops accepted");
DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(5_000)
.bytes_per_sec(10 * 1024 * 1024)
.bytes_burst_capacity(50 * 1024 * 1024)
.throttle
.validate()
.expect("burst > rate accepted");
DiskConfig::default()
.throttle
.validate()
.expect("no throttle accepted");
DiskConfig::default()
.iops(1_000)
.bytes_per_sec(1_000_000)
.throttle
.validate()
.expect("rate without burst accepted");
}
#[test]
fn validate_rejects_burst_below_rate() {
let err = DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(500)
.throttle
.validate()
.expect_err("burst < iops rejected");
assert_eq!(
err,
DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Iops,
burst: 500,
rate: 1_000,
},
"unexpected error variant",
);
let msg = err.to_string();
assert!(
msg.contains("iops_burst_capacity") && msg.contains("must be >="),
"unexpected error message: {msg}",
);
assert!(
msg.contains("pass 0 to clear"),
"remediation hint missing: {msg}",
);
let err = DiskConfig::default()
.bytes_per_sec(10_000)
.bytes_burst_capacity(5_000)
.throttle
.validate()
.expect_err("burst < bytes_per_sec rejected");
assert_eq!(
err,
DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Bytes,
burst: 5_000,
rate: 10_000,
},
"unexpected error variant",
);
let msg = err.to_string();
assert!(
msg.contains("bytes_burst_capacity") && msg.contains("must be >="),
"unexpected error message: {msg}",
);
assert!(
msg.contains("pass 0 to clear"),
"remediation hint missing: {msg}",
);
}
#[test]
fn validate_rejects_burst_one_below_rate() {
let err = DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(999)
.throttle
.validate()
.expect_err("iops burst one below rate must be rejected");
assert_eq!(
err,
DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Iops,
burst: 999,
rate: 1_000,
},
);
let msg = err.to_string();
assert!(
msg.contains("iops_burst_capacity") && msg.contains("must be >="),
"unexpected error message: {msg}",
);
let err = DiskConfig::default()
.bytes_per_sec(1_000)
.bytes_burst_capacity(999)
.throttle
.validate()
.expect_err("bytes burst one below rate must be rejected");
assert_eq!(
err,
DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Bytes,
burst: 999,
rate: 1_000,
},
);
let msg = err.to_string();
assert!(
msg.contains("bytes_burst_capacity") && msg.contains("must be >="),
"unexpected error message: {msg}",
);
}
#[test]
fn iops_clear_after_burst_set_validates_clean() {
DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(5_000)
.iops(0)
.throttle
.validate()
.expect("iops-cleared throttle must validate clean");
}
#[test]
fn validate_rejects_burst_without_rate() {
let err = DiskConfig::default()
.iops_burst_capacity(5_000)
.throttle
.validate()
.expect_err("burst without iops rejected");
assert_eq!(
err,
DiskThrottleValidationError::BurstWithoutRate {
dimension: ThrottleDimension::Iops,
},
);
let msg = err.to_string();
assert!(
msg.contains("iops_burst_capacity") && msg.contains("without iops"),
"unexpected error message: {msg}",
);
assert!(
msg.contains("pass 0 to clear"),
"remediation hint missing: {msg}",
);
let err = DiskConfig::default()
.bytes_burst_capacity(5_000)
.throttle
.validate()
.expect_err("burst without bytes_per_sec rejected");
assert_eq!(
err,
DiskThrottleValidationError::BurstWithoutRate {
dimension: ThrottleDimension::Bytes,
},
);
let msg = err.to_string();
assert!(
msg.contains("bytes_burst_capacity") && msg.contains("without bytes_per_sec"),
"unexpected error message: {msg}",
);
assert!(
msg.contains("pass 0 to clear"),
"remediation hint missing: {msg}",
);
}
#[test]
fn validation_error_dimension_accessor() {
let err = DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Iops,
burst: 500,
rate: 1_000,
};
assert_eq!(err.dimension(), ThrottleDimension::Iops);
let err = DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Bytes,
burst: 500,
rate: 1_000,
};
assert_eq!(err.dimension(), ThrottleDimension::Bytes);
let err = DiskThrottleValidationError::BurstWithoutRate {
dimension: ThrottleDimension::Iops,
};
assert_eq!(err.dimension(), ThrottleDimension::Iops);
let err = DiskThrottleValidationError::BurstWithoutRate {
dimension: ThrottleDimension::Bytes,
};
assert_eq!(err.dimension(), ThrottleDimension::Bytes);
}
#[test]
fn throttle_dimension_field_names() {
assert_eq!(ThrottleDimension::Iops.burst_field(), "iops_burst_capacity");
assert_eq!(ThrottleDimension::Iops.rate_field(), "iops");
assert_eq!(
ThrottleDimension::Bytes.burst_field(),
"bytes_burst_capacity",
);
assert_eq!(ThrottleDimension::Bytes.rate_field(), "bytes_per_sec");
}
#[test]
fn disk_throttle_validation_error_downcasts_through_anyhow() {
let typed = DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(500)
.throttle
.validate()
.expect_err("burst < iops rejected");
let wrapped = anyhow::anyhow!(typed).context("invalid disk throttle");
let recovered = wrapped
.chain()
.find_map(|c| c.downcast_ref::<DiskThrottleValidationError>())
.expect(
"DiskThrottleValidationError must remain downcastable through \
the production anyhow wrap; lost typing means library \
consumers cannot route programmatic recovery",
);
assert_eq!(
*recovered,
DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Iops,
burst: 500,
rate: 1_000,
},
);
let rendered = format!("{wrapped:#}");
assert!(
rendered.contains("invalid disk throttle"),
"anyhow context must survive the wrap: {rendered}",
);
}
#[test]
fn validate_first_failure_wins_iops_before_bytes() {
let throttle = DiskConfig::default()
.iops(1_000)
.iops_burst_capacity(500) .bytes_per_sec(10_000)
.bytes_burst_capacity(5_000) .throttle;
let err = throttle
.validate()
.expect_err("both-dimensions-bad must reject");
assert_eq!(
err,
DiskThrottleValidationError::BurstBelowRate {
dimension: ThrottleDimension::Iops,
burst: 500,
rate: 1_000,
},
"iops violation must surface first; refactor that aggregates \
or reverses the check order would change this",
);
assert_eq!(err.dimension(), ThrottleDimension::Iops);
let throttle = DiskConfig::default()
.iops_burst_capacity(5_000)
.bytes_burst_capacity(8_000)
.throttle;
let err = throttle
.validate()
.expect_err("both-without-rate must reject");
assert_eq!(
err,
DiskThrottleValidationError::BurstWithoutRate {
dimension: ThrottleDimension::Iops,
},
"iops violation must surface first across both \
BurstBelowRate and BurstWithoutRate variants",
);
}
#[test]
fn disk_config_burst_serde_roundtrip() {
let original = DiskConfig::default()
.iops(2_500)
.iops_burst_capacity(10_000)
.bytes_per_sec(50 * 1024 * 1024)
.bytes_burst_capacity(200 * 1024 * 1024);
let json = serde_json::to_string(&original).expect("serialize burst DiskConfig");
let parsed: DiskConfig = serde_json::from_str(&json).expect("deserialize burst DiskConfig");
assert_eq!(parsed, original);
assert_eq!(parsed.throttle.iops, NonZeroU64::new(2_500));
assert_eq!(parsed.throttle.iops_burst_capacity, NonZeroU64::new(10_000));
assert_eq!(
parsed.throttle.bytes_per_sec,
NonZeroU64::new(50 * 1024 * 1024)
);
assert_eq!(
parsed.throttle.bytes_burst_capacity,
NonZeroU64::new(200 * 1024 * 1024)
);
}
}