rm unneeded function, better cleanup

stephen-derosa · stephen-derosa · commit 90788da8c0cf · 2026-07-06T15:37:36.000-06:00
diff --git a/include/livekit/subscription_thread_dispatcher.h b/include/livekit/subscription_thread_dispatcher.h
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <cstdint>
 #include <functional>
 #include <memory>
@@ -259,6 +260,9 @@ class LIVEKIT_API SubscriptionThreadDispatcher {
     std::shared_ptr<AudioStream> audio_stream;
     std::shared_ptr<VideoStream> video_stream;
     std::thread thread;
+    /// SID of the subscribed track backing this reader, used to skip redundant
+    /// reader restarts when the same publication is re-subscribed.
+    std::string track_sid;
   };
 
   /// Compound lookup key for a remote participant identity and data track name.
@@ -289,6 +293,9 @@ class LIVEKIT_API SubscriptionThreadDispatcher {
   /// Active read-side resources for one data track stream subscription.
   struct ActiveDataReader {
     std::shared_ptr<RemoteDataTrack> remote_track;
+    /// Set true when this reader is being replaced or torn down so the reader
+    /// thread can abort a subscription that is still in flight.
+    std::atomic<bool> cancelled{false};
     std::mutex sub_mutex;
     std::shared_ptr<DataTrackStream> stream; // guarded by sub_mutex
     std::thread thread;
@@ -333,18 +340,21 @@ class LIVEKIT_API SubscriptionThreadDispatcher {
                                      const RegisteredVideoCallback& callback);
 
   /// Extract and close the data reader for a given callback ID, returning its
-  /// thread.  Must be called with @ref lock_ held.
+  /// thread.  Marks the reader cancelled so a subscription still in flight is
+  /// aborted.  Must be called with @ref lock_ held.
   std::thread extractDataReaderThreadLocked(DataFrameCallbackId id);
 
-  /// Extract and close the data reader for a given (participant, track_name)
-  /// key, returning its thread.  Must be called with @ref lock_ held.
-  std::thread extractDataReaderThreadLocked(const DataCallbackKey& key);
-
   /// Start a data reader thread for the given callback ID, key, and track.
   /// Must be called with @ref lock_ held.
   std::thread startDataReaderLocked(DataFrameCallbackId id, const DataCallbackKey& key,
                                     const std::shared_ptr<RemoteDataTrack>& track, const DataFrameCallback& cb);
 
+  /// Remove @p reader from @ref active_data_readers_ if the slot for @p id
+  /// still refers to it.  Called by the reader thread itself when it exits
+  /// after a failed or cancelled subscription so it does not leave a stale
+  /// entry behind.  Acquires @ref lock_.
+  void eraseDataReaderIfCurrent(DataFrameCallbackId id, const std::shared_ptr<ActiveDataReader>& reader);
+
   /// Protects callback registration maps and active reader state.
   mutable std::mutex lock_;
 
diff --git a/src/subscription_thread_dispatcher.cpp b/src/subscription_thread_dispatcher.cpp
@@ -259,6 +259,9 @@ void SubscriptionThreadDispatcher::handleDataTrackUnpublished(const std::string&
     for (auto it = active_data_readers_.begin(); it != active_data_readers_.end();) {
       auto& reader = it->second;
       if (reader->remote_track && reader->remote_track->info().sid == sid) {
+        // Mark cancelled before closing so a subscription still in flight aborts
+        // even if the stream has not been created yet.
+        reader->cancelled = true;
         {
           const std::scoped_lock<std::mutex> sub_guard(reader->sub_mutex);
           if (reader->stream) {
@@ -312,6 +315,10 @@ void SubscriptionThreadDispatcher::stopAll() {
     video_callbacks_.clear();
 
     for (auto& [id, reader] : active_data_readers_) {
+      // Mark cancelled before closing so a subscription still in flight aborts
+      // even if the stream has not been created yet, avoiding a join that never
+      // completes during teardown.
+      reader->cancelled = true;
       {
         const std::scoped_lock<std::mutex> sub_guard(reader->sub_mutex);
         if (reader->stream) {
@@ -397,6 +404,16 @@ std::thread SubscriptionThreadDispatcher::startAudioReaderLocked(const CallbackK
                                                                  const AudioFrameCallback& cb,
                                                                  const AudioStream::Options& opts) {
   LK_LOG_DEBUG("Starting audio reader for participant={} track_name={}", key.participant_identity, key.track_name);
+
+  auto existing = active_readers_.find(key);
+  if (existing != active_readers_.end() && existing->second.track_sid == track->sid()) {
+    LK_LOG_DEBUG(
+        "Skipping audio reader start for participant={} track_name={} because a "
+        "reader for sid={} is already active",
+        key.participant_identity, key.track_name, track->sid());
+    return {};
+  }
+
   auto old_thread = extractReaderThreadLocked(key);
 
   if (static_cast<int>(active_readers_.size()) >= kMaxActiveReaders) {
@@ -415,6 +432,7 @@ std::thread SubscriptionThreadDispatcher::startAudioReaderLocked(const CallbackK
 
   ActiveReader reader;
   reader.audio_stream = stream;
+  reader.track_sid = track->sid();
   const std::string participant_identity = key.participant_identity;
   const std::string track_name = key.track_name;
   // NOLINTBEGIN(bugprone-lambda-function-name,bugprone-exception-escape)
@@ -454,6 +472,16 @@ std::thread SubscriptionThreadDispatcher::startVideoReaderLocked(const CallbackK
                                                                  const std::shared_ptr<Track>& track,
                                                                  const RegisteredVideoCallback& callback) {
   LK_LOG_DEBUG("Starting video reader for participant={} track_name={}", key.participant_identity, key.track_name);
+
+  auto existing = active_readers_.find(key);
+  if (existing != active_readers_.end() && existing->second.track_sid == track->sid()) {
+    LK_LOG_DEBUG(
+        "Skipping video reader start for participant={} track_name={} because a "
+        "reader for sid={} is already active",
+        key.participant_identity, key.track_name, track->sid());
+    return {};
+  }
+
   auto old_thread = extractReaderThreadLocked(key);
 
   if (static_cast<int>(active_readers_.size()) >= kMaxActiveReaders) {
@@ -472,6 +500,7 @@ std::thread SubscriptionThreadDispatcher::startVideoReaderLocked(const CallbackK
 
   ActiveReader reader;
   reader.video_stream = stream;
+  reader.track_sid = track->sid();
   auto legacy_cb = callback.legacy_callback;
   auto event_cb = callback.event_callback;
   const std::string participant_identity = key.participant_identity;
@@ -522,6 +551,9 @@ std::thread SubscriptionThreadDispatcher::extractDataReaderThreadLocked(DataFram
   }
   auto reader = std::move(it->second);
   active_data_readers_.erase(it);
+  // Mark cancelled before closing so a subscription still in flight aborts even
+  // if the stream has not been created yet.
+  reader->cancelled = true;
   {
     const std::scoped_lock<std::mutex> guard(reader->sub_mutex);
     if (reader->stream) {
@@ -531,28 +563,36 @@ std::thread SubscriptionThreadDispatcher::extractDataReaderThreadLocked(DataFram
   return std::move(reader->thread);
 }
 
-std::thread SubscriptionThreadDispatcher::extractDataReaderThreadLocked(const DataCallbackKey& key) {
-  for (auto it = active_data_readers_.begin(); it != active_data_readers_.end(); ++it) {
-    if (it->second && it->second->remote_track &&
-        it->second->remote_track->publisherIdentity() == key.participant_identity &&
-        it->second->remote_track->info().name == key.track_name) {
-      auto reader = std::move(it->second);
-      active_data_readers_.erase(it);
-      {
-        const std::scoped_lock<std::mutex> guard(reader->sub_mutex);
-        if (reader->stream) {
-          reader->stream->close();
-        }
-      }
-      return std::move(reader->thread);
-    }
+void SubscriptionThreadDispatcher::eraseDataReaderIfCurrent(DataFrameCallbackId id,
+                                                            const std::shared_ptr<ActiveDataReader>& reader) {
+  const std::scoped_lock<std::mutex> lock(lock_);
+  auto it = active_data_readers_.find(id);
+  if (it == active_data_readers_.end() || it->second != reader) {
+    // The slot was already extracted or replaced; the owner joins that thread.
+    return;
   }
-  return {};
+  // Detach so the reader can be destroyed by its own thread's final shared_ptr
+  // release without tripping std::thread's joinable-at-destruction check. The
+  // thread is already returning when this runs, so nothing is left to join.
+  if (reader->thread.joinable()) {
+    reader->thread.detach();
+  }
+  active_data_readers_.erase(it);
 }
 
 std::thread SubscriptionThreadDispatcher::startDataReaderLocked(DataFrameCallbackId id, const DataCallbackKey& key,
                                                                 const std::shared_ptr<RemoteDataTrack>& track,
                                                                 const DataFrameCallback& cb) {
+  auto existing = active_data_readers_.find(id);
+  if (existing != active_data_readers_.end() && existing->second->remote_track &&
+      existing->second->remote_track->info().sid == track->info().sid) {
+    LK_LOG_DEBUG(
+        "Skipping data reader start for \"{}\" track=\"{}\" because a reader for "
+        "sid={} is already active",
+        key.participant_identity, key.track_name, track->info().sid);
+    return {};
+  }
+
   auto old_thread = extractDataReaderThreadLocked(id);
 
   const int total_active = static_cast<int>(active_readers_.size()) + static_cast<int>(active_data_readers_.size());
@@ -571,7 +611,7 @@ std::thread SubscriptionThreadDispatcher::startDataReaderLocked(DataFrameCallbac
   auto identity = key.participant_identity;
   auto track_name = key.track_name;
   // NOLINTBEGIN(bugprone-lambda-function-name)
-  reader->thread = std::thread([reader, track, cb, identity, track_name]() {
+  reader->thread = std::thread([this, id, reader, track, cb, identity, track_name]() {
     LK_LOG_INFO("Data reader thread: subscribing to \"{}\" track=\"{}\"", identity, track_name);
     std::shared_ptr<DataTrackStream> stream;
     auto subscribe_result = track->subscribe();
@@ -581,13 +621,21 @@ std::thread SubscriptionThreadDispatcher::startDataReaderLocked(DataFrameCallbac
           "Failed to subscribe to data track \"{}\" from \"{}\": code={} "
           "message={}",
           track_name, identity, static_cast<std::uint32_t>(error.code), error.message);
+      eraseDataReaderIfCurrent(id, reader);
       return;
     }
     stream = subscribe_result.value();
     LK_LOG_INFO("Data reader thread: subscribed to \"{}\" track=\"{}\"", identity, track_name);
 
     {
       const std::scoped_lock<std::mutex> guard(reader->sub_mutex);
+      // A replacement or teardown may have cancelled this reader while the
+      // subscribe was in flight. Close the fresh stream and bail so we do not
+      // leave a second live subscription behind.
+      if (reader->cancelled.load()) {
+        stream->close();
+        return;
+      }
       reader->stream = stream;
     }
 
@@ -606,6 +654,9 @@ std::thread SubscriptionThreadDispatcher::startDataReaderLocked(DataFrameCallbac
           "\"{}\": code={} message={}",
           track_name, identity, static_cast<std::uint32_t>(error->code), error->message);
     }
+    // Clean our own slot if the stream ended on its own (server EOS) and no
+    // extract/teardown already claimed it. A no-op when we were extracted.
+    eraseDataReaderIfCurrent(id, reader);
     LK_LOG_INFO("Data reader thread exiting for \"{}\" track=\"{}\"", identity, track_name);
   });
   // NOLINTEND(bugprone-lambda-function-name)
diff --git a/src/tests/unit/test_subscription_thread_dispatcher.cpp b/src/tests/unit/test_subscription_thread_dispatcher.cpp
@@ -20,6 +20,8 @@
 #include <livekit/livekit.h>
 
 #include <atomic>
+#include <memory>
+#include <mutex>
 #include <thread>
 #include <unordered_map>
 #include <vector>
@@ -37,13 +39,25 @@ class SubscriptionThreadDispatcherTest : public ::testing::Test {
   using DataCallbackKey = SubscriptionThreadDispatcher::DataCallbackKey;
   using DataCallbackKeyHash = SubscriptionThreadDispatcher::DataCallbackKeyHash;
 
+  using ActiveDataReader = SubscriptionThreadDispatcher::ActiveDataReader;
+
   static auto& audioCallbacks(SubscriptionThreadDispatcher& dispatcher) { return dispatcher.audio_callbacks_; }
   static auto& videoCallbacks(SubscriptionThreadDispatcher& dispatcher) { return dispatcher.video_callbacks_; }
   static auto& activeReaders(SubscriptionThreadDispatcher& dispatcher) { return dispatcher.active_readers_; }
   static auto& dataCallbacks(SubscriptionThreadDispatcher& dispatcher) { return dispatcher.data_callbacks_; }
   static auto& activeDataReaders(SubscriptionThreadDispatcher& dispatcher) { return dispatcher.active_data_readers_; }
   static auto& remoteDataTracks(SubscriptionThreadDispatcher& dispatcher) { return dispatcher.remote_data_tracks_; }
   static int maxActiveReaders() { return SubscriptionThreadDispatcher::kMaxActiveReaders; }
+
+  static std::thread extractDataReader(SubscriptionThreadDispatcher& dispatcher, DataFrameCallbackId id) {
+    const std::scoped_lock<std::mutex> lock(dispatcher.lock_);
+    return dispatcher.extractDataReaderThreadLocked(id);
+  }
+
+  static void eraseDataReaderIfCurrent(SubscriptionThreadDispatcher& dispatcher, DataFrameCallbackId id,
+                                       const std::shared_ptr<ActiveDataReader>& reader) {
+    dispatcher.eraseDataReaderIfCurrent(id, reader);
+  }
 };
 
 // ============================================================================
@@ -478,6 +492,57 @@ TEST_F(SubscriptionThreadDispatcherTest, NoRemoteDataTracksInitially) {
   EXPECT_TRUE(remoteDataTracks(dispatcher).empty());
 }
 
+// ============================================================================
+// Data reader replacement: cancellation and self-erase
+// ============================================================================
+
+TEST_F(SubscriptionThreadDispatcherTest, ActiveDataReaderNotCancelledByDefault) {
+  auto reader = std::make_shared<ActiveDataReader>();
+  EXPECT_FALSE(reader->cancelled.load());
+}
+
+TEST_F(SubscriptionThreadDispatcherTest, ExtractDataReaderMarksCancelledAndRemovesEntry) {
+  SubscriptionThreadDispatcher dispatcher;
+  auto reader = std::make_shared<ActiveDataReader>();
+  activeDataReaders(dispatcher)[0] = reader;
+
+  auto extracted = extractDataReader(dispatcher, 0);
+
+  EXPECT_TRUE(reader->cancelled.load()) << "Extract must cancel so an in-flight subscribe aborts";
+  EXPECT_FALSE(extracted.joinable()) << "No real thread was attached to the seeded reader";
+  EXPECT_TRUE(activeDataReaders(dispatcher).empty());
+}
+
+TEST_F(SubscriptionThreadDispatcherTest, ExtractMissingDataReaderIsNoOp) {
+  SubscriptionThreadDispatcher dispatcher;
+  auto extracted = extractDataReader(dispatcher, 42);
+  EXPECT_FALSE(extracted.joinable());
+}
+
+TEST_F(SubscriptionThreadDispatcherTest, EraseDataReaderIfCurrentRemovesMatchingEntry) {
+  SubscriptionThreadDispatcher dispatcher;
+  auto reader = std::make_shared<ActiveDataReader>();
+  activeDataReaders(dispatcher)[0] = reader;
+
+  eraseDataReaderIfCurrent(dispatcher, 0, reader);
+
+  EXPECT_TRUE(activeDataReaders(dispatcher).empty());
+}
+
+TEST_F(SubscriptionThreadDispatcherTest, EraseDataReaderIfCurrentLeavesReplacedEntry) {
+  SubscriptionThreadDispatcher dispatcher;
+  auto original = std::make_shared<ActiveDataReader>();
+  auto replacement = std::make_shared<ActiveDataReader>();
+  activeDataReaders(dispatcher)[0] = replacement;
+
+  // The original reader exited after being replaced; it must not evict the
+  // newer reader that now owns the same callback id.
+  eraseDataReaderIfCurrent(dispatcher, 0, original);
+
+  ASSERT_EQ(activeDataReaders(dispatcher).size(), 1u);
+  EXPECT_EQ(activeDataReaders(dispatcher)[0], replacement);
+}
+
 // ============================================================================
 // Data track destruction safety
 // ============================================================================
