Use timeval to set up SO_RCVTIMEO (dmlc#1734)

Refactor code to use timeval struct for setting up the SO_RCVTIMEO variable timeout value for sockets in case of Linux. Changed SetTimeout API timeout input to seconds. Co-authored-by: Chao Ma <[email protected]>
action-wei · Jul 3, 2020 · f80fc9b · f80fc9b
1 parent 4ddd477
commit f80fc9b
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 5 deletions.
diff --git a/src/rpc/network/socket_communicator.cc b/src/rpc/network/socket_communicator.cc
@@ -190,7 +190,7 @@ bool SocketReceiver::Wait(const char* addr, int num_sender) {
   }
   // Initialize socket and socket-thread
   server_socket_ = new TCPSocket();
-  server_socket_->SetTimeout(kTimeOut * 60 * 1000);  // millsec
+  server_socket_->SetTimeout(kTimeOut);  // seconds
   // Bind socket
   if (server_socket_->Bind(ip.c_str(), port) == false) {
     LOG(FATAL) << "Cannot bind to " << ip << ":" << port;

diff --git a/src/rpc/network/socket_communicator.h b/src/rpc/network/socket_communicator.h
@@ -21,7 +21,7 @@ namespace dgl {
 namespace network {
 
 static constexpr int kMaxTryCount = 1024;    // maximal connection: 1024
-static constexpr int kTimeOut = 10;          // 10 minutes for socket timeout
+static constexpr int kTimeOut = 10 * 60;     // 10 minutes (in seconds) for socket timeout
 static constexpr int kMaxConnection = 1024;  // maximal connection: 1024
 
 /*!

diff --git a/src/rpc/network/tcp_socket.cc b/src/rpc/network/tcp_socket.cc
@@ -137,8 +137,17 @@ bool TCPSocket::SetBlocking(bool flag) {
 #endif  // _WIN32
 
 void TCPSocket::SetTimeout(int timeout) {
-  setsockopt(socket_, SOL_SOCKET, SO_RCVTIMEO,
-    reinterpret_cast<char*>(&timeout), sizeof(timeout));
+  #ifdef _WIN32
+    timeout = timeout * 1000;  // WIN API accepts millsec
+    setsockopt(socket_, SOL_SOCKET, SO_RCVTIMEO,
+               reinterpret_cast<char*>(&timeout), sizeof(timeout));
+  #else  // !_WIN32
+    struct timeval tv;
+    tv.tv_sec = timeout;
+    tv.tv_usec = 0;
+    setsockopt(socket_, SOL_SOCKET, SO_RCVTIMEO,
+               &tv, sizeof(tv));
+  #endif  // _WIN32
 }
 
 bool TCPSocket::ShutDown(int ways) {

diff --git a/src/rpc/network/tcp_socket.h b/src/rpc/network/tcp_socket.h
@@ -79,7 +79,7 @@ class TCPSocket {
 
   /*!
    * \brief Set timeout for socket
-   * \param timeout millsec timeout
+   * \param timeout seconds timeout
    */
   void SetTimeout(int timeout);