haematite-cli 0.21.8

use std::collections::BTreeMap;
use std::fmt;
use std::num::NonZeroU32;

use haematite_query::handle::{Error, StartupStage};
use haematite_query::{Ready, Refusal};

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

/// What the service knows about its query worker after the probe at start.
///
/// The settings are kept whether or not the probe passed: a service whose
/// worker cannot start still serves every other verb, and the refusal kept
/// here names why a query cannot run.
#[derive(Debug)]
pub struct Prepared {
    settings: Settings,
    outcome: Result<Ready, Failure>,
}

impl Prepared {
    #[must_use]
    pub const fn settings(&self) -> &Settings {
        &self.settings
    }
    /// What the worker stated at the probe, or the named refusal.
    ///
    /// # Errors
    /// Returns the named failure retained by the startup probe.
    pub const fn ready(&self) -> Result<&Ready, &Failure> {
        self.outcome.as_ref()
    }
}

#[derive(Debug)]
pub struct ProbeError {
    pub failure: Failure,
    pub ending: Option<Box<Error>>,
}

impl fmt::Display for ProbeError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(&self.failure, formatter)
    }
}
impl std::error::Error for ProbeError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        self.ending
            .as_deref()
            .map(|error| error as &dyn std::error::Error)
    }
}

pub(super) fn classified_name(error: &Error) -> &'static str {
    match error.refusal.code {
        Refusal::QueryTimeLimit => return "query_worker_start_timeout",
        Refusal::QueryMemoryLimit => return "query_worker_start_memory_limit",
        _ => {}
    }
    let Some(startup) = error.startup else {
        return "query_worker_provenance_missing";
    };
    match startup.stage {
        StartupStage::Spawn => "query_worker_not_executable",
        StartupStage::Identity => "query_worker_identity_mismatch",
        StartupStage::Stopping => "query_worker_stop_failed",
        StartupStage::UnexpectedFrame
        | StartupStage::Failed(_)
        | StartupStage::UnofferedReading => "query_worker_not_ready",
        StartupStage::Owner => "query_worker_owner_failed",
        // During the exchange itself: a frame that is not the protocol, a
        // worker that exited (its status was waited for), or neither.
        StartupStage::SendStart | StartupStage::ReadReady => {
            if error.refusal.code == Refusal::QueryProtocol {
                "query_worker_not_ready"
            } else if error.status.is_some() {
                "query_worker_exited_before_ready"
            } else {
                "query_worker_owner_failed"
            }
        }
    }
}

fn classify(error: Error, settings: &Settings) -> ProbeError {
    let code = classified_name(&error);
    let mut message = format!("worker {}: {error}", settings.worker.display());
    if let Some(startup) = error.startup
        && let StartupStage::Failed(code) = startup.stage
    {
        match serde_json::to_string(&code) {
            Ok(code) => message = format!("{message}; Failed code {code}"),
            Err(cause) => {
                return ProbeError {
                    failure: Failure::new(
                        "query_worker_diagnostic_failed",
                        format!("cannot encode worker Failed code: {cause}; {message}"),
                    ),
                    ending: Some(Box::new(error)),
                };
            }
        }
    }
    if error.refusal.code == Refusal::QueryTimeLimit {
        message = format!(
            "{message}; startup allowance is query_time_ms={} ms",
            settings.time_ms
        );
    }
    if let Some(cause) = &error.stderr_error {
        message = format!("{message}; stderr reader: {cause}");
    }
    ProbeError {
        failure: Failure::new(code, message),
        ending: Some(Box::new(error)),
    }
}

/// Start, verify and stop the configured query worker once.
///
/// # Errors
/// Refuses invalid or inaccessible worker paths, an expired startup allowance, a failed startup exchange or a failed stop, preserving worker provenance.
pub fn run(settings: &Settings, frame_bytes: NonZeroU32) -> Result<Ready, ProbeError> {
    let boundary = |failure| ProbeError {
        failure,
        ending: None,
    };
    let deadline = start::deadline(settings).map_err(boundary)?;
    start::check(settings, deadline).map_err(boundary)?;
    let request = start::request(settings, frame_bytes, ".".to_owned(), None, BTreeMap::new());
    let (mut worker, ready) =
        start::connect(settings, &request, deadline).map_err(|error| classify(error, settings))?;
    let pid = worker.pid();
    worker
        .stop()
        .map_err(|error| classify(error.at_start(StartupStage::Stopping, Some(pid)), settings))?;
    Ok(ready)
}

/// Probe the configured worker once. A failed probe is kept, not returned:
/// the caller starts the service either way.
#[must_use]
pub fn prepare(settings: Settings, frame_bytes: NonZeroU32) -> Prepared {
    let outcome = run(&settings, frame_bytes).map_err(|error| error.failure);
    Prepared { settings, outcome }
}

#[cfg(all(test, unix))]
#[path = "probe_tests.rs"]
mod tests;