129 lines
3.8 KiB
Rust
129 lines
3.8 KiB
Rust
// Copied from tower/src/semaphore.rs, commit:
|
|
// d4d1c67 hedge: use auto-resizing histograms (#484)
|
|
//
|
|
// When we upgrade to tower 0.4, we can use tokio's PollSemaphore, like tower's:
|
|
// ccfaffc buffer, limit: use `tokio-util`'s `PollSemaphore` (#556)
|
|
|
|
// Ignore lints on this copied code
|
|
#![allow(dead_code)]
|
|
|
|
pub(crate) use self::sync::OwnedSemaphorePermit as Permit;
|
|
use futures::FutureExt;
|
|
use futures_core::ready;
|
|
use std::{
|
|
fmt,
|
|
future::Future,
|
|
mem,
|
|
pin::Pin,
|
|
sync::{Arc, Weak},
|
|
task::{Context, Poll},
|
|
};
|
|
use tokio::sync;
|
|
|
|
#[derive(Debug)]
|
|
pub(crate) struct Semaphore {
|
|
semaphore: Arc<sync::Semaphore>,
|
|
state: State,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub(crate) struct Close {
|
|
semaphore: Weak<sync::Semaphore>,
|
|
permits: usize,
|
|
}
|
|
|
|
enum State {
|
|
Waiting(Pin<Box<dyn Future<Output = Permit> + Send + Sync + 'static>>),
|
|
Ready(Permit),
|
|
Empty,
|
|
}
|
|
|
|
impl Semaphore {
|
|
pub(crate) fn new_with_close(permits: usize) -> (Self, Close) {
|
|
let semaphore = Arc::new(sync::Semaphore::new(permits));
|
|
let close = Close {
|
|
semaphore: Arc::downgrade(&semaphore),
|
|
permits,
|
|
};
|
|
let semaphore = Self {
|
|
semaphore,
|
|
state: State::Empty,
|
|
};
|
|
(semaphore, close)
|
|
}
|
|
|
|
pub(crate) fn new(permits: usize) -> Self {
|
|
Self {
|
|
semaphore: Arc::new(sync::Semaphore::new(permits)),
|
|
state: State::Empty,
|
|
}
|
|
}
|
|
|
|
pub(crate) fn poll_acquire(&mut self, cx: &mut Context<'_>) -> Poll<()> {
|
|
loop {
|
|
self.state = match self.state {
|
|
State::Ready(_) => return Poll::Ready(()),
|
|
State::Waiting(ref mut fut) => {
|
|
let permit = ready!(Pin::new(fut).poll(cx));
|
|
State::Ready(permit)
|
|
}
|
|
State::Empty => State::Waiting(Box::pin(
|
|
self.semaphore
|
|
.clone()
|
|
.acquire_owned()
|
|
.map(|result| result.expect("internal semaphore is never closed")),
|
|
)),
|
|
};
|
|
}
|
|
}
|
|
|
|
pub(crate) fn take_permit(&mut self) -> Option<Permit> {
|
|
if let State::Ready(permit) = mem::replace(&mut self.state, State::Empty) {
|
|
return Some(permit);
|
|
}
|
|
None
|
|
}
|
|
}
|
|
|
|
impl Clone for Semaphore {
|
|
fn clone(&self) -> Self {
|
|
Self {
|
|
semaphore: self.semaphore.clone(),
|
|
state: State::Empty,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl fmt::Debug for State {
|
|
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
|
match self {
|
|
State::Waiting(_) => f
|
|
.debug_tuple("State::Waiting")
|
|
.field(&format_args!("..."))
|
|
.finish(),
|
|
State::Ready(ref r) => f.debug_tuple("State::Ready").field(&r).finish(),
|
|
State::Empty => f.debug_tuple("State::Empty").finish(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Close {
|
|
/// Close the semaphore, waking any remaining tasks currently awaiting a permit.
|
|
pub(crate) fn close(self) {
|
|
// The maximum number of permits that a `tokio::sync::Semaphore`
|
|
// can hold is usize::MAX >> 3. If we attempt to add more than that
|
|
// number of permits, the semaphore will panic.
|
|
// XXX(eliza): another shift is kinda janky but if we add (usize::MAX
|
|
// > 3 - initial permits) the semaphore impl panics (I think due to a
|
|
// bug in tokio?).
|
|
// TODO(eliza): Tokio should _really_ just expose `Semaphore::close`
|
|
// publicly so we don't have to do this nonsense...
|
|
const MAX: usize = std::usize::MAX >> 4;
|
|
if let Some(semaphore) = self.semaphore.upgrade() {
|
|
// If we added `MAX - available_permits`, any tasks that are
|
|
// currently holding permits could drop them, overflowing the max.
|
|
semaphore.add_permits(MAX - self.permits);
|
|
}
|
|
}
|
|
}
|