add token from config and change the signal handler to SIGTERM

3 files changed, 31 insertions, 6 deletions

diff --git a/libs/leash/Cargo.toml b/libs/leash/Cargo.toml
index 5cd54a5..32f385c 100644
--- a/libs/leash/Cargo.toml
+++ b/libs/leash/Cargo.toml
@@ -9,5 +9,5 @@ edition = "2021"
 shared = { path = "../shared" }
 anyhow = "1.0.68"
 tokio = { version = "1.23.0", features = ["full"] }
-
+pretty_env_logger = "0.4"
 serde = "1.0.152"
\ No newline at end of file
diff --git a/libs/leash/src/lib.rs b/libs/leash/src/lib.rs
index 360db12..1de7687 100644
--- a/libs/leash/src/lib.rs
+++ b/libs/leash/src/lib.rs
@@ -1,19 +1,29 @@
 use anyhow::Result;
 use serde::de::DeserializeOwned;
-use shared::config::Settings;
+use shared::{
+    config::Settings,
+    log::{error, info},
+};
 use std::{future::Future, pin::Pin};
+use tokio::{signal::{unix::SignalKind}, sync::oneshot};
 
 pub type AnyhowResultFuture<T> = Pin<Box<dyn Future<Output = Result<T>>>>;
 pub trait Component: Send + Sync + 'static + Sized {
     type Config: Default + Clone + DeserializeOwned;
 
     const SERVICE_NAME: &'static str;
-    fn start(&self, settings: Settings<Self::Config>) -> AnyhowResultFuture<()>;
+    fn start(
+        &self,
+        settings: Settings<Self::Config>,
+        stop: oneshot::Receiver<()>,
+    ) -> AnyhowResultFuture<()>;
     fn new() -> Self;
 
     fn _internal_start(self) -> AnyhowResultFuture<()> {
         Box::pin(async move {
+            pretty_env_logger::init();
             let settings = Settings::<Self::Config>::new(Self::SERVICE_NAME);
+            let (stop, stop_channel) = oneshot::channel();
 
             // Start the grpc healthcheck
             tokio::spawn(async move {});
@@ -21,7 +31,21 @@ pub trait Component: Send + Sync + 'static + Sized {
             // Start the prometheus monitoring job
             tokio::spawn(async move {});
 
-            self.start(settings?).await
+            tokio::spawn(async move {
+                match tokio::signal::unix::signal(SignalKind::terminate()).unwrap().recv().await {
+                    Some(()) => {
+                        info!("Stopping program.");
+
+                        stop.send(()).unwrap();
+                    }
+                    None => {
+                        error!("Unable to listen for shutdown signal");
+                        // we also shut down in case of error
+                    }
+                }
+            });
+
+            self.start(settings?, stop_channel).await
         })
     }
 }
@@ -41,6 +65,7 @@ macro_rules! ignite {
 #[cfg(test)]
 mod test {
     use serde::Deserialize;
+    use tokio::sync::oneshot;
 
     use crate::Component;
 
@@ -57,6 +82,7 @@ mod test {
         fn start(
             &self,
             _settings: shared::config::Settings<Self::Config>,
+            _stop: oneshot::Receiver<()>,
         ) -> crate::AnyhowResultFuture<()> {
             Box::pin(async move { Ok(()) })
         }
@@ -65,6 +91,6 @@ mod test {
             Self {}
         }
     }
-    
+
     ignite!(TestComponent);
 }
diff --git a/libs/shared/Cargo.toml b/libs/shared/Cargo.toml
index ab19ce8..ce08fbc 100644
--- a/libs/shared/Cargo.toml
+++ b/libs/shared/Cargo.toml
@@ -4,7 +4,6 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-pretty_env_logger = "0.4"
 log = { version = "0.4", features = ["std"] }
 serde = { version = "1.0.8", features = ["derive"] }
 serde_repr = "0.1"