use std::collections::{HashMap, HashSet};
use crate::http::Request;
pub const DEFAULT_INLINE_BUDGET_THRESHOLD_BYTES: usize = 102_400;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Decision {
Inline,
Preload(String),
}
#[derive(Default)]
pub(crate) struct InlineBudgetState {
pub(crate) cumulative: HashMap<String, usize>,
pub(crate) warned: HashSet<String>,
}
impl InlineBudgetState {
pub(crate) fn record_and_decide(
&mut self,
key: &str,
bytes: usize,
threshold: usize,
fallback_url: &str,
route_pattern: &str,
) -> Decision {
let entry = self.cumulative.entry(key.to_string()).or_insert(0);
*entry = entry.saturating_add(bytes);
let cumulative = *entry;
if cumulative <= threshold {
return Decision::Inline;
}
if self.warned.insert(key.to_string()) {
tracing::warn!(
key = %key,
cumulative_bytes = cumulative,
threshold_bytes = threshold,
fallback_url = %fallback_url,
route_pattern = %route_pattern,
"inline_budget: threshold crossed; flipping to Preload"
);
}
Decision::Preload(fallback_url.to_string())
}
}
pub(crate) fn decide(req: &mut Request, key: &str, bytes: usize, fallback_url: &str) -> Decision {
let threshold = crate::Config::get::<crate::AppConfig>()
.map(|c| c.inline_budget_threshold_bytes)
.unwrap_or(DEFAULT_INLINE_BUDGET_THRESHOLD_BYTES);
let route_pattern = req.route_pattern().unwrap_or_default();
if req.get::<InlineBudgetState>().is_none() {
req.insert(InlineBudgetState::default());
}
let state = req
.get_mut::<InlineBudgetState>()
.expect("InlineBudgetState was just inserted above");
state.record_and_decide(key, bytes, threshold, fallback_url, &route_pattern)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn decision_enum_is_clone_and_eq() {
let a = Decision::Inline;
let b = a.clone();
assert_eq!(a, b);
let c = Decision::Preload("/x".to_string());
let d = c.clone();
assert_eq!(c, d);
assert_ne!(a, c);
}
#[test]
fn inline_budget_state_default_is_empty() {
let s = InlineBudgetState::default();
assert!(s.cumulative.is_empty());
assert!(s.warned.is_empty());
}
#[test]
fn decides_inline_below_threshold() {
let mut state = InlineBudgetState::default();
let d = state.record_and_decide("k", 50_000, 102_400, "/fb", "");
assert_eq!(d, Decision::Inline);
assert_eq!(*state.cumulative.get("k").unwrap(), 50_000);
assert!(state.warned.is_empty());
}
#[test]
fn decides_inline_at_exact_threshold() {
let mut state = InlineBudgetState::default();
let d = state.record_and_decide("k", 102_400, 102_400, "/fb", "");
assert_eq!(d, Decision::Inline);
assert_eq!(*state.cumulative.get("k").unwrap(), 102_400);
assert!(state.warned.is_empty());
}
#[test]
fn decides_preload_above_threshold() {
let mut state = InlineBudgetState::default();
let d = state.record_and_decide("k", 102_401, 102_400, "/fb", "");
assert_eq!(d, Decision::Preload("/fb".to_string()));
assert_eq!(*state.cumulative.get("k").unwrap(), 102_401);
assert!(state.warned.contains("k"));
assert_eq!(state.warned.len(), 1);
}
#[test]
fn decides_preload_after_accumulation() {
let mut state = InlineBudgetState::default();
let d1 = state.record_and_decide("k", 40_000, 102_400, "/fb", "");
assert_eq!(d1, Decision::Inline);
assert_eq!(*state.cumulative.get("k").unwrap(), 40_000);
assert!(state.warned.is_empty());
let d2 = state.record_and_decide("k", 40_000, 102_400, "/fb", "");
assert_eq!(d2, Decision::Inline);
assert_eq!(*state.cumulative.get("k").unwrap(), 80_000);
assert!(state.warned.is_empty());
let d3 = state.record_and_decide("k", 40_000, 102_400, "/fb", "");
assert_eq!(d3, Decision::Preload("/fb".to_string()));
assert_eq!(*state.cumulative.get("k").unwrap(), 120_000);
assert!(state.warned.contains("k"));
assert_eq!(state.warned.len(), 1);
}
#[test]
fn warn_fires_once_per_key() {
let mut state = InlineBudgetState::default();
let d1 = state.record_and_decide("k", 200_000, 102_400, "/fb", "");
assert!(matches!(d1, Decision::Preload(_)));
assert!(state.warned.contains("k"));
assert_eq!(state.warned.len(), 1);
let d2 = state.record_and_decide("k", 1, 102_400, "/fb", "");
assert!(matches!(d2, Decision::Preload(_)));
assert!(state.warned.contains("k"));
assert_eq!(state.warned.len(), 1); }
#[test]
fn warn_independent_per_key() {
let mut state = InlineBudgetState::default();
let _ = state.record_and_decide("key_a", 200_000, 102_400, "/fb-a", "");
let _ = state.record_and_decide("key_b", 200_000, 102_400, "/fb-b", "");
assert!(state.warned.contains("key_a"));
assert!(state.warned.contains("key_b"));
assert_eq!(state.warned.len(), 2);
}
}