feat: craned health check

etc/config.yaml

@@ -162,6 +162,12 @@ Container:
   RuntimeDelete: /usr/bin/runc --rootless=true --root=/run/user/%U/ delete --force %u.%U.%j.%x
   RuntimeRun: /usr/bin/runc --rootless=true --root=/run/user/%U/ run -b %b %u.%U.%j.%x
 
+HealthCheck:
+  Program: /usr/sbin/nhc
+  Interval: 60
+  NodeState: ANY # IDLE, ALLOC, MIXED, ANY
+  Cycle: false
+
 Plugin:
   # Toggle the plugin module in CraneSched
   Enabled: false

protos/Crane.proto

@@ -922,6 +922,21 @@ message QueryTxnLogReply {
   repeated Txn txn_log_list = 3;
 }
 
+message SendHealthCheckResultRequest {
+  string craned_id = 1;
+  bool healthy = 2;
+}
+
+message QueryNodeStateRequest {
+  string craned_id = 1;
+}
+
+message QueryNodeStateReply {
+  CranedResourceState state = 1;
+  bool drain = 2;
+  bool ok = 3;
+}
+
 // Todo: Divide service into two parts: one for Craned and one for Crun
 //  We need to distinguish the message sender
 //  and have some kind of authentication
@@ -988,6 +1003,9 @@ service CraneCtldForInternal {
   rpc CranedRegister(CranedRegisterRequest) returns (CranedRegisterReply);
   rpc CranedPing(CranedPingRequest) returns (CranedPingReply);
 
+  rpc SendHealthCheckResult(SendHealthCheckResultRequest) returns (google.protobuf.Empty);
+  rpc QueryNodeState(QueryNodeStateRequest) returns (QueryNodeStateReply);
+
   /* RPCs called from Cfored */
   rpc CforedStream(stream StreamCforedRequest) returns (stream StreamCtldReply);
 }

src/CraneCtld/CranedMetaContainer.cpp

@@ -825,6 +825,51 @@ crane::grpc::ModifyCranedStateReply CranedMetaContainer::ChangeNodeState(
   return reply;
 }
 
