From ff242af2ee3cca874c918ffc3e621d36b2118ef6 Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Sun, 29 Jan 2023 13:36:24 -0600
Subject: [PATCH] Add healthcheck for db

---
 src/lib.rs   |  6 ++++++
 src/store.rs | 27 +++++++++++++++++++++++++--
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 8eafef0..634a523 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -285,6 +285,7 @@ pub fn configure(cfg: &mut web::ServiceConfig, state: State, client: Client) {
         client,
     )))
     .app_data(web::Data::new(state))
+    .route("/healthz", web::get().to(healthz))
     .service(web::resource("/static/{filename}").route(web::get().to(static_files)))
     .service(web::resource("/404").route(web::get().to(not_found)))
     .service(
@@ -355,6 +356,11 @@ where
     }
 }
 
+async fn healthz(state: web::Data<State>) -> Result<HttpResponse, StateError> {
+    state.store.check_health().await.stateful(&state)?;
+    Ok(HttpResponse::Ok().finish())
+}
+
 #[tracing::instrument(name = "Static files")]
 async fn static_files(filename: web::Path<String>, state: web::Data<State>) -> HttpResponse {
     let filename = filename.into_inner();
diff --git a/src/store.rs b/src/store.rs
index 41538d3..32569f5 100644
--- a/src/store.rs
+++ b/src/store.rs
@@ -1,12 +1,20 @@
 use crate::{Collection, CollectionPath, Direction, Entry, EntryPath, MoveEntryPath, Token};
 use actix_web::web;
 use sled::{Db, Tree};
-use std::collections::BTreeMap;
-use std::ops::RangeInclusive;
+use std::{
+    collections::BTreeMap,
+    ops::RangeInclusive,
+    sync::{
+        atomic::{AtomicU64, Ordering},
+        Arc,
+    },
+};
 use uuid::Uuid;
 
 #[derive(Clone)]
 pub(crate) struct Store {
+    healthz: Tree,
+    healthz_counter: Arc<AtomicU64>,
     tree: Tree,
 }
 
@@ -104,10 +112,25 @@ impl<'a> DeleteEntry<'a> {
 impl Store {
     pub(crate) fn new(db: &Db) -> Result<Self, sled::Error> {
         Ok(Store {
+            healthz: db.open_tree("healthz")?,
+            healthz_counter: Arc::new(AtomicU64::new(0)),
             tree: db.open_tree("collections")?,
         })
     }
 
+    pub(crate) async fn check_health(&self) -> Result<(), Error> {
+        let next = self.healthz_counter.fetch_add(1, Ordering::Relaxed);
+
+        let healthz = self.healthz.clone();
+        web::block(move || healthz.insert("healthz", &next.to_be_bytes()[..])).await??;
+
+        self.healthz.flush_async().await?;
+
+        let healthz = self.healthz.clone();
+        web::block(move || healthz.get("healthz")).await??;
+        Ok(())
+    }
+
     async fn create_collection(&self, config: CreateCollection<'_>) -> Result<(), Error> {
         let collection_key = config.collection_path.key();
         let collection_value = serde_json::to_string(&config.collection)?;