THRIFT-4618: Use poll() instead of select() in C++ TNonblockingServer if available (#1580)

4 files changed, 48 insertions, 4 deletions

diff --git a/build/cmake/ConfigureChecks.cmake b/build/cmake/ConfigureChecks.cmake
index 6b9c6a32f..457bfe05d 100644
--- a/build/cmake/ConfigureChecks.cmake
+++ b/build/cmake/ConfigureChecks.cmake
@@ -45,6 +45,7 @@ check_include_file(sys/socket.h HAVE_SYS_SOCKET_H)
 check_include_file(sys/stat.h HAVE_SYS_STAT_H)
 check_include_file(sys/time.h HAVE_SYS_TIME_H)
 check_include_file(sys/un.h HAVE_SYS_UN_H)
+check_include_file(poll.h HAVE_POLL_H)
 check_include_file(sys/poll.h HAVE_SYS_POLL_H)
 check_include_file(sys/select.h HAVE_SYS_SELECT_H)
 check_include_file(sched.h HAVE_SCHED_H)
diff --git a/build/cmake/config.h.in b/build/cmake/config.h.in
index c5d4d307d..39d8270aa 100644
--- a/build/cmake/config.h.in
+++ b/build/cmake/config.h.in
@@ -121,6 +121,9 @@
 /* Define to 1 if you have the <sys/un.h> header file. */
 #cmakedefine HAVE_SYS_UN_H 1
 
+/* Define to 1 if you have the <poll.h> header file. */
+#cmakedefine HAVE_POLL_H 1
+
 /* Define to 1 if you have the <sys/poll.h> header file. */
 #cmakedefine HAVE_SYS_POLL_H 1
 
diff --git a/configure.ac b/configure.ac
index 917f6fa22..14a8c2e0b 100755
--- a/configure.ac
+++ b/configure.ac
@@ -640,6 +640,7 @@ AC_CHECK_HEADERS([sys/ioctl.h])
 AC_CHECK_HEADERS([sys/socket.h])
 AC_CHECK_HEADERS([sys/time.h])
 AC_CHECK_HEADERS([sys/un.h])
+AC_CHECK_HEADERS([poll.h])
 AC_CHECK_HEADERS([sys/poll.h])
 AC_CHECK_HEADERS([sys/resource.h])
 AC_CHECK_HEADERS([unistd.h])
diff --git a/lib/cpp/src/thrift/server/TNonblockingServer.cpp b/lib/cpp/src/thrift/server/TNonblockingServer.cpp
index f89b5f793..194d59fa0 100644
--- a/lib/cpp/src/thrift/server/TNonblockingServer.cpp
+++ b/lib/cpp/src/thrift/server/TNonblockingServer.cpp
@@ -28,7 +28,11 @@
 #include <algorithm>
 #include <iostream>
 
-#ifdef HAVE_SYS_SELECT_H
+#ifdef HAVE_POLL_H
+#include <poll.h>
+#elif HAVE_SYS_POLL_H
+#include <sys/poll.h>
+#elif HAVE_SYS_SELECT_H
 #include <sys/select.h>
 #endif
 
@@ -1291,10 +1295,44 @@ bool TNonblockingIOThread::notify(TNonblockingServer::TConnection* conn) {
     return false;
   }
 
-  fd_set wfds, efds;
-  long ret = -1;
+  int ret = -1;
   long kSize = sizeof(conn);
-  const char* pos = reinterpret_cast<const char*>(&conn);
+  const char * pos = (const char *)const_cast_sockopt(&conn);
+
+#if defined(HAVE_POLL_H) || defined(HAVE_SYS_POLL_H)
+  struct pollfd pfd = {fd, POLLOUT, 0};
+
+  while (kSize > 0) {
+    pfd.revents = 0;
+    ret = poll(&pfd, 1, -1);
+    if (ret < 0) {
+      return false;
+    } else if (ret == 0) {
+      continue;
+    }
+
+    if (pfd.revents & POLLHUP || pfd.revents & POLLERR) {
+      ::THRIFT_CLOSESOCKET(fd);
+      return false;
+    }
+
+    if (pfd.revents & POLLOUT) {
+      ret = send(fd, pos, kSize, 0);
+      if (ret < 0) {
+        if (errno == EAGAIN) {
+          continue;
+        }
+
+        ::THRIFT_CLOSESOCKET(fd);
+        return false;
+      }
+
+      kSize -= ret;
+      pos += ret;
+    }
+  }
+#else
+  fd_set wfds, efds;
 
   while (kSize > 0) {
     FD_ZERO(&wfds);
@@ -1328,6 +1366,7 @@ bool TNonblockingIOThread::notify(TNonblockingServer::TConnection* conn) {
       pos += ret;
     }
   }
+#endif
 
   return true;
 }