Add max retries and timeout to heartbeat constructor

jupyter-xeus · Apr 4, 2024 · 3bb4e84 · 3bb4e84
1 parent 63bceb4
commit 3bb4e84
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 13 deletions.
diff --git a/src/xclient_zmq_impl.cpp b/src/xclient_zmq_impl.cpp
@@ -23,7 +23,7 @@ namespace xeus
         , m_shell_client(context, config.m_transport, config.m_ip, config.m_shell_port)
         , m_control_client(context, config.m_transport, config.m_ip, config.m_control_port)
         , m_iopub_client(context, config)
-        , m_heartbeat_client(context, config)
+        , m_heartbeat_client(context, config, m_max_retry, m_heartbeat_timeout)
         , p_messenger(context)
         , m_error_handler(eh)
     {
@@ -177,7 +177,6 @@ namespace xeus
     void xclient_zmq_impl::start()
     {
         start_iopub_thread();
-        // pass in desired timeout value below
         start_heartbeat_thread();
     }
 
@@ -186,9 +185,9 @@ namespace xeus
         m_iopub_thread = std::move(xthread(&xiopub_client::run, p_iopub_client.get()));
     }
 
-    void xclient_zmq_impl::start_heartbeat_thread(long timeout)
+    void xclient_zmq_impl::start_heartbeat_thread()
     {
-        m_heartbeat_thread = std::move(xthread(&xheartbeat_client::run, p_heartbeat_client.get(), timeout));
+        m_heartbeat_thread = std::move(xthread(&xheartbeat_client::run, p_heartbeat_client.get()));
     }
 
     xmessage xclient_zmq_impl::deserialize(zmq::multipart_t& wire_msg) const

diff --git a/src/xclient_zmq_impl.hpp b/src/xclient_zmq_impl.hpp
@@ -82,7 +82,7 @@ namespace xeus
 
     private:
         void start_iopub_thread();
-        void start_heartbeat_thread(long timeout);
+        void start_heartbeat_thread();
         void poll(long timeout);
 
         using authentication_ptr = std::unique_ptr<xauthentication>;
@@ -93,6 +93,9 @@ namespace xeus
         xiopub_client m_iopub_client;
         xheartbeat_client m_heartbeat_client;
 
+        const std::size_t m_max_retry = 3;
+        const long m_heartbeat_timeout = std::chrono::milliseconds(90).count();
+
         xclient_messenger p_messenger;
 
         nl::json::error_handler_t m_error_handler;

diff --git a/src/xheartbeat_client.cpp b/src/xheartbeat_client.cpp
@@ -16,9 +16,13 @@ namespace xeus
 {
 
     xheartbeat_client::xheartbeat_client(zmq::context_t& context,
-                                 const xeus::xconfiguration& config)
+                                 const xeus::xconfiguration& config,
+                                 const std::size_t max_retry,
+                                 const long timeout)
         : m_heartbeat(context, zmq::socket_type::req)
         , m_controller(context, zmq::socket_type::rep)
+        , m_max_retry(max_retry)
+        , m_heartbeat_timeout(timeout)
     {
         m_heartbeat.connect(get_end_point(config.m_transport, config.m_ip, config.m_hb_port));
         init_socket(m_controller, get_controller_end_point("heartbeat"));
@@ -51,18 +55,17 @@ namespace xeus
         m_kernel_status_listener(status);
     }
 
-    void xheartbeat_client::run(long timeout)
+    void xheartbeat_client::run()
     {
         bool stop = false;
-        int retry_count = 0;
-        const int max_retries = 3;
+        std::size_t retry_count = 0;
 
         while(!stop)
         {
             send_heartbeat_message();
-            if(!wait_for_answer(timeout))
+            if(!wait_for_answer(m_heartbeat_timeout))
             {
-                if (retry_count < max_retries)
+                if (retry_count < m_max_retry)
                 {
                     ++retry_count;
                 } else {

diff --git a/src/xheartbeat_client.hpp b/src/xheartbeat_client.hpp
@@ -25,11 +25,13 @@ namespace xeus
         using kernel_status_listener = std::function<void(bool)>;
 
         xheartbeat_client(zmq::context_t& context,
-                        const xeus::xconfiguration& config);
+                        const xeus::xconfiguration& config,
+                        const std::size_t max_retry,
+                        const long timeout);
 
         ~xheartbeat_client();
 
-        void run(long timeout);
+        void run();
 
         void register_kernel_status_listener(const kernel_status_listener& l);
         void notify_kernel_dead(bool status);
@@ -42,6 +44,8 @@ namespace xeus
         zmq::socket_t m_controller;
 
         kernel_status_listener m_kernel_status_listener;
+        const std::size_t m_max_retry;
+        const long m_heartbeat_timeout;
     };
 }