feat: Introduce independent health service for more policies

Author: zuston

src/app.rs

@@ -766,6 +766,10 @@ impl AppManager {
         self.apps.get(app_id).map(|v| v.value().clone())
     }
 
+    pub fn get_alive_app_number(&self) -> usize {
+        self.apps.len()
+    }
+
     pub fn register(
         &self,
         app_id: String,

src/config.rs

@@ -183,6 +183,13 @@ impl Default for RuntimeConfig {
 
 // =========================================================
 
+#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Default)]
+pub struct HealthServiceConfig {
+    pub alive_app_number_max_limit: Option<usize>,
+}
+
+// =========================================================
+
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
 pub struct HybridStoreConfig {
     #[serde(default = "as_default_memory_spill_high_watermark")]
@@ -280,9 +287,15 @@ pub struct Config {
     pub http_monitor_service_port: u16,
 
     pub tracing: Option<TracingConfig>,
+
+    #[serde(default = "as_default_health_service_config")]
+    pub health_service_config: HealthServiceConfig,
 }
 
 // ====
+fn as_default_health_service_config() -> HealthServiceConfig {
+    Default::default()
+}
 fn as_default_hybrid_store_config() -> HybridStoreConfig {
     HybridStoreConfig::default()
 }

src/health_service.rs

@@ -0,0 +1,46 @@
+use crate::app::AppManagerRef;
+use crate::config::HealthServiceConfig;
+use crate::storage::HybridStorage;
+use anyhow::Result;
+
+#[derive(Clone)]
+pub struct HealthService {
+    app_manager_ref: AppManagerRef,
+    hybrid_storage: HybridStorage,
+
+    alive_app_number_limit: Option<usize>,
+}
+
+impl HealthService {
+    pub fn new(
+        app_manager: &AppManagerRef,
+        storage: &HybridStorage,
+        conf: &HealthServiceConfig,
+    ) -> Self {
+        Self {
+            app_manager_ref: app_manager.clone(),
+            hybrid_storage: storage.clone(),
+            alive_app_number_limit: conf.alive_app_number_max_limit,
+        }
+    }
+
+    pub async fn is_healthy(&self) -> Result<bool> {
+        if !self.app_manager_ref.store_is_healthy().await? {
+            return Ok(false);
+        }
+
+        // for the initial deploy, to ensure the service stable.
+        // this could be removed in the future.
+        // case1: app number limit
+        // case2: once disk unhealthy, mark the service unhealthy
+
+        if let Some(limit) = self.alive_app_number_limit {
+            let alive_app_number = self.app_manager_ref.get_alive_app_number();
+            if alive_app_number > limit {
+                return Ok(false);
+            }
+        }
+
+        Ok(true)
+    }
+}

src/heartbeat.rs

@@ -2,8 +2,8 @@ use crate::app::{AppManagerRef, SHUFFLE_SERVER_ID, SHUFFLE_SERVER_IP};
 use crate::config::Config;
 use crate::grpc::protobuf::uniffle::coordinator_server_client::CoordinatorServerClient;
 use crate::grpc::protobuf::uniffle::{ShuffleServerHeartBeatRequest, ShuffleServerId};
+use crate::health_service::HealthService;
 use crate::runtime::manager::RuntimeManager;
-use crate::util::get_local_ip;
 use log::info;
 use std::time::Duration;
 use tonic::transport::Channel;
@@ -13,7 +13,16 @@ const DEFAULT_SHUFFLE_SERVER_TAG: &str = "ss_v4";
 pub struct HeartbeatTask;
 
 impl HeartbeatTask {
-    pub fn init(config: &Config, runtime_manager: RuntimeManager, app_manager: AppManagerRef) {
+    pub fn init(
+        config: &Config,
+        runtime_manager: &RuntimeManager,
+        app_manager: &AppManagerRef,
+        health_service: &HealthService,
+    ) {
+        let runtime_manager = runtime_manager.clone();
+        let app_manager = app_manager.clone();
+        let health_service = health_service.clone();
+
         let coordinator_quorum = config.coordinator_quorum.clone();
         let tags = config.tags.clone().unwrap_or(vec![]);
 
@@ -48,7 +57,7 @@ impl HeartbeatTask {
                 all_tags.push(DEFAULT_SHUFFLE_SERVER_TAG.to_string());
                 all_tags.extend_from_slice(&*tags);
 
-                let healthy = app_manager.store_is_healthy().await.unwrap_or(false);
+                let healthy = health_service.is_healthy().await.unwrap_or(false);
                 let memory_snapshot = app_manager
                     .store_memory_snapshot()
                     .await

src/lib.rs

@@ -40,6 +40,7 @@ pub mod urpc;
 pub mod util;
 
 pub mod event_bus;
+mod health_service;
 mod kerberos;
 mod reject;
 pub mod semaphore_with_index;

src/main.rs

@@ -20,6 +20,7 @@
 use crate::app::AppManager;
 use crate::common::init_global_variable;
 use crate::config::Config;
+use crate::health_service::HealthService;
 use crate::heartbeat::HeartbeatTask;
 use crate::http::{HTTPServer, HttpMonitorService};
 use crate::log_service::LogService;
@@ -46,6 +47,7 @@ pub mod constant;
 mod error;
 pub mod event_bus;
 pub mod grpc;
+pub mod health_service;
 pub mod heartbeat;
 mod http;
 pub mod kerberos;
@@ -96,9 +98,12 @@ fn main() -> Result<()> {
     let app_manager_ref = AppManager::get_ref(runtime_manager.clone(), config.clone(), &storage);
     storage.with_app_manager(&app_manager_ref);
 
+    let health_service =
+        HealthService::new(&app_manager_ref, &storage, &config.health_service_config);
+
     MetricService::init(&config, runtime_manager.clone());
     FastraceWrapper::init(config.clone());
-    HeartbeatTask::init(&config, runtime_manager.clone(), app_manager_ref.clone());
+    HeartbeatTask::init(&config, &runtime_manager, &app_manager_ref, &health_service);
     HttpMonitorService::init(&config, runtime_manager.clone());
 
     DefaultRpcService {}.start(&config, runtime_manager, app_manager_ref)?;

src/store/hybrid.rs

@@ -536,7 +536,7 @@ impl Store for HybridStore {
         let cold = check_healthy(self.cold_store.as_ref())
             .await
             .unwrap_or(false);
-        Ok(self.hot_store.is_healthy().await? && (warm || cold))
+        Ok(self.hot_store.is_healthy().await? && warm && cold)
     }
 
     #[trace]