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) {}
 };