Skip to main content

haematite_cli/serve/query/
probe.rs

1use std::collections::BTreeMap;
2use std::fmt;
3use std::num::NonZeroU32;
4
5use haematite_query::handle::{Error, StartupStage};
6use haematite_query::{Ready, Refusal};
7
8use super::{Settings, start};
9use crate::serve::storage::Failure;
10
11/// What the service knows about its query worker after the probe at start.
12///
13/// The settings are kept whether or not the probe passed: a service whose
14/// worker cannot start still serves every other verb, and the refusal kept
15/// here names why a query cannot run.
16#[derive(Debug)]
17pub struct Prepared {
18    settings: Settings,
19    outcome: Result<Ready, Failure>,
20}
21
22impl Prepared {
23    #[must_use]
24    pub const fn settings(&self) -> &Settings {
25        &self.settings
26    }
27    /// What the worker stated at the probe, or the named refusal.
28    ///
29    /// # Errors
30    /// Returns the named failure retained by the startup probe.
31    pub const fn ready(&self) -> Result<&Ready, &Failure> {
32        self.outcome.as_ref()
33    }
34}
35
36#[derive(Debug)]
37pub struct ProbeError {
38    pub failure: Failure,
39    pub ending: Option<Box<Error>>,
40}
41
42impl fmt::Display for ProbeError {
43    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
44        fmt::Display::fmt(&self.failure, formatter)
45    }
46}
47impl std::error::Error for ProbeError {
48    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
49        self.ending
50            .as_deref()
51            .map(|error| error as &dyn std::error::Error)
52    }
53}
54
55pub(super) fn classified_name(error: &Error) -> &'static str {
56    match error.refusal.code {
57        Refusal::QueryTimeLimit => return "query_worker_start_timeout",
58        Refusal::QueryMemoryLimit => return "query_worker_start_memory_limit",
59        _ => {}
60    }
61    let Some(startup) = error.startup else {
62        return "query_worker_provenance_missing";
63    };
64    match startup.stage {
65        StartupStage::Spawn => "query_worker_not_executable",
66        StartupStage::Identity => "query_worker_identity_mismatch",
67        StartupStage::Stopping => "query_worker_stop_failed",
68        StartupStage::UnexpectedFrame
69        | StartupStage::Failed(_)
70        | StartupStage::UnofferedReading => "query_worker_not_ready",
71        StartupStage::Owner => "query_worker_owner_failed",
72        // During the exchange itself: a frame that is not the protocol, a
73        // worker that exited (its status was waited for), or neither.
74        StartupStage::SendStart | StartupStage::ReadReady => {
75            if error.refusal.code == Refusal::QueryProtocol {
76                "query_worker_not_ready"
77            } else if error.status.is_some() {
78                "query_worker_exited_before_ready"
79            } else {
80                "query_worker_owner_failed"
81            }
82        }
83    }
84}
85
86fn classify(error: Error, settings: &Settings) -> ProbeError {
87    let code = classified_name(&error);
88    let mut message = format!("worker {}: {error}", settings.worker.display());
89    if let Some(startup) = error.startup
90        && let StartupStage::Failed(code) = startup.stage
91    {
92        match serde_json::to_string(&code) {
93            Ok(code) => message = format!("{message}; Failed code {code}"),
94            Err(cause) => {
95                return ProbeError {
96                    failure: Failure::new(
97                        "query_worker_diagnostic_failed",
98                        format!("cannot encode worker Failed code: {cause}; {message}"),
99                    ),
100                    ending: Some(Box::new(error)),
101                };
102            }
103        }
104    }
105    if error.refusal.code == Refusal::QueryTimeLimit {
106        message = format!(
107            "{message}; startup allowance is query_time_ms={} ms",
108            settings.time_ms
109        );
110    }
111    if let Some(cause) = &error.stderr_error {
112        message = format!("{message}; stderr reader: {cause}");
113    }
114    ProbeError {
115        failure: Failure::new(code, message),
116        ending: Some(Box::new(error)),
117    }
118}
119
120/// Start, verify and stop the configured query worker once.
121///
122/// # Errors
123/// Refuses invalid or inaccessible worker paths, an expired startup allowance, a failed startup exchange or a failed stop, preserving worker provenance.
124pub fn run(settings: &Settings, frame_bytes: NonZeroU32) -> Result<Ready, ProbeError> {
125    let boundary = |failure| ProbeError {
126        failure,
127        ending: None,
128    };
129    let deadline = start::deadline(settings).map_err(boundary)?;
130    start::check(settings, deadline).map_err(boundary)?;
131    let request = start::request(settings, frame_bytes, ".".to_owned(), None, BTreeMap::new());
132    let (mut worker, ready) =
133        start::connect(settings, &request, deadline).map_err(|error| classify(error, settings))?;
134    let pid = worker.pid();
135    worker
136        .stop()
137        .map_err(|error| classify(error.at_start(StartupStage::Stopping, Some(pid)), settings))?;
138    Ok(ready)
139}
140
141/// Probe the configured worker once. A failed probe is kept, not returned:
142/// the caller starts the service either way.
143#[must_use]
144pub fn prepare(settings: Settings, frame_bytes: NonZeroU32) -> Prepared {
145    let outcome = run(&settings, frame_bytes).map_err(|error| error.failure);
146    Prepared { settings, outcome }
147}
148
149#[cfg(all(test, unix))]
150#[path = "probe_tests.rs"]
151mod tests;