wireshift_core/

ring.rs

//! Ring manager and completion routing.
//!
//! This module defines [`Ring`], the primary interface for submitting operations
//! and receiving completions, as well as [`ChainContext`] for dependent operation
//! execution.
//!
//! # Ring Lifecycle
//!
//! ```text
//! 1. Create: Ring::new(config) -> Result<Ring>
//!    - Validates configuration
//!    - Creates backend via BackendFactory
//!    - Spawns backend worker threads if needed
//!
//! 2. Submit: ring.submit(op) -> Result<Request<Output>>
//!    - Assigns unique request id
//!    - Converts operation to descriptor
//!    - Sends to backend
//!
//! 3. Complete: request.wait(timeout) -> Result<Output>
//!    - Waits for backend completion
//!    - Maps completion payload to typed output
//!
//! 4. Shutdown: (implicit on Drop)
//!    - Signals backend to shut down
//!    - Waits for in-flight operations
//!    - Joins worker threads
//! ```
//!
//! # Fork Safety
//!
//! Rings are **not fork-safe**. The ring records the process ID at creation
//! and will return an error if used from a forked child process. Create a new
//! ring after forking.

use std::borrow::Cow;
use std::collections::VecDeque;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::mpsc::{self, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use crate::backend::{
    AnyBackend, Backend, BackendCompletion, BackendFactory, BackendKind, BackendSubmission,
    CancellationHandle,
};
use crate::completion::{CompletionEvent, CompletionKind, Request};
use crate::config::RingConfig;
use crate::error::{Error, Result};
use crate::op::{boxed_mapper, BoxedMapper, Op};

pub(crate) struct InflightEntry {
    id: u64,
    op_name: &'static str,
    mapper: BoxedMapper,
    responder: Sender<Result<Box<dyn std::any::Any + Send>>>,
}

struct InflightTable {
    slots: Vec<Option<InflightEntry>>,
    generations: Vec<u64>,
}

impl InflightTable {
    fn new(capacity: usize) -> Self {
        Self {
            slots: (0..capacity).map(|_| None).collect(),
            generations: vec![0; capacity],
        }
    }

    fn insert(
        &mut self,
        start_slot: usize,
        op_name: &'static str,
        mapper: BoxedMapper,
        responder: Sender<Result<Box<dyn std::any::Any + Send>>>,
    ) -> Result<u64> {
        let capacity = self.slots.len();
        for offset in 0..capacity {
            let slot = (start_slot + offset) % capacity;
            if self.slots[slot].is_none() {
                let id = self.generations[slot]
                    .saturating_mul(capacity as u64)
                    .saturating_add(slot as u64)
                    .saturating_add(1);
                self.slots[slot] = Some(InflightEntry {
                    id,
                    op_name,
                    mapper,
                    responder,
                });
                return Ok(id);
            }
        }
        Err(Error::submission(
            "ring submission failed: no inflight slot available",
            "drain completions before submitting more work",
        ))
    }

    fn remove(&mut self, id: u64) -> Option<InflightEntry> {
        let slot = inflight_slot(id, self.slots.len());
        let matches = self
            .slots
            .get(slot)
            .and_then(Option::as_ref)
            .is_some_and(|entry| entry.id == id);
        if matches {
            self.generations[slot] = self.generations[slot].saturating_add(1);
            self.slots[slot].take()
        } else {
            None
        }
    }
}

#[allow(clippy::cast_possible_truncation)]
fn inflight_slot(id: u64, capacity: usize) -> usize {
    // id starts at 1 (0 is reserved/invalid). Saturating sub prevents underflow
    // if id is ever erroneously 0  -  maps to slot 0 instead of wrapping to MAX.
    (id.saturating_sub(1) % capacity as u64) as usize
}

/// Shared ring state used by request handles.
pub struct RingInner {
    pub(crate) config: RingConfig,
    completion_rx: Mutex<Receiver<BackendCompletion>>,
    inflight: Mutex<InflightTable>,
    inflight_count: AtomicUsize,
    backend: AnyBackend,
    next_slot: AtomicUsize,
    pid: u32,
    closed: AtomicBool,
}

impl RingInner {
    pub(crate) fn pump_completion(&self, timeout: Option<Duration>) -> Result<CompletionEvent> {
        self.check_fork()?;
        let completion = {
            let receiver = self.completion_rx.lock().map_err(|_| {
                Error::completion(
                    "completion receiver mutex was poisoned",
                    "avoid panicking while waiting for ring completions",
                )
            })?;
            if let Some(timeout) = timeout {
                match receiver.recv_timeout(timeout) {
                    Ok(value) => value,
                    Err(mpsc::RecvTimeoutError::Timeout) => {
                        return Err(Error::timeout(
                            "completion wait timed out",
                            "increase the wait timeout or submit less work concurrently",
                        ));
                    }
                    Err(mpsc::RecvTimeoutError::Disconnected) => {
                        return Err(Error::completion(
                            "backend completion channel disconnected",
                            "keep the ring alive until the backend thread is done",
                        ));
                    }
                }
            } else {
                receiver.recv().map_err(|_| {
                    Error::completion(
                        "backend completion channel disconnected",
                        "keep the ring alive until the backend thread is done",
                    )
                })?
            }
        };
        self.route_completion(completion)
    }

    fn route_completion(&self, completion: BackendCompletion) -> Result<CompletionEvent> {
        let entry = {
            let mut table = self.inflight.lock().map_err(|_| {
                Error::completion(
                    "inflight table mutex was poisoned",
                    "avoid panicking while routing operation completions",
                )
            })?;
            let entry = table.remove(completion.id).ok_or_else(|| {
                Error::completion(
                    "completion for unknown request id",
                    "avoid completing the same request twice and keep request ids unique",
                )
            })?;
            self.inflight_count.fetch_sub(1, Ordering::Release);
            entry
        };
        let kind = match &completion.result {
            Ok(_) => CompletionKind::Completed,
            Err(Error::Canceled { .. }) => CompletionKind::Canceled,
            Err(_) => CompletionKind::Failed,
        };
        let result = completion.result.and_then(entry.mapper);
        let send_result = entry.responder.send(result).map_err(|_| {
            Error::completion(
                "request receiver dropped before completion",
                "keep the request handle alive until the result is consumed",
            )
        });
        if let Err(error) = send_result {
            tracing::warn!("{error}");
        }
        Ok(CompletionEvent::new(completion.id, entry.op_name, kind))
    }

    fn check_fork(&self) -> Result<()> {
        let pid = std::process::id();
        if pid != self.pid {
            return Err(Error::ProcessState {
                message: Cow::Owned(format!(
                    "ring was created in process {} but used from forked process {}",
                    self.pid, pid
                )),
                fix: Cow::Borrowed(
                    "create a new ring after fork; io_uring and fallback worker state are not fork-safe",
                ),
            });
        }
        Ok(())
    }
}

impl Drop for RingInner {
    fn drop(&mut self) {
        if !self.closed.swap(true, Ordering::AcqRel) {
            if let Err(error) = self.backend.shutdown() {
                tracing::warn!("{error}");
            }
        }
    }
}

/// Submission and completion router.
#[derive(Clone)]
pub struct Ring<F: BackendFactory> {
    inner: Arc<RingInner>,
    factory: PhantomData<F>,
}

/// Sequential chain execution context for dependent operations.
pub struct ChainContext<'a, F: BackendFactory> {
    ring: &'a Ring<F>,
}

