From 524b7080797a5cedb0a9f5572e31bb38e90350dd Mon Sep 17 00:00:00 2001
From: yihaoDeng <yhdeng@taosdata.com>
Date: Thu, 1 Apr 2021 02:06:48 +0800
Subject: [PATCH] add connection timeout

---
 src/os/inc/osArm32.h         |  1 +
 src/os/inc/osArm64.h         |  1 +
 src/os/inc/osLinux32.h       |  1 +
 src/os/inc/osLinux64.h       |  1 +
 src/os/inc/osSocket.h        |  2 ++
 src/os/src/detail/osSocket.c |  3 +++
 src/util/src/tsocket.c       | 43 ++++++++++++++++++++++++++++++++++--
 7 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/src/os/inc/osArm32.h b/src/os/inc/osArm32.h
index 24ff95522e..54835a1ca8 100644
--- a/src/os/inc/osArm32.h
+++ b/src/os/inc/osArm32.h
@@ -77,6 +77,7 @@ extern "C" {
 #include <sys/resource.h>
 #include <error.h>
 #include <math.h>
+#include <poll.h>
 
 #define TAOS_OS_FUNC_LZ4
   #define BUILDIN_CLZL(val) __builtin_clzll(val)
diff --git a/src/os/inc/osArm64.h b/src/os/inc/osArm64.h
index 22f0000e96..76098f6846 100644
--- a/src/os/inc/osArm64.h
+++ b/src/os/inc/osArm64.h
@@ -78,6 +78,7 @@ extern "C" {
 #include <error.h>
 #include <linux/sysctl.h>
 #include <math.h>
+#include <poll.h>
 
 #ifdef __cplusplus
 }
diff --git a/src/os/inc/osLinux32.h b/src/os/inc/osLinux32.h
index cfef05368f..eee4999399 100644
--- a/src/os/inc/osLinux32.h
+++ b/src/os/inc/osLinux32.h
@@ -77,6 +77,7 @@ extern "C" {
 #include <sys/resource.h>
 #include <error.h>
 #include <math.h>
+#include <poll.h>
 
 #define TAOS_OS_FUNC_LZ4
 #define BUILDIN_CLZL(val) __builtin_clzll(val)
diff --git a/src/os/inc/osLinux64.h b/src/os/inc/osLinux64.h
index a2febd51b7..0dfc819da3 100644
--- a/src/os/inc/osLinux64.h
+++ b/src/os/inc/osLinux64.h
@@ -75,6 +75,7 @@ extern "C" {
 #include <fcntl.h>
 #include <sys/utsname.h>
 #include <sys/resource.h>
+#include <poll.h>
 #ifndef _ALPINE
 #include <error.h>
 #endif
diff --git a/src/os/inc/osSocket.h b/src/os/inc/osSocket.h
index 111fca712c..cf75b74a9a 100644
--- a/src/os/inc/osSocket.h
+++ b/src/os/inc/osSocket.h
@@ -68,6 +68,8 @@ void    taosSetMaskSIGPIPE();
 // TAOS_OS_FUNC_SOCKET_SETSOCKETOPT
 int32_t taosSetSockOpt(SOCKET socketfd, int32_t level, int32_t optname, void *optval, int32_t optlen);
 
+int32_t taosGetSockOpt(SOCKET socketfd, int32_t level, int32_t optname, void *optval, int32_t* optlen);
+
 // TAOS_OS_FUNC_SOCKET_INET
 uint32_t    taosInetAddr(char *ipAddr);
 const char *taosInetNtoa(struct in_addr ipInt);
diff --git a/src/os/src/detail/osSocket.c b/src/os/src/detail/osSocket.c
index 1186a6dd0a..db976acccd 100644
--- a/src/os/src/detail/osSocket.c
+++ b/src/os/src/detail/osSocket.c
@@ -71,6 +71,9 @@ int32_t taosSetSockOpt(SOCKET socketfd, int32_t level, int32_t optname, void *op
   return setsockopt(socketfd, level, optname, optval, (socklen_t)optlen);
 }
 
+int32_t taosGetSockOpt(SOCKET socketfd, int32_t level, int32_t optname, void *optval, int32_t* optlen) {
+  return getsockopt(socketfd, level, optname, optval, (socklen_t *)optlen);
+} 
 #endif
 
 #ifndef TAOS_OS_FUNC_SOCKET_INET
diff --git a/src/util/src/tsocket.c b/src/util/src/tsocket.c
index b33cdb8b80..ec4918e3e3 100644
--- a/src/util/src/tsocket.c
+++ b/src/util/src/tsocket.c
@@ -22,6 +22,8 @@
   #define SIGPIPE EPIPE
 #endif
 
+#define TCP_CONN_TIMEOUT 3  // conn timeout
+
 int32_t taosGetFqdn(char *fqdn) {
   char hostname[1024];
   hostname[1023] = '\0';
@@ -346,10 +348,47 @@ SOCKET taosOpenTcpClientSocket(uint32_t destIp, uint16_t destPort, uint32_t clie
   serverAddr.sin_addr.s_addr = destIp;
   serverAddr.sin_port = (uint16_t)htons((uint16_t)destPort);
 
-  ret = connect(sockFd, (struct sockaddr *)&serverAddr, sizeof(serverAddr));
+#ifdef _TD_LINUX 
+  taosSetNonblocking(sockFd, 1);   
+  ret = connect(sockFd, (struct sockaddr *)&serverAddr, sizeof(serverAddr)); 
+  if (ret == -1) {
+    if (errno == EHOSTUNREACH) {
+      uError("failed to connect socket, ip:0x%x, port:%hu(%s)", destIp, destPort, strerror(errno));
+      taosCloseSocket(sockFd);
+      return -1; 
+    } else if (errno == EINPROGRESS || errno == EAGAIN || errno == EWOULDBLOCK) {
+      struct pollfd wfd[1]; 
+
+      wfd[0].fd = sockFd;
+      wfd[0].events = POLLOUT;
+    
+      int res = poll(wfd, 1, TCP_CONN_TIMEOUT);
+      if (res == -1 || res == 0) {
+        uError("failed to connect socket, ip:0x%x, port:%hu(poll error or conn timeout)", destIp, destPort);
+        taosCloseSocket(sockFd); //  
+        return -1;
+      }
+      int optVal = -1, optLen = sizeof(int);  
+      if ((0 != taosGetSockOpt(sockFd, SOL_SOCKET, SO_ERROR, &optVal, &optLen)) || (optVal != 0)) {
+        uError("failed to connect socket, ip:0x%x, port:%hu(connect host error)", destIp, destPort);
+        taosCloseSocket(sockFd); //  
+        return -1;
+      }
+      ret = 0;
+    } else { // Other error
+      uError("failed to connect socket, ip:0x%x, port:%hu(target host cannot be reached)", destIp, destPort);
+      taosCloseSocket(sockFd); //  
+      return -1; 
+    } 
+  }
+  taosSetNonblocking(sockFd, 0);   
+
+#else
+  ret = connect(sockFd, (struct sockaddr *)&serverAddr, sizeof(serverAddr)); 
+#endif
 
   if (ret != 0) {
-    // uError("failed to connect socket, ip:0x%x, port:%hu(%s)", destIp, destPort, strerror(errno));
+    uError("failed to connect socket, ip:0x%x, port:%hu(%s)", destIp, destPort, strerror(errno));
     taosCloseSocket(sockFd);
     sockFd = -1;
   } else {
-- 
GitLab