[lldb] Add a callback version of TCPSocket::Accept

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;