[lldb] Fix race condition in Process::WaitForProcessToStop()

lldb/include/lldb/Utility/Listener.h

@@ -53,6 +53,13 @@ class Listener : public std::enable_shared_from_this<Listener> {
 
   void AddEvent(lldb::EventSP &event);
 
+  /// Transfers all events matching the specified broadcaster and type to the
+  /// destination listener. This can be useful when setting up a hijack listener,
+  /// to ensure that no relevant events are missed.
+  void MoveEvents(lldb::ListenerSP destination,
+                  Broadcaster *broadcaster, // nullptr for any broadcaster
+                  uint32_t event_type_mask);
+
   void Clear();
 
   const char *GetName() { return m_name.c_str(); }

lldb/source/Utility/Broadcaster.cpp

@@ -335,6 +335,18 @@ bool Broadcaster::BroadcasterImpl::HijackBroadcaster(
       "{0} Broadcaster(\"{1}\")::HijackBroadcaster (listener(\"{2}\")={3})",
       static_cast<void *>(this), GetBroadcasterName(),
       listener_sp->m_name.c_str(), static_cast<void *>(listener_sp.get()));
+
+  // Move pending events from the previous listener to the hijack listener.
+  // This ensures that no relevant event queued before the transition is missed
+  // by the hijack listener.
+  ListenerSP prev_listener;
+  if (!m_hijacking_listeners.empty())
+    prev_listener = m_hijacking_listeners.back();
+  else if (m_primary_listener_sp)
+    prev_listener = m_primary_listener_sp;
+  if (prev_listener && listener_sp)
+    prev_listener->MoveEvents(listener_sp, &m_broadcaster, event_mask);
+
   m_hijacking_listeners.push_back(listener_sp);
   m_hijacking_masks.push_back(event_mask);
   return true;
@@ -367,6 +379,19 @@ void Broadcaster::BroadcasterImpl::RestoreBroadcaster() {
              static_cast<void *>(this), GetBroadcasterName(),
              listener_sp->m_name.c_str(),
              static_cast<void *>(listener_sp.get()));
+
+    // Move any remaining events from the hijack listener back to
+    // the previous listener. This ensures that no events are dropped when
+    // restoring the original listener.
+    ListenerSP prev_listener;
+    if (m_hijacking_listeners.size() > 1)
+      prev_listener = m_hijacking_listeners[m_hijacking_listeners.size() - 2];
+    else if (m_primary_listener_sp)
+      prev_listener = m_primary_listener_sp;
+    if (listener_sp && prev_listener && !m_hijacking_masks.empty())
+      listener_sp->MoveEvents(prev_listener, &m_broadcaster,
+                              m_hijacking_masks.back());
+
     m_hijacking_listeners.pop_back();
   }
   if (!m_hijacking_masks.empty())

lldb/source/Utility/Listener.cpp

@@ -176,6 +176,33 @@ void Listener::AddEvent(EventSP &event_sp) {
   m_events_condition.notify_all();
 }
 
