feat: wait until connect to master or timeout

facebookincubator · Feb 18, 2025 · 337f3c5 · 337f3c5
1 parent 7e19c28
commit 337f3c5
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 7 deletions.
diff --git a/gloo/rendezvous/tcp_store.cc b/gloo/rendezvous/tcp_store.cc
@@ -16,6 +16,7 @@
 #include <arpa/inet.h>
 #include <iostream>
 #include <stdexcept>
+#include <signal.h>
 
 #ifndef _WIN32
 #include <unistd.h>
@@ -209,11 +210,26 @@ namespace gloo
         GLOO_THROW(err);
       }
 
-      // connect to server
-      if (connect(new_server_fd, (struct sockaddr *)&server_address, sizeof(server_address)) < 0)
+      const auto start = std::chrono::steady_clock::now();
+      auto timeout = std::chrono::seconds(timeout_);
+      while (true)
       {
-        auto err = std::string("Connection to server failed: ") + strerror(errno);
-        GLOO_THROW(err);
+        // connect to server
+        if (connect(new_server_fd, (struct sockaddr *)&server_address, sizeof(server_address)) == 0)
+        {
+          break;
+        }
+
+        // check timeout
+        const auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(
+            std::chrono::steady_clock::now() - start);
+        if (timeout != kNoTimeout && elapsed > timeout)
+        {
+          GLOO_THROW_IO_EXCEPTION(GLOO_ERROR_MSG(
+              "Connection to master timeout for " + std::to_string(timeout_) + " seconds"));
+        }
+        /* sleep override */
+        std::this_thread::sleep_for(std::chrono::milliseconds(10));
       }
 
       return new_server_fd;
@@ -358,7 +374,6 @@ namespace gloo
         }
         /* sleep override */
         std::this_thread::sleep_for(std::chrono::milliseconds(10));
-        std::this_thread::sleep_for(std::chrono::milliseconds(1000));
       }
     }
 

diff --git a/gloo/rendezvous/tcp_store.h b/gloo/rendezvous/tcp_store.h
@@ -16,6 +16,8 @@
 #include <netdb.h>
 #include <arpa/inet.h>
 
+#define SOCKET_INIT_TIMEOUT_SECONDS 30
+
 namespace gloo
 {
   namespace rendezvous
@@ -24,7 +26,7 @@ namespace gloo
     class TCPStore : public Store
     {
     public:
-      explicit TCPStore(const std::string &hostname, int port, int world_size, bool is_master, int timeout = 30);
+      explicit TCPStore(const std::string &hostname, int port, int world_size, bool is_master, int timeout = SOCKET_INIT_TIMEOUT_SECONDS);
       virtual ~TCPStore();
 
       virtual void set(const std::string &key, const std::vector<char> &data)
@@ -64,7 +66,7 @@ namespace gloo
 
       std::mutex mtx;
 
-      int server_fd;
+      int server_fd = -1;
       std::map<std::string, std::vector<char>> data_;
     };