[lldb] Add a callback version of TCPSocket::Accept (#106955)

The existing function already used the MainLoop class, which allows one
to wait on multiple events at once. It needed to do this in order to
wait for v4 and v6 connections simultaneously. However, since it was
creating its own instance of MainLoop, this meant that it was impossible
to multiplex these sockets with anything else.

This patch simply adds a version of this function which uses an
externally provided main loop instance, which allows the caller to add
any events it deems necessary. The previous function becomes a very thin
wrapper over the new one.

Author: labath

lldb/include/lldb/Host/common/TCPSocket.h

@@ -9,6 +9,7 @@
 #ifndef LLDB_HOST_COMMON_TCPSOCKET_H
 #define LLDB_HOST_COMMON_TCPSOCKET_H
 
+#include "lldb/Host/MainLoopBase.h"
 #include "lldb/Host/Socket.h"
 #include "lldb/Host/SocketAddress.h"
 #include <map>
@@ -40,6 +41,16 @@ class TCPSocket : public Socket {
 
   Status Connect(llvm::StringRef name) override;
   Status Listen(llvm::StringRef name, int backlog) override;
+
+  // Use the provided main loop instance to accept new connections. The callback
+  // will be called (from MainLoop::Run) for each new connection. This function
+  // does not block.
+  llvm::Expected<std::vector<MainLoopBase::ReadHandleUP>>
+  Accept(MainLoopBase &loop,
+         std::function<void(std::unique_ptr<TCPSocket> socket)> sock_cb);
+
+  // Accept a single connection and "return" it in the pointer argument. This
+  // function blocks until the connection arrives.
   Status Accept(Socket *&conn_socket) override;
 
   Status CreateSocket(int domain);

lldb/source/Host/common/TCPSocket.cpp

@@ -19,6 +19,7 @@
 
 #include "llvm/Config/llvm-config.h"
 #include "llvm/Support/Errno.h"
+#include "llvm/Support/Error.h"
 #include "llvm/Support/WindowsError.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -254,67 +255,63 @@ void TCPSocket::CloseListenSockets() {
   m_listen_sockets.clear();
 }
 
-Status TCPSocket::Accept(Socket *&conn_socket) {
-  Status error;
-  if (m_listen_sockets.size() == 0) {
-    error = Status::FromErrorString("No open listening sockets!");
-    return error;
-  }
+llvm::Expected<std::vector<MainLoopBase::ReadHandleUP>> TCPSocket::Accept(
+    MainLoopBase &loop,
+    std::function<void(std::unique_ptr<TCPSocket> socket)> sock_cb) {
+  if (m_listen_sockets.size() == 0)
+    return llvm::createStringError("No open listening sockets!");
 
-  NativeSocket sock = kInvalidSocketValue;
-  NativeSocket listen_sock = kInvalidSocketValue;
-  lldb_private::SocketAddress AcceptAddr;
-  MainLoop accept_loop;
   std::vector<MainLoopBase::ReadHandleUP> handles;
   for (auto socket : m_listen_sockets) {
     auto fd = socket.first;
-    auto inherit = this->m_child_processes_inherit;
-    auto io_sp = IOObjectSP(new TCPSocket(socket.first, false, inherit));
-    handles.emplace_back(accept_loop.RegisterReadObject(
-        io_sp, [fd, inherit, &sock, &AcceptAddr, &error,
-                        &listen_sock](MainLoopBase &loop) {
-          socklen_t sa_len = AcceptAddr.GetMaxLength();
-          sock = AcceptSocket(fd, &AcceptAddr.sockaddr(), &sa_len, inherit,
-                              error);
-          listen_sock = fd;
-          loop.RequestTermination();
-        }, error));
-    if (error.Fail())
-      return error;
-  }
-
-  bool accept_connection = false;
-  std::unique_ptr<TCPSocket> accepted_socket;
-  // Loop until we are happy with our connection
-  while (!accept_connection) {
-    accept_loop.Run();
-
-    if (error.Fail())
-        return error;
+    auto io_sp =
+        std::make_shared<TCPSocket>(fd, false, this->m_child_processes_inherit);
+    auto cb = [this, fd, sock_cb](MainLoopBase &loop) {
+      lldb_private::SocketAddress AcceptAddr;
+      socklen_t sa_len = AcceptAddr.GetMaxLength();
+      Status error;
+      NativeSocket sock = AcceptSocket(fd, &AcceptAddr.sockaddr(), &sa_len,
+                                       m_child_processes_inherit, error);
+      Log *log = GetLog(LLDBLog::Host);
+      if (error.Fail()) {
+        LLDB_LOG(log, "AcceptSocket({0}): {1}", fd, error);
+        return;
+      }
 
-    lldb_private::SocketAddress &AddrIn = m_listen_sockets[listen_sock];
-    if (!AddrIn.IsAnyAddr() && AcceptAddr != AddrIn) {
-      if (sock != kInvalidSocketValue) {
+      const lldb_private::SocketAddress &AddrIn = m_listen_sockets[fd];
+      if (!AddrIn.IsAnyAddr() && AcceptAddr != AddrIn) {
         CLOSE_SOCKET(sock);
-        sock = kInvalidSocketValue;
+        LLDB_LOG(log, "rejecting incoming connection from {0} (expecting {1})",
+                 AcceptAddr.GetIPAddress(), AddrIn.GetIPAddress());
+        return;
       }
-      llvm::errs() << llvm::formatv(
-          "error: rejecting incoming connection from {0} (expecting {1})",
-          AcceptAddr.GetIPAddress(), AddrIn.GetIPAddress());
-      continue;
-    }
-    accept_connection = true;
-    accepted_socket.reset(new TCPSocket(sock, *this));
+      std::unique_ptr<TCPSocket> sock_up(new TCPSocket(sock, *this));
+
+      // Keep our TCP packets coming without any delays.
+      sock_up->SetOptionNoDelay();
+
+      sock_cb(std::move(sock_up));
+    };
+    Status error;
+    handles.emplace_back(loop.RegisterReadObject(io_sp, cb, error));
+    if (error.Fail())
+      return error.ToError();
   }
 
-  if (!accepted_socket)
-    return error;
+  return handles;
+}
 
-  // Keep our TCP packets coming without any delays.
-  accepted_socket->SetOptionNoDelay();
-  error.Clear();
-  conn_socket = accepted_socket.release();
-  return error;
+Status TCPSocket::Accept(Socket *&conn_socket) {
+  MainLoop accept_loop;
+  llvm::Expected<std::vector<MainLoopBase::ReadHandleUP>> expected_handles =
+      Accept(accept_loop,
+             [&accept_loop, &conn_socket](std::unique_ptr<TCPSocket> sock) {
+               conn_socket = sock.release();
+               accept_loop.RequestTermination();
+             });
+  if (!expected_handles)
+    return Status(expected_handles.takeError());
+  return accept_loop.Run();
 }
 
 int TCPSocket::SetOptionNoDelay() {

lldb/unittests/Host/SocketTest.cpp

@@ -9,6 +9,7 @@
 #include "TestingSupport/Host/SocketTestUtilities.h"
 #include "TestingSupport/SubsystemRAII.h"
 #include "lldb/Host/Config.h"
+#include "lldb/Host/MainLoop.h"
 #include "lldb/Utility/UriParser.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
@@ -95,6 +96,43 @@ TEST_P(SocketTest, TCPListen0ConnectAccept) {
                             &socket_b_up);
 }
 
+TEST_P(SocketTest, TCPMainLoopAccept) {
+  if (!HostSupportsProtocol())
+    return;
+
+  const bool child_processes_inherit = false;
+  auto listen_socket_up =
+      std::make_unique<TCPSocket>(true, child_processes_inherit);
+  Status error = listen_socket_up->Listen(
+      llvm::formatv("[{0}]:0", GetParam().localhost_ip).str(), 5);
+  ASSERT_THAT_ERROR(error.ToError(), llvm::Succeeded());
+  ASSERT_TRUE(listen_socket_up->IsValid());
+
+  MainLoop loop;
+  std::unique_ptr<TCPSocket> accepted_socket_up;
+  auto expected_handles = listen_socket_up->Accept(
+      loop, [&accepted_socket_up, &loop](std::unique_ptr<TCPSocket> sock_up) {
+        accepted_socket_up = std::move(sock_up);
+        loop.RequestTermination();
+      });
+  ASSERT_THAT_EXPECTED(expected_handles, llvm::Succeeded());
+
+  std::unique_ptr<TCPSocket> connect_socket_up(
+      new TCPSocket(true, child_processes_inherit));
+  ASSERT_THAT_ERROR(
+      connect_socket_up
+          ->Connect(llvm::formatv("[{0}]:{1}", GetParam().localhost_ip,
+                                  listen_socket_up->GetLocalPortNumber())
+                        .str())
+          .ToError(),
+      llvm::Succeeded());
+  ASSERT_TRUE(connect_socket_up->IsValid());
+
+  loop.Run();
+  ASSERT_TRUE(accepted_socket_up);
+  ASSERT_TRUE(accepted_socket_up->IsValid());
+}
+
 TEST_P(SocketTest, TCPGetAddress) {
   std::unique_ptr<TCPSocket> socket_a_up;
   std::unique_ptr<TCPSocket> socket_b_up;