wireshift-core 0.1.1

Core typed operations, buffers, and backend traits for wireshift
Documentation
//! Callback logic via an async runtime `Waker`.
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>>,
}

/// Bridges ring completions into a task waker.
///
/// ```no_run
/// use wireshift::{Ring, RingConfig, strategy::WakerDrain};
///
/// let ring = Ring::new(RingConfig::default())?;
/// let drain = WakerDrain::new(ring);
/// let _ = drain.drain_ready();
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
pub struct WakerDrain {
    state: Arc<WakerState>,
    running: Arc<AtomicBool>,
    worker: Option<thread::JoinHandle<()>>,
}

impl WakerDrain {
    /// Creates a waker-backed completion bridge.
    #[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),
        }
    }

    /// Registers the most recent waker to notify when new completions arrive.
    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(())
    }

    /// Drains all queued completion notifications accumulated by the worker.
    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();
        }
    }
}