[lldb] Add timeout argument to Socket::Accept (#117691)

Allows us to stop waiting for a connection if it doesn't come in a
certain amount of time. Right now, I'm keeping the status quo (infitnite
wait) in the "production" code, but using smaller (finite) values in
tests. (A lot of these tests create "loopback" connections, where a
really short wait is sufficient: on linux at least even a poll (0s wait)
is sufficient if the other end has connect()ed already, but this doesn't
seem to be the case on Windows, so I'm using a 1s wait in these cases).

Author: labath

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)));