From f80fc9bdfcb6a7ced9555a7aa0683b922dfacdf0 Mon Sep 17 00:00:00 2001
From: mszarma <michal.szarmach@intel.com>
Date: Fri, 3 Jul 2020 12:02:44 +0200
Subject: [PATCH] Use timeval to set up SO_RCVTIMEO (#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 <mctt90@gmail.com>
---
 src/rpc/network/socket_communicator.cc |  2 +-
 src/rpc/network/socket_communicator.h  |  2 +-
 src/rpc/network/tcp_socket.cc          | 13 +++++++++++--
 src/rpc/network/tcp_socket.h           |  2 +-
 4 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/rpc/network/socket_communicator.cc b/src/rpc/network/socket_communicator.cc
index 4bb4d9d3db63..c0dc8a807668 100644
--- 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
index 3932c0f0f335..cf05dba71a01 100644
--- 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
index 59827b5bd185..4d75a2d7e313 100644
--- 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
index 125809dc1dff..6649689c0715 100644
--- 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);