crank/crates/crank-runtime/src/limits.rs

use std::{env, num::ParseIntError};

use thiserror::Error;

const DEFAULT_MAX_CONCURRENT_UNARY: usize = 64;
const DEFAULT_MAX_CONCURRENT_WINDOW: usize = 16;
const DEFAULT_MAX_CONCURRENT_SESSIONS: usize = 16;
const DEFAULT_MAX_CONCURRENT_JOBS: usize = 16;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RuntimeLimits {
    pub max_concurrent_unary: usize,
    pub max_concurrent_window: usize,
    pub max_concurrent_sessions: usize,
    pub max_concurrent_jobs: usize,
}

impl Default for RuntimeLimits {
    fn default() -> Self {
        Self {
            max_concurrent_unary: DEFAULT_MAX_CONCURRENT_UNARY,
            max_concurrent_window: DEFAULT_MAX_CONCURRENT_WINDOW,
            max_concurrent_sessions: DEFAULT_MAX_CONCURRENT_SESSIONS,
            max_concurrent_jobs: DEFAULT_MAX_CONCURRENT_JOBS,
        }
    }
}

impl RuntimeLimits {
    pub fn from_env() -> Result<Self, RuntimeLimitsConfigError> {
        Ok(Self {
            max_concurrent_unary: parse_limit(
                "CRANK_RUNTIME_MAX_CONCURRENT_UNARY",
                DEFAULT_MAX_CONCURRENT_UNARY,
            )?,
            max_concurrent_window: parse_limit(
                "CRANK_RUNTIME_MAX_CONCURRENT_WINDOW",
                DEFAULT_MAX_CONCURRENT_WINDOW,
            )?,
            max_concurrent_sessions: parse_limit(
                "CRANK_RUNTIME_MAX_CONCURRENT_SESSIONS",
                DEFAULT_MAX_CONCURRENT_SESSIONS,
            )?,
            max_concurrent_jobs: parse_limit(
                "CRANK_RUNTIME_MAX_CONCURRENT_JOBS",
                DEFAULT_MAX_CONCURRENT_JOBS,
            )?,
        })
    }
}

fn parse_limit(name: &'static str, default: usize) -> Result<usize, RuntimeLimitsConfigError> {
    match env::var(name) {
        Ok(raw) => {
            let value =
                raw.parse::<usize>()
                    .map_err(|source| RuntimeLimitsConfigError::InvalidValue {
                        name,
                        value: raw,
                        source,
                    })?;
            if value == 0 {
                return Err(RuntimeLimitsConfigError::ZeroValue { name });
            }
            Ok(value)
        }
        Err(env::VarError::NotPresent) => Ok(default),
        Err(env::VarError::NotUnicode(_)) => Err(RuntimeLimitsConfigError::InvalidUnicode { name }),
    }
}

#[derive(Debug, Error)]
pub enum RuntimeLimitsConfigError {
    #[error("{name} must contain valid UTF-8")]
    InvalidUnicode { name: &'static str },
    #[error("{name} must be a positive integer, got {value}")]
    InvalidValue {
        name: &'static str,
        value: String,
        source: ParseIntError,
    },
    #[error("{name} must be greater than zero")]
    ZeroValue { name: &'static str },
}

#[cfg(test)]
mod tests {
    use super::{RuntimeLimits, RuntimeLimitsConfigError};

    #[test]
    fn defaults_are_positive() {
        let limits = RuntimeLimits::default();

        assert!(limits.max_concurrent_unary > 0);
        assert!(limits.max_concurrent_window > 0);
        assert!(limits.max_concurrent_sessions > 0);
        assert!(limits.max_concurrent_jobs > 0);
    }

    #[test]
    fn rejects_zero_limit_values() {
        unsafe {
            std::env::set_var("CRANK_RUNTIME_MAX_CONCURRENT_UNARY", "0");
        }

        let error = RuntimeLimits::from_env().unwrap_err();

        assert!(matches!(
            error,
            RuntimeLimitsConfigError::ZeroValue {
                name: "CRANK_RUNTIME_MAX_CONCURRENT_UNARY"
            }
        ));

        unsafe {
            std::env::remove_var("CRANK_RUNTIME_MAX_CONCURRENT_UNARY");
        }
    }
}