use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::task::{Context, Poll, Waker};
use std::time::Duration;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SaturationPolicy {
Wait {
max_queued: usize,
queue_timeout: Option<Duration>,
},
Reject,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum Scope {
Operations,
Attempts,
}
#[derive(Debug, Clone)]
pub struct ConcurrencyLimit {
permits: usize,
scope: Scope,
saturation: SaturationPolicy,
}
impl ConcurrencyLimit {
pub fn operations(n: usize) -> Self {
Self::new(n, Scope::Operations)
}
pub fn attempts(n: usize) -> Self {
Self::new(n, Scope::Attempts)
}
fn new(permits: usize, scope: Scope) -> Self {
Self {
permits: permits.max(1),
scope,
saturation: SaturationPolicy::Wait {
max_queued: usize::MAX,
queue_timeout: None,
},
}
}
pub fn max_queued(mut self, n: usize) -> Self {
self.saturation = match self.saturation {
SaturationPolicy::Wait { queue_timeout, .. } => SaturationPolicy::Wait {
max_queued: n,
queue_timeout,
},
SaturationPolicy::Reject => SaturationPolicy::Reject,
};
self
}
pub fn queue_timeout(mut self, d: Duration) -> Self {
self.saturation = match self.saturation {
SaturationPolicy::Wait { max_queued, .. } => SaturationPolicy::Wait {
max_queued,
queue_timeout: Some(d),
},
SaturationPolicy::Reject => SaturationPolicy::Reject,
};
self
}
pub fn reject(mut self) -> Self {
self.saturation = SaturationPolicy::Reject;
self
}
pub(crate) fn build(&self) -> Arc<Semaphore> {
Semaphore::new(self.permits)
}
pub(crate) fn compile(&self) -> CompiledConcurrency {
CompiledConcurrency {
sem: self.build(),
saturation: self.saturation,
scope: self.scope,
}
}
}
#[derive(Debug)]
pub(crate) struct CompiledConcurrency {
pub(crate) sem: Arc<Semaphore>,
pub(crate) saturation: SaturationPolicy,
pub(crate) scope: Scope,
}
impl From<usize> for ConcurrencyLimit {
fn from(n: usize) -> Self {
ConcurrencyLimit::operations(n)
}
}
impl From<u32> for ConcurrencyLimit {
fn from(n: u32) -> Self {
ConcurrencyLimit::operations(n as usize)
}
}
#[derive(Debug)]
struct SemState {
permits: usize,
waiters: VecDeque<Waker>,
}
#[derive(Debug)]
pub(crate) struct Semaphore {
state: Mutex<SemState>,
}
impl Semaphore {
fn new(permits: usize) -> Arc<Self> {
Arc::new(Self {
state: Mutex::new(SemState {
permits,
waiters: VecDeque::new(),
}),
})
}
pub(crate) fn queued(self: &Arc<Self>) -> usize {
self.state.lock().unwrap().waiters.len()
}
pub(crate) fn try_acquire(self: &Arc<Self>) -> Option<Permit> {
let mut st = self.state.lock().unwrap();
if st.permits > 0 {
st.permits -= 1;
Some(Permit {
sem: Arc::clone(self),
})
} else {
None
}
}
pub(crate) fn acquire(self: &Arc<Self>) -> Acquire {
Acquire {
sem: Arc::clone(self),
}
}
fn release(&self) {
let mut st = self.state.lock().unwrap();
st.permits += 1;
if let Some(w) = st.waiters.pop_front() {
w.wake();
}
}
}
#[derive(Debug)]
pub(crate) struct Permit {
sem: Arc<Semaphore>,
}
impl Drop for Permit {
fn drop(&mut self) {
self.sem.release();
}
}
pub(crate) struct Acquire {
sem: Arc<Semaphore>,
}
impl Future for Acquire {
type Output = Permit;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Permit> {
let mut st = self.sem.state.lock().unwrap();
if st.permits > 0 {
st.permits -= 1;
return Poll::Ready(Permit {
sem: Arc::clone(&self.sem),
});
}
st.waiters.push_back(cx.waker().clone());
Poll::Pending
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn permit_release_lets_next_in() {
let cfg = ConcurrencyLimit::operations(1);
let sem = cfg.build();
let p1 = sem.try_acquire().expect("first permit");
assert!(sem.try_acquire().is_none(), "saturated");
let acq = sem.acquire();
tokio::pin!(acq);
let waiter = async { (&mut acq).await };
tokio::pin!(waiter);
let releaser = async {
tokio::task::yield_now().await;
drop(p1); };
let (_p2, ()) = tokio::join!(waiter, releaser);
assert!(sem.try_acquire().is_none());
}
#[test]
fn queued_counts_waiters() {
let sem = ConcurrencyLimit::operations(1).build();
let _p = sem.try_acquire().unwrap();
assert_eq!(sem.queued(), 0);
}
#[test]
fn saturation_builders() {
let c = ConcurrencyLimit::attempts(8)
.max_queued(4)
.queue_timeout(Duration::from_millis(10))
.compile();
assert_eq!(c.scope, Scope::Attempts);
match c.saturation {
SaturationPolicy::Wait {
max_queued,
queue_timeout,
} => {
assert_eq!(max_queued, 4);
assert_eq!(queue_timeout, Some(Duration::from_millis(10)));
}
_ => panic!("expected Wait"),
}
assert_eq!(
ConcurrencyLimit::operations(2)
.reject()
.compile()
.saturation,
SaturationPolicy::Reject
);
}
}