117 lines
2.8 KiB
Rust
117 lines
2.8 KiB
Rust
use std::collections::VecDeque;
|
|
|
|
use crate::notify::Notify;
|
|
use crate::sync::{Arc, Mutex};
|
|
|
|
pub(super) fn bounded<T>(capacity: usize) -> Queue<T> {
|
|
Queue::bounded(capacity)
|
|
}
|
|
|
|
pub(crate) struct Queue<T> {
|
|
inner: Arc<QueueState<T>>,
|
|
capacity: usize,
|
|
}
|
|
|
|
struct QueueState<T> {
|
|
queue: Mutex<VecDeque<T>>,
|
|
push_notify: Notify,
|
|
pop_notify: Notify,
|
|
}
|
|
|
|
impl<T> Queue<T> {
|
|
pub(super) fn bounded(capacity: usize) -> Self {
|
|
metrics::counter!("async-cpupool.queue.created").increment(1);
|
|
Self {
|
|
inner: Arc::new(QueueState {
|
|
queue: Mutex::new(VecDeque::new()),
|
|
push_notify: Notify::new(),
|
|
pop_notify: Notify::new(),
|
|
}),
|
|
capacity,
|
|
}
|
|
}
|
|
|
|
pub(super) fn len(&self) -> usize {
|
|
self.inner.queue.lock().expect("not poisoned").len()
|
|
}
|
|
|
|
pub(super) fn is_full_or(&self) -> Result<(), usize> {
|
|
let len = self.len();
|
|
|
|
if len >= self.capacity {
|
|
Ok(())
|
|
} else {
|
|
Err(len)
|
|
}
|
|
}
|
|
|
|
pub(super) async fn push(&self, mut item: T) {
|
|
loop {
|
|
let listener = self.inner.push_notify.listen().await;
|
|
|
|
if let Some(returned_item) = self.try_push(item) {
|
|
item = returned_item;
|
|
|
|
listener.await;
|
|
} else {
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(super) fn try_push(&self, item: T) -> Option<T> {
|
|
match self.try_push_impl(item) {
|
|
Some(item) => Some(item),
|
|
None => {
|
|
self.inner.pop_notify.notify_one();
|
|
metrics::counter!("async-cpupool.queue.pushed").increment(1);
|
|
None
|
|
}
|
|
}
|
|
}
|
|
|
|
fn try_push_impl(&self, item: T) -> Option<T> {
|
|
let mut guard = self.inner.queue.lock().expect("not poisoned");
|
|
|
|
if self.capacity <= guard.len() {
|
|
Some(item)
|
|
} else {
|
|
guard.push_back(item);
|
|
None
|
|
}
|
|
}
|
|
|
|
pub(super) async fn pop(&self) -> T {
|
|
loop {
|
|
let listener = self.inner.pop_notify.listen().await;
|
|
|
|
if let Some(item) = self.try_pop() {
|
|
self.inner.push_notify.notify_one();
|
|
metrics::counter!("async-cpupool.queue.popped").increment(1);
|
|
return item;
|
|
}
|
|
|
|
listener.await;
|
|
}
|
|
}
|
|
|
|
fn try_pop(&self) -> Option<T> {
|
|
self.inner.queue.lock().expect("not poisoned").pop_front()
|
|
}
|
|
}
|
|
|
|
impl<T> Clone for Queue<T> {
|
|
fn clone(&self) -> Self {
|
|
Self {
|
|
inner: Arc::clone(&self.inner),
|
|
capacity: self.capacity,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<T> Drop for Queue<T> {
|
|
fn drop(&mut self) {
|
|
metrics::counter!("async-cpupool.queue.dropped").increment(1);
|
|
}
|
|
}
|