From a3638a29831244c77d20b2faf03a842d24e5946ab28fc9e3cd80b4a261a4a70d Mon Sep 17 00:00:00 2001
From: raven <7156279+RavenX8@users.noreply.github.com>
Date: Tue, 26 Nov 2024 10:28:30 -0500
Subject: [PATCH] - add: signal handler to deregister from consul

---
 auth-service/src/main.rs     | 21 +++++++++++----------
 database-service/src/main.rs | 20 +++++++++++---------
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/auth-service/src/main.rs b/auth-service/src/main.rs
index 5a8cbbb..e0aaa2d 100644
--- a/auth-service/src/main.rs
+++ b/auth-service/src/main.rs
@@ -6,7 +6,9 @@ use auth_service::grpc::MyAuthService;
 use dotenv::dotenv;
 use std::env;
 use std::net::ToSocketAddrs;
+use tokio::{select, signal};
 use tonic::transport::Server;
+use tracing::info;
 use warp::Filter;
 
 mod consul_registration;
@@ -33,6 +35,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let service_port = port.clone();
     let health_check_url = format!("http://{}:{}/health", service_address, service_port);
     let health_check_endpoint_addr = format!("{}:8081", service_address);
+    let db_address = get_service_address(&consul_url, "database-service").await?;
 
     // Register service with Consul
     consul_registration::register_service(
@@ -50,13 +53,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     tokio::spawn(warp::serve(health_route).run(health_check_endpoint_addr.to_socket_addrs()?.next().unwrap()));
 
-    let db_address = get_service_address(&consul_url, "database-service").await?;
-
-    tokio::spawn(async move {
-        tokio::signal::ctrl_c().await.unwrap();
-        consul_registration::deregister_service(&consul_url, service_name.as_str()).await.expect("");
-    });
-
     let db_url = format!("http://{}:{}", db_address.Address, db_address.Port);
     let database_client = DatabaseClient::connect(&db_url).await?;
 
@@ -69,10 +65,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     println!("Authentication Service running on {}", addr);
 
     // Start the gRPC server
-    Server::builder()
+    tokio::spawn(Server::builder()
         .add_service(AuthServiceServer::new(auth_service))
-        .serve(address)
-        .await?;
+        .serve(address));
+    
+    select! {
+        _ = signal::ctrl_c() => {},
+    }
 
+    consul_registration::deregister_service(&consul_url, service_name.as_str()).await.expect("");
+    info!("service {} deregistered", service_name);
     Ok(())
 }
diff --git a/database-service/src/main.rs b/database-service/src/main.rs
index f7704c2..3c34bfc 100644
--- a/database-service/src/main.rs
+++ b/database-service/src/main.rs
@@ -8,7 +8,9 @@ use sqlx::postgres::PgPoolOptions;
 use std::env;
 use std::net::ToSocketAddrs;
 use std::sync::Arc;
+ use tokio::{select, signal};
 use tonic::transport::Server;
+use tracing::info;
 use warp::Filter;
 
 mod consul_registration;
@@ -51,11 +53,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     tokio::spawn(warp::serve(health_route).run(health_check_endpoint_addr.to_socket_addrs()?.next().unwrap()));
 
-    tokio::spawn(async move {
-        tokio::signal::ctrl_c().await.unwrap();
-        consul_registration::deregister_service(&consul_url, service_name.as_str()).await.expect("");
-    });
-
     let full_addr = format!("{}:{}", &addr, port);
     let address = full_addr.parse().expect("Invalid address");
     let pool = PgPoolOptions::new()
@@ -73,10 +70,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     // Pass `shared_cache` into services as needed
     println!("Database Service running on {}", address);
-    Server::builder()
-        .add_service(DatabaseServiceServer::new(database_service))
-        .serve(address)
-        .await?;
+    tokio::spawn(Server::builder()
+            .add_service(DatabaseServiceServer::new(database_service))
+            .serve(address));
 
+    select! {
+        _ = signal::ctrl_c() => {},
+    }
+    
+    consul_registration::deregister_service(&consul_url, service_name.as_str()).await.expect("");
+    info!("service {} deregistered", service_name);
     Ok(())
 }