haematite-cli 0.21.7

use std::collections::BTreeMap;
use std::fmt;
use std::num::{NonZeroU32, NonZeroUsize};
use std::sync::atomic::{AtomicUsize, Ordering};

use haematite_query::Ready;
use haematite_query::handle::{Error, Handle};
use serde_json::value::RawValue;

use super::{Settings, start};
use crate::serve::storage::Failure;

pub struct Slots {
    running: AtomicUsize,
    limit: NonZeroUsize,
}

impl Slots {
    #[must_use]
    pub const fn new(limit: NonZeroUsize) -> Self {
        Self {
            running: AtomicUsize::new(0),
            limit,
        }
    }

    /// Reserve one query process slot.
    ///
    /// # Errors
    /// Returns `query_busy` when the configured process limit is already reached.
    pub fn try_acquire(&self) -> Result<Slot<'_>, Failure> {
        let mut running = self.running.load(Ordering::Acquire);
        loop {
            if running >= self.limit.get() {
                return Err(Failure::new(
                    "query_busy",
                    format!(
                        "query_max_running={} queries are already running",
                        self.limit
                    ),
                ));
            }
            match self.running.compare_exchange_weak(
                running,
                running + 1,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => return Ok(Slot { slots: self }),
                Err(observed) => running = observed,
            }
        }
    }
}

#[must_use]
pub struct Slot<'a> {
    slots: &'a Slots,
}

impl Drop for Slot<'_> {
    fn drop(&mut self) {
        self.slots.running.fetch_sub(1, Ordering::Release);
    }
}

#[derive(Debug)]
pub enum StartError {
    Admission(Failure),
    Worker(Box<Error>),
}

impl fmt::Display for StartError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Admission(error) => fmt::Display::fmt(error, formatter),
            Self::Worker(error) => fmt::Display::fmt(error, formatter),
        }
    }
}

impl std::error::Error for StartError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        Some(match self {
            Self::Admission(error) => error,
            Self::Worker(error) => error.as_ref(),
        })
    }
}

/// The process is dropped and reaped before its admission slot is released.
#[must_use]
pub struct Running<'a> {
    worker: Option<Handle>,
    ready: Ready,
    slot: Option<Slot<'a>>,
}

impl<'a> Running<'a> {
    /// Reserve a slot and start one query worker.
    ///
    /// # Errors
    /// Returns an admission refusal when settings are absent or slots are full, or the worker's startup failure.
    pub fn start(
        slots: &'a Slots,
        settings: Option<&Settings>,
        frame_bytes: NonZeroU32,
        filter: String,
        alternate: Option<String>,
        parameters: BTreeMap<String, Box<RawValue>>,
    ) -> Result<Self, StartError> {
        let settings = settings.ok_or_else(|| {
            StartError::Admission(Failure::new(
                "query_engine_unavailable",
                "query settings are not configured",
            ))
        })?;
        let slot = slots.try_acquire().map_err(StartError::Admission)?;
        let (worker, ready) = start::open(settings, frame_bytes, filter, alternate, parameters)
            .map_err(|error| StartError::Worker(Box::new(error)))?;
        Ok(Self {
            worker: Some(worker),
            ready,
            slot: Some(slot),
        })
    }

    #[must_use]
    pub const fn ready(&self) -> &Ready {
        &self.ready
    }

    /// Borrow the worker while this running query still owns it.
    ///
    /// # Errors
    /// Returns `query_stopped` after the worker has been released.
    pub fn worker(&mut self) -> Result<&mut Handle, Error> {
        self.worker.as_mut().ok_or_else(|| {
            Error::new(
                haematite_query::Refusal::QueryStopped,
                "query worker already released",
            )
        })
    }

    /// Stop and reap the worker before releasing its slot.
    ///
    /// # Errors
    /// Propagates a missing worker or the worker's stop failure.
    pub fn stop(mut self) -> Result<(), Error> {
        self.worker()?.stop()
    }
}

impl Drop for Running<'_> {
    fn drop(&mut self) {
        drop(self.worker.take());
        drop(self.slot.take());
    }
}

#[cfg(test)]
#[path = "admission_tests.rs"]
mod tests;