shep_core/

values.rs

//! Config value newtypes: memory sizes and durations

use core::fmt;
use core::str::FromStr;

use serde::{Deserialize, Deserializer, Serialize, Serializer};

// Binary units per the Flockfile grammar `^\d+(G|M|K)?$`: K/M/G are
// KiB/MiB/GiB, not decimal. Unit definitions, not tuning thresholds.
const KIB: u64 = 1024;
const MIB: u64 = 1024 * KIB;
const GIB: u64 = 1024 * MIB;

/// A memory quantity in bytes, used for memory-limit thresholds
///
/// Parses the Flockfile grammar `^\d+(G|M|K)?$` (binary units; plain digits
/// are bytes). Ordering compares byte counts, so a configured limit compares
/// directly against a sampled RSS wrapped with [`MemSize::from_bytes`].
///
/// # Example
/// ```
/// use shep_core::values::MemSize;
///
/// let limit: MemSize = "512M".parse()?;
/// assert_eq!(limit.bytes(), 512 << 20);
/// assert!("512MB".parse::<MemSize>().is_err()); // strict grammar
/// # Ok::<(), shep_core::values::ParseMemSizeError>(())
/// ```
// wire format: changing this is a breaking change (serialized as its string
// form inside AppConfig, which travels over the client<->daemon socket)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct MemSize(u64);

impl MemSize {
    /// Wraps a raw byte count, e.g. an RSS sample
    #[inline]
    #[must_use]
    pub const fn from_bytes(bytes: u64) -> Self {
        Self(bytes)
    }

    /// Returns the quantity in bytes
    #[inline]
    #[must_use]
    pub const fn bytes(self) -> u64 {
        self.0
    }
}

impl FromStr for MemSize {
    type Err = ParseMemSizeError;

    /// Parses `^\d+(G|M|K)?$`: binary units, plain digits = bytes
    ///
    /// # Errors
    ///
    /// - [`ParseMemSizeError::Empty`]: empty input.
    /// - [`ParseMemSizeError::MissingDigits`]: unit suffix with no digits.
    /// - [`ParseMemSizeError::InvalidCharacter`]: anything outside ASCII
    ///   digits plus one trailing `G`/`M`/`K` (lowercase, whitespace,
    ///   fractions, multi-letter suffixes all land here).
    /// - [`ParseMemSizeError::Overflow`]: byte count exceeds `u64::MAX`.
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.is_empty() {
            return Err(ParseMemSizeError::Empty);
        }
        let (digits, multiplier) = match s.as_bytes()[s.len() - 1] {
            b'G' => (&s[..s.len() - 1], GIB),
            b'M' => (&s[..s.len() - 1], MIB),
            b'K' => (&s[..s.len() - 1], KIB),
            _ => (s, 1),
        };
        if digits.is_empty() {
            return Err(ParseMemSizeError::MissingDigits);
        }
        if !digits.bytes().all(|b| b.is_ascii_digit()) {
            return Err(ParseMemSizeError::InvalidCharacter);
        }
        let value: u64 = digits.parse().map_err(|_| ParseMemSizeError::Overflow)?;
        value
            .checked_mul(multiplier)
            .map(Self)
            .ok_or(ParseMemSizeError::Overflow)
    }
}

/// Formats with the largest binary unit dividing the value exactly;
/// output always re-parses to the same value
impl fmt::Display for MemSize {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.0 {
            0 => f.write_str("0"),
            b if b % GIB == 0 => write!(f, "{}G", b / GIB),
            b if b % MIB == 0 => write!(f, "{}M", b / MIB),
            b if b % KIB == 0 => write!(f, "{}K", b / KIB),
            b => write!(f, "{b}"),
        }
    }
}

impl Serialize for MemSize {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.collect_str(self)
    }
}

impl<'de> Deserialize<'de> for MemSize {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        // String, not &str: the toml deserializer cannot always borrow
        let s = String::deserialize(deserializer)?;
        s.parse().map_err(serde::de::Error::custom)
    }
}

