From 51525108c23b18fabf72b238209997d9ea4d6971 Mon Sep 17 00:00:00 2001
From: Kishore Nallan <kishorenc@gmail.com>
Date: Sun, 3 Mar 2024 17:46:21 +0530
Subject: [PATCH] Add health with rusage.

---
 include/core_api.h            |  2 ++
 src/core_api.cpp              | 29 +++++++++++++++++++++++++++++
 src/main/typesense_server.cpp |  2 +-
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/include/core_api.h b/include/core_api.h
index f6ebca2c..08df91b9 100644
--- a/include/core_api.h
+++ b/include/core_api.h
@@ -113,6 +113,8 @@ bool get_debug(const std::shared_ptr<http_req>& req, const std::shared_ptr<http_
 
 bool get_health(const std::shared_ptr<http_req>& req, const std::shared_ptr<http_res>& res);
 
+bool get_health_with_resource_usage(const std::shared_ptr<http_req>& req, const std::shared_ptr<http_res>& res);
+
 bool post_health(const std::shared_ptr<http_req>& req, const std::shared_ptr<http_res>& res);
 
 bool get_metrics_json(const std::shared_ptr<http_req>& req, const std::shared_ptr<http_res>& res);
diff --git a/src/core_api.cpp b/src/core_api.cpp
index 04167bea..dceb32c5 100644
--- a/src/core_api.cpp
+++ b/src/core_api.cpp
@@ -311,6 +311,35 @@ bool get_debug(const std::shared_ptr<http_req>& req, const std::shared_ptr<http_
     return true;
 }
 
+bool get_health_with_resource_usage(const std::shared_ptr<http_req>& req, const std::shared_ptr<http_res>& res) {
+    nlohmann::json result;
+    bool alive = server->is_alive();
+
+    auto resource_error = cached_resource_stat_t::get_instance().get_out_of_resource_error();
+    if (resource_error != cached_resource_stat_t::resource_check_t::OK) {
+        result["resource_error"] = std::string(magic_enum::enum_name(resource_error));
+    }
+
+    if(req->params.count("cpu_threshold") != 0 && StringUtils::is_float(req->params["cpu_threshold"])) {
+        float cpu_threshold = std::stof(req->params["cpu_threshold"]);
+        SystemMetrics sys_metrics;
+        std::vector<cpu_stat_t> cpu_stats = sys_metrics.get_cpu_stats();
+        if(!cpu_stats.empty() && StringUtils::is_float(cpu_stats[0].active)) {
+            alive = alive && (std::stof(cpu_stats[0].active) < cpu_threshold);
+        }
+    }
+
+    result["ok"] = alive;
+
+    if(alive) {
+        res->set_body(200, result.dump());
+    } else {
+        res->set_body(503, result.dump());
+    }
+
+    return alive;
+}
+
 bool get_health(const std::shared_ptr<http_req>& req, const std::shared_ptr<http_res>& res) {
     nlohmann::json result;
     bool alive = server->is_alive();
diff --git a/src/main/typesense_server.cpp b/src/main/typesense_server.cpp
index e495b58c..65f9094e 100644
--- a/src/main/typesense_server.cpp
+++ b/src/main/typesense_server.cpp
@@ -86,6 +86,7 @@ void master_server_routes() {
     server->get("/stats.json", get_stats_json);
     server->get("/debug", get_debug);
     server->get("/health", get_health);
+    server->get("/health_with_rusage", get_health_with_resource_usage);
     server->post("/health", post_health);
     server->get("/status", get_status);
 
@@ -100,7 +101,6 @@ void master_server_routes() {
     server->get("/conversations/models/:id", get_conversation_model);
     server->del("/conversations/models/:id", del_conversation_model);
 
-
     server->get("/conversations", get_conversations);
     server->get("/conversations/:id", get_conversation);
     server->del("/conversations/:id", del_conversation);