[llvm] Add raw_socket_stream (PR #73603)

[via llvm-commits]

https://github.com/criis updated https://github.com/llvm/llvm-project/pull/73603

>From abe3119e8cf0423974acbf2d2db5919ae2f00b9d Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Tue, 14 Nov 2023 14:01:17 -0800
Subject: [PATCH 01/14] Add raw_socket_stream

---
 llvm/include/llvm/Support/raw_ostream.h       |  26 +++
 llvm/lib/Support/raw_ostream.cpp              | 153 +++++++++++++++++-
 llvm/unittests/Support/CMakeLists.txt         |   1 +
 .../Support/raw_socket_stream_test.cpp        |  55 +++++++
 4 files changed, 233 insertions(+), 2 deletions(-)
 create mode 100644 llvm/unittests/Support/raw_socket_stream_test.cpp

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 1e01eb9ea19c41..f135ee5435379f 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -630,6 +630,32 @@ class raw_fd_stream : public raw_fd_ostream {
   static bool classof(const raw_ostream *OS);
 };
 
+//===----------------------------------------------------------------------===//
+// Socket Streams
+//===----------------------------------------------------------------------===//
+
+/// A raw stream for sockets reading/writing
+
+class raw_socket_stream : public raw_fd_ostream {
+  StringRef SocketPath;
+  bool ShouldUnlink;
+
+  uint64_t current_pos() const override { return 0; }
+  
+public:
+  int get_socket() {
+    return get_fd();
+  }
+  
+  static int MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog, std::error_code &EC);
+
+  raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC);
+  raw_socket_stream(StringRef SockPath, std::error_code &EC);
+  ~raw_socket_stream();
+
+  Expected<std::string> read_impl();
+};
+
 //===----------------------------------------------------------------------===//
 // Output Stream Adaptors
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 8908e7b6a150ca..ff15b39ec69d8b 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -1,4 +1,4 @@
-//===--- raw_ostream.cpp - Implement the raw_ostream classes --------------===//
+ //===--- raw_ostream.cpp - Implement the raw_ostream classes --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -23,11 +23,17 @@
 #include "llvm/Support/NativeFormatting.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/Program.h"
+#include "llvm/Support/Threading.h"
+#include "llvm/Support/Error.h"
 #include <algorithm>
 #include <cerrno>
 #include <cstdio>
 #include <sys/stat.h>
 
+#include <sys/socket.h>
+#include <sys/un.h>
+#include <iostream>
+
 // <fcntl.h> may provide O_BINARY.
 #if defined(HAVE_FCNTL_H)
 # include <fcntl.h>
@@ -58,6 +64,9 @@
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
+#include "raw_ostream.h"
+#include <afunix.h>
+#include <io.h>
 #endif
 
 using namespace llvm;
@@ -644,7 +653,7 @@ raw_fd_ostream::raw_fd_ostream(int fd, bool shouldClose, bool unbuffered,
   // Check if this is a console device. This is not equivalent to isatty.
   IsWindowsConsole =
       ::GetFileType((HANDLE)::_get_osfhandle(fd)) == FILE_TYPE_CHAR;
-#endif
+#endif // _WIN32
 
   // Get the starting position.
   off_t loc = ::lseek(FD, 0, SEEK_CUR);
@@ -942,6 +951,146 @@ bool raw_fd_stream::classof(const raw_ostream *OS) {
   return OS->get_kind() == OStreamKind::OK_FDStream;
 }
 
+//===----------------------------------------------------------------------===//
+//  raw_socket_stream
+//===----------------------------------------------------------------------===//
+
+int raw_socket_stream::MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog, std::error_code &EC) {
+
+#ifdef _WIN32
+  SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+#else
+  int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+#endif // defined(_WIN32)
+
+#ifdef _WIN32
+  if (MaybeWinsocket == INVALID_SOCKET) {
+#else
+  if (MaybeWinsocket == -1) {
+#endif // _WIN32
+    std::string Msg = "socket create error" + std::string(strerror(errno));
+    std::perror(Msg.c_str());
+    std::cout << Msg << std::endl;
+    EC = std::make_error_code(std::errc::connection_aborted);
+    return -1;  
+  }
+
+  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) {
+    if (errno == EADDRINUSE) {
+      ::close(MaybeWinsocket);
+      EC = std::make_error_code(std::errc::address_in_use);
+    } else {
+      EC = std::make_error_code(std::errc::inappropriate_io_control_operation);
+    }
+    return -1;
+  }
+
+  if (listen(MaybeWinsocket, MaxBacklog) == -1) {
+    EC = std::make_error_code(std::errc::address_not_available);
+    return -1;
+  }
+#ifdef _WIN32
+  return _open_osfhandle(MaybeWinsocket, 0); // flags? https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/open-osfhandle?view=msvc-170
+#else
+  return MaybeWinsocket;
+#endif // _WIN32
+}
+
+int GetSocketFD(StringRef SocketPath, std::error_code &EC) {
+#ifdef _WIN32
+  SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (MaybeWinsocket == INVALID_SOCKET) {
+#else
+  int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (MaybeWinsocket == -1) {
+#endif // _WIN32
+    std::string Msg = "socket create error" + std::string(strerror(errno));
+    std::perror(Msg.c_str());
+    EC = std::make_error_code(std::errc::connection_aborted);
+    return -1;
+  }
+
+  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);
+
+  int status = connect(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr));
+  if (status == -1) {
+    std::string Msg = "socket connect error" + std::string(strerror(errno));
+    std::perror(Msg.c_str());
+    EC = std::make_error_code(std::errc::connection_aborted);
+    return -1;
+  }
+#ifdef _WIN32
+  return _open_osfhandle(MaybeWinsocket, 0);
+#else
+  return MaybeWinsocket;
+#endif // _WIN32
+}
+
+static int ServerAccept(int FD) {
+  int AcceptFD;
+#ifdef _WIN32
+  SOCKET WinServerSock = _get_osfhandle(FD);
+  SOCKET WinAcceptSock = ::accept(WinServerSock, NULL, NULL);
+  AcceptFD = _open_osfhandle(WinAcceptSock, 0); // flags?
+#else
+  AcceptFD = ::accept(FD, NULL, NULL);
+#endif //_WIN32
+  return AcceptFD;
+}
+
+// Server
+// Call raw_fd_ostream with ShouldClose=false
+raw_socket_stream::raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC) : raw_fd_ostream(ServerAccept(SocketFD), true) {
+  SocketPath = SockPath;
+  ShouldUnlink = true;
+}
+
+// Client
+raw_socket_stream::raw_socket_stream(StringRef SockPath, std::error_code &EC) : raw_fd_ostream(GetSocketFD(SockPath, EC), true, true, OStreamKind::OK_OStream ) {
+  SocketPath = SockPath;
+  ShouldUnlink = false;
+}
+
+raw_socket_stream::~raw_socket_stream() {
+  if (ShouldUnlink) {
+    unlink(SocketPath.str().c_str());
+  }
+}
+
+Expected<std::string> raw_socket_stream::read_impl() {
+  const size_t BUFFER_SIZE = 4096;
+  std::vector<char> Buffer(BUFFER_SIZE);
+
+  int Socket = get_socket();
+  assert(Socket >= 0 && "Socket not found.");
+
+  ssize_t n;
+#ifdef _WIN32
+  SOCKET MaybeWinsocket = _get_osfhandle(Socket);
+#else
+  int MaybeWinsocket = Socket;
+#endif // _WIN32
+  n = ::read(MaybeWinsocket, Buffer.data(), Buffer.size());
+
+  if (n < 0) {
+      std::string Msg = "Buffer read error: " + std::string(strerror(errno));
+      return llvm::make_error<StringError>(Msg, inconvertibleErrorCode());
+  }
+
+  if (n == 0) {
+      return llvm::make_error<StringError>("EOF", inconvertibleErrorCode());
+  }
+  return std::string(Buffer.data());
+}
+
 //===----------------------------------------------------------------------===//
 //  raw_string_ostream
 //===----------------------------------------------------------------------===//
diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
index e1bf793536b686..df35a7b7f3626a 100644
--- a/llvm/unittests/Support/CMakeLists.txt
+++ b/llvm/unittests/Support/CMakeLists.txt
@@ -103,6 +103,7 @@ add_llvm_unittest(SupportTests
   raw_ostream_test.cpp
   raw_pwrite_stream_test.cpp
   raw_sha1_ostream_test.cpp
+  raw_socket_stream_test.cpp
   xxhashTest.cpp
 
   DEPENDS
diff --git a/llvm/unittests/Support/raw_socket_stream_test.cpp b/llvm/unittests/Support/raw_socket_stream_test.cpp
new file mode 100644
index 00000000000000..b2f756b792289b
--- /dev/null
+++ b/llvm/unittests/Support/raw_socket_stream_test.cpp
@@ -0,0 +1,55 @@
+#include <stdlib.h>
+#include <iostream>
+#include <future>
+#include "llvm/ADT/SmallString.h"
+#include "llvm/Config/llvm-config.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/FileUtilities.h"
+#include "llvm/Support/raw_ostream.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+TEST(raw_socket_streamTest, CLIENT_TO_SERVER_AND_SERVER_TO_CLIENT) {
+
+  SmallString<100> SocketPath("/tmp/test_raw_socket_stream.sock");
+  std::error_code ECServer, ECClient;
+
+  int ServerFD = raw_socket_stream::MakeServerSocket(SocketPath, 3, ECServer);
+
+  raw_socket_stream Client(SocketPath, ECClient);
+  EXPECT_TRUE(!ECClient);
+
+  raw_socket_stream Client2(SocketPath, ECClient);
+
+  raw_socket_stream Server(ServerFD, SocketPath, ECServer);
+  EXPECT_TRUE(!ECServer);
+
+  Client << "01234567";
+  Client.flush();
+
+  Client2 << "abcdefgh";
+  Client2.flush();
+
+  Expected<std::string> from_client = Server.read_impl();
+
+  if (auto E = from_client.takeError()) {
+    return; // FIXME: Do something.
+  }
+  EXPECT_EQ("01234567", (*from_client));
+
+  Server << "76543210";
+  Server.flush();
+
+  Expected<std::string> from_server = Client.read_impl();
+    if (auto E = from_server.takeError()) {
+    return;
+    // YIKES! 😩
+  }
+  EXPECT_EQ("76543210", (*from_server));
+
+}
+} // namespace
\ No newline at end of file

>From 67a0ed599be8843e027419881c2eab4456f988ef Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Mon, 27 Nov 2023 18:12:18 -0800
Subject: [PATCH 02/14] clang-format

---
 llvm/include/llvm/Support/raw_ostream.h       | 11 +++---
 llvm/lib/Support/raw_ostream.cpp              | 35 ++++++++++++-------
 .../Support/raw_socket_stream_test.cpp        |  9 +++--
 3 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index f135ee5435379f..c126abace7c25f 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -641,13 +641,12 @@ class raw_socket_stream : public raw_fd_ostream {
   bool ShouldUnlink;
 
   uint64_t current_pos() const override { return 0; }
-  
+
 public:
-  int get_socket() {
-    return get_fd();
-  }
-  
-  static int MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog, std::error_code &EC);
+  int get_socket() { return get_fd(); }
+
+  static int MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog,
+                              std::error_code &EC);
 
   raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC);
   raw_socket_stream(StringRef SockPath, std::error_code &EC);
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index ff15b39ec69d8b..505b62e33b3888 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -1,4 +1,4 @@
- //===--- raw_ostream.cpp - Implement the raw_ostream classes --------------===//
+//===--- raw_ostream.cpp - Implement the raw_ostream classes --------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -15,6 +15,7 @@
 #include "llvm/Config/config.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/Duration.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Format.h"
@@ -24,15 +25,14 @@
 #include "llvm/Support/Process.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/Threading.h"
-#include "llvm/Support/Error.h"
 #include <algorithm>
 #include <cerrno>
 #include <cstdio>
 #include <sys/stat.h>
 
+#include <iostream>
 #include <sys/socket.h>
 #include <sys/un.h>
-#include <iostream>
 
 // <fcntl.h> may provide O_BINARY.
 #if defined(HAVE_FCNTL_H)
@@ -61,10 +61,10 @@
 #endif
 
 #ifdef _WIN32
+#include "raw_ostream.h"
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
-#include "raw_ostream.h"
 #include <afunix.h>
 #include <io.h>
 #endif