/// Failure to parse a [`MemSize`] from the grammar `^\d+(G|M|K)?$`
///
/// `#[non_exhaustive]`: a future grammar revision, such as fractional sizes
/// (`1.5G`) or a binary-vs-decimal distinction, would want its own variant
/// rather than folding into [`Self::InvalidCharacter`]'s catch-all, and
/// shep-core is a published library an out-of-tree matcher should not
/// break for.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParseMemSizeError {
    /// The input string was empty
    Empty,
    /// A unit suffix with no digits before it (`"M"`)
    MissingDigits,
    /// A character outside ASCII digits plus one optional trailing
    /// `G`/`M`/`K`: covers lowercase units, whitespace, signs, fractions,
    /// and multi-letter suffixes such as `"MB"`
    InvalidCharacter,
    /// The quantity in bytes does not fit in `u64`
    Overflow,
}

impl fmt::Display for ParseMemSizeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Self::Empty => "memory size is empty",
            Self::MissingDigits => "memory size has a unit suffix but no digits",
            Self::InvalidCharacter => {
                "memory size must be ASCII digits with an optional trailing G, M, or K"
            }
            Self::Overflow => "memory size in bytes overflows u64",
        })
    }
}

impl core::error::Error for ParseMemSizeError {}

/// String-shaped, matching this type's `Serialize`/`Deserialize`, which go
/// through `Display`/`FromStr` rather than the wrapped `u64`. A derive here
/// would emit `{"type":"integer"}` and describe a wire form that does not
/// exist.
///
/// The pattern is `FromStr`'s own grammar, lifted from its doc comment
/// above. If you change one, change the other: the paired test below is
/// what catches it.
#[cfg(feature = "schema")]
impl schemars::JsonSchema for MemSize {
    fn schema_name() -> std::borrow::Cow<'static, str> {
        "MemSize".into()
    }

    fn json_schema(_: &mut schemars::SchemaGenerator) -> schemars::Schema {
        schemars::json_schema!({
            "type": "string",
            "pattern": r"^\d+(G|M|K)?$",
            "description": "A byte quantity: digits, optionally suffixed G, M or K (binary units).",
        })
    }
}

/// A duration from the Flockfile grammar `^\d+(ms|h|m|s)?$`
///
/// Plain digits are milliseconds; `ms`/`s`/`m`/`h` are
/// milliseconds/seconds/minutes/hours. Used for `min_uptime`,
/// `kill_timeout`, and the other lifecycle timers.
///
/// `ms` is checked before the single-letter suffixes, so `m` still means
/// minutes and only a trailing `ms` means milliseconds: `5m` and `5ms`
/// differ by a factor of sixty thousand.
///
/// # Example
/// ```
/// use shep_core::values::UpDuration;
///
/// assert_eq!("30s".parse::<UpDuration>()?.as_millis(), 30_000);
/// assert_eq!("500ms".parse::<UpDuration>()?.as_millis(), 500);
/// assert_eq!("5m".parse::<UpDuration>()?.as_millis(), 300_000);
/// assert!("30S".parse::<UpDuration>().is_err()); // lowercase units only
/// # Ok::<(), shep_core::values::ParseUpDurationError>(())
/// ```
// wire format: changing this is a breaking change (string form in AppConfig)
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct UpDuration(core::time::Duration);

impl UpDuration {
    /// Wraps a raw millisecond count
    #[inline]
    #[must_use]
    pub const fn from_millis(ms: u64) -> Self {
        Self(core::time::Duration::from_millis(ms))
    }

    /// Returns the wrapped [`core::time::Duration`]
    #[inline]
    #[must_use]
    pub const fn as_duration(self) -> core::time::Duration {
        self.0
    }

    /// Returns the duration in whole milliseconds
    #[inline]
    #[must_use]
    pub const fn as_millis(self) -> u64 {
        // Sound: every constructor bounds millis to u64 (`from_millis`
        // stores its argument directly; `FromStr` reaches this type only
        // via a `checked_mul` that already fits in u64). Revisit if a raw
        // `Duration` constructor is ever added.
        self.0.as_millis() as u64
    }
}

impl FromStr for UpDuration {
    type Err = ParseUpDurationError;