impl<F: BackendFactory> ChainContext<'_, F> {
    /// Submits one operation and waits for its typed result.
    pub fn run<O: Op>(&self, op: O) -> Result<O::Output> {
        self.ring
            .submit(op)?
            .wait(Some(std::time::Duration::from_secs(30)))
    }

    /// Runs a bounded batch within the surrounding chain.
    pub fn batch<I, O>(&self, ops: I) -> Result<Vec<O::Output>>
    where
        I: IntoIterator<Item = O>,
        O: Op,
    {
        self.ring.batch(ops)
    }
}

impl<F: BackendFactory> Ring<F> {
    /// Creates a ring using the selected backend policy.
    pub fn new(config: RingConfig) -> Result<Self> {
        config.validate()?;
        let (completion_tx, completion_rx) = mpsc::channel();
        let backend = F::default().create_enum(&config, completion_tx)?;
        Ok(Self {
            inner: Arc::new(RingInner {
                completion_rx: Mutex::new(completion_rx),
                inflight: Mutex::new(InflightTable::new(config.queue_depth as usize)),
                config,
                inflight_count: AtomicUsize::new(0),
                backend,
                next_slot: AtomicUsize::new(0),
                pid: std::process::id(),
                closed: AtomicBool::new(false),
            }),
            factory: PhantomData,
        })
    }

