Add healthcheck for db

This commit is contained in:
asonix 2023-01-29 11:36:09 -06:00
parent 5e5dd27a05
commit f0f40db8c3
3 changed files with 34 additions and 3 deletions

View file

@ -985,6 +985,11 @@ async fn identifier<R: FullRepo, S: Store>(
})))
}
async fn healthz<R: FullRepo>(repo: web::Data<R>) -> Result<HttpResponse, Error> {
repo.health_check().await?;
Ok(HttpResponse::Ok().finish())
}
fn transform_error(error: actix_form_data::Error) -> actix_web::Error {
let error: Error = error.into();
let error: actix_web::Error = error.into();
@ -1039,6 +1044,7 @@ async fn launch<R: FullRepo + 'static, SC: StoreConfig + 'static>(
.app_data(web::Data::new(repo))
.app_data(web::Data::new(store))
.app_data(web::Data::new(build_client()))
.route("/healthz", web::get().to(healthz::<R>))
.service(
web::scope("/image")
.service(

View file

@ -60,6 +60,8 @@ pub(crate) trait FullRepo:
+ Clone
+ Debug
{
async fn health_check(&self) -> Result<(), Error>;
#[tracing::instrument(skip(self))]
async fn identifier_from_alias<I: Identifier + 'static>(
&self,
@ -103,7 +105,12 @@ pub(crate) trait FullRepo:
}
}
impl<T> FullRepo for actix_web::web::Data<T> where T: FullRepo {}
#[async_trait::async_trait(?Send)]
impl<T> FullRepo for actix_web::web::Data<T> where T: FullRepo {
async fn health_check(&self) -> Result<(), Error> {
T::health_check(self).await
}
}
pub(crate) trait BaseRepo {
type Bytes: AsRef<[u8]> + From<Vec<u8>> + Clone;

View file

@ -13,7 +13,10 @@ use sled::{CompareAndSwapError, Db, IVec, Tree};
use std::{
collections::HashMap,
pin::Pin,
sync::{Arc, RwLock},
sync::{
atomic::{AtomicU64, Ordering},
Arc, RwLock,
},
};
use tokio::sync::Notify;
@ -52,6 +55,8 @@ pub(crate) enum SledError {
#[derive(Clone)]
pub(crate) struct SledRepo {
healthz_count: Arc<AtomicU64>,
healthz: Tree,
settings: Tree,
identifier_details: Tree,
hashes: Tree,
@ -74,6 +79,8 @@ pub(crate) struct SledRepo {
impl SledRepo {
pub(crate) fn new(db: Db) -> Result<Self, SledError> {
Ok(SledRepo {
healthz_count: Arc::new(AtomicU64::new(0)),
healthz: db.open_tree("pict-rs-healthz-tree")?,
settings: db.open_tree("pict-rs-settings-tree")?,
identifier_details: db.open_tree("pict-rs-identifier-details-tree")?,
hashes: db.open_tree("pict-rs-hashes-tree")?,
@ -99,7 +106,18 @@ impl BaseRepo for SledRepo {
type Bytes = IVec;
}
impl FullRepo for SledRepo {}
#[async_trait::async_trait(?Send)]
impl FullRepo for SledRepo {
async fn health_check(&self) -> Result<(), Error> {
let next = self.healthz_count.fetch_add(1, Ordering::Relaxed);
b!(self.healthz, {
healthz.insert("healthz", &next.to_be_bytes()[..])
});
self.healthz.flush_async().await?;
b!(self.healthz, healthz.get("healthz"));
Ok(())
}
}
#[derive(serde::Deserialize, serde::Serialize)]
enum InnerUploadResult {