    /// Parses `^\d+(ms|h|m|s)?$`: plain digits are milliseconds
    ///
    /// `ms` is matched before the single-letter suffixes below, so a
    /// trailing `m` alone still means minutes.
    ///
    /// # Errors
    ///
    /// - [`ParseUpDurationError::Empty`]: empty input.
    /// - [`ParseUpDurationError::MissingDigits`]: unit with no digits.
    /// - [`ParseUpDurationError::InvalidCharacter`]: anything outside ASCII
    ///   digits plus one trailing lowercase `h`/`m`/`s`/`ms`.
    /// - [`ParseUpDurationError::Overflow`]: milliseconds overflow `u64`.
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.is_empty() {
            return Err(ParseUpDurationError::Empty);
        }
        // `ms` first: it shares its trailing `m` with the minutes suffix, so
        // checking the single-letter match first would parse "5ms" as "5m"
        // followed by a stray "s" and reject it, or worse, alias the two
        // suffixes if that stray byte were ever tolerated.
        let (digits, ms_per_unit) = if let Some(rest) = s.strip_suffix("ms") {
            (rest, 1)
        } else {
            match s.as_bytes()[s.len() - 1] {
                b'h' => (&s[..s.len() - 1], 3_600_000),
                b'm' => (&s[..s.len() - 1], 60_000),
                b's' => (&s[..s.len() - 1], 1_000),
                _ => (s, 1),
            }
        };
        if digits.is_empty() {
            return Err(ParseUpDurationError::MissingDigits);
        }
        if !digits.bytes().all(|b| b.is_ascii_digit()) {
            return Err(ParseUpDurationError::InvalidCharacter);
        }
        let value: u64 = digits.parse().map_err(|_| ParseUpDurationError::Overflow)?;
        value
            .checked_mul(ms_per_unit)
            .map(Self::from_millis)
            .ok_or(ParseUpDurationError::Overflow)
    }
}

/// Formats with the largest unit dividing the value exactly (ms as digits)
impl fmt::Display for UpDuration {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let ms = self.as_millis();
        match ms {
            0 => f.write_str("0"),
            v if v % 3_600_000 == 0 => write!(f, "{}h", v / 3_600_000),
            v if v % 60_000 == 0 => write!(f, "{}m", v / 60_000),
            v if v % 1_000 == 0 => write!(f, "{}s", v / 1_000),
            v => write!(f, "{v}"),
        }
    }
}

impl Serialize for UpDuration {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.collect_str(self)
    }
}

impl<'de> Deserialize<'de> for UpDuration {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        // String, not &str: the toml deserializer cannot always borrow
        let s = String::deserialize(deserializer)?;
        s.parse().map_err(serde::de::Error::custom)
    }
}

/// Failure to parse an [`UpDuration`] from the grammar `^\d+(ms|h|m|s)?$`
///
/// `#[non_exhaustive]`, for the same reason as [`ParseMemSizeError`]: a
/// future grammar revision, such as fractional durations or a `d`/`w`
/// unit, would want its own variant rather than folding into
/// [`Self::InvalidCharacter`]'s catch-all.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParseUpDurationError {
    /// The input string was empty
    Empty,
    /// A unit suffix with no digits before it (`"s"`)
    MissingDigits,
    /// A character outside ASCII digits plus one optional trailing
    /// lowercase `h`/`m`/`s`/`ms`
    InvalidCharacter,
    /// The duration in milliseconds does not fit in `u64`
    Overflow,
}

impl fmt::Display for ParseUpDurationError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Self::Empty => "duration is empty",
            Self::MissingDigits => "duration has a unit suffix but no digits",
            Self::InvalidCharacter => {
                "duration must be ASCII digits with an optional trailing h, m, s, or ms"
            }
            Self::Overflow => "duration in milliseconds overflows u64",
        })
    }
}

impl core::error::Error for ParseUpDurationError {}

/// String-shaped, matching this type's `Serialize`/`Deserialize`: see
/// [`MemSize`]'s own `JsonSchema` impl for the full reasoning, which applies
/// here unchanged.
#[cfg(feature = "schema")]
impl schemars::JsonSchema for UpDuration {
    fn schema_name() -> std::borrow::Cow<'static, str> {
        "UpDuration".into()
    }

    fn json_schema(_: &mut schemars::SchemaGenerator) -> schemars::Schema {
        schemars::json_schema!({
            "type": "string",
            "pattern": r"^\d+(ms|h|m|s)?$",
            "description": "A duration: digits, optionally suffixed ms, h, m or s. Plain digits are milliseconds.",
        })
    }
}

