From 7274538b7007cee84207063dd551723473f66eac Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Fri, 7 Jul 2023 12:05:42 -0500
Subject: [PATCH] Try adding store healthcheck

---
 src/lib.rs                |  8 ++++++--
 src/store.rs              | 10 ++++++++++
 src/store/file_store.rs   | 12 ++++++++++++
 src/store/object_store.rs | 21 +++++++++++++++++++++
 4 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index d0413c6..3bd6bfe 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1060,8 +1060,12 @@ async fn identifier<R: FullRepo, S: Store>(
     })))
 }
 
-async fn healthz<R: FullRepo>(repo: web::Data<R>) -> Result<HttpResponse, Error> {
+async fn healthz<R: FullRepo, S: Store>(
+    repo: web::Data<R>,
+    store: web::Data<S>,
+) -> Result<HttpResponse, Error> {
     repo.health_check().await?;
+    store.health_check().await?;
     Ok(HttpResponse::Ok().finish())
 }
 
@@ -1104,7 +1108,7 @@ fn configure_endpoints<R: FullRepo + 'static, S: Store + 'static>(
         .app_data(web::Data::new(repo))
         .app_data(web::Data::new(store))
         .app_data(web::Data::new(client))
-        .route("/healthz", web::get().to(healthz::<R>))
+        .route("/healthz", web::get().to(healthz::<R, S>))
         .service(
             web::scope("/image")
                 .service(
diff --git a/src/store.rs b/src/store.rs
index 550b531..467d105 100644
--- a/src/store.rs
+++ b/src/store.rs
@@ -77,6 +77,8 @@ pub(crate) trait Store: Clone + Debug {
     type Identifier: Identifier + 'static;
     type Stream: Stream<Item = std::io::Result<Bytes>> + Unpin + 'static;
 
+    async fn health_check(&self) -> Result<(), StoreError>;
+
     async fn save_async_read<Reader>(&self, reader: Reader) -> Result<Self::Identifier, StoreError>
     where
         Reader: AsyncRead + Unpin + 'static;
@@ -115,6 +117,10 @@ where
     type Identifier = T::Identifier;
     type Stream = T::Stream;
 
+    async fn health_check(&self) -> Result<(), StoreError> {
+        T::health_check(self).await
+    }
+
     async fn save_async_read<Reader>(&self, reader: Reader) -> Result<Self::Identifier, StoreError>
     where
         Reader: AsyncRead + Unpin + 'static,
@@ -170,6 +176,10 @@ where
     type Identifier = T::Identifier;
     type Stream = T::Stream;
 
+    async fn health_check(&self) -> Result<(), StoreError> {
+        T::health_check(self).await
+    }
+
     async fn save_async_read<Reader>(&self, reader: Reader) -> Result<Self::Identifier, StoreError>
     where
         Reader: AsyncRead + Unpin + 'static,
diff --git a/src/store/file_store.rs b/src/store/file_store.rs
index 90be306..d091477 100644
--- a/src/store/file_store.rs
+++ b/src/store/file_store.rs
@@ -54,6 +54,14 @@ impl Store for FileStore {
     type Identifier = FileId;
     type Stream = Pin<Box<dyn Stream<Item = std::io::Result<Bytes>>>>;
 
+    async fn health_check(&self) -> Result<(), StoreError> {
+        tokio::fs::metadata(&self.root_dir)
+            .await
+            .map_err(FileError::from)?;
+
+        Ok(())
+    }
+
     #[tracing::instrument(skip(reader))]
     async fn save_async_read<Reader>(
         &self,
@@ -157,6 +165,10 @@ impl FileStore {
     pub(crate) async fn build(root_dir: PathBuf, repo: Repo) -> Result<Self, StoreError> {
         let path_gen = init_generator(&repo).await?;
 
+        tokio::fs::create_dir_all(&root_dir)
+            .await
+            .map_err(FileError::from)?;
+
         Ok(FileStore {
             root_dir,
             path_gen,
diff --git a/src/store/object_store.rs b/src/store/object_store.rs
index b849f4a..1b5bdc9 100644
--- a/src/store/object_store.rs
+++ b/src/store/object_store.rs
@@ -169,6 +169,21 @@ impl Store for ObjectStore {
     type Identifier = ObjectId;
     type Stream = Pin<Box<dyn Stream<Item = std::io::Result<Bytes>>>>;
 
+    async fn health_check(&self) -> Result<(), StoreError> {
+        let response = self
+            .head_bucket_request()
+            .await?
+            .send()
+            .await
+            .map_err(ObjectError::from)?;
+
+        if !response.status().is_success() {
+            return Err(status_error(response).await);
+        }
+
+        Ok(())
+    }
+
     async fn save_async_read<Reader>(&self, reader: Reader) -> Result<Self::Identifier, StoreError>
     where
         Reader: AsyncRead + Unpin + 'static,
@@ -438,6 +453,12 @@ impl ObjectStore {
         })
     }
 
+    async fn head_bucket_request(&self) -> Result<ClientRequest, StoreError> {
+        let action = self.bucket.head_bucket(Some(&self.credentials));
+
+        Ok(self.build_request(action))
+    }
+
     async fn put_object_request(&self) -> Result<(ClientRequest, ObjectId), StoreError> {
         let path = self.next_file().await?;