    /// Returns the active backend kind.
    #[must_use]
    pub fn backend_kind(&self) -> BackendKind {
        self.inner.backend.kind()
    }

    /// Returns the configured queue depth.
    #[must_use]
    pub fn queue_depth(&self) -> u32 {
        self.inner.config.queue_depth
    }

    /// Submits a batch of operations and returns ordered results.
    ///
    /// The ring keeps at most `queue_depth` requests in flight at once, so
    /// callers can batch arbitrarily large iterators without manually chunking.
    pub fn batch<I, O>(&self, ops: I) -> Result<Vec<O::Output>>
    where
        I: IntoIterator<Item = O>,
        O: Op,
    {
        let mut results = Vec::new();
        let mut inflight: VecDeque<Request<O::Output>> = VecDeque::new();
        let max_inflight = self.inner.config.queue_depth as usize;

        for op in ops {
            if inflight.len() == max_inflight {
                let request = inflight.pop_front().ok_or_else(|| {
                    Error::completion(
                        "batch bookkeeping lost an in-flight request",
                        "keep the batch queue intact while draining completed operations",
                    )
                })?;
                results.push(request.wait(Some(std::time::Duration::from_secs(2)))?);
            }
            inflight.push_back(self.submit(op)?);
        }

        while let Some(request) = inflight.pop_front() {
            results.push(request.wait(Some(std::time::Duration::from_secs(2)))?);
        }

        Ok(results)
    }

    /// Runs a dependent chain of operations in one call.
    ///
    /// This is the high-level API for flows where later operations require
    /// outputs from earlier ones, such as `connect -> send -> recv`.
    pub fn chain<T>(&self, run: impl FnOnce(ChainContext<'_, F>) -> Result<T>) -> Result<T> {
        run(ChainContext { ring: self })
    }

    /// Submits an operation and returns a typed request handle.
    pub fn submit<O: Op>(&self, op: O) -> Result<Request<O::Output>> {
        self.inner.check_fork()?;
        let op_name = op.name();
        let descriptor = op.into_descriptor()?;
        let (sender, receiver) = mpsc::channel();
        let id = {
            let mut table = self.inner.inflight.lock().map_err(|_| {
                Error::submission(
                    "inflight table mutex was poisoned",
                    "avoid panicking while submitting operations",
                )
            })?;
            let id = table.insert(
                self.inner.next_slot.fetch_add(1, Ordering::Relaxed)
                    % self.inner.config.queue_depth as usize,
                op_name,
                boxed_mapper::<O>(),
                sender,
                )
                .map_err(|_| {
                    Error::submission(
                        format!(
                            "ring submission failed: SQ full (capacity: {}, in-flight: {})",
                            self.inner.config.queue_depth,
                            self.inner.inflight_count.load(Ordering::Acquire)
                        ),
                        "increase submission_queue_depth in RingConfig or drain completions before submitting more",
                    )
                })?;
            self.inner.inflight_count.fetch_add(1, Ordering::Release);
            id
        };
        let submission =
            BackendSubmission::new(id, descriptor).with_timeout(self.inner.config.default_timeout);
        if let Err(error) = self.inner.backend.submit(submission) {
            let mut table = self
                .inner
                .inflight
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            table.remove(id);
            self.inner.inflight_count.fetch_sub(1, Ordering::Release);
            return Err(error);
        }
        Ok(Request {
            id,
            ring: self.inner.clone(),
            receiver,
            marker: std::marker::PhantomData,
        })
    }

    /// Blocks until one completion is routed.
    pub fn complete(&self, timeout: Option<Duration>) -> Result<CompletionEvent> {
        self.inner.pump_completion(timeout)
    }

    /// Attempts to cancel a submitted operation.
    pub fn cancel(&self, id: u64) -> Result<()> {
        self.inner.check_fork()?;
        self.inner.backend.cancel(CancellationHandle { target: id })
    }

    /// Returns the number of tracked in-flight operations.
    pub fn in_flight(&self) -> Result<usize> {
        Ok(self.inner.inflight_count.load(Ordering::Acquire))
    }
}

impl<F: BackendFactory> Drop for Ring<F> {
    fn drop(&mut self) {
        if Arc::strong_count(&self.inner) == 1 && !self.inner.closed.swap(true, Ordering::Relaxed) {
            if let Err(error) = self.inner.backend.shutdown() {
                tracing::error!(%error, "wireshift backend shutdown failed during Ring drop");
            }
        }
    }
}