From c19a8ff8f98871a111b251803768870f802925cf Mon Sep 17 00:00:00 2001
From: Niladri Halder <niladri.halder26@gmail.com>
Date: Tue, 3 Dec 2024 21:28:15 +0000
Subject: [PATCH] feat(rest): merge health service into main REST service

Signed-off-by: Niladri Halder <niladri.halder26@gmail.com>
---
 control-plane/rest/Cargo.toml          |  1 -
 control-plane/rest/service/src/main.rs | 51 ++++++++------------------
 2 files changed, 16 insertions(+), 36 deletions(-)

diff --git a/control-plane/rest/Cargo.toml b/control-plane/rest/Cargo.toml
index a7c7de0d6..a3f26f862 100644
--- a/control-plane/rest/Cargo.toml
+++ b/control-plane/rest/Cargo.toml
@@ -19,7 +19,6 @@ path = "./src/lib.rs"
 rustls = { version = "0.23.19", default-features = false }
 rustls-pemfile = "2.2.0"
 actix-web = { version = "4.9.0", features = ["rustls-0_23"] }
-tokio = { version = "1.41.0", default-features = false, features = ["macros"] }
 actix-service = "2.0.2"
 opentelemetry = { version = "0.26.0" }
 tracing-actix-web = { version = "0.7.14", features = ["opentelemetry_0_26"] }
diff --git a/control-plane/rest/service/src/main.rs b/control-plane/rest/service/src/main.rs
index 78e0b1dd3..f25238a37 100644
--- a/control-plane/rest/service/src/main.rs
+++ b/control-plane/rest/service/src/main.rs
@@ -22,9 +22,8 @@ use grpc::{client::CoreClient, operations::jsongrpc::client::JsonGrpcClient};
 use http::Uri;
 use rustls::{pki_types::PrivateKeyDer, ServerConfig};
 use rustls_pemfile::{certs, rsa_private_keys};
-use std::{fs::File, io::BufReader, net::SocketAddr, time::Duration};
+use std::{fs::File, io::BufReader, time::Duration};
 use stor_port::transport_api::{RequestMinTimeout, TimeoutOptions};
-use tokio::try_join;
 use utils::{
     tracing_telemetry::{FmtLayer, FmtStyle, KeyValue},
     DEFAULT_GRPC_CLIENT_ADDR,
@@ -40,10 +39,6 @@ pub(crate) struct CliArgs {
     #[clap(long)]
     http: Option<String>,
 
-    /// The bind address for the health REST server.
-    #[clap(long, default_value = "[::]:9091")]
-    health_endpoint: SocketAddr,
-
     /// The CORE gRPC Server URL or address to connect to the services.
     #[clap(long, short = 'z', default_value = DEFAULT_GRPC_CLIENT_ADDR)]
     core_grpc: Uri,
@@ -96,10 +91,6 @@ pub(crate) struct CliArgs {
     #[clap(long, short, default_value_t = num_cpus::get_physical())]
     workers: usize,
 
-    /// Set the number of health service workers. Uses a minimum of 1 worker.
-    #[clap(long, default_value_t = 1)]
-    health_workers: usize,
-
     /// Set the max number of workers to start.
     /// The value 0 means the number of available physical CPUs is used.
     #[clap(long, short, default_value = utils::DEFAULT_REST_MAX_WORKER_THREADS)]
@@ -241,8 +232,14 @@ async fn main() -> anyhow::Result<()> {
         .with_tracing_tags(cli_args.tracing_tags.clone())
         .init("rest-server");
 
+    let cached_core_state =
+        Data::new(CachedCoreState::new(cli_args.core_liveness_check_frequency).await);
+
     let app = move || {
         actix_web::App::new()
+            .app_data(cached_core_state.clone())
+            .service(liveness)
+            .service(readiness)
             .wrap(tracing_actix_web::TracingLogger::default())
             .wrap(middleware::Logger::default())
             .app_data(authentication::init(get_jwk_path()))
@@ -251,7 +248,7 @@ async fn main() -> anyhow::Result<()> {
 
     // Initialize the core client to be used in rest
     CORE_CLIENT
-        .set(CoreClient::new(cli_args.core_grpc.clone(), timeout_opts()).await)
+        .set(CoreClient::new(cli_args.core_grpc, timeout_opts()).await)
         .ok()
         .expect("Expect to be initialised only once");
 
@@ -263,34 +260,18 @@ async fn main() -> anyhow::Result<()> {
             .expect("Expect to be initialised only once");
     }
 
-    let main_server =
+    let server =
         HttpServer::new(app).bind_rustls_0_23(CliArgs::args().https, get_certificates()?)?;
-    let main_server = if let Some(http) = CliArgs::args().http {
-        main_server.bind(http).map_err(anyhow::Error::from)?
+    let result = if let Some(http) = CliArgs::args().http {
+        server.bind(http).map_err(anyhow::Error::from)?
     } else {
-        main_server
+        server
     }
-    .workers(workers(&CliArgs::args()));
-
-    let cached_core_state =
-        Data::new(CachedCoreState::new(cli_args.core_liveness_check_frequency).await);
-    let health_server = HttpServer::new(move || {
-        actix_web::App::new()
-            .app_data(cached_core_state.clone())
-            .service(liveness)
-            .service(readiness)
-            .wrap(tracing_actix_web::TracingLogger::default())
-            .wrap(middleware::Logger::default())
-    })
-    .bind(cli_args.health_endpoint)?
-    // Use a minimum of 1 worker.
-    .workers(cli_args.health_workers.max(1));
-
-    let result = try_join!(main_server.run(), health_server.run());
+    .workers(workers(&CliArgs::args()))
+    .run()
+    .await;
 
     utils::tracing_telemetry::flush_traces();
 
-    result?;
-
-    Ok(())
+    result.map_err(|e| e.into())
 }