[lldb] Add timeout argument to Socket::Accept

lldb/include/lldb/Host/Socket.h

@@ -13,6 +13,7 @@
 #include <string>
 
 #include "lldb/Host/MainLoopBase.h"
+#include "lldb/Utility/Timeout.h"
 #include "lldb/lldb-private.h"
 
 #include "lldb/Host/SocketAddress.h"
@@ -108,7 +109,7 @@ class Socket : public IOObject {
 
   // Accept a single connection and "return" it in the pointer argument. This
   // function blocks until the connection arrives.
-  virtual Status Accept(Socket *&socket);
+  virtual Status Accept(const Timeout<std::micro> &timeout, Socket *&socket);
 
   // Initialize a Tcp Socket object in listening mode.  listen and accept are
   // implemented separately because the caller may wish to manipulate or query

lldb/source/Host/common/Socket.cpp

@@ -460,7 +460,8 @@ NativeSocket Socket::CreateSocket(const int domain, const int type,
   return sock;
 }
 
-Status Socket::Accept(Socket *&socket) {
+Status Socket::Accept(const Timeout<std::micro> &timeout, Socket *&socket) {
+  socket = nullptr;
   MainLoop accept_loop;
   llvm::Expected<std::vector<MainLoopBase::ReadHandleUP>> expected_handles =
       Accept(accept_loop,
@@ -470,7 +471,15 @@ Status Socket::Accept(Socket *&socket) {
              });
   if (!expected_handles)
     return Status::FromError(expected_handles.takeError());
-  return accept_loop.Run();
+  if (timeout) {
+    accept_loop.AddCallback(
+        [](MainLoopBase &loop) { loop.RequestTermination(); }, *timeout);
+  }
+  if (Status status = accept_loop.Run(); status.Fail())
+    return status;
+  if (socket)
+    return Status();
+  return Status(std::make_error_code(std::errc::timed_out));
 }
 
 NativeSocket Socket::AcceptSocket(NativeSocket sockfd, struct sockaddr *addr,

lldb/source/Host/posix/ConnectionFileDescriptorPosix.cpp

@@ -543,7 +543,7 @@ lldb::ConnectionStatus ConnectionFileDescriptor::AcceptSocket(
 
   if (!error.Fail()) {
     post_listen_callback(*listening_socket);
-    error = listening_socket->Accept(accepted_socket);
+    error = listening_socket->Accept(/*timeout=*/std::nullopt, accepted_socket);
   }
 
   if (!error.Fail()) {

lldb/source/Plugins/Process/gdb-remote/GDBRemoteCommunication.cpp

@@ -1223,10 +1223,6 @@ GDBRemoteCommunication::ConnectLocally(GDBRemoteCommunication &client,
           listen_socket.Listen("localhost:0", backlog).ToError())
     return error;
 
-  Socket *accept_socket = nullptr;
-  std::future<Status> accept_status = std::async(
-      std::launch::async, [&] { return listen_socket.Accept(accept_socket); });
-
   llvm::SmallString<32> remote_addr;
   llvm::raw_svector_ostream(remote_addr)
       << "connect://localhost:" << listen_socket.GetLocalPortNumber();
@@ -1238,10 +1234,15 @@ GDBRemoteCommunication::ConnectLocally(GDBRemoteCommunication &client,
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "Unable to connect: %s", status.AsCString());
 
-  client.SetConnection(std::move(conn_up));
-  if (llvm::Error error = accept_status.get().ToError())
-    return error;
+  // The connection was already established above, so a short timeout is
+  // sufficient.
+  Socket *accept_socket = nullptr;
+  if (Status accept_status =
+          listen_socket.Accept(std::chrono::seconds(1), accept_socket);
+      accept_status.Fail())
+    return accept_status.takeError();
 
+  client.SetConnection(std::move(conn_up));
   server.SetConnection(
       std::make_unique<ConnectionFileDescriptor>(accept_socket));
   return llvm::Error::success();

lldb/unittests/Host/MainLoopTest.cpp

@@ -42,7 +42,8 @@ class MainLoopTest : public testing::Test {
         llvm::formatv("localhost:{0}", listen_socket_up->GetLocalPortNumber())
             .str());
     ASSERT_TRUE(error.Success());
-    ASSERT_TRUE(listen_socket_up->Accept(accept_socket).Success());
+    ASSERT_TRUE(listen_socket_up->Accept(std::chrono::seconds(1), accept_socket)
+                    .Success());
 
     callback_count = 0;
     socketpair[0] = std::move(connect_socket_up);

lldb/unittests/Host/SocketTest.cpp

@@ -12,7 +12,9 @@
 #include "lldb/Host/MainLoop.h"
 #include "lldb/Utility/UriParser.h"
 #include "llvm/Testing/Support/Error.h"
+#include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <chrono>
 
 using namespace lldb_private;
 
@@ -133,6 +135,27 @@ TEST_P(SocketTest, TCPListen0ConnectAccept) {
                             &socket_b_up);
 }
 
+TEST_P(SocketTest, TCPAcceptTimeout) {
+  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());
+
+  Socket *socket;
+  ASSERT_THAT_ERROR(
+      listen_socket_up->Accept(std::chrono::milliseconds(10), socket)
+          .takeError(),
+      llvm::Failed<llvm::ErrorInfoBase>(
+          testing::Property(&llvm::ErrorInfoBase::convertToErrorCode,
+                            std::make_error_code(std::errc::timed_out))));
+}
+
 TEST_P(SocketTest, TCPMainLoopAccept) {
   if (!HostSupportsProtocol())
     return;

lldb/unittests/TestingSupport/Host/SocketTestUtilities.cpp

@@ -19,11 +19,6 @@
 
 using namespace lldb_private;
 
-static void AcceptThread(Socket *listen_socket, bool child_processes_inherit,
-                         Socket **accept_socket, Status *error) {
-  *error = listen_socket->Accept(*accept_socket);
-}
-
 template <typename SocketType>
 void lldb_private::CreateConnectedSockets(
     llvm::StringRef listen_remote_address,
@@ -38,12 +33,6 @@ void lldb_private::CreateConnectedSockets(
   ASSERT_THAT_ERROR(error.ToError(), llvm::Succeeded());
   ASSERT_TRUE(listen_socket_up->IsValid());
 
-  Status accept_error;
-  Socket *accept_socket;
-  std::thread accept_thread(AcceptThread, listen_socket_up.get(),
-                            child_processes_inherit, &accept_socket,
-                            &accept_error);
-
   std::string connect_remote_address = get_connect_addr(*listen_socket_up);
   std::unique_ptr<SocketType> connect_socket_up(
       new SocketType(true, child_processes_inherit));
@@ -55,9 +44,13 @@ void lldb_private::CreateConnectedSockets(
   a_up->swap(connect_socket_up);
   ASSERT_TRUE((*a_up)->IsValid());
 
-  accept_thread.join();
+  Socket *accept_socket;
+  ASSERT_THAT_ERROR(
+      listen_socket_up->Accept(std::chrono::seconds(1), accept_socket)
+          .takeError(),
+      llvm::Succeeded());
+
   b_up->reset(static_cast<SocketType *>(accept_socket));
-  ASSERT_THAT_ERROR(accept_error.ToError(), llvm::Succeeded());
   ASSERT_NE(nullptr, b_up->get());
   ASSERT_TRUE((*b_up)->IsValid());
 

lldb/unittests/tools/lldb-server/tests/TestClient.cpp

@@ -26,9 +26,13 @@ using namespace lldb_private;
 using namespace llvm;
 using namespace llgs_tests;
 
+static std::chrono::seconds GetDefaultTimeout() {
+  return std::chrono::seconds{10};
+}
+
 TestClient::TestClient(std::unique_ptr<Connection> Conn) {
   SetConnection(std::move(Conn));
-  SetPacketTimeout(std::chrono::seconds(10));
+  SetPacketTimeout(GetDefaultTimeout());
 }
 
 TestClient::~TestClient() {
@@ -117,7 +121,10 @@ TestClient::launchCustom(StringRef Log, bool disable_stdio,
     return status.ToError();
 
   Socket *accept_socket;
-  listen_socket.Accept(accept_socket);
+  if (llvm::Error E =
+          listen_socket.Accept(2 * GetDefaultTimeout(), accept_socket)
+              .takeError())
+    return E;
   auto Conn = std::make_unique<ConnectionFileDescriptor>(accept_socket);
   auto Client = std::unique_ptr<TestClient>(new TestClient(std::move(Conn)));