[Support] Add Interrupt Handling in Pipe

src/support/error_handling.h

@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file err_handling.h
+ * \brief Common error handling functions for socket.h and pipe.h
+ */
+#ifndef TVM_SUPPORT_ERR_HANDLING_H_
+#define TVM_SUPPORT_ERR_HANDLING_H_
+#include <errno.h>
+
+namespace tvm {
+namespace support {
+   /*!
+   * \return last error of socket operation
+   */
+  static int GetLastErrorCode() {
+#ifdef _WIN32
+    return WSAGetLastError();
+#else
+    return errno;
+#endif
+  }
+  /*!
+   * \brief Call a function and retry if an EINTR error is encountered.
+   *
+   *  Socket operations can return EINTR when the interrupt handler
+   *  is registered by the execution environment(e.g. python).
+   *  We should retry if there is no KeyboardInterrupt recorded in
+   *  the environment.
+   *
+   * \note This function is needed to avoid rare interrupt event
+   *       in long running server code.
+   *
+   * \param func The function to retry.
+   * \return The return code returned by function f or error_value on retry failure.
+   */
+  template <typename FuncType>
+  inline ssize_t RetryCallOnEINTR(FuncType func) {
+    ssize_t ret = func();
+    // common path
+    if (ret != -1) return ret;
+    // less common path
+    do {
+      if (GetLastErrorCode() == EINTR) {
+        // Call into env check signals to see if there are
+        // environment specific(e.g. python) signal exceptions.
+        // This function will throw an exception if there is
+        // if the process received a signal that requires TVM to return immediately (e.g. SIGINT).
+        runtime::EnvCheckSignals();
+      } else {
+        // other errors
+        return ret;
+      }
+      ret = func();
+    } while (ret == -1);
+    return ret;
+  }
+}  // namespace support
+}  // namespace tvm
+#endif  // TVM_SUPPORT_ERR_HANDLING_H_

src/support/pipe.h

@@ -36,6 +36,7 @@
 #include <cstdlib>
 #include <cstring>
 #endif
+#include "error_handling.h"
 
 namespace tvm {
 namespace support {
@@ -52,6 +53,7 @@ class Pipe : public dmlc::Stream {
 #endif
   /*! \brief destructor */
   ~Pipe() { Flush(); }
+
   using Stream::Read;
   using Stream::Write;
   /*!
@@ -67,8 +69,8 @@ class Pipe : public dmlc::Stream {
     ICHECK(ReadFile(handle_, static_cast<TCHAR*>(ptr), size, &nread, nullptr))
         << "Read Error: " << GetLastError();
 #else
-    ssize_t nread;
-    nread = read(handle_, ptr, size);
+    ssize_t nread = RetryCallOnEINTR(
+          [&]() { return read(handle_, ptr, size); });
     ICHECK_GE(nread, 0) << "Write Error: " << strerror(errno);
 #endif
     return static_cast<size_t>(nread);
@@ -87,8 +89,8 @@ class Pipe : public dmlc::Stream {
            static_cast<size_t>(nwrite) == size)
         << "Write Error: " << GetLastError();
 #else
-    ssize_t nwrite;
-    nwrite = write(handle_, ptr, size);
+    ssize_t nwrite = RetryCallOnEINTR(
+          [&]() { return write(handle_, ptr, size); });
     ICHECK_EQ(static_cast<size_t>(nwrite), size) << "Write Error: " << strerror(errno);
 #endif
   }

src/support/socket.h

@@ -39,7 +39,6 @@
 #endif
 #else
 #include <arpa/inet.h>
-#include <errno.h>
 #include <fcntl.h>
 #include <netdb.h>
 #include <netinet/in.h>
@@ -56,8 +55,9 @@
 #include <unordered_map>
 #include <vector>
 
-#include "../support/ssize.h"
-#include "../support/utils.h"
+#include "ssize.h"
+#include "utils.h"
+#include "error_handling.h"
 
 #if defined(_WIN32)
 static inline int poll(struct pollfd* pfd, int nfds, int timeout) {
@@ -307,19 +307,10 @@ class Socket {
       Error("Socket::Close double close the socket or close without create");
     }
   }
-  /*!
-   * \return last error of socket operation
-   */
-  static int GetLastError() {
-#ifdef _WIN32
-    return WSAGetLastError();
-#else
-    return errno;
-#endif
-  }
+
   /*! \return whether last error was would block */
   static bool LastErrorWouldBlock() {
-    int errsv = GetLastError();
+    int errsv = GetLastErrorCode();
 #ifdef _WIN32
     return errsv == WSAEWOULDBLOCK;
 #else
@@ -355,50 +346,14 @@ class Socket {
    * \param msg The error message.
    */
   static void Error(const char* msg) {
-    int errsv = GetLastError();
+    int errsv = GetLastErrorCode();
 #ifdef _WIN32
     LOG(FATAL) << "Socket " << msg << " Error:WSAError-code=" << errsv;
 #else
     LOG(FATAL) << "Socket " << msg << " Error:" << strerror(errsv);
 #endif
   }
 
-  /*!
-   * \brief Call a function and retry if an EINTR error is encountered.
-   *
-   *  Socket operations can return EINTR when the interrupt handler
-   *  is registered by the execution environment(e.g. python).
-   *  We should retry if there is no KeyboardInterrupt recorded in
-   *  the environment.
-   *
-   * \note This function is needed to avoid rare interrupt event
-   *       in long running server code.
-   *
-   * \param func The function to retry.
-   * \return The return code returned by function f or error_value on retry failure.
-   */
-  template <typename FuncType>
-  ssize_t RetryCallOnEINTR(FuncType func) {
-    ssize_t ret = func();
-    // common path
-    if (ret != -1) return ret;
-    // less common path
-    do {
-      if (GetLastError() == EINTR) {
-        // Call into env check signals to see if there are
-        // environment specific(e.g. python) signal exceptions.
-        // This function will throw an exception if there is
-        // if the process received a signal that requires TVM to return immediately (e.g. SIGINT).
-        runtime::EnvCheckSignals();
-      } else {
-        // other errors
-        return ret;
-      }
-      ret = func();
-    } while (ret == -1);
-    return ret;
-  }
-
  protected:
   explicit Socket(SockType sockfd) : sockfd(sockfd) {}
 };