+void CranedMetaContainer::UpdateNodeStateWithHealthCheck_(
+    const CranedId& craned_id, bool is_health) {
+  if (!craned_meta_map_.Contains(craned_id)) {
+    CRANE_ERROR(
+        "Health check: unknown craned_id '{}', cannot update drain state.",
+        craned_id);
+    return;
+  }
+
+  auto craned_meta = craned_meta_map_[craned_id];
+  // only update drain state when the reason is 'Node failed health check'
+  if (craned_meta->drain &&
+      craned_meta->state_reason == HealthCheckFailedReason) {
+    craned_meta->drain = !is_health;
+    if (is_health) craned_meta->state_reason.clear();
+  } else if (!craned_meta->drain && !is_health) {
+    craned_meta->drain = true;
+    craned_meta->state_reason = HealthCheckFailedReason;
+  }
+  CRANE_TRACE("Health check: craned_id '{}' drain state changed to {}.",
+              craned_id, craned_meta->drain);
+}
+
+void CranedMetaContainer::QueryNodeState(
+    const CranedId& craned_id, crane::grpc::QueryNodeStateReply* response) {
+  if (!craned_meta_map_.Contains(craned_id)) {
+    CRANE_ERROR(
+        "Health check: unknown craned_id '{}', cannot update drain state.",
+        craned_id);
+    response->set_ok(false);
+    return;
+  }
+
+  auto craned_meta = craned_meta_map_[craned_id];
+  if (craned_meta->res_in_use.IsZero())
+    response->set_state(crane::grpc::CranedResourceState::CRANE_IDLE);
+  else if (craned_meta->res_avail.allocatable_res.IsAnyZero())
+    response->set_state(crane::grpc::CranedResourceState::CRANE_ALLOC);
+  else
+    response->set_state(crane::grpc::CranedResourceState::CRANE_MIX);
+
+  response->set_drain(craned_meta->drain);
+  response->set_ok(true);
+}
+
 CraneExpected<void> CranedMetaContainer::ModifyPartitionAcl(
     const std::string& partition_name, bool is_allowed_list,
     std::unordered_set<std::string>&& accounts) {

src/CraneCtld/CranedMetaContainer.h

@@ -89,6 +89,12 @@ class CranedMetaContainer final {
   crane::grpc::ModifyCranedStateReply ChangeNodeState(
       const crane::grpc::ModifyCranedStateRequest& request);
 
+  void UpdateNodeStateWithHealthCheck_(const CranedId& craned_id,
+                                       bool is_health);
+
+  void QueryNodeState(const CranedId& craned_id,
+                      crane::grpc::QueryNodeStateReply* response);
+
   CraneExpected<void> ModifyPartitionAcl(
       const std::string& partition_name, bool is_allowed_list,
       std::unordered_set<std::string>&& accounts);

src/CraneCtld/RpcService/CtldGrpcServer.cpp

@@ -168,6 +168,25 @@ grpc::Status CtldForInternalServiceImpl::CranedPing(
   return grpc::Status::OK;
 }
 
+grpc::Status CtldForInternalServiceImpl::SendHealthCheckResult(
+    grpc::ServerContext *context,
+    const crane::grpc::SendHealthCheckResultRequest *request,
+    google::protobuf::Empty *response) {
+  g_meta_container->UpdateNodeStateWithHealthCheck_(request->craned_id(),
+                                                    request->healthy());
+
+  return grpc::Status::OK;
+}
+
+grpc::Status CtldForInternalServiceImpl::QueryNodeState(
+    grpc::ServerContext *context,
+    const crane::grpc::QueryNodeStateRequest *request,
+    crane::grpc::QueryNodeStateReply *response) {
+  g_meta_container->QueryNodeState(request->craned_id(), response);
+
+  return grpc::Status::OK;
+}
+
 grpc::Status CtldForInternalServiceImpl::CforedStream(
     grpc::ServerContext *context,
     grpc::ServerReaderWriter<crane::grpc::StreamCtldReply,

src/CraneCtld/RpcService/CtldGrpcServer.h

@@ -192,6 +192,16 @@ class CtldForInternalServiceImpl final
                           const crane::grpc::CranedPingRequest *request,
                           crane::grpc::CranedPingReply *response) override;
 
+  grpc::Status SendHealthCheckResult(
+      grpc::ServerContext *context,
+      const crane::grpc::SendHealthCheckResultRequest *request,
+      google::protobuf::Empty *response) override;
+
+  grpc::Status QueryNodeState(
+      grpc::ServerContext *context,
+      const crane::grpc::QueryNodeStateRequest *request,
+      crane::grpc::QueryNodeStateReply *response) override;
+
   grpc::Status CforedStream(
       grpc::ServerContext *context,
       grpc::ServerReaderWriter<crane::grpc::StreamCtldReply,

src/Craned/Core/Craned.cpp

@@ -648,6 +648,41 @@ void ParseConfig(int argc, char** argv) {
           }
         }
 
+        if (config["HealthCheck"]) {
+          const auto& health_check_config = config["HealthCheck"];
+          g_config.HealthCheck.Program =
+              YamlValueOr(health_check_config["Program"], "");
+          if (g_config.HealthCheck.Program.empty()) {
+            CRANE_ERROR("HealthCheckProgram is not configured");
+            std::exit(1);
+          }
+          g_config.HealthCheck.Interval =
+              YamlValueOr<uint64_t>(health_check_config["Interval"], 0L);
+          std::string node_state;
+          node_state = absl::StripAsciiWhitespace(absl::AsciiStrToLower(
+              YamlValueOr(health_check_config["NodeState"], "any")));
+          if (node_state != "any" && node_state != "idle" &&
+              node_state != "mixed" && node_state != "alloc") {
+            CRANE_WARN("HealthCheckNodeState is not valid, reset to any");
+            node_state = "any";
+          }
+          if (node_state == "any") {
+            g_config.HealthCheck.NodeState =
+                Craned::Config::HealthCheckConfig::ANY;
+          } else if (node_state == "idle") {
+            g_config.HealthCheck.NodeState =
+                Craned::Config::HealthCheckConfig::IDLE;
+          } else if (node_state == "mixed") {
+            g_config.HealthCheck.NodeState =
+                Craned::Config::HealthCheckConfig::MIXED;
+          } else if (node_state == "alloc") {
+            g_config.HealthCheck.NodeState =
+                Craned::Config::HealthCheckConfig::ALLOC;
+          }
+          g_config.HealthCheck.Cycle =
+              YamlValueOr<bool>(health_check_config["Cycle"], false);
+        }
+
         if (config["Plugin"]) {
           const auto& plugin_config = config["Plugin"];
           g_config.Plugin.Enabled =

src/Craned/Core/CranedPublicDefs.h

@@ -102,6 +102,15 @@ struct Config {
   };
   ContainerConfig Container;
 
+  struct HealthCheckConfig {
+    enum NodeStateEnum { IDLE, ALLOC, MIXED, ANY, NONDRAINED_IDLE };
+    std::string Program;
+    uint64_t Interval{0};
+    NodeStateEnum NodeState;
+    bool Cycle;
+  };
+  HealthCheckConfig HealthCheck;
+
   struct PluginConfig {
     bool Enabled{false};
     std::string PlugindSockPath;

src/Craned/Core/CtldClient.cpp

@@ -380,6 +380,7 @@ CtldClient::~CtldClient() {
   CRANE_TRACE("Waiting for CtldClient thread to finish.");
   if (m_async_send_thread_.joinable()) m_async_send_thread_.join();
   if (m_uvw_thread_.joinable()) m_uvw_thread_.join();
+  if (m_health_check_thread_.joinable()) m_health_check_thread_.join();
 }
 
 void CtldClient::Init() {
@@ -460,6 +461,24 @@ void CtldClient::InitGrpcChannel(const std::string& server_address) {
   // std::unique_ptr will automatically release the dangling stub.
   m_stub_ = CraneCtldForInternal::NewStub(m_ctld_channel_);
 
+  if (g_config.HealthCheck.Interval > 0L) {
+    HealthCheck_();
+    m_health_check_thread_ = std::thread([this] {
+      std::mt19937 rng{std::random_device{}()};
+      do {
+        uint64_t interval = g_config.HealthCheck.Interval;
+        int delay = interval;
+        if (g_config.HealthCheck.Cycle) {
+          std::uniform_int_distribution<int> dist(1, interval);
+          delay = dist(rng);
+        }
+        std::this_thread::sleep_for(std::chrono::seconds(delay));
+        if (m_stopping_ || !m_stub_) return;
+        if (CheckNodeState_()) HealthCheck_();
+      } while (true);
+    });
+  }
+
   m_async_send_thread_ = std::thread([this] { AsyncSendThread_(); });
 }
 
@@ -703,6 +722,119 @@ void CtldClient::SendStatusChanges_() {
   }
 }
 
+void CtldClient::SendHealthCheckResult_(bool is_health) const {
+  if (m_stopping_ || !m_stub_) return;
+
+  grpc::ClientContext context;
+  crane::grpc::SendHealthCheckResultRequest request;
+  google::protobuf::Empty reply;
+
+  request.set_craned_id(g_config.CranedIdOfThisNode);
+  request.set_healthy(is_health);
+
+  auto result = m_stub_->SendHealthCheckResult(&context, request, &reply);
+  if (!result.ok()) {
+    CRANE_ERROR("SendHealthCheckResult failed: is_health={}", is_health);
+  }
+}
+
+void CtldClient::HealthCheck_() {
+  if (!g_server->ReadyFor(RequestSource::CTLD)) return;
+
+  CRANE_DEBUG("Health checking.....");
+
+  subprocess_s subprocess{};
+  std::vector<const char*> argv = {g_config.HealthCheck.Program.c_str(),
+                                   nullptr};
+
+  if (subprocess_create(argv.data(), 0, &subprocess) != 0) {
+    CRANE_ERROR(
+        "[Craned Subprocess] HealthCheck subprocess creation failed: {}.",
+        strerror(errno));
+    SendHealthCheckResult_(false);
+    return;
+  }
+
+  pid_t pid = subprocess.child;
+  int result = 0;
+
+  auto fut = std::async(std::launch::async,
+                        [pid, &result]() { return waitpid(pid, &result, 0); });
+
+  bool child_exited = false;
+  if (fut.wait_for(std::chrono::milliseconds(MaxHealthCheckWaitTime)) ==
+      std::future_status::ready) {
+    if (fut.get() == pid) {
+      child_exited = true;
+    }
+  }
+
+  auto read_stream = [](std::FILE* f) {
+    std::string out;
+    char buf[4096];
+    while (std::fgets(buf, sizeof(buf), f)) out.append(buf);
+    return out;
+  };
+
+  if (!child_exited) {
+    kill(pid, SIGKILL);
+    waitpid(pid, &result, 0);
+    std::string stdout_str = read_stream(subprocess_stdout(&subprocess));
+    std::string stderr_str = read_stream(subprocess_stderr(&subprocess));
+    CRANE_WARN("HealthCheck: Timeout. stdout: {}, stderr: {}", stdout_str,
+               stderr_str);
+    SendHealthCheckResult_(false);
+    subprocess_destroy(&subprocess);
+    return;
+  }
+
+  if (subprocess_destroy(&subprocess) != 0)
+    CRANE_ERROR("[Craned Subprocess] HealthCheck destroy failed.");
+
+  if (result != 0) {
+    std::string stdout_str = read_stream(subprocess_stdout(&subprocess));
+    std::string stderr_str = read_stream(subprocess_stderr(&subprocess));
+    CRANE_WARN("HealthCheck: Failed (exit code:{}). stdout: {}, stderr: {}",
+               result, stdout_str, stderr_str);
+    SendHealthCheckResult_(false);
+    return;
+  }
+
+  CRANE_DEBUG("Health check success.");
+  SendHealthCheckResult_(true);
+}
+
+bool CtldClient::CheckNodeState_() {
+  if (g_config.HealthCheck.NodeState == Config::HealthCheckConfig::ANY)
+    return true;
+
+  grpc::ClientContext context;
+  crane::grpc::QueryNodeStateRequest req;
+  crane::grpc::QueryNodeStateReply reply;
+  req.set_craned_id(g_config.CranedIdOfThisNode);
+  auto result = m_stub_->QueryNodeState(&context, req, &reply);
+  if (!result.ok() || !reply.ok()) {
+    CRANE_ERROR("QueryNodeState failed");
+    return false;
+  }
+
+  switch (g_config.HealthCheck.NodeState) {
+  case Config::HealthCheckConfig::NONDRAINED_IDLE:
+    return !reply.drain() &&
+           reply.state() == crane::grpc::CranedResourceState::CRANE_IDLE;
+  case Config::HealthCheckConfig::IDLE:
+    return reply.state() == crane::grpc::CranedResourceState::CRANE_IDLE;
+  case Config::HealthCheckConfig::MIXED:
+    return reply.state() == crane::grpc::CranedResourceState::CRANE_MIX;
+  case Config::HealthCheckConfig::ALLOC:
+    return reply.state() == crane::grpc::CranedResourceState::CRANE_ALLOC;
+  case Config::HealthCheckConfig::ANY:
+    break;
+  }
+
+  return false;
+}
+
 bool CtldClient::Ping_() {
   grpc::ClientContext context;
   context.set_deadline(std::chrono::system_clock::now() +

src/Craned/Core/CtldClient.h

@@ -209,6 +209,12 @@ class CtldClient {
 
   void SendStatusChanges_();
 
+  void SendHealthCheckResult_(bool is_health) const;
+
+  void HealthCheck_();
+
+  bool CheckNodeState_();
+
   absl::Mutex m_step_status_change_mtx_;
 
   std::list<TaskStatusChangeQueueElem> m_step_status_change_list_
@@ -237,6 +243,8 @@ class CtldClient {
   std::shared_ptr<uvw::timer_handle> m_ping_handle_;
   std::atomic_bool m_ping_ctld_{false};
   std::atomic<std::chrono::steady_clock::time_point> m_last_active_time_;
+
+  std::thread m_health_check_thread_;
 };
 
 }  // namespace Craned

src/Utilities/PublicHeader/include/crane/PublicHeader.h

@@ -49,6 +49,9 @@ using CraneExpectedRich = std::expected<T, CraneRichError>;
 constexpr const char* kLogPattern =
     "[%^%L%$ %C-%m-%d %H:%M:%S.%e %s:%#][%n] %v";
 
+constexpr const char* HealthCheckFailedReason = "Node failed health check";
+constexpr int MaxHealthCheckWaitTime = 60000;
+
 inline const char* const kDefaultHost = "0.0.0.0";
 
 inline const char* kCtldDefaultPort = "10011";