#[cfg(test)]
mod mem_size_tests {
    use super::*;

    #[test]
    fn plain_digits_parse_as_bytes() {
        assert_eq!("123".parse::<MemSize>().unwrap().bytes(), 123);
    }

    #[test]
    fn units_are_binary() {
        assert_eq!("7K".parse::<MemSize>().unwrap().bytes(), 7 * 1024);
        assert_eq!("512M".parse::<MemSize>().unwrap().bytes(), 512 << 20);
        assert_eq!("3G".parse::<MemSize>().unwrap().bytes(), 3 << 30);
    }

    #[test]
    fn rejects_spec_violations() {
        use ParseMemSizeError::*;
        assert_eq!("".parse::<MemSize>(), Err(Empty));
        assert_eq!("G".parse::<MemSize>(), Err(MissingDigits));
        assert_eq!("512m".parse::<MemSize>(), Err(InvalidCharacter)); // lowercase
        assert_eq!(" 512M".parse::<MemSize>(), Err(InvalidCharacter)); // whitespace
        assert_eq!("1.5G".parse::<MemSize>(), Err(InvalidCharacter)); // fraction
        assert_eq!("512MB".parse::<MemSize>(), Err(InvalidCharacter)); // multi-letter
        assert_eq!("18446744073709551616".parse::<MemSize>(), Err(Overflow));
        assert_eq!("17179869184G".parse::<MemSize>(), Err(Overflow));
    }

    #[test]
    fn display_uses_largest_exact_unit_and_round_trips() {
        for bytes in [
            0u64,
            1,
            1023,
            1024,
            1536,
            1 << 20,
            (1 << 30) + 1024,
            u64::MAX,
        ] {
            let size = MemSize::from_bytes(bytes);
            let reparsed: MemSize = size.to_string().parse().unwrap();
            assert_eq!(reparsed, size, "display of {bytes} bytes must reparse");
        }
        assert_eq!(MemSize::from_bytes(3 << 30).to_string(), "3G");
        assert_eq!(MemSize::from_bytes(1536).to_string(), "1536");
    }

    #[test]
    fn serde_uses_string_form() {
        let size: MemSize = serde_json::from_str("\"512M\"").unwrap();
        assert_eq!(size.bytes(), 512 << 20);
        assert_eq!(serde_json::to_string(&size).unwrap(), "\"512M\"");
        assert!(serde_json::from_str::<MemSize>("\"512MB\"").is_err());
    }

    /// The pattern must agree with `FromStr`, not just be self-consistent.
    /// `512T` and `1P` are in the reject list because a widened suffix set
    /// is the way this pattern most plausibly drifts.
    #[cfg(feature = "schema")]
    #[test]
    fn the_schema_pattern_agrees_with_from_str() {
        let schema = serde_json::to_value(schemars::schema_for!(MemSize)).unwrap();
        let pattern = schema["pattern"].as_str().unwrap();
        let re = regex::Regex::new(pattern).unwrap();
        for accepted in ["512M", "1G", "4096", "7K"] {
            assert!(re.is_match(accepted), "pattern rejects {accepted}");
            assert!(
                accepted.parse::<MemSize>().is_ok(),
                "FromStr rejects {accepted}"
            );
        }
        for rejected in ["512MB", "512m", "1.5G", "", "M", "512T", "1P", "512g"] {
            assert!(!re.is_match(rejected), "pattern accepts {rejected}");
            assert!(
                rejected.parse::<MemSize>().is_err(),
                "FromStr accepts {rejected}"
            );
        }
    }
}

#[cfg(test)]
mod up_duration_tests {
    use super::*;

    #[test]
    fn plain_digits_are_milliseconds() {
        assert_eq!("1600".parse::<UpDuration>().unwrap().as_millis(), 1600);
    }

