use std::sync::Arc;
use parking_lot::{Condvar, Mutex};
pub const DEFAULT_LLM_PERMITS: usize = 1;
#[derive(Clone)]
pub struct LlmGate {
inner: Arc<Inner>,
}
struct Inner {
available: Mutex<usize>,
cond: Condvar,
}
pub struct Permit {
inner: Arc<Inner>,
}
impl LlmGate {
pub fn new(permits: usize) -> Self {
Self {
inner: Arc::new(Inner {
available: Mutex::new(permits.max(1)),
cond: Condvar::new(),
}),
}
}
pub fn try_acquire(&self) -> Option<Permit> {
let mut avail = self.inner.available.lock();
if *avail == 0 {
return None;
}
*avail -= 1;
Some(Permit {
inner: self.inner.clone(),
})
}
pub fn acquire(&self) -> Permit {
let mut avail = self.inner.available.lock();
while *avail == 0 {
self.inner.cond.wait(&mut avail);
}
*avail -= 1;
Permit {
inner: self.inner.clone(),
}
}
pub fn available_permits(&self) -> usize {
*self.inner.available.lock()
}
}
impl Default for LlmGate {
fn default() -> Self {
Self::new(DEFAULT_LLM_PERMITS)
}
}
impl Drop for Permit {
fn drop(&mut self) {
let mut avail = self.inner.available.lock();
*avail += 1;
self.inner.cond.notify_one();
}
}
#[cfg(test)]
#[path = "gate_tests.rs"]
mod tests;