[Lldb-commits] [lldb] [lldb] Add a callback version of TCPSocket::Accept (PR #106955)

Pavel Labath via lldb-commits lldb-commits at lists.llvm.org
Tue Sep 3 00:31:25 PDT 2024


https://github.com/labath updated https://github.com/llvm/llvm-project/pull/106955

>From 50e1aeec73b5c5ad3571d040c2ecc04c8d7e4009 Mon Sep 17 00:00:00 2001
From: Pavel Labath <pavel at labath.sk>
Date: Mon, 2 Sep 2024 11:18:50 +0200
Subject: [PATCH 1/2] [lldb] Add a callback version of TCPSocket::Accept

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.
---
 lldb/include/lldb/Host/common/TCPSocket.h | 11 +++
 lldb/source/Host/common/TCPSocket.cpp     | 96 +++++++++++------------
 lldb/unittests/Host/SocketTest.cpp        | 35 +++++++++
 3 files changed, 91 insertions(+), 51 deletions(-)

diff --git a/lldb/include/lldb/Host/common/TCPSocket.h b/lldb/include/lldb/Host/common/TCPSocket.h
index b782c9e6096c44..78e80568e39967 100644
--- a/lldb/include/lldb/Host/common/TCPSocket.h
+++ b/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);
diff --git a/lldb/source/Host/common/TCPSocket.cpp b/lldb/source/Host/common/TCPSocket.cpp
index ea26d8433c370a..4699e34bff924c 100644
--- a/lldb/source/Host/common/TCPSocket.cpp
+++ b/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,60 @@ 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;
-  }
+    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);
+
+      const lldb_private::SocketAddress &AddrIn = m_listen_sockets[fd];
+      if (!AddrIn.IsAnyAddr() && AcceptAddr != AddrIn) {
+        CLOSE_SOCKET(sock);
+        LLDB_LOG(log, "rejecting incoming connection from {0} (expecting {1})",
+                 AcceptAddr.GetIPAddress(), AddrIn.GetIPAddress());
+      }
+      std::unique_ptr<TCPSocket> sock_up(new TCPSocket(sock, *this));
 
-  bool accept_connection = false;
-  std::unique_ptr<TCPSocket> accepted_socket;
-  // Loop until we are happy with our connection
-  while (!accept_connection) {
-    accept_loop.Run();
+      // 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;
-
-    lldb_private::SocketAddress &AddrIn = m_listen_sockets[listen_sock];
-    if (!AddrIn.IsAnyAddr() && AcceptAddr != AddrIn) {
-      if (sock != kInvalidSocketValue) {
-        CLOSE_SOCKET(sock);
-        sock = kInvalidSocketValue;
-      }
-      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));
+      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() {
diff --git a/lldb/unittests/Host/SocketTest.cpp b/lldb/unittests/Host/SocketTest.cpp
index 78e1e11df06562..43650a4e3b8d2c 100644
--- a/lldb/unittests/Host/SocketTest.cpp
+++ b/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,40 @@ TEST_P(SocketTest, TCPListen0ConnectAccept) {
                             &socket_b_up);
 }
 
+TEST_P(SocketTest, TCPMainLoopAccept) {
+  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;

>From eb48e73d5b77680733d795ccd35cf1e571231942 Mon Sep 17 00:00:00 2001
From: Pavel Labath <pavel at labath.sk>
Date: Tue, 3 Sep 2024 09:30:51 +0200
Subject: [PATCH 2/2] some (not all) fixes

---
 lldb/source/Host/common/TCPSocket.cpp | 5 ++++-
 lldb/unittests/Host/SocketTest.cpp    | 3 +++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/lldb/source/Host/common/TCPSocket.cpp b/lldb/source/Host/common/TCPSocket.cpp
index 4699e34bff924c..fc005814308d90 100644
--- a/lldb/source/Host/common/TCPSocket.cpp
+++ b/lldb/source/Host/common/TCPSocket.cpp
@@ -273,14 +273,17 @@ llvm::Expected<std::vector<MainLoopBase::ReadHandleUP>> TCPSocket::Accept(
       NativeSocket sock = AcceptSocket(fd, &AcceptAddr.sockaddr(), &sa_len,
                                        m_child_processes_inherit, error);
       Log *log = GetLog(LLDBLog::Host);
-      if (error.Fail())
+      if (error.Fail()) {
         LLDB_LOG(log, "AcceptSocket({0}): {1}", fd, error);
+        return;
+      }
 
       const lldb_private::SocketAddress &AddrIn = m_listen_sockets[fd];
       if (!AddrIn.IsAnyAddr() && AcceptAddr != AddrIn) {
         CLOSE_SOCKET(sock);
         LLDB_LOG(log, "rejecting incoming connection from {0} (expecting {1})",
                  AcceptAddr.GetIPAddress(), AddrIn.GetIPAddress());
+        return;
       }
       std::unique_ptr<TCPSocket> sock_up(new TCPSocket(sock, *this));
 
diff --git a/lldb/unittests/Host/SocketTest.cpp b/lldb/unittests/Host/SocketTest.cpp
index 43650a4e3b8d2c..3a356d11ba1a51 100644
--- a/lldb/unittests/Host/SocketTest.cpp
+++ b/lldb/unittests/Host/SocketTest.cpp
@@ -97,6 +97,9 @@ TEST_P(SocketTest, TCPListen0ConnectAccept) {
 }
 
 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);



More information about the lldb-commits mailing list