Fix Mock GPU Plugin disconnect issue

Author: qxy11

lldb/tools/lldb-server/Plugins/MockGPU/LLDBServerPluginMockGPU.cpp

@@ -18,7 +18,6 @@
 #include <sys/types.h>
 #include <sys/uio.h>
 #include <unistd.h>
-#include <thread>
 
 using namespace lldb;
 using namespace lldb_private;
@@ -95,41 +94,6 @@ bool LLDBServerPluginMockGPU::HandleEventFileDescriptorEvent(int fd) {
   return false;
 }
 
-void LLDBServerPluginMockGPU::AcceptAndMainLoopThread(
-    std::unique_ptr<TCPSocket> listen_socket_up) {
-  Log *log = GetLog(GDBRLog::Plugin);
-  LLDB_LOGF(log, "%s spawned", __PRETTY_FUNCTION__);
-  Socket *socket = nullptr;
-  Status error = listen_socket_up->Accept(std::chrono::seconds(30), socket);
-  // Scope for lock guard.
-  {
-    // Protect access to m_is_listening and m_is_connected.
-    std::lock_guard<std::mutex> guard(m_connect_mutex);
-    m_is_listening = false;
-    if (error.Fail()) {
-      LLDB_LOGF(log, "%s error returned from Accept(): %s", __PRETTY_FUNCTION__, 
-                error.AsCString());  
-      return;
-    }
-    m_is_connected = true;
-  }
-
-  LLDB_LOGF(log, "%s initializing connection", __PRETTY_FUNCTION__);
-  std::unique_ptr<Connection> connection_up(
-      new ConnectionFileDescriptor(std::unique_ptr<Socket>(socket)));
-  m_gdb_server->InitializeConnection(std::move(connection_up));
-  LLDB_LOGF(log, "%s running main loop", __PRETTY_FUNCTION__);
-  m_main_loop_status = m_main_loop.Run();
-  LLDB_LOGF(log, "%s main loop exited!", __PRETTY_FUNCTION__);
-  if (m_main_loop_status.Fail()) {
-    LLDB_LOGF(log, "%s main loop exited with an error: %s", __PRETTY_FUNCTION__,
-              m_main_loop_status.AsCString());
-
-  }
-  // Protect access to m_is_connected.
-  std::lock_guard<std::mutex> guard(m_connect_mutex);
-  m_is_connected = false;
-}
 
 std::optional<GPUPluginConnectionInfo> LLDBServerPluginMockGPU::CreateConnection() {
   std::lock_guard<std::mutex> guard(m_connect_mutex);
@@ -153,9 +117,20 @@ std::optional<GPUPluginConnectionInfo> LLDBServerPluginMockGPU::CreateConnection
     connection_info.connect_url = llvm::formatv("connect://localhost:{}", 
                                                 listen_port);
     LLDB_LOGF(log, "%s listening to %u", __PRETTY_FUNCTION__, listen_port);
-    std::thread t(&LLDBServerPluginMockGPU::AcceptAndMainLoopThread, this, 
-                  std::move(*sock));
-    t.detach();
+    
+    m_listen_socket = std::move(*sock);
+    m_is_connected = false;
+    auto res = m_listen_socket->Accept(m_main_loop, [this](std::unique_ptr<Socket> socket) {
+      std::unique_ptr<Connection> connection_up(new ConnectionFileDescriptor(std::move(socket)));
+      this->m_gdb_server->InitializeConnection(std::move(connection_up));
+      m_is_connected = true;
+    });
+    if (res) {
+      m_read_handles = std::move(*res);
+    } else {
+      LLDB_LOGF(log, "MockGPU failed to accept: %s", llvm::toString(res.takeError()).c_str());
+    }
+    
     return connection_info;
   } else {
     std::string error = llvm::toString(sock.takeError());

lldb/tools/lldb-server/Plugins/MockGPU/LLDBServerPluginMockGPU.h

@@ -10,7 +10,9 @@
 #define LLDB_TOOLS_LLDB_SERVER_LLDBSERVERPLUGINMOCKGPU_H
 
 #include "Plugins/Process/gdb-remote/LLDBServerPlugin.h"
+#include "lldb/Host/MainLoopBase.h"
 #include "lldb/Utility/Status.h"
+#include <mutex>
 
 // This is a mock GPU plugin that is used for testing the LLDBServerPlugin. It
 // should be run with the following code as the main binary:
@@ -61,7 +63,8 @@ namespace lldb_server {
 
 class LLDBServerPluginMockGPU : public LLDBServerPlugin {
 public:
-  LLDBServerPluginMockGPU(LLDBServerPlugin::GDBServer &native_process, MainLoop &main_loop);
+  LLDBServerPluginMockGPU(LLDBServerPlugin::GDBServer &native_process,
+                          MainLoop &main_loop);
   ~LLDBServerPluginMockGPU() override;
   llvm::StringRef GetPluginName() override;
   int GetEventFileDescriptorAtIndex(size_t idx) override;
@@ -74,11 +77,14 @@ class LLDBServerPluginMockGPU : public LLDBServerPlugin {
 private:
   std::optional<GPUPluginConnectionInfo> CreateConnection();
   void CloseFDs();
-  void AcceptAndMainLoopThread(std::unique_ptr<TCPSocket> listen_socket_up);
 
   // Used with a socketpair to get events on the native ptrace event queue.
   int m_fds[2] = {-1, -1};
   Status m_main_loop_status;
+
+  // Connection state tracking
+  std::unique_ptr<TCPSocket> m_listen_socket;
+  std::vector<MainLoopBase::ReadHandleUP> m_read_handles;
 };
 
 } // namespace lldb_server