WIP (not protable) - Add self-pipe trick to ListeningServer

Author: cpsughrue

llvm/include/llvm/Support/raw_socket_stream.h

@@ -58,8 +58,15 @@ class WSABalancer {
 /// \endcode
 ///
 class ListeningSocket {
+
+  /// If ListeningSocket::shutdown is used by a signal handler to clean up
+  /// ListeningSocket resources FD may be closed while ::poll is waiting for
+  /// FD to become ready to perform I/O. When FD is closed ::poll will
+  /// continue to block so use the self-pipe trick to get ::poll to return
+  int PipeFD[2];
+  std::mutex PipeMutex;
   std::atomic<int> FD;
-  std::string SocketPath;
+  std::string SocketPath; // Never modified
   ListeningSocket(int SocketFD, StringRef SocketPath);
 
 #ifdef _WIN32

llvm/lib/Support/raw_socket_stream.cpp

@@ -17,7 +17,9 @@
 #include "llvm/Support/FileSystem.h"
 
 #include <atomic>
+#include <fcntl.h>
 #include <poll.h>
+#include <thread>
 
 #ifndef _WIN32
 #include <sys/socket.h>
@@ -168,7 +170,7 @@ ListeningSocket::createListeningUnixSocket(StringRef SocketPath,
 Expected<std::unique_ptr<raw_socket_stream>>
 ListeningSocket::accept(std::optional<std::chrono::milliseconds> Timeout) {
 
-  struct pollfd FDs[1];
+  struct pollfd FDs[2];
   FDs[0].events = POLLIN;
 #ifdef _WIN32
   SOCKET WinServerSock = _get_osfhandle(FD);
@@ -177,8 +179,16 @@ ListeningSocket::accept(std::optional<std::chrono::milliseconds> Timeout) {
   FDs[0].fd = FD;
 #endif
 
+  FDs[1].events = POLLIN;
+  PipeMutex.lock();
+  if (::pipe(PipeFD) == -1)
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "pipe failed");
+  FDs[1].fd = PipeFD[0];
+  PipeMutex.unlock();
+
   int TimeoutCount = Timeout.value_or(std::chrono::milliseconds(-1)).count();
-  int PollStatus = ::poll(FDs, 1, TimeoutCount);
+  int PollStatus = ::poll(FDs, 2, TimeoutCount);
 
   if (PollStatus == -1)
     return llvm::make_error<StringError>(getLastSocketErrorCode(),
@@ -212,6 +222,11 @@ void ListeningSocket::shutdown() {
     return;
   ::close(FD);
   ::unlink(SocketPath.c_str());
+
+  char Byte = 'A';
+  PipeMutex.lock();
+  write(PipeFD[1], &Byte, 1);
+  PipeMutex.unlock();
   FD = -1;
 }
 

llvm/unittests/Support/raw_socket_stream_test.cpp

@@ -97,4 +97,40 @@ TEST(raw_socket_streamTest, TIMEOUT_PROVIDED) {
     ASSERT_EQ(EC, std::errc::timed_out);
   });
 }
+
+TEST(raw_socket_streamTest, FILE_DESCRIPTOR_CLOSED) {
+  if (!hasUnixSocketSupport())
+    GTEST_SKIP();
+
+  SmallString<100> SocketPath;
+  llvm::sys::fs::createUniquePath("fd_closed.sock", SocketPath, true);
+
+  // Make sure socket file does not exist. May still be there from the last test
+  std::remove(SocketPath.c_str());
+
+  Expected<ListeningSocket> MaybeServerListener =
+      ListeningSocket::createListeningUnixSocket(SocketPath);
+  ASSERT_THAT_EXPECTED(MaybeServerListener, llvm::Succeeded());
+  ListeningSocket ServerListener = std::move(*MaybeServerListener);
+
+  // Create a separate thread to close the socket after a delay. Simulates a
+  // signal handler calling ServerListener::shutdown
+  std::thread CloseThread([&]() {
+    std::this_thread::sleep_for(std::chrono::seconds(2));
+    ServerListener.shutdown();
+  });
+
+  Expected<std::unique_ptr<raw_socket_stream>> MaybeServer =
+      ServerListener.accept();
+
+  // Wait for the CloseThread to finish
+  CloseThread.join();
+
+  ASSERT_THAT_EXPECTED(MaybeServer, Failed());
+  llvm::Error Err = MaybeServer.takeError();
+  llvm::handleAllErrors(std::move(Err), [&](const llvm::StringError &SE) {
+    std::error_code EC = SE.convertToErrorCode();
+    ASSERT_EQ(EC, std::errc::bad_file_descriptor);
+  });
+}
 } // namespace