use scc::HashMap;
use std::sync::Arc;
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use url::Url;
#[derive(Clone)]
pub(super) struct BackfillAdmission {
gates: Arc<HashMap<Url, Arc<Semaphore>>>,
per_pds_limit: usize,
}
impl BackfillAdmission {
pub(super) fn new(global_limit: usize, per_pds_limit: usize) -> Self {
// leave room for another PDS whenever the global limit allows it
let fair_limit = global_limit.saturating_sub(1).max(1);
Self {
gates: Arc::new(HashMap::new()),
per_pds_limit: per_pds_limit.min(fair_limit).max(1),
}
}
pub(super) async fn try_acquire(&self, pds: &Url) -> Option<OwnedSemaphorePermit> {
let limit = self.per_pds_limit;
let gate = self
.gates
.entry_async(pds.clone())
.await
.or_insert_with(|| Arc::new(Semaphore::new(limit)))
.get()
.clone();
gate.try_acquire_owned().ok()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn saturated_pds_does_not_block_another_pds() {
let admission = BackfillAdmission::new(4, 4);
let first = Url::parse("https://first.example").unwrap();
let second = Url::parse("https://second.example").unwrap();
let permits = futures::future::join_all((0..3).map(|_| admission.try_acquire(&first)))
.await
.into_iter()
.collect::<Option<Vec<_>>>()
.expect("first PDS should receive its fair share");
assert!(admission.try_acquire(&first).await.is_none());
assert!(admission.try_acquire(&second).await.is_some());
drop(permits);
assert!(admission.try_acquire(&first).await.is_some());
}
#[tokio::test]
async fn single_global_slot_remains_usable() {
let admission = BackfillAdmission::new(1, 10);
let pds = Url::parse("https://pds.example").unwrap();
assert!(admission.try_acquire(&pds).await.is_some());
}
}