    #[test]
    fn units_seconds_minutes_hours() {
        assert_eq!("30s".parse::<UpDuration>().unwrap().as_millis(), 30_000);
        assert_eq!("5m".parse::<UpDuration>().unwrap().as_millis(), 300_000);
        assert_eq!("2h".parse::<UpDuration>().unwrap().as_millis(), 7_200_000);
    }

    /// `ms` and `m` share a trailing byte, so a naive last-byte match would
    /// parse "5ms" as "5m" (a 60,000x error) or reject it. Pinned adjacently
    /// so a regression shows as a wrong multiplier, not a rejected string.
    #[test]
    fn milliseconds_do_not_alias_minutes() {
        assert_eq!("500ms".parse::<UpDuration>().unwrap().as_millis(), 500);
        assert_eq!("5ms".parse::<UpDuration>().unwrap().as_millis(), 5);
        assert_eq!("5m".parse::<UpDuration>().unwrap().as_millis(), 300_000);
        // A bare trailing `m` at end of input is still minutes.
        assert_eq!("1m".parse::<UpDuration>().unwrap().as_millis(), 60_000);
    }

    #[test]
    fn rejects_spec_violations() {
        use ParseUpDurationError::*;
        assert_eq!("".parse::<UpDuration>(), Err(Empty));
        assert_eq!("s".parse::<UpDuration>(), Err(MissingDigits));
        assert_eq!("ms".parse::<UpDuration>(), Err(MissingDigits));
        assert_eq!("30S".parse::<UpDuration>(), Err(InvalidCharacter)); // uppercase
        assert_eq!("1.5s".parse::<UpDuration>(), Err(InvalidCharacter));
        assert_eq!("30 s".parse::<UpDuration>(), Err(InvalidCharacter));
        assert_eq!("30MS".parse::<UpDuration>(), Err(InvalidCharacter)); // uppercase ms
        // Digit string itself overflows u64 before any unit multiplication.
        assert_eq!("99999999999999999999h".parse::<UpDuration>(), Err(Overflow));
        // Digit string fits u64 on its own, but overflows on the ×3_600_000
        // (hours-to-ms) multiplication.
        assert_eq!("9999999999999999h".parse::<UpDuration>(), Err(Overflow));
    }

    #[test]
    fn display_round_trips() {
        for ms in [
            0u64, 1, 999, 1000, 1600, 30_000, 300_000, 7_200_000, 3_601_000,
        ] {
            let d = UpDuration::from_millis(ms);
            assert_eq!(d.to_string().parse::<UpDuration>().unwrap(), d, "{ms}ms");
        }
        assert_eq!(UpDuration::from_millis(30_000).to_string(), "30s");
        assert_eq!(UpDuration::from_millis(1600).to_string(), "1600");
        assert_eq!(UpDuration::from_millis(7_200_000).to_string(), "2h");
    }

    #[test]
    fn serde_uses_string_form() {
        let d: UpDuration = serde_json::from_str("\"30s\"").unwrap();
        assert_eq!(d.as_millis(), 30_000);
        assert_eq!(serde_json::to_string(&d).unwrap(), "\"30s\"");
    }

    /// Same requirement as [`MemSize`]'s schema test: the pattern must
    /// agree with `FromStr`, not just be self-consistent.
    #[cfg(feature = "schema")]
    #[test]
    fn the_schema_pattern_agrees_with_from_str() {
        let schema = serde_json::to_value(schemars::schema_for!(UpDuration)).unwrap();
        let pattern = schema["pattern"].as_str().unwrap();
        let re = regex::Regex::new(pattern).unwrap();
        for accepted in ["1600", "30s", "5m", "2h", "500ms"] {
            assert!(re.is_match(accepted), "pattern rejects {accepted}");
            assert!(
                accepted.parse::<UpDuration>().is_ok(),
                "FromStr rejects {accepted}"
            );
        }
        for rejected in ["30S", "1.5s", "30 s", "", "s", "30d", "30w", "30MS"] {
            assert!(!re.is_match(rejected), "pattern accepts {rejected}");
            assert!(
                rejected.parse::<UpDuration>().is_err(),
                "FromStr accepts {rejected}"
            );
        }
    }
}