use crossbeam_queue::SegQueue;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::task::Waker;
use std::thread;
use std::time::Duration;
use crate::completion::CompletionEvent;
use crate::error::{Error, Result};
use crate::{BackendFactory, Ring};
struct WakerState {
queued: SegQueue<Result<CompletionEvent>>,
waker: Mutex<Option<Waker>>,
}
pub struct WakerDrain {
state: Arc<WakerState>,
running: Arc<AtomicBool>,
worker: Option<thread::JoinHandle<()>>,
}
impl WakerDrain {
#[must_use]
pub fn new<F: BackendFactory>(ring: Ring<F>) -> Self {
let state = Arc::new(WakerState {
queued: SegQueue::new(),
waker: Mutex::new(None),
});
let running = Arc::new(AtomicBool::new(true));
let worker_state = state.clone();
let worker_running = running.clone();
let worker = thread::spawn(move || {
while worker_running.load(Ordering::Acquire) {
match ring.complete(Some(Duration::from_millis(50))) {
Ok(event) => {
worker_state.queued.push(Ok(event));
wake_registered(&worker_state);
}
Err(Error::Timeout { .. }) => {}
Err(error) => {
worker_state.queued.push(Err(error));
wake_registered(&worker_state);
break;
}
}
}
});
Self {
state,
running,
worker: Some(worker),
}
}
pub fn register(&self, waker: &Waker) -> Result<()> {
let mut slot = self.state.waker.lock().map_err(|_| {
Error::completion(
"waker registry mutex was poisoned",
"avoid panicking while registering a completion waker",
)
})?;
*slot = Some(waker.clone());
Ok(())
}
pub fn drain_ready(&self) -> Result<Vec<Result<CompletionEvent>>> {
let mut results = Vec::new();
while let Some(item) = self.state.queued.pop() {
results.push(item);
}
Ok(results)
}
}
impl Drop for WakerDrain {
fn drop(&mut self) {
self.running.store(false, Ordering::Release);
if let Some(worker) = self.worker.take() {
let _ = worker.join();
}
}
}
fn wake_registered(state: &WakerState) {
if let Ok(slot) = state.waker.lock() {
if let Some(waker) = slot.as_ref() {
waker.wake_by_ref();
}
}
}