+void Listener::MoveEvents(
+    ListenerSP destination,
+    Broadcaster *broadcaster, // nullptr for any broadcaster
+    uint32_t event_type_mask) {
+  Log *log = GetLog(LLDBLog::Events);
+
+  std::lock_guard<std::mutex> guard(m_events_mutex);
+  auto pos = m_events.begin();
+  while (pos != m_events.end()) {
+    EventSP &event_sp = *pos;
+    if (event_sp &&
+        ((broadcaster == nullptr) || event_sp->BroadcasterIs(broadcaster)) &&
+        (event_type_mask == 0 || event_type_mask & event_sp->GetType())) {
+      LLDB_LOGF(
+          log, "%p Listener('%s')::MoveEvents moving event %p to %p('%s')",
+          static_cast<void *>(this), m_name.c_str(),
+          static_cast<void *>(event_sp.get()),
+          static_cast<void *>(destination.get()), destination->GetName());
+
+      destination->AddEvent(event_sp);
+      pos = m_events.erase(pos);
+    } else {
+      ++pos;
+    }
+  }
+}
+
 bool Listener::FindNextEventInternal(
     std::unique_lock<std::mutex> &lock,
     Broadcaster *broadcaster, // nullptr for any broadcaster

lldb/unittests/Utility/ListenerTest.cpp

@@ -150,3 +150,84 @@ TEST(ListenerTest, StartStopListeningForEventSpec) {
   ASSERT_EQ(event_sp->GetBroadcaster(), &broadcaster1);
   ASSERT_FALSE(listener_sp->GetEvent(event_sp, std::chrono::seconds(0)));
 }
+
+TEST(ListenerTest, MoveEventsOnHijackAndRestore) {
+  Broadcaster broadcaster(nullptr, "test-broadcaster");
+  const uint32_t event_mask = 1;
+  EventSP event_sp;
+
+  // Create the original listener and start listening.
+  ListenerSP original_listener = Listener::MakeListener("original-listener");
+  ASSERT_EQ(event_mask, original_listener->StartListeningForEvents(&broadcaster,
+                                                                   event_mask));
+  broadcaster.SetPrimaryListener(original_listener);
+
+  // Queue two events to original listener, but do not consume them yet.
+  broadcaster.BroadcastEvent(event_mask, nullptr);
+  broadcaster.BroadcastEvent(event_mask, nullptr);
+
+  // Hijack.
+  ListenerSP hijack_listener = Listener::MakeListener("hijack-listener");
+  broadcaster.HijackBroadcaster(hijack_listener, event_mask);
+
+  // The events should have been moved to the hijack listener.
+  EXPECT_FALSE(original_listener->GetEvent(event_sp, std::chrono::seconds(0)));
+  EXPECT_TRUE(hijack_listener->GetEvent(event_sp, std::chrono::seconds(0)));
+  EXPECT_TRUE(hijack_listener->GetEvent(event_sp, std::chrono::seconds(0)));
+
+  // Queue two events while hijacked, but do not consume them yet.
+  broadcaster.BroadcastEvent(event_mask, nullptr);
+  broadcaster.BroadcastEvent(event_mask, nullptr);
+
+  // Restore the original listener.
+  broadcaster.RestoreBroadcaster();
+
+  // The events queued while hijacked should have been moved back to the
+  // original listener.
+  EXPECT_FALSE(hijack_listener->GetEvent(event_sp, std::chrono::seconds(0)));
+  EXPECT_TRUE(original_listener->GetEvent(event_sp, std::chrono::seconds(0)));
+  EXPECT_TRUE(original_listener->GetEvent(event_sp, std::chrono::seconds(0)));
+}
+
+TEST(ListenerTest, MoveEventsBetweenHijackers) {
+  Broadcaster broadcaster(nullptr, "test-broadcaster");
+  const uint32_t event_mask = 1;
+  EventSP event_sp;
+
+  // Create the original listener and start listening.
+  ListenerSP original_listener = Listener::MakeListener("original-listener");
+  ASSERT_EQ(event_mask, original_listener->StartListeningForEvents(&broadcaster,
+                                                                   event_mask));
+  broadcaster.SetPrimaryListener(original_listener);
+
+  // First hijack.
+  ListenerSP hijack_listener1 = Listener::MakeListener("hijack-listener1");
+  broadcaster.HijackBroadcaster(hijack_listener1, event_mask);
+
+  // Queue two events while hijacked, but do not consume
+  // them yet.
+  broadcaster.BroadcastEvent(event_mask, nullptr);
+  broadcaster.BroadcastEvent(event_mask, nullptr);
+
+  // Second hijack.
+  ListenerSP hijack_listener2 = Listener::MakeListener("hijack-listener2");
+  broadcaster.HijackBroadcaster(hijack_listener2, event_mask);
+
+  // The second hijacker should have the events now.
+  EXPECT_FALSE(hijack_listener1->GetEvent(event_sp, std::chrono::seconds(0)));
+  EXPECT_TRUE(hijack_listener2->GetEvent(event_sp, std::chrono::seconds(0)));
+  EXPECT_TRUE(hijack_listener2->GetEvent(event_sp, std::chrono::seconds(0)));
+
+  // Queue two events while hijacked with second hijacker, but do not consume
+  // them yet.
+  broadcaster.BroadcastEvent(event_mask, nullptr);
+  broadcaster.BroadcastEvent(event_mask, nullptr);
+
+  // Restore the previous hijacker.
+  broadcaster.RestoreBroadcaster();
+
+  // The first hijacker should now have the events.
+  EXPECT_FALSE(hijack_listener2->GetEvent(event_sp, std::chrono::seconds(0)));
+  EXPECT_TRUE(hijack_listener1->GetEvent(event_sp, std::chrono::seconds(0)));
+  EXPECT_TRUE(hijack_listener1->GetEvent(event_sp, std::chrono::seconds(0)));
+}