[llvm] [llvm][Support] Improvements to raw_socket_stream functionality and documentation (PR #84710)

[Michael Spencer via llvm-commits]

================
@@ -56,117 +61,210 @@ static std::error_code getLastSocketErrorCode() {
 #endif
 }
 
-ListeningSocket::ListeningSocket(int SocketFD, StringRef SocketPath)
-    : FD(SocketFD), SocketPath(SocketPath) {}
+static sockaddr_un setSocketAddr(StringRef SocketPath) {
+  struct sockaddr_un Addr;
+  memset(&Addr, 0, sizeof(Addr));
+  Addr.sun_family = AF_UNIX;
+  strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
+  return Addr;
+}
+
+static Expected<int> getSocketFD(StringRef SocketPath) {
+#ifdef _WIN32
+  SOCKET Socket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (Socket == INVALID_SOCKET) {
+#else
+  int Socket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (Socket == -1) {
+#endif // _WIN32
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "Create socket failed");
+  }
+
+  struct sockaddr_un Addr = setSocketAddr(SocketPath);
+  if (::connect(Socket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1)
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "Connect socket failed");
+
+#ifdef _WIN32
+  return _open_osfhandle(Socket, 0);
+#else
+  return Socket;
+#endif // _WIN32
+}
+
+ListeningSocket::ListeningSocket(int SocketFD, StringRef SocketPath,
+                                 int PipeFD[2])
+    : FD(SocketFD), SocketPath(SocketPath), PipeFD{PipeFD[0], PipeFD[1]} {}
 
 ListeningSocket::ListeningSocket(ListeningSocket &&LS)
-    : FD(LS.FD), SocketPath(LS.SocketPath) {
+    : FD(LS.FD.load()), SocketPath(LS.SocketPath),
+      PipeFD{LS.PipeFD[0], LS.PipeFD[1]} {
+
   LS.FD = -1;
+  LS.SocketPath.clear();
+  LS.PipeFD[0] = -1;
+  LS.PipeFD[1] = -1;
 }
 
-Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath,
-                                                      int MaxBacklog) {
+Expected<ListeningSocket>
+ListeningSocket::createListeningUnixSocket(StringRef SocketPath,
+                                           int MaxBacklog) {
+
+  // Handle instances where the target socket address already exists and
+  // differentiate between a preexisting file with and without a bound socket
+  //
+  // ::bind will return std::errc:address_in_use if a file at the socket address
+  // already exists (e.g., the file was not properly unlinked due to a crash)
+  // even if another socket has not yet binded to that address
+  if (llvm::sys::fs::exists(SocketPath)) {
+    Expected<int> MaybeFD = getSocketFD(SocketPath);
+    if (!MaybeFD) {
+
+      // Regardless of the error, notify the caller that a file already exists
+      // at the desired socket address and that there is no bound socket at that
+      // address. The file must be removed before ::bind can use the address
+      consumeError(MaybeFD.takeError());
+      return llvm::make_error<StringError>(
+          std::make_error_code(std::errc::file_exists),
+          "Socket address unavailable");
+    }
+    ::close(std::move(*MaybeFD));
+
+    // Notify caller that the provided socket address already has a bound socket
+    return llvm::make_error<StringError>(
+        std::make_error_code(std::errc::address_in_use),
+        "Socket address unavailable");
+  }
 
 #ifdef _WIN32
   WSABalancer _;
-  SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
-  if (MaybeWinsocket == INVALID_SOCKET) {
+  SOCKET Socket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (Socket == INVALID_SOCKET)
 #else
-  int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
-  if (MaybeWinsocket == -1) {
+  int Socket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (Socket == -1)
 #endif
     return llvm::make_error<StringError>(getLastSocketErrorCode(),
                                          "socket create failed");
-  }
 
-  struct sockaddr_un Addr;
-  memset(&Addr, 0, sizeof(Addr));
-  Addr.sun_family = AF_UNIX;
-  strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
-
-  if (bind(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1) {
-    std::error_code Err = getLastSocketErrorCode();
-    if (Err == std::errc::address_in_use)
-      ::close(MaybeWinsocket);
-    return llvm::make_error<StringError>(Err, "Bind error");
+  struct sockaddr_un Addr = setSocketAddr(SocketPath);
+  if (::bind(Socket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1) {
+    // Grab error code from call to ::bind before calling ::close
+    std::error_code EC = getLastSocketErrorCode();
+    ::close(Socket);
+    return llvm::make_error<StringError>(EC, "Bind error");
   }
-  if (listen(MaybeWinsocket, MaxBacklog) == -1) {
+
+  // Mark socket as passive so incoming connections can be accepted
+  if (::listen(Socket, MaxBacklog) == -1)
     return llvm::make_error<StringError>(getLastSocketErrorCode(),
                                          "Listen error");
-  }
-  int UnixSocket;
+
+  int PipeFD[2];
+#ifdef _WIN32
+  // Reserve 1 byte for the pipe and use default textmode
+  if (::_pipe(PipeFD, 1, 0) == -1)
+#else
+  if (::pipe(PipeFD) == -1)
+#endif // _WIN32
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "pipe failed");
+
 #ifdef _WIN32
-  UnixSocket = _open_osfhandle(MaybeWinsocket, 0);
+  return ListeningSocket{_open_osfhandle(Socket, 0), SocketPath, PipeFD};
 #else
-  UnixSocket = MaybeWinsocket;
+  return ListeningSocket{Socket, SocketPath, PipeFD};
 #endif // _WIN32
-  return ListeningSocket{UnixSocket, SocketPath};
 }
 
-Expected<std::unique_ptr<raw_socket_stream>> ListeningSocket::accept() {
-  int AcceptFD;
+Expected<std::unique_ptr<raw_socket_stream>>
+ListeningSocket::accept(std::optional<std::chrono::milliseconds> Timeout) {
+
+  struct pollfd FDs[2];
+  FDs[0].events = POLLIN;
 #ifdef _WIN32
   SOCKET WinServerSock = _get_osfhandle(FD);
+  FDs[0].fd = WinServerSock;
+#else
+  FDs[0].fd = FD;
+#endif
+  FDs[1].events = POLLIN;
+  FDs[1].fd = PipeFD[0];
+
+  int TimeoutCount = Timeout.value_or(std::chrono::milliseconds(-1)).count();
+#ifdef _WIN32
+  int PollStatus = WSAPoll(FDs, 2, TimeoutCount);
+  if (PollStatus == SOCKET_ERROR)
+#else
+  int PollStatus = ::poll(FDs, 2, TimeoutCount);
+  if (PollStatus == -1)
+#endif
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "poll failed");
----------------
Bigcheese wrote:

This will return early if poll is interrupted by a signal, you need to handle `std::errc::interrupted` and retry after first updating the timeout to account for the already elapsed time. 

https://github.com/llvm/llvm-project/pull/84710