use std::{collections::HashMap, path::PathBuf, sync::Arc, time::Duration};
use bestool_canopy::schema::CheckSeverity;
use futures::{StreamExt, future::BoxFuture, stream::BoxStream};
use jiff::Timestamp;
use miette::{Result, miette};
use serde_json::{Value, json};
use tokio::sync::{Mutex, mpsc};
use tracing::warn;
use crate::doctor::{
self,
check::{Check, CheckStatus},
progress::DoctorEvent,
stat::{MetricsSnapshot, StatusCounts},
};
use crate::tasks::TaskEndpointHandler;
use crate::{BackgroundTask, TaskContext, TaskEndpoint, TaskEndpointResponse};
const DOCTOR_INTERVAL: Duration = Duration::from_secs(60);
pub type BackupDispatch = Arc<dyn Fn(Vec<String>) + Send + Sync>;
fn cap_check(check: Check, severities: Option<&HashMap<String, CheckSeverity>>) -> Check {
match severities {
Some(map) => {
let ceiling = doctor::sweep::severity_ceiling(map, check.name);
Check {
status: check.status.cap_to(ceiling),
..check
}
}
None => check,
}
}
pub struct DoctorTask {
inner: Arc<DoctorTaskInner>,
}
enum TamanuSource {
Fixed,
Discover { root: Option<PathBuf> },
}
struct DoctorTaskInner {
binary_version: String,
tamanu: Mutex<Option<doctor::SweepTamanu>>,
tamanu_source: TamanuSource,
pg_version_cache: Mutex<Option<String>>,
latest: Mutex<Option<LatestSweep>>,
check_severities: Mutex<Option<HashMap<String, CheckSeverity>>>,
backup_dispatch: Option<BackupDispatch>,
}
#[derive(Clone)]
struct LatestSweep {
computed_at: Timestamp,
sweep: doctor::SweepResult,
}
impl DoctorTask {
pub fn new(binary_version: String, tamanu: Option<doctor::SweepTamanu>) -> Self {
Self {
inner: Arc::new(DoctorTaskInner {
binary_version,
tamanu: Mutex::new(tamanu),
tamanu_source: TamanuSource::Fixed,
pg_version_cache: Mutex::new(None),
latest: Mutex::new(None),
check_severities: Mutex::new(None),
backup_dispatch: None,
}),
}
}
pub fn with_tamanu_discovery(self, root: Option<PathBuf>) -> Self {
let mut inner =
Arc::try_unwrap(self.inner).unwrap_or_else(|_| panic!("DoctorTask already shared"));
inner.tamanu_source = TamanuSource::Discover { root };
Self {
inner: Arc::new(inner),
}
}
pub fn with_backup_dispatch(self, dispatch: BackupDispatch) -> Self {
let mut inner =
Arc::try_unwrap(self.inner).unwrap_or_else(|_| panic!("DoctorTask already shared"));
inner.backup_dispatch = Some(dispatch);
Self {
inner: Arc::new(inner),
}
}
pub fn metrics_handle(&self) -> DoctorMetricsHandle {
DoctorMetricsHandle {
inner: self.inner.clone(),
}
}
}
#[derive(Clone)]
pub struct DoctorMetricsHandle {
inner: Arc<DoctorTaskInner>,
}
impl DoctorMetricsHandle {
pub async fn snapshot(&self) -> Option<MetricsSnapshot> {
let latest = self.inner.latest.lock().await.clone()?;
let sweep = self.inner.capped(latest.sweep).await;
let counts = census(&sweep.results);
let stats = sweep
.results
.iter()
.flat_map(|(check, _)| check.stats.iter().map(|stat| (check.name, stat.clone())))
.collect();
Some(MetricsSnapshot {
computed_at: latest.computed_at,
stats,
counts,
})
}
}
fn census(results: &[(Check, bool)]) -> StatusCounts {
let mut counts = StatusCounts::default();
for (check, _) in results {
match &check.status {
CheckStatus::Pass => counts.passing += 1,
CheckStatus::Warning(_) => counts.warning += 1,
CheckStatus::Fail(_) => counts.failing += 1,
CheckStatus::Skip(_) => counts.skipped += 1,
CheckStatus::Broken(_) => counts.broken += 1,
}
}
counts
}
impl DoctorTaskInner {
async fn resolve_tamanu(&self) -> Option<doctor::SweepTamanu> {
let TamanuSource::Discover { root } = &self.tamanu_source else {
return self.tamanu.lock().await.clone();
};
self.apply_discovery(doctor::discover_sweep_tamanu(root.as_deref()).await)
.await
}
async fn apply_discovery(
&self,
discovered: Result<Option<doctor::SweepTamanu>>,
) -> Option<doctor::SweepTamanu> {
let mut guard = self.tamanu.lock().await;
match discovered {
Ok(resolved) => *guard = resolved,
Err(err) => warn!(
%err,
"could not resolve the Tamanu install; sweeping against the last known context"
),
}
guard.clone()
}
async fn run_sweep(
self: &Arc<Self>,
ctx: &TaskContext,
progress: Option<doctor::progress::ProgressSender>,
enable_heal: bool,
) -> Result<doctor::SweepResult> {
let cached = self.pg_version_cache.lock().await.clone();
let tamanu = self.resolve_tamanu().await;
let sweep = doctor::perform_sweep(
&self.binary_version,
tamanu,
ctx.http_client.clone(),
&[],
&[],
cached,
progress,
ctx.canopy_client.clone(),
enable_heal,
)
.await?;
if let Some(ref version) = sweep.pg_version {
let mut guard = self.pg_version_cache.lock().await;
if guard.is_none() {
*guard = Some(version.clone());
}
}
let latest = LatestSweep {
computed_at: Timestamp::now(),
sweep: sweep.clone(),
};
*self.latest.lock().await = Some(latest);
Ok(sweep)
}
async fn severities_snapshot(&self) -> Option<HashMap<String, CheckSeverity>> {
self.check_severities.lock().await.clone()
}
async fn capped(&self, mut sweep: doctor::SweepResult) -> doctor::SweepResult {
if let Some(severities) = self.severities_snapshot().await {
sweep.apply_severities(&severities);
}
sweep
}
async fn tick(self: &Arc<Self>, ctx: &TaskContext) -> Result<()> {
let sweep = self.run_sweep(ctx, None, true).await?;
let Some(server_id) = sweep.server_id else {
warn!("no metaServerId available; skipping canopy status push");
return Ok(());
};
let Some(canopy) = ctx.canopy_client.as_ref() else {
warn!("no canopy client available; skipping canopy status push");
return Ok(());
};
let response = canopy
.status(&server_id, &sweep.payload)
.await
.map_err(|err| miette!("posting doctor status to canopy: {err}"))?;
*self.check_severities.lock().await = Some(response.check_severities);
let tags = response.tags.0.into_iter().collect();
if let Err(err) = bestool_tamanu::server_info::save_cached_tags(&tags) {
warn!(%err, "could not refresh tags cache from status response");
}
let backup_now = response.backup_now;
if !backup_now.is_empty() {
match &self.backup_dispatch {
Some(dispatch) => dispatch(backup_now),
None => warn!(
?backup_now,
"canopy requested a backup but no backup dispatcher is configured"
),
}
}
Ok(())
}
async fn endpoint_latest(self: Arc<Self>) -> TaskEndpointResponse {
let snapshot = self.latest.lock().await.clone();
match snapshot {
Some(s) => {
let sweep = self.capped(s.sweep).await;
TaskEndpointResponse::Json(json!({
"computedAt": s.computed_at.to_string(),
"serverId": sweep.server_id,
"payload": sweep.payload,
}))
}
None => TaskEndpointResponse::Error {
status: 503,
message: "no doctor sweep cached yet (daemon may have just started)".into(),
},
}
}
async fn endpoint_recompute(self: Arc<Self>, ctx: TaskContext) -> TaskEndpointResponse {
let (progress_tx, mut progress_rx) = mpsc::unbounded_channel::<DoctorEvent>();
let (out_tx, out_rx) = mpsc::unbounded_channel::<Value>();
let severities = self.severities_snapshot().await;
let task_self = self.clone();
tokio::spawn(async move {
let progress_forward_tx = out_tx.clone();
let stream_severities = severities.clone();
let forwarder = tokio::spawn(async move {
while let Some(event) = progress_rx.recv().await {
let DoctorEvent::Completed(check) = event;
let check = cap_check(check, stream_severities.as_ref());
let _ = progress_forward_tx.send(json!({
"event": "check",
"check": check.to_streaming_json(),
}));
}
});
match task_self.run_sweep(&ctx, Some(progress_tx), false).await {
Ok(mut sweep) => {
if let Some(severities) = &severities {
sweep.apply_severities(severities);
}
let _ = forwarder.await;
let _ = out_tx.send(json!({
"event": "done",
"computedAt": Timestamp::now().to_string(),
"serverId": sweep.server_id,
"payload": sweep.payload,
}));
}
Err(err) => {
let _ = forwarder.await;
let _ = out_tx.send(json!({
"event": "error",
"message": format!("{err:?}"),
}));
}
}
});
let stream: BoxStream<'static, Value> =
Box::pin(tokio_stream::wrappers::UnboundedReceiverStream::new(out_rx).map(|v| v));
TaskEndpointResponse::JsonLines(stream)
}
}
impl BackgroundTask for DoctorTask {
fn name(&self) -> &'static str {
"doctor"
}
fn interval(&self) -> Duration {
DOCTOR_INTERVAL
}
fn run<'a>(&'a self, ctx: &'a TaskContext) -> BoxFuture<'a, Result<()>> {
let inner = self.inner.clone();
Box::pin(async move { inner.tick(ctx).await })
}
fn http_endpoints(&self) -> Vec<TaskEndpoint> {
let latest_handler: TaskEndpointHandler = {
let inner = self.inner.clone();
Arc::new(move |_ctx| {
let inner = inner.clone();
Box::pin(async move { inner.endpoint_latest().await })
})
};
let recompute_handler: TaskEndpointHandler = {
let inner = self.inner.clone();
Arc::new(move |ctx| {
let inner = inner.clone();
Box::pin(async move { inner.endpoint_recompute(ctx).await })
})
};
vec![
TaskEndpoint {
name: "latest",
handler: latest_handler,
},
TaskEndpoint {
name: "recompute",
handler: recompute_handler,
},
]
}
}
#[cfg(test)]
mod tests {
use node_semver::Version;
use bestool_tamanu::config::{Database, TamanuConfig};
use super::*;
use crate::doctor::check::CheckStatus;
const DB_URL: &str = "postgres://u:p@localhost/tamanu";
fn sweep_tamanu(version: &str) -> doctor::SweepTamanu {
doctor::SweepTamanu {
version: Version::parse(version).unwrap(),
root: PathBuf::from("/opt/tamanu"),
config: Arc::new(TamanuConfig::from_database(
Database::from_url(DB_URL).unwrap(),
)),
database_url: DB_URL.into(),
has_install: true,
is_tamanu: true,
}
}
fn inner(tamanu: Option<doctor::SweepTamanu>, tamanu_source: TamanuSource) -> DoctorTaskInner {
DoctorTaskInner {
binary_version: "0.0.0-test".into(),
tamanu: Mutex::new(tamanu),
tamanu_source,
pg_version_cache: Mutex::new(None),
latest: Mutex::new(None),
check_severities: Mutex::new(None),
backup_dispatch: None,
}
}
#[tokio::test]
async fn discovery_replaces_the_previous_tamanu_context() {
let inner = inner(Some(sweep_tamanu("2.54.0")), TamanuSource::Fixed);
let resolved = inner
.apply_discovery(Ok(Some(sweep_tamanu("2.55.0"))))
.await
.expect("a context");
assert_eq!(resolved.version, Version::parse("2.55.0").unwrap());
assert_eq!(
inner.tamanu.lock().await.as_ref().unwrap().version,
Version::parse("2.55.0").unwrap()
);
}
#[tokio::test]
async fn discovery_failure_keeps_the_last_known_context() {
let inner = inner(Some(sweep_tamanu("2.54.0")), TamanuSource::Fixed);
let resolved = inner
.apply_discovery(Err(miette!("no tamanu discovered")))
.await
.expect("the last known context");
assert_eq!(resolved.version, Version::parse("2.54.0").unwrap());
}
#[tokio::test]
async fn discovery_clears_the_context_when_tamanu_is_gone() {
let inner = inner(Some(sweep_tamanu("2.54.0")), TamanuSource::Fixed);
assert!(inner.apply_discovery(Ok(None)).await.is_none());
assert!(inner.tamanu.lock().await.is_none());
}
#[tokio::test]
async fn fixed_source_reuses_the_context_it_was_given() {
let inner = inner(Some(sweep_tamanu("2.54.0")), TamanuSource::Fixed);
let resolved = inner.resolve_tamanu().await.expect("a context");
assert_eq!(resolved.version, Version::parse("2.54.0").unwrap());
}
#[test]
fn cap_check_applies_ceiling_when_present() {
let mut severities = HashMap::new();
severities.insert("disk_free".to_string(), CheckSeverity::Warn);
let check = Check::fail("disk_free", "1% free", "out of space");
let capped = cap_check(check, Some(&severities));
match capped.status {
CheckStatus::Warning(r) => assert_eq!(r, "out of space"),
other => panic!("expected Warning, got {other:?}"),
}
}
#[test]
fn cap_check_absent_check_defaults_to_warn() {
let check = Check::fail("brand_new", "bad", "reason");
let capped = cap_check(check, Some(&HashMap::new()));
assert!(matches!(capped.status, CheckStatus::Warning(_)));
}
#[test]
fn cap_check_no_mapping_is_a_noop() {
let check = Check::fail("disk_free", "1% free", "out of space");
let capped = cap_check(check, None);
assert!(matches!(capped.status, CheckStatus::Fail(_)));
}
#[test]
fn census_counts_each_status() {
let results = vec![
(Check::pass("a", ""), true),
(Check::pass("b", ""), true),
(Check::warning("c", "", "w"), true),
(Check::fail("d", "", "f"), true),
(Check::skip("e", "", "s"), true),
(Check::broken("g", "", "b"), true),
];
let c = census(&results);
assert_eq!(c.passing, 2);
assert_eq!(c.warning, 1);
assert_eq!(c.failing, 1);
assert_eq!(c.skipped, 1);
assert_eq!(c.broken, 1);
assert_eq!(c.total(), 6);
assert_eq!(c.active(), 5);
}
#[test]
fn census_reflects_severity_capping() {
let mut sweep = doctor::SweepResult {
server_id: None,
results: vec![(Check::fail("disk_free", "1% free", "out of space"), true)],
overall: doctor::check::OverallResult::Failing,
payload: json!({}),
pg_version: None,
};
let mut severities = HashMap::new();
severities.insert("disk_free".to_string(), CheckSeverity::Warn);
sweep.apply_severities(&severities);
let c = census(&sweep.results);
assert_eq!(c.failing, 0);
assert_eq!(c.warning, 1);
}
}