@@ -955,7 +955,9 @@ bool raw_fd_stream::classof(const raw_ostream *OS) {
 //  raw_socket_stream
 //===----------------------------------------------------------------------===//
 
-int raw_socket_stream::MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog, std::error_code &EC) {
+int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
+                                        unsigned int MaxBacklog,
+                                        std::error_code &EC) {
 
 #ifdef _WIN32
   SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
@@ -972,14 +974,14 @@ int raw_socket_stream::MakeServerSocket(StringRef SocketPath, unsigned int MaxBa
     std::perror(Msg.c_str());
     std::cout << Msg << std::endl;
     EC = std::make_error_code(std::errc::connection_aborted);
-    return -1;  
+    return -1;
   }
 
   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) {
     if (errno == EADDRINUSE) {
       ::close(MaybeWinsocket);
@@ -995,7 +997,10 @@ int raw_socket_stream::MakeServerSocket(StringRef SocketPath, unsigned int MaxBa
     return -1;
   }
 #ifdef _WIN32
-  return _open_osfhandle(MaybeWinsocket, 0); // flags? https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/open-osfhandle?view=msvc-170
+  return _open_osfhandle(
+      MaybeWinsocket,
+      0); // flags?
+          // https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/open-osfhandle?view=msvc-170
 #else
   return MaybeWinsocket;
 #endif // _WIN32
@@ -1048,13 +1053,17 @@ static int ServerAccept(int FD) {
 
 // Server
 // Call raw_fd_ostream with ShouldClose=false
-raw_socket_stream::raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC) : raw_fd_ostream(ServerAccept(SocketFD), true) {
+raw_socket_stream::raw_socket_stream(int SocketFD, StringRef SockPath,
+                                     std::error_code &EC)
+    : raw_fd_ostream(ServerAccept(SocketFD), true) {
   SocketPath = SockPath;
   ShouldUnlink = true;
 }
 
 // Client
-raw_socket_stream::raw_socket_stream(StringRef SockPath, std::error_code &EC) : raw_fd_ostream(GetSocketFD(SockPath, EC), true, true, OStreamKind::OK_OStream ) {
+raw_socket_stream::raw_socket_stream(StringRef SockPath, std::error_code &EC)
+    : raw_fd_ostream(GetSocketFD(SockPath, EC), true, true,
+                     OStreamKind::OK_OStream) {
   SocketPath = SockPath;
   ShouldUnlink = false;
 }
@@ -1081,12 +1090,12 @@ Expected<std::string> raw_socket_stream::read_impl() {
   n = ::read(MaybeWinsocket, Buffer.data(), Buffer.size());
 
   if (n < 0) {
-      std::string Msg = "Buffer read error: " + std::string(strerror(errno));
-      return llvm::make_error<StringError>(Msg, inconvertibleErrorCode());
+    std::string Msg = "Buffer read error: " + std::string(strerror(errno));
+    return llvm::make_error<StringError>(Msg, inconvertibleErrorCode());
   }
 
   if (n == 0) {
-      return llvm::make_error<StringError>("EOF", inconvertibleErrorCode());
+    return llvm::make_error<StringError>("EOF", inconvertibleErrorCode());
   }
   return std::string(Buffer.data());
 }
diff --git a/llvm/unittests/Support/raw_socket_stream_test.cpp b/llvm/unittests/Support/raw_socket_stream_test.cpp
index b2f756b792289b..c2bc9c4b3e1b2b 100644
--- a/llvm/unittests/Support/raw_socket_stream_test.cpp
+++ b/llvm/unittests/Support/raw_socket_stream_test.cpp
@@ -1,6 +1,3 @@
-#include <stdlib.h>
-#include <iostream>
-#include <future>
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Casting.h"
@@ -8,6 +5,9 @@
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/raw_ostream.h"
 #include "gtest/gtest.h"
+#include <future>
+#include <iostream>
+#include <stdlib.h>
 
 using namespace llvm;
 
@@ -45,11 +45,10 @@ TEST(raw_socket_streamTest, CLIENT_TO_SERVER_AND_SERVER_TO_CLIENT) {
   Server.flush();
 
   Expected<std::string> from_server = Client.read_impl();
-    if (auto E = from_server.takeError()) {
+  if (auto E = from_server.takeError()) {
     return;
     // YIKES! 😩
   }
   EXPECT_EQ("76543210", (*from_server));
-
 }
 } // namespace
\ No newline at end of file

>From 91a587464093ca8e2452ab6954711d7314434404 Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Tue, 28 Nov 2023 09:58:38 -0800
Subject: [PATCH 03/14] Windows

---
 llvm/lib/Support/raw_ostream.cpp | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 505b62e33b3888..3ea7c2cc4c974c 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -31,8 +31,13 @@
 #include <sys/stat.h>
 
 #include <iostream>
+#ifdef _WIN32
+#include <afunix.h>
+#include <winsock2.h>
+#else
 #include <sys/socket.h>
 #include <sys/un.h>
+#endif // _WIN32
 
 // <fcntl.h> may provide O_BINARY.
 #if defined(HAVE_FCNTL_H)
@@ -61,7 +66,6 @@
 #endif
 
 #ifdef _WIN32
-#include "raw_ostream.h"
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
@@ -979,7 +983,9 @@ int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
 
   struct sockaddr_un Addr;
   memset(&Addr, 0, sizeof(Addr));
+#ifndef _WIN32
   Addr.sun_family = AF_UNIX;
+#endif // _WIN32
   strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
 
   if (bind(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1) {
@@ -997,10 +1003,7 @@ int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
     return -1;
   }
 #ifdef _WIN32
-  return _open_osfhandle(
-      MaybeWinsocket,
-      0); // flags?
-          // https://learn.microsoft.com/en-us/cpp/c-runtime-library/reference/open-osfhandle?view=msvc-170
+  return _open_osfhandle(MaybeWinsocket, 0);
 #else
   return MaybeWinsocket;
 #endif // _WIN32
@@ -1044,7 +1047,7 @@ static int ServerAccept(int FD) {
 #ifdef _WIN32
   SOCKET WinServerSock = _get_osfhandle(FD);
   SOCKET WinAcceptSock = ::accept(WinServerSock, NULL, NULL);
-  AcceptFD = _open_osfhandle(WinAcceptSock, 0); // flags?
+  AcceptFD = _open_osfhandle(WinAcceptSock, 0);
 #else
   AcceptFD = ::accept(FD, NULL, NULL);
 #endif //_WIN32

>From 835930307d9f0fd59f2af7bd220caf5c4995974b Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Thu, 30 Nov 2023 09:13:31 -0800
Subject: [PATCH 04/14] Link against winsock and move #include

---
 llvm/lib/Support/CMakeLists.txt  | 2 +-
 llvm/lib/Support/raw_ostream.cpp | 7 ++-----
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Support/CMakeLists.txt b/llvm/lib/Support/CMakeLists.txt
index b96d62c7a6224d..80854d1d09d985 100644
--- a/llvm/lib/Support/CMakeLists.txt
+++ b/llvm/lib/Support/CMakeLists.txt
@@ -40,7 +40,7 @@ endif()
 if( MSVC OR MINGW )
   # libuuid required for FOLDERID_Profile usage in lib/Support/Windows/Path.inc.
   # advapi32 required for CryptAcquireContextW in lib/Support/Windows/Path.inc.
-  set(system_libs ${system_libs} psapi shell32 ole32 uuid advapi32)
+  set(system_libs ${system_libs} psapi shell32 ole32 uuid advapi32 Ws2_32)
 elseif( CMAKE_HOST_UNIX )
   if( HAVE_LIBRT )
     set(system_libs ${system_libs} rt)
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 3ea7c2cc4c974c..21bdc616a6c87c 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -30,11 +30,7 @@
 #include <cstdio>
 #include <sys/stat.h>
 
-#include <iostream>
-#ifdef _WIN32
-#include <afunix.h>
-#include <winsock2.h>
-#else
+#ifndef _WIN32
 #include <sys/socket.h>
 #include <sys/un.h>
 #endif // _WIN32
@@ -69,6 +65,7 @@
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
+#include <winsock2.h>
 #include <afunix.h>
 #include <io.h>
 #endif

>From b71591921a64fe35cf9b605c6daeb724e33ffa47 Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Thu, 30 Nov 2023 10:59:13 -0800
Subject: [PATCH 05/14] Don't cout

---
 llvm/lib/Support/raw_ostream.cpp | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 21bdc616a6c87c..6e7654543735c2 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -971,9 +971,6 @@ int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
 #else
   if (MaybeWinsocket == -1) {
 #endif // _WIN32
-    std::string Msg = "socket create error" + std::string(strerror(errno));
-    std::perror(Msg.c_str());
-    std::cout << Msg << std::endl;
     EC = std::make_error_code(std::errc::connection_aborted);
     return -1;
   }
@@ -1014,8 +1011,6 @@ int GetSocketFD(StringRef SocketPath, std::error_code &EC) {
   int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
   if (MaybeWinsocket == -1) {
 #endif // _WIN32
-    std::string Msg = "socket create error" + std::string(strerror(errno));
-    std::perror(Msg.c_str());
     EC = std::make_error_code(std::errc::connection_aborted);
     return -1;
   }
@@ -1027,8 +1022,6 @@ int GetSocketFD(StringRef SocketPath, std::error_code &EC) {
 
   int status = connect(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr));
   if (status == -1) {
-    std::string Msg = "socket connect error" + std::string(strerror(errno));
-    std::perror(Msg.c_str());
     EC = std::make_error_code(std::errc::connection_aborted);
     return -1;
   }

>From 36053900af980ff799c11ad87c6a5154f22dfb04 Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Thu, 30 Nov 2023 11:05:17 -0800
Subject: [PATCH 06/14] Ignore clang format

---
 llvm/lib/Support/raw_ostream.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 6e7654543735c2..86bd0e561ac97a 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -65,9 +65,11 @@
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
+// clang-format off
 #include <winsock2.h>
 #include <afunix.h>
 #include <io.h>
+// clang-format on
 #endif
 
 using namespace llvm;

>From a43bed950f9062b191c6768243906145aca2abb8 Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Tue, 5 Dec 2023 09:27:51 -0800
Subject: [PATCH 07/14] WIP

---
 llvm/include/llvm/Support/raw_ostream.h |  50 +++++++++---
 llvm/lib/Support/raw_ostream.cpp        | 102 +++++++++++++++++++++++-
 2 files changed, 138 insertions(+), 14 deletions(-)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index c126abace7c25f..532ab0c73252ad 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -615,6 +615,8 @@ class raw_fd_stream : public raw_fd_ostream {
   /// immediately destroyed.
   raw_fd_stream(StringRef Filename, std::error_code &EC);
 
+  raw_fd_stream(int fd, bool shouldClose);
+
   /// This reads the \p Size bytes into a buffer pointed by \p Ptr.
   ///
   /// \param Ptr The start of the buffer to hold data to be read.
@@ -636,25 +638,47 @@ class raw_fd_stream : public raw_fd_ostream {
 
 /// A raw stream for sockets reading/writing
 
-class raw_socket_stream : public raw_fd_ostream {
-  StringRef SocketPath;
-  bool ShouldUnlink;
-
+class raw_socket_stream;
+class ListeningSocket {
+  int FD;
+  ListeningSocket(int SocketFD);
+  // ...
+public:
+  static Expected<ListeningSocket> createUnix(StringRef SocketPath);
+  Expected<raw_socket_stream> accept();
+  ~ListeningSocket();
+};
+class raw_socket_stream : public raw_fd_stream {
   uint64_t current_pos() const override { return 0; }
-
 public:
-  int get_socket() { return get_fd(); }
+  raw_socket_stream(int SocketFD);
 
-  static int MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog,
-                              std::error_code &EC);
-
-  raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC);
-  raw_socket_stream(StringRef SockPath, std::error_code &EC);
+  /// Create a \p raw_socket_stream connected to the Unix domain socket at \p
+  /// SocketPath.
+  static Expected<raw_socket_stream> createConnectedUnix(StringRef SocketPath);
   ~raw_socket_stream();
-
-  Expected<std::string> read_impl();
 };
 
+
+// class raw_socket_stream : public raw_fd_ostream {
+//   StringRef SocketPath;
+//   bool ShouldUnlink;
+
+//   uint64_t current_pos() const override { return 0; }
+
+// public:
+//   int get_socket() { return get_fd(); }
+
+//   static int MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog,
+//                               std::error_code &EC);
+
+//   raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC);
+//   raw_socket_stream(StringRef SockPath, std::error_code &EC);
+//   ~raw_socket_stream();
+
+//   Expected<std::string> read_impl();
+// };
+
 //===----------------------------------------------------------------------===//
 // Output Stream Adaptors
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 86bd0e561ac97a..c71a8b7c87f6be 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -65,11 +65,12 @@
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
+// winsock2.h must be included before afunix.h. Briefly turn off clang-format to avoid error.
 // clang-format off
 #include <winsock2.h>
 #include <afunix.h>
-#include <io.h>
 // clang-format on
+#include <io.h>
 #endif
 
 using namespace llvm;
@@ -940,6 +941,10 @@ raw_fd_stream::raw_fd_stream(StringRef Filename, std::error_code &EC)
     EC = std::make_error_code(std::errc::invalid_argument);
 }
 
+raw_fd_stream::raw_fd_stream(int fd, bool shouldClose)
+    : raw_fd_ostream(fd, shouldClose, false, OStreamKind::OK_FDStream) {}
+
+
 ssize_t raw_fd_stream::read(char *Ptr, size_t Size) {
   assert(get_fd() >= 0 && "File already closed.");
   ssize_t Ret = ::read(get_fd(), (void *)Ptr, Size);
@@ -958,6 +963,99 @@ bool raw_fd_stream::classof(const raw_ostream *OS) {
 //  raw_socket_stream
 //===----------------------------------------------------------------------===//
 
+ListeningSocket::ListeningSocket(int SocketFD) : FD(SocketFD) {}
+
+Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
+
+  int MaxBacklog = 3;
+#ifdef _WIN32
+  SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (MaybeWinsocket == INVALID_SOCKET) {
+#else
+  int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (MaybeWinsocket == -1) {
+#endif
+    return llvm::make_error<StringError>("socket create failed", inconvertibleErrorCode());
+  }
+    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::string Msg;
+    if (errno == EADDRINUSE) {
+      ::close(MaybeWinsocket);
+      return llvm::make_error<StringError>("Address in use", inconvertibleErrorCode());
+    } else {
+      return llvm::make_error<StringError>("Unknown error", inconvertibleErrorCode());
+    }
+  }
+  if (listen(MaybeWinsocket, MaxBacklog) == -1) {
+    return llvm::make_error<StringError>("Listen error", inconvertibleErrorCode());
+  }
+  int UnixSocket;
+#ifdef _WIN32
+  UnixSocket = _open_osfhandle(MaybeWinsocket, 0);
+#else
+  UnixSocket = MaybeWinsocket;
+#endif // _WIN32
+  ListeningSocket ListenSocket(UnixSocket);
+  return ListenSocket;
+}
+
+Expected<raw_socket_stream> ListeningSocket::accept() {
+  int AcceptFD;
+#ifdef _WIN32
+  SOCKET WinServerSock = _get_osfhandle(FD);
+  SOCKET WinAcceptSock = ::accept(WinServerSock, NULL, NULL);
+  AcceptFD = _open_osfhandle(WinAcceptSock, 0);
+#else
+  AcceptFD = ::accept(FD, NULL, NULL);
+#endif //_WIN32
+  if (AcceptFD == -1) {
+    return llvm::make_error<StringError>("Accept failed", inconvertibleErrorCode());
+  }
+  Expected<raw_socket_stream> MaybeSocketStream(AcceptFD);
+  return MaybeSocketStream;
+}
+
+// Expected?
+int GetSocketFD(StringRef SocketPath) {
+#ifdef _WIN32
+  SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (MaybeWinsocket == INVALID_SOCKET) {
+#else
+  int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
+  if (MaybeWinsocket == -1) {
+#endif // _WIN32
+    // return llvm::make_error<StringError>("Create socket failed", inconvertibleErrorCode());
+  }
+
+  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);
+
+  int status = connect(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr));
+  if (status == -1) {
+    // return llvm::make_error<StringError>("Connect socket failed", inconvertibleErrorCode());
+  }
+#ifdef _WIN32
+  return _open_osfhandle(MaybeWinsocket, 0);
+#else
+  return MaybeWinsocket;
+#endif // _WIN32
+}
+
+raw_socket_stream::raw_socket_stream(int AcceptFD): raw_fd_stream(AcceptFD, false) {}
+
+Expected<raw_socket_stream> createConnectedUnix(StringRef SocketPath) {
+  Expected<raw_socket_stream> MaybeSocketStream(GetSocketFD(SocketPath));
+  return MaybeSocketStream;
+}
+
+/*
 int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
                                         unsigned int MaxBacklog,
                                         std::error_code &EC) {
@@ -1095,6 +1193,8 @@ Expected<std::string> raw_socket_stream::read_impl() {
   return std::string(Buffer.data());
 }
 
+*/
+
 //===----------------------------------------------------------------------===//
 //  raw_string_ostream
 //===----------------------------------------------------------------------===//

>From 169f2ef074756f461b14408a0fb5bddd6bf957ec Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Tue, 5 Dec 2023 15:49:08 -0800
Subject: [PATCH 08/14] WIP

---
 llvm/include/llvm/Support/raw_ostream.h       | 25 ++-----------
 llvm/lib/Support/raw_ostream.cpp              | 20 ++++++++--
 .../Support/raw_socket_stream_test.cpp        | 37 +++++++------------
 3 files changed, 34 insertions(+), 48 deletions(-)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 532ab0c73252ad..93b4731bc8d1af 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -641,17 +641,20 @@ class raw_fd_stream : public raw_fd_ostream {
 class raw_socket_stream;
 class ListeningSocket {
   int FD;
-  ListeningSocket(int SocketFD);
+  StringRef SocketPath;
+  ListeningSocket(int SocketFD, StringRef SocketPath);
   // ...
 public:
   static Expected<ListeningSocket> createUnix(StringRef SocketPath);
   Expected<raw_socket_stream> accept();
+  ListeningSocket(ListeningSocket &&LS);
   ~ListeningSocket();
 };
 class raw_socket_stream : public raw_fd_stream {
   uint64_t current_pos() const override { return 0; }
 public:
   raw_socket_stream(int SocketFD);
+  raw_socket_stream(raw_socket_stream &&socket_stream);
 
   /// Create a \p raw_socket_stream connected to the Unix domain socket at \p
   /// SocketPath.
@@ -659,26 +662,6 @@ class raw_socket_stream : public raw_fd_stream {
   ~raw_socket_stream();
 };
 
-
-// class raw_socket_stream : public raw_fd_ostream {
-//   StringRef SocketPath;
-//   bool ShouldUnlink;
-
-//   uint64_t current_pos() const override { return 0; }
-
-// public:
-//   int get_socket() { return get_fd(); }
-
-//   static int MakeServerSocket(StringRef SocketPath, unsigned int MaxBacklog,
-//                               std::error_code &EC);
-
-//   raw_socket_stream(int SocketFD, StringRef SockPath, std::error_code &EC);
-//   raw_socket_stream(StringRef SockPath, std::error_code &EC);
-//   ~raw_socket_stream();
-
-//   Expected<std::string> read_impl();
-// };
-
 //===----------------------------------------------------------------------===//
 // Output Stream Adaptors
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index c71a8b7c87f6be..8fa4eb720bd933 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -963,7 +963,9 @@ bool raw_fd_stream::classof(const raw_ostream *OS) {
 //  raw_socket_stream
 //===----------------------------------------------------------------------===//
 
-ListeningSocket::ListeningSocket(int SocketFD) : FD(SocketFD) {}
+ListeningSocket::ListeningSocket(int SocketFD, StringRef SocketPath) : FD(SocketFD), SocketPath(SocketPath) {}
+
+ListeningSocket::ListeningSocket(ListeningSocket &&LS) : FD(LS.FD), SocketPath(LS.SocketPath) {}
 
 Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
 
@@ -977,7 +979,8 @@ Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
 #endif
     return llvm::make_error<StringError>("socket create failed", inconvertibleErrorCode());
   }
-    struct sockaddr_un Addr;
+
+  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);
@@ -1000,7 +1003,7 @@ Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
 #else
   UnixSocket = MaybeWinsocket;
 #endif // _WIN32
-  ListeningSocket ListenSocket(UnixSocket);
+  ListeningSocket ListenSocket(UnixSocket, SocketPath);
   return ListenSocket;
 }
 
@@ -1020,6 +1023,11 @@ Expected<raw_socket_stream> ListeningSocket::accept() {
   return MaybeSocketStream;
 }
 
+ListeningSocket::~ListeningSocket() {
+  ::close(FD);
+  unlink(SocketPath.str().c_str());
+}
+
 // Expected?
 int GetSocketFD(StringRef SocketPath) {
 #ifdef _WIN32
@@ -1050,11 +1058,15 @@ int GetSocketFD(StringRef SocketPath) {
 
 raw_socket_stream::raw_socket_stream(int AcceptFD): raw_fd_stream(AcceptFD, false) {}
 
-Expected<raw_socket_stream> createConnectedUnix(StringRef SocketPath) {
+raw_socket_stream::raw_socket_stream(raw_socket_stream &&socket_stream) : raw_fd_stream(socket_stream.get_fd(), false) {}
+
+Expected<raw_socket_stream> raw_socket_stream::createConnectedUnix(StringRef SocketPath) {
   Expected<raw_socket_stream> MaybeSocketStream(GetSocketFD(SocketPath));
   return MaybeSocketStream;
 }
 
+raw_socket_stream::~raw_socket_stream() {}
+
 /*
 int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
                                         unsigned int MaxBacklog,
diff --git a/llvm/unittests/Support/raw_socket_stream_test.cpp b/llvm/unittests/Support/raw_socket_stream_test.cpp
index c2bc9c4b3e1b2b..18e705fe3091b7 100644
--- a/llvm/unittests/Support/raw_socket_stream_test.cpp
+++ b/llvm/unittests/Support/raw_socket_stream_test.cpp
@@ -18,37 +18,28 @@ TEST(raw_socket_streamTest, CLIENT_TO_SERVER_AND_SERVER_TO_CLIENT) {
   SmallString<100> SocketPath("/tmp/test_raw_socket_stream.sock");
   std::error_code ECServer, ECClient;
 
-  int ServerFD = raw_socket_stream::MakeServerSocket(SocketPath, 3, ECServer);
+  char Bytes[8];
 
-  raw_socket_stream Client(SocketPath, ECClient);
-  EXPECT_TRUE(!ECClient);
+  Expected<ListeningSocket> MaybeServerListener = ListeningSocket::createUnix(SocketPath);
 
-  raw_socket_stream Client2(SocketPath, ECClient);
+  ListeningSocket ServerListener = std::move(*MaybeServerListener);
 
-  raw_socket_stream Server(ServerFD, SocketPath, ECServer);
-  EXPECT_TRUE(!ECServer);
+  Expected<raw_socket_stream> MaybeClient = raw_socket_stream::createConnectedUnix(SocketPath);
 
-  Client << "01234567";
-  Client.flush();
+  raw_socket_stream Client = std::move(*MaybeClient);
+
+  Expected<raw_socket_stream> MaybeServer = ServerListener.accept();
 
-  Client2 << "abcdefgh";
-  Client2.flush();
+  raw_socket_stream Server = std::move(*MaybeServer);
 
-  Expected<std::string> from_client = Server.read_impl();
+  Client << "01234567";
+  Client.flush();
 
-  if (auto E = from_client.takeError()) {
-    return; // FIXME: Do something.
-  }
-  EXPECT_EQ("01234567", (*from_client));
+  ssize_t BytesRead = Server.read(Bytes, 8);
 
-  Server << "76543210";
-  Server.flush();
+  std::string string(reinterpret_cast<char *>(Bytes), 8);
 
-  Expected<std::string> from_server = Client.read_impl();
-  if (auto E = from_server.takeError()) {
-    return;
-    // YIKES! 😩
-  }
-  EXPECT_EQ("76543210", (*from_server));
+  ASSERT_EQ(8, BytesRead);
+  ASSERT_EQ("01234567", string);
 }
 } // namespace
\ No newline at end of file

>From 9f0f71695ee179153dd24f59050e7faa7bcc4f7e Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Wed, 6 Dec 2023 16:34:18 -0800
Subject: [PATCH 09/14] Michael's diff

---
 llvm/include/llvm/Support/raw_ostream.h       | 12 ++--
 llvm/lib/Support/raw_ostream.cpp              | 65 ++++++++++++-------
 .../Support/raw_socket_stream_test.cpp        | 21 +++---
 3 files changed, 59 insertions(+), 39 deletions(-)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 93b4731bc8d1af..ee00198d9e4dbc 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -641,24 +641,24 @@ class raw_fd_stream : public raw_fd_ostream {
 class raw_socket_stream;
 class ListeningSocket {
   int FD;
-  StringRef SocketPath;
+  std::string SocketPath;
   ListeningSocket(int SocketFD, StringRef SocketPath);
-  // ...
+
 public:
   static Expected<ListeningSocket> createUnix(StringRef SocketPath);
-  Expected<raw_socket_stream> accept();
+  Expected<std::unique_ptr<raw_socket_stream>> accept();
   ListeningSocket(ListeningSocket &&LS);
   ~ListeningSocket();
 };
 class raw_socket_stream : public raw_fd_stream {
   uint64_t current_pos() const override { return 0; }
+
 public:
   raw_socket_stream(int SocketFD);
-  raw_socket_stream(raw_socket_stream &&socket_stream);
-
   /// Create a \p raw_socket_stream connected to the Unix domain socket at \p
   /// SocketPath.
-  static Expected<raw_socket_stream> createConnectedUnix(StringRef SocketPath);
+  static Expected<std::unique_ptr<raw_socket_stream>>
+  createConnectedUnix(StringRef SocketPath);
   ~raw_socket_stream();
 };
 
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 8fa4eb720bd933..0dfba93d73946f 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -963,9 +963,20 @@ bool raw_fd_stream::classof(const raw_ostream *OS) {
 //  raw_socket_stream
 //===----------------------------------------------------------------------===//
 
+std::error_code getLastSocketErrorCode() {
+#ifdef _WIN32
+  return std::error_code(::WSAGetLastError(), std::system_category());
+#else
+  return std::error_code(errno, std::system_category());
+#endif
+}
+
 ListeningSocket::ListeningSocket(int SocketFD, StringRef SocketPath) : FD(SocketFD), SocketPath(SocketPath) {}
 
-ListeningSocket::ListeningSocket(ListeningSocket &&LS) : FD(LS.FD), SocketPath(LS.SocketPath) {}
+ListeningSocket::ListeningSocket(ListeningSocket &&LS)
+    : FD(LS.FD), SocketPath(LS.SocketPath) {
+  LS.FD = -1;
+}
 
 Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
 
@@ -977,7 +988,8 @@ Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
   int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
   if (MaybeWinsocket == -1) {
 #endif
-    return llvm::make_error<StringError>("socket create failed", inconvertibleErrorCode());
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "socket create failed");
   }
 
   struct sockaddr_un Addr;
@@ -986,16 +998,14 @@ Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
   strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
 
   if (bind(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1) {
-    std::string Msg;
-    if (errno == EADDRINUSE) {
+    std::error_code Err = getLastSocketErrorCode();
+    if (Err == std::errc::address_in_use)
       ::close(MaybeWinsocket);
-      return llvm::make_error<StringError>("Address in use", inconvertibleErrorCode());
-    } else {
-      return llvm::make_error<StringError>("Unknown error", inconvertibleErrorCode());
-    }
+    return llvm::make_error<StringError>(Err, "Bind error");
   }
   if (listen(MaybeWinsocket, MaxBacklog) == -1) {
-    return llvm::make_error<StringError>("Listen error", inconvertibleErrorCode());
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "Listen error");
   }
   int UnixSocket;
 #ifdef _WIN32
@@ -1007,7 +1017,7 @@ Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
   return ListenSocket;
 }
 
-Expected<raw_socket_stream> ListeningSocket::accept() {
+Expected<std::unique_ptr<raw_socket_stream>> ListeningSocket::accept() {
   int AcceptFD;
 #ifdef _WIN32
   SOCKET WinServerSock = _get_osfhandle(FD);
@@ -1016,20 +1026,21 @@ Expected<raw_socket_stream> ListeningSocket::accept() {
 #else
   AcceptFD = ::accept(FD, NULL, NULL);
 #endif //_WIN32
-  if (AcceptFD == -1) {
-    return llvm::make_error<StringError>("Accept failed", inconvertibleErrorCode());
-  }
-  Expected<raw_socket_stream> MaybeSocketStream(AcceptFD);
-  return MaybeSocketStream;
+  if (AcceptFD == -1)
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "Accept failed");
+  return std::make_unique<raw_socket_stream>(AcceptFD);
 }
 
 ListeningSocket::~ListeningSocket() {
+  if (FD == -1)
+    return;
   ::close(FD);
-  unlink(SocketPath.str().c_str());
+  unlink(SocketPath.c_str());
 }
 
 // Expected?
-int GetSocketFD(StringRef SocketPath) {
+Expected<int> GetSocketFD(StringRef SocketPath) {
 #ifdef _WIN32
   SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
   if (MaybeWinsocket == INVALID_SOCKET) {
@@ -1037,7 +1048,8 @@ int GetSocketFD(StringRef SocketPath) {
   int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
   if (MaybeWinsocket == -1) {
 #endif // _WIN32
-    // return llvm::make_error<StringError>("Create socket failed", inconvertibleErrorCode());
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "Create socket failed");
   }
 
   struct sockaddr_un Addr;
@@ -1047,7 +1059,8 @@ int GetSocketFD(StringRef SocketPath) {
 
   int status = connect(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr));
   if (status == -1) {
-    // return llvm::make_error<StringError>("Connect socket failed", inconvertibleErrorCode());
+    return llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                         "Connect socket failed");
   }
 #ifdef _WIN32
   return _open_osfhandle(MaybeWinsocket, 0);
@@ -1056,13 +1069,15 @@ int GetSocketFD(StringRef SocketPath) {
 #endif // _WIN32
 }
 
-raw_socket_stream::raw_socket_stream(int AcceptFD): raw_fd_stream(AcceptFD, false) {}
-
-raw_socket_stream::raw_socket_stream(raw_socket_stream &&socket_stream) : raw_fd_stream(socket_stream.get_fd(), false) {}
+raw_socket_stream::raw_socket_stream(int SocketFD)
+    : raw_fd_stream(SocketFD, true) {}
 
-Expected<raw_socket_stream> raw_socket_stream::createConnectedUnix(StringRef SocketPath) {
-  Expected<raw_socket_stream> MaybeSocketStream(GetSocketFD(SocketPath));
-  return MaybeSocketStream;
+Expected<std::unique_ptr<raw_socket_stream>>
+raw_socket_stream::createConnectedUnix(StringRef SocketPath) {
+  Expected<int> FD = GetSocketFD(SocketPath);
+  if (!FD)
+    return FD.takeError();
+  return std::make_unique<raw_socket_stream>(*FD);
 }
 
 raw_socket_stream::~raw_socket_stream() {}
diff --git a/llvm/unittests/Support/raw_socket_stream_test.cpp b/llvm/unittests/Support/raw_socket_stream_test.cpp
index 18e705fe3091b7..943bc24bca9631 100644
--- a/llvm/unittests/Support/raw_socket_stream_test.cpp
+++ b/llvm/unittests/Support/raw_socket_stream_test.cpp
@@ -4,6 +4,7 @@
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 #include <future>
 #include <iostream>
@@ -14,30 +15,34 @@ using namespace llvm;
 namespace {
 
 TEST(raw_socket_streamTest, CLIENT_TO_SERVER_AND_SERVER_TO_CLIENT) {
-
-  SmallString<100> SocketPath("/tmp/test_raw_socket_stream.sock");
-  std::error_code ECServer, ECClient;
+  SmallString<100> SocketPath;
+  llvm::sys::fs::createUniquePath("test_raw_socket_stream.sock", SocketPath, true);
 
   char Bytes[8];
 
   Expected<ListeningSocket> MaybeServerListener = ListeningSocket::createUnix(SocketPath);
+  ASSERT_THAT_EXPECTED(MaybeServerListener, llvm::Succeeded());
 
   ListeningSocket ServerListener = std::move(*MaybeServerListener);
 
-  Expected<raw_socket_stream> MaybeClient = raw_socket_stream::createConnectedUnix(SocketPath);
+  Expected<std::unique_ptr<raw_socket_stream>> MaybeClient =
+      raw_socket_stream::createConnectedUnix(SocketPath);
+  ASSERT_THAT_EXPECTED(MaybeClient, llvm::Succeeded());
 
-  raw_socket_stream Client = std::move(*MaybeClient);
+  raw_socket_stream &Client = **MaybeClient;
 
-  Expected<raw_socket_stream> MaybeServer = ServerListener.accept();
+  Expected<std::unique_ptr<raw_socket_stream>> MaybeServer =
+      ServerListener.accept();
+  ASSERT_THAT_EXPECTED(MaybeServer, llvm::Succeeded());
 
-  raw_socket_stream Server = std::move(*MaybeServer);
+  raw_socket_stream &Server = **MaybeServer;
 
   Client << "01234567";
   Client.flush();
 
   ssize_t BytesRead = Server.read(Bytes, 8);
 
-  std::string string(reinterpret_cast<char *>(Bytes), 8);
+  std::string string(Bytes, 8);
 
   ASSERT_EQ(8, BytesRead);
   ASSERT_EQ("01234567", string);

>From 6e2f3fd514d91ab3349f6c88f391e769a3405b88 Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Wed, 6 Dec 2023 16:41:35 -0800
Subject: [PATCH 10/14] Clang format

---
 llvm/lib/Support/raw_ostream.cpp                  | 7 ++++---
 llvm/unittests/Support/raw_socket_stream_test.cpp | 6 ++++--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 0dfba93d73946f..ae7a6dd782eef8 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -65,7 +65,8 @@
 #include "llvm/Support/ConvertUTF.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/Windows/WindowsSupport.h"
-// winsock2.h must be included before afunix.h. Briefly turn off clang-format to avoid error.
+// winsock2.h must be included before afunix.h. Briefly turn off clang-format to
+// avoid error.
 // clang-format off
 #include <winsock2.h>
 #include <afunix.h>
@@ -944,7 +945,6 @@ raw_fd_stream::raw_fd_stream(StringRef Filename, std::error_code &EC)
 raw_fd_stream::raw_fd_stream(int fd, bool shouldClose)
     : raw_fd_ostream(fd, shouldClose, false, OStreamKind::OK_FDStream) {}
 
-
 ssize_t raw_fd_stream::read(char *Ptr, size_t Size) {
   assert(get_fd() >= 0 && "File already closed.");
   ssize_t Ret = ::read(get_fd(), (void *)Ptr, Size);
@@ -971,7 +971,8 @@ std::error_code getLastSocketErrorCode() {
 #endif
 }
 
-ListeningSocket::ListeningSocket(int SocketFD, StringRef SocketPath) : FD(SocketFD), SocketPath(SocketPath) {}
+ListeningSocket::ListeningSocket(int SocketFD, StringRef SocketPath)
+    : FD(SocketFD), SocketPath(SocketPath) {}
 
 ListeningSocket::ListeningSocket(ListeningSocket &&LS)
     : FD(LS.FD), SocketPath(LS.SocketPath) {
diff --git a/llvm/unittests/Support/raw_socket_stream_test.cpp b/llvm/unittests/Support/raw_socket_stream_test.cpp
index 943bc24bca9631..53eb86ae45d29d 100644
--- a/llvm/unittests/Support/raw_socket_stream_test.cpp
+++ b/llvm/unittests/Support/raw_socket_stream_test.cpp
@@ -16,11 +16,13 @@ namespace {
 
 TEST(raw_socket_streamTest, CLIENT_TO_SERVER_AND_SERVER_TO_CLIENT) {
   SmallString<100> SocketPath;
-  llvm::sys::fs::createUniquePath("test_raw_socket_stream.sock", SocketPath, true);
+  llvm::sys::fs::createUniquePath("test_raw_socket_stream.sock", SocketPath,
+                                  true);
 
   char Bytes[8];
 
-  Expected<ListeningSocket> MaybeServerListener = ListeningSocket::createUnix(SocketPath);
+  Expected<ListeningSocket> MaybeServerListener =
+      ListeningSocket::createUnix(SocketPath);
   ASSERT_THAT_EXPECTED(MaybeServerListener, llvm::Succeeded());
 
   ListeningSocket ServerListener = std::move(*MaybeServerListener);

>From a6472b631ecdc837809b9cea0b131a54ebb37632 Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Thu, 7 Dec 2023 13:58:34 -0800
Subject: [PATCH 11/14] Call WSAStartup/Cleanup

---
 llvm/lib/Support/raw_ostream.cpp | 158 ++++---------------------------
 1 file changed, 18 insertions(+), 140 deletions(-)

diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index ae7a6dd782eef8..04349c3beb547a 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -983,6 +983,13 @@ Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
 
   int MaxBacklog = 3;
 #ifdef _WIN32
+  WSADATA WsaData = {0};
+  if (WSAStartup(MAKEWORD(2, 2), &WsaData) != 0) {
+    llvm::Error E = llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                                  "WSAStartup failed");
+    WSACleanup();
+    return E;
+  }
   SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
   if (MaybeWinsocket == INVALID_SOCKET) {
 #else
@@ -1038,11 +1045,22 @@ ListeningSocket::~ListeningSocket() {
     return;
   ::close(FD);
   unlink(SocketPath.c_str());
+#ifdef _WIN32
+  WSACleanup();
+#endif // _WIN32
 }
 
 // Expected?
 Expected<int> GetSocketFD(StringRef SocketPath) {
 #ifdef _WIN32
+  WSADATA WsaData = {0};
+  if (WSAStartup(MAKEWORD(2, 2), &WsaData) != 0) {
+    llvm::Error E = llvm::make_error<StringError>(getLastSocketErrorCode(),
+                                                  "WSAStartup failed");
+    WSACleanup();
+    return E;
+  }
+
   SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
   if (MaybeWinsocket == INVALID_SOCKET) {
 #else
@@ -1083,146 +1101,6 @@ raw_socket_stream::createConnectedUnix(StringRef SocketPath) {
 
 raw_socket_stream::~raw_socket_stream() {}
 
-/*
-int raw_socket_stream::MakeServerSocket(StringRef SocketPath,
-                                        unsigned int MaxBacklog,
-                                        std::error_code &EC) {
-
-#ifdef _WIN32
-  SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
-#else
-  int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
-#endif // defined(_WIN32)
-
-#ifdef _WIN32
-  if (MaybeWinsocket == INVALID_SOCKET) {
-#else
-  if (MaybeWinsocket == -1) {
-#endif // _WIN32
-    EC = std::make_error_code(std::errc::connection_aborted);
-    return -1;
-  }
-
-  struct sockaddr_un Addr;
-  memset(&Addr, 0, sizeof(Addr));
-#ifndef _WIN32
-  Addr.sun_family = AF_UNIX;
-#endif // _WIN32
-  strncpy(Addr.sun_path, SocketPath.str().c_str(), sizeof(Addr.sun_path) - 1);
-
-  if (bind(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr)) == -1) {
-    if (errno == EADDRINUSE) {
-      ::close(MaybeWinsocket);
-      EC = std::make_error_code(std::errc::address_in_use);
-    } else {
-      EC = std::make_error_code(std::errc::inappropriate_io_control_operation);
-    }
-    return -1;
-  }
-
-  if (listen(MaybeWinsocket, MaxBacklog) == -1) {
-    EC = std::make_error_code(std::errc::address_not_available);
-    return -1;
-  }
-#ifdef _WIN32
-  return _open_osfhandle(MaybeWinsocket, 0);
-#else
-  return MaybeWinsocket;
-#endif // _WIN32
-}
-
-int GetSocketFD(StringRef SocketPath, std::error_code &EC) {
-#ifdef _WIN32
-  SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
-  if (MaybeWinsocket == INVALID_SOCKET) {
-#else
-  int MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
-  if (MaybeWinsocket == -1) {
-#endif // _WIN32
-    EC = std::make_error_code(std::errc::connection_aborted);
-    return -1;
-  }
-
-  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);
-
-  int status = connect(MaybeWinsocket, (struct sockaddr *)&Addr, sizeof(Addr));
-  if (status == -1) {
-    EC = std::make_error_code(std::errc::connection_aborted);
-    return -1;
-  }
-#ifdef _WIN32
-  return _open_osfhandle(MaybeWinsocket, 0);
-#else
-  return MaybeWinsocket;
-#endif // _WIN32
-}
-
-static int ServerAccept(int FD) {
-  int AcceptFD;
-#ifdef _WIN32
-  SOCKET WinServerSock = _get_osfhandle(FD);
-  SOCKET WinAcceptSock = ::accept(WinServerSock, NULL, NULL);
-  AcceptFD = _open_osfhandle(WinAcceptSock, 0);
-#else
-  AcceptFD = ::accept(FD, NULL, NULL);
-#endif //_WIN32
-  return AcceptFD;
-}
-
-// Server
-// Call raw_fd_ostream with ShouldClose=false
-raw_socket_stream::raw_socket_stream(int SocketFD, StringRef SockPath,
-                                     std::error_code &EC)
-    : raw_fd_ostream(ServerAccept(SocketFD), true) {
-  SocketPath = SockPath;
-  ShouldUnlink = true;
-}
-
-// Client
-raw_socket_stream::raw_socket_stream(StringRef SockPath, std::error_code &EC)
-    : raw_fd_ostream(GetSocketFD(SockPath, EC), true, true,
-                     OStreamKind::OK_OStream) {
-  SocketPath = SockPath;
-  ShouldUnlink = false;
-}
-
-raw_socket_stream::~raw_socket_stream() {
-  if (ShouldUnlink) {
-    unlink(SocketPath.str().c_str());
-  }
-}
-
-Expected<std::string> raw_socket_stream::read_impl() {
-  const size_t BUFFER_SIZE = 4096;
-  std::vector<char> Buffer(BUFFER_SIZE);
-
-  int Socket = get_socket();
-  assert(Socket >= 0 && "Socket not found.");
-
-  ssize_t n;
-#ifdef _WIN32
-  SOCKET MaybeWinsocket = _get_osfhandle(Socket);
-#else
-  int MaybeWinsocket = Socket;
-#endif // _WIN32
-  n = ::read(MaybeWinsocket, Buffer.data(), Buffer.size());
-
-  if (n < 0) {
-    std::string Msg = "Buffer read error: " + std::string(strerror(errno));
-    return llvm::make_error<StringError>(Msg, inconvertibleErrorCode());
-  }
-
-  if (n == 0) {
-    return llvm::make_error<StringError>("EOF", inconvertibleErrorCode());
-  }
-  return std::string(Buffer.data());
-}
-
-*/
-
 //===----------------------------------------------------------------------===//
 //  raw_string_ostream
 //===----------------------------------------------------------------------===//

>From 738a4835ba2c30691f2141e889ce467603b573d7 Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Mon, 11 Dec 2023 14:55:07 -0800
Subject: [PATCH 12/14] Default for MaxBacklog

---
 llvm/include/llvm/Support/raw_ostream.h | 5 ++++-
 llvm/lib/Support/raw_ostream.cpp        | 4 ++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index ee00198d9e4dbc..0262bfdbb11073 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -16,6 +16,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/DataTypes.h"
+#include "llvm/Support/Threading.h"
 #include <cassert>
 #include <cstddef>
 #include <cstdint>
@@ -645,7 +646,9 @@ class ListeningSocket {
   ListeningSocket(int SocketFD, StringRef SocketPath);
 
 public:
-  static Expected<ListeningSocket> createUnix(StringRef SocketPath);
+  static Expected<ListeningSocket> createUnix(
+      StringRef SocketPath,
+      int MaxBacklog = llvm::hardware_concurrency().compute_thread_count());
   Expected<std::unique_ptr<raw_socket_stream>> accept();
   ListeningSocket(ListeningSocket &&LS);
   ~ListeningSocket();
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 04349c3beb547a..97851c228c9875 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -979,9 +979,9 @@ ListeningSocket::ListeningSocket(ListeningSocket &&LS)
   LS.FD = -1;
 }
 
-Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath) {
+Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath,
+                                                      int MaxBacklog) {
 
-  int MaxBacklog = 3;
 #ifdef _WIN32
   WSADATA WsaData = {0};
   if (WSAStartup(MAKEWORD(2, 2), &WsaData) != 0) {

>From 8f106fd5a0cfaf21b2d1e40ec6dde1771dbb39db Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Mon, 11 Dec 2023 16:38:18 -0800
Subject: [PATCH 13/14] WSABAlancer

---
 llvm/include/llvm/Support/raw_ostream.h | 16 +++++++++++
 llvm/lib/Support/raw_ostream.cpp        | 36 ++++++++++++-------------
 2 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 0262bfdbb11073..3f49918f7bfaa6 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -640,10 +640,23 @@ class raw_fd_stream : public raw_fd_ostream {
 /// A raw stream for sockets reading/writing
 
 class raw_socket_stream;
+
+// Make sure that calls to WSAStartup and WSACleanup are balanced.
+#ifdef _WIN32
+class WSABalancer {
+public:
+  WSABalancer();
+  ~WSABalancer();
+}
+#endif // _WIN32
+
 class ListeningSocket {
   int FD;
   std::string SocketPath;
   ListeningSocket(int SocketFD, StringRef SocketPath);
+#ifdef _WIN32
+  WSABalancer _;
+#endif // _WIN32
 
 public:
   static Expected<ListeningSocket> createUnix(
@@ -655,6 +668,9 @@ class ListeningSocket {
 };
 class raw_socket_stream : public raw_fd_stream {
   uint64_t current_pos() const override { return 0; }
+#ifdef _WIN32
+  WSABalancer _;
+#endif // _WIN32
 
 public:
   raw_socket_stream(int SocketFD);
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 97851c228c9875..ffb83889ec9afd 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -963,6 +963,20 @@ bool raw_fd_stream::classof(const raw_ostream *OS) {
 //  raw_socket_stream
 //===----------------------------------------------------------------------===//
 
+#ifdef _WIN32
+WSABalancer::WSABalancer() {
+  WSADATA WsaData = {0};
+  if (WSAStartup(MAKEWORD(2, 2), &WsaData) != 0) {
+    llvm::report_fatal_error("WSAStartup failed");
+  }
+}
+
+WSABalancer::~WSABalancer() {
+  WSACleanup();
+}
+
+#endif // _WIN32
+
 std::error_code getLastSocketErrorCode() {
 #ifdef _WIN32
   return std::error_code(::WSAGetLastError(), std::system_category());
@@ -983,13 +997,7 @@ Expected<ListeningSocket> ListeningSocket::createUnix(StringRef SocketPath,
                                                       int MaxBacklog) {
 
 #ifdef _WIN32
-  WSADATA WsaData = {0};
-  if (WSAStartup(MAKEWORD(2, 2), &WsaData) != 0) {
-    llvm::Error E = llvm::make_error<StringError>(getLastSocketErrorCode(),
-                                                  "WSAStartup failed");
-    WSACleanup();
-    return E;
-  }
+  WSABalancer _;
   SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
   if (MaybeWinsocket == INVALID_SOCKET) {
 #else
@@ -1045,22 +1053,11 @@ ListeningSocket::~ListeningSocket() {
     return;
   ::close(FD);
   unlink(SocketPath.c_str());
-#ifdef _WIN32
-  WSACleanup();
-#endif // _WIN32
 }
 
 // Expected?
 Expected<int> GetSocketFD(StringRef SocketPath) {
 #ifdef _WIN32
-  WSADATA WsaData = {0};
-  if (WSAStartup(MAKEWORD(2, 2), &WsaData) != 0) {
-    llvm::Error E = llvm::make_error<StringError>(getLastSocketErrorCode(),
-                                                  "WSAStartup failed");
-    WSACleanup();
-    return E;
-  }
-
   SOCKET MaybeWinsocket = socket(AF_UNIX, SOCK_STREAM, 0);
   if (MaybeWinsocket == INVALID_SOCKET) {
 #else
@@ -1093,6 +1090,9 @@ raw_socket_stream::raw_socket_stream(int SocketFD)
 
 Expected<std::unique_ptr<raw_socket_stream>>
 raw_socket_stream::createConnectedUnix(StringRef SocketPath) {
+#ifdef _WIN32
+  WSABalancer _;
+#endif // _WIN32
   Expected<int> FD = GetSocketFD(SocketPath);
   if (!FD)
     return FD.takeError();

>From 9a4e084c5fb21e5501803077d640b68effb10b23 Mon Sep 17 00:00:00 2001
From: Christian Riis <criis at apple.com>
Date: Mon, 11 Dec 2023 21:54:16 -0800
Subject: [PATCH 14/14] Format and missing semi

---
 llvm/include/llvm/Support/raw_ostream.h | 2 +-
 llvm/lib/Support/raw_ostream.cpp        | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 3f49918f7bfaa6..7c8d264afeff2f 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -647,7 +647,7 @@ class WSABalancer {
 public:
   WSABalancer();
   ~WSABalancer();
-}
+};
 #endif // _WIN32
 
 class ListeningSocket {
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index ffb83889ec9afd..0868b97010de06 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -971,9 +971,7 @@ WSABalancer::WSABalancer() {
   }
 }
 
-WSABalancer::~WSABalancer() {
-  WSACleanup();
-}
+WSABalancer::~WSABalancer() { WSACleanup(); }
 
 #endif // _WIN32