use std::sync::Arc;
use oxirs_core::sla::SlaClass;
use super::admission::{ClusterAdmissionController, SlaError};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProposerOutcome {
Admitted,
Rejected(SlaError),
}
#[derive(Debug, Clone)]
pub struct SlaProposerGate {
controller: Arc<ClusterAdmissionController>,
}
impl SlaProposerGate {
pub fn new(controller: Arc<ClusterAdmissionController>) -> Self {
Self { controller }
}
pub fn controller(&self) -> &Arc<ClusterAdmissionController> {
&self.controller
}
pub fn try_acquire(&self, class: SlaClass) -> Result<(), SlaError> {
self.controller.acquire_permit(class)
}
pub fn admit(&self, class: SlaClass) -> ProposerOutcome {
match self.controller.acquire_permit(class) {
Ok(()) => ProposerOutcome::Admitted,
Err(e) => ProposerOutcome::Rejected(e),
}
}
pub fn release(&self, class: SlaClass) -> Result<(), SlaError> {
self.controller.release_permit(class)
}
pub fn scoped<F, T>(&self, class: SlaClass, f: F) -> Result<T, SlaError>
where
F: FnOnce() -> T,
{
self.controller.scoped_admit(class, f)
}
pub fn in_flight(&self, class: SlaClass) -> Result<usize, SlaError> {
self.controller.in_flight(class)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sla::admission::{SlaAdmissionConfig, SlaClassQuota};
use std::collections::HashMap;
fn small_gate() -> SlaProposerGate {
let mut quotas = HashMap::new();
quotas.insert(
SlaClass::Bronze,
SlaClassQuota {
max_qps: None,
max_concurrent: 1,
token_cost: 1.0,
},
);
let controller = Arc::new(ClusterAdmissionController::new(SlaAdmissionConfig {
quotas,
}));
SlaProposerGate::new(controller)
}
#[test]
fn admitted_outcome_for_first_request() {
let gate = small_gate();
assert_eq!(gate.admit(SlaClass::Bronze), ProposerOutcome::Admitted);
}
#[test]
fn second_request_blocked_by_concurrency_cap() {
let gate = small_gate();
gate.try_acquire(SlaClass::Bronze).expect("first");
let outcome = gate.admit(SlaClass::Bronze);
match outcome {
ProposerOutcome::Rejected(SlaError::ConcurrencyCapExceeded { limit, .. }) => {
assert_eq!(limit, 1);
}
other => panic!("unexpected outcome: {:?}", other),
}
}
#[test]
fn release_frees_slot() {
let gate = small_gate();
gate.try_acquire(SlaClass::Bronze).expect("first");
gate.release(SlaClass::Bronze).expect("release");
gate.try_acquire(SlaClass::Bronze).expect("second");
}
#[test]
fn scoped_runs_callback_and_releases() {
let gate = small_gate();
let result = gate.scoped(SlaClass::Bronze, || 42).expect("scoped");
assert_eq!(result, 42);
assert_eq!(gate.in_flight(SlaClass::Bronze).expect("count"), 0);
}
#[test]
fn unregistered_class_rejected() {
let gate = small_gate();
let outcome = gate.admit(SlaClass::Gold);
assert!(matches!(
outcome,
ProposerOutcome::Rejected(SlaError::ClassNotRegistered { .. })
));
}
}