|
|
|
@ -4,6 +4,9 @@ use tokio::{
|
|
|
|
|
sync::{
|
|
|
|
|
Semaphore,
|
|
|
|
|
SemaphorePermit,
|
|
|
|
|
RwLock,
|
|
|
|
|
RwLockReadGuard,
|
|
|
|
|
RwLockWriteGuard,
|
|
|
|
|
},
|
|
|
|
|
};
|
|
|
|
|
use futures::{
|
|
|
|
@ -16,11 +19,12 @@ pub struct State
|
|
|
|
|
{
|
|
|
|
|
cfg: config::Config,
|
|
|
|
|
mtx: Option<Semaphore>,
|
|
|
|
|
stop: RwLock<()>,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Guard for operations inside state
|
|
|
|
|
#[derive(Debug)]
|
|
|
|
|
pub struct Permit<'a>(Option<SemaphorePermit<'a>>);
|
|
|
|
|
pub struct Permit<'a>(RwLockReadGuard<'a, ()>, Option<SemaphorePermit<'a>>);
|
|
|
|
|
|
|
|
|
|
impl State
|
|
|
|
|
{
|
|
|
|
@ -30,13 +34,25 @@ impl State
|
|
|
|
|
Self {
|
|
|
|
|
mtx: cfg.limit.map(|x| Semaphore::new(x.into())),
|
|
|
|
|
cfg,
|
|
|
|
|
stop: RwLock::new(()),
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Acquire exclusive lock, preventing other tasks from acquiring through either `lock` or `lockout`.
|
|
|
|
|
///
|
|
|
|
|
/// This function will yield until all other `lock` or `lockout` guards currently acquired (if any) are released.
|
|
|
|
|
pub async fn lockout(&self) -> RwLockWriteGuard<'_, ()>
|
|
|
|
|
{
|
|
|
|
|
self.stop.write().await
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// Acquire a permit for concurrent work, yielding the task if needed.
|
|
|
|
|
///
|
|
|
|
|
/// If there are currently `config().limit` premits acquired, this will yield until at least one if released.
|
|
|
|
|
/// If there is a `lockout` guard acquired, this will yield until it is released too.
|
|
|
|
|
pub async fn lock(&self) -> Permit<'_>
|
|
|
|
|
{
|
|
|
|
|
Permit(OptionFuture::from(self.mtx.as_ref().map(Semaphore::acquire)).await)
|
|
|
|
|
Permit(self.stop.read().await, OptionFuture::from(self.mtx.as_ref().map(Semaphore::acquire)).await)
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
/// The config object
|
|
|
|
|