Wake up the heartbeat thread immediately (#5081)

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/core/control/ctrl_client.cpp

@@ -40,9 +40,12 @@ GrpcCtrlClient::GrpcCtrlClient(const ProcessCtx& process_ctx) : process_ctx_(pro
     LoadServerRequest request;
     LoadServerResponse response;
     while (true) {
+      const auto wait_duration = std::chrono::seconds(sleep_second_dis(gen));
       {
         std::unique_lock<std::mutex> lck(need_heartbeat_thread_stop_mtx_);
-        if (need_heartbeat_thread_stop_) { break; }
+        const bool stopped = need_heartbeat_thread_stop_cv_.wait_for(
+            lck, wait_duration, [&]() { return need_heartbeat_thread_stop_; });
+        if (stopped) { break; }
       }
       for (size_t i = 0; i < rpc_client_.GetStubSize(); ++i) {
         grpc::ClientContext client_ctx;
@@ -51,7 +54,6 @@ GrpcCtrlClient::GrpcCtrlClient(const ProcessCtx& process_ctx) : process_ctx_(pro
             &client_ctx, request, &response))
             << "Machine " << i << " lost";
       }
-      std::this_thread::sleep_for(std::chrono::seconds(sleep_second_dis(gen)));
     }
   });
 }
@@ -118,6 +120,7 @@ void GrpcCtrlClient::StopHeartbeat() {
     std::unique_lock<std::mutex> lck(need_heartbeat_thread_stop_mtx_);
     already_stopped = need_heartbeat_thread_stop_;
     need_heartbeat_thread_stop_ = true;
+    need_heartbeat_thread_stop_cv_.notify_all();
   }
   if (!already_stopped) { heartbeat_thread_.join(); }
 }

oneflow/core/rpc/include/grpc.h

@@ -58,6 +58,7 @@ class GrpcCtrlClient final : public CtrlClient {
   ProcessCtx process_ctx_;
   bool need_heartbeat_thread_stop_;
   std::mutex need_heartbeat_thread_stop_mtx_;
+  std::condition_variable need_heartbeat_thread_stop_cv_;
   std::thread heartbeat_thread_;
   RpcClient rpc_client_;
 };