From 1a33452e95b950fa14db8d48d2d4a3bbaa5a3031 Mon Sep 17 00:00:00 2001
From: winlin <winlin@vip.126.com>
Date: Sun, 20 Sep 2020 22:09:03 +0800
Subject: [PATCH] Refine resource manager, fix loop and context switching bug

---
 trunk/src/app/srs_app_conn.cpp    | 29 +++++++++++++-
 trunk/src/app/srs_app_conn.hpp    |  2 +
 trunk/src/utest/srs_utest_rtc.cpp | 64 +++++++++++++++++++++++++++++++
 3 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/trunk/src/app/srs_app_conn.cpp b/trunk/src/app/srs_app_conn.cpp
index a1a9485cd..720ef21f8 100644
--- a/trunk/src/app/srs_app_conn.cpp
+++ b/trunk/src/app/srs_app_conn.cpp
@@ -47,6 +47,7 @@ SrsResourceManager::SrsResourceManager(const std::string& label, bool verbose)
     label_ = label;
     cond = srs_cond_new();
     trd = NULL;
+    p_disposing_ = NULL;
 }
 
 SrsResourceManager::~SrsResourceManager()
@@ -168,15 +169,30 @@ void SrsResourceManager::remove(ISrsResource* c)
         srs_trace("before dispose resource(%s), zombies=%d", c->desc().c_str(), (int)zombies_.size());
     }
 
-    if (std::find(zombies_.begin(), zombies_.end(), c) == zombies_.end()) {
-        zombies_.push_back(c);
+    // Only notify when not removed(in zombies_).
+    vector<ISrsResource*>::iterator it = std::find(zombies_.begin(), zombies_.end(), c);
+    if (it != zombies_.end()) {
+        return;
     }
 
+    // Also ignore when we are disposing it.
+    if (p_disposing_) {
+        it = std::find(p_disposing_->begin(), p_disposing_->end(), c);
+        if (it != p_disposing_->end()) {
+            return;
+        }
+    }
+
+    // Push to zombies, we will free it in another coroutine.
+    zombies_.push_back(c);
+
+    // Notify other handlers to handle the before-dispose event.
     for (int i = 0; i < (int)handlers_.size(); i++) {
         ISrsDisposingHandler* h = handlers_.at(i);
         h->on_before_dispose(c);
     }
 
+    // Notify the coroutine to free it.
     srs_cond_signal(cond);
 }
 
@@ -191,10 +207,19 @@ void SrsResourceManager::clear()
         srs_trace("clear zombies=%d connections", (int)zombies_.size());
     }
 
+    do_clear();
+
+    // Reset it for it points to a local object.
+    p_disposing_ = NULL;
+}
+
+void SrsResourceManager::do_clear()
+{
     // To prevent thread switch when delete connection,
     // we copy all connections then free one by one.
     vector<ISrsResource*> copy;
     copy.swap(zombies_);
+    p_disposing_ = &copy;
 
     vector<ISrsResource*>::iterator it;
     for (it = copy.begin(); it != copy.end(); ++it) {
diff --git a/trunk/src/app/srs_app_conn.hpp b/trunk/src/app/srs_app_conn.hpp
index 7ae7fc610..f9706ef8d 100644
--- a/trunk/src/app/srs_app_conn.hpp
+++ b/trunk/src/app/srs_app_conn.hpp
@@ -64,6 +64,7 @@ private:
     std::vector<ISrsDisposingHandler*> handlers_;
     // The zombie connections, we will delete it asynchronously.
     std::vector<ISrsResource*> zombies_;
+    std::vector<ISrsResource*>* p_disposing_;
 private:
     // The connections without any id.
     std::vector<ISrsResource*> conns_;
@@ -96,6 +97,7 @@ public:
     virtual void remove(ISrsResource* c);
 private:
     void clear();
+    void do_clear();
     void dispose(ISrsResource* c);
 };
 
diff --git a/trunk/src/utest/srs_utest_rtc.cpp b/trunk/src/utest/srs_utest_rtc.cpp
index 291becfcb..40fedaada 100644
--- a/trunk/src/utest/srs_utest_rtc.cpp
+++ b/trunk/src/utest/srs_utest_rtc.cpp
@@ -65,10 +65,74 @@ public:
     }
 };
 
+class MockResourceSelf : public ISrsResource, public ISrsDisposingHandler
+{
+public:
+    bool remove_in_before_dispose;
+    bool remove_in_disposing;
+    SrsResourceManager* manager_;
+    MockResourceSelf(SrsResourceManager* manager) {
+        remove_in_before_dispose = remove_in_disposing = false;
+        manager_ = manager;
+        manager_->subscribe(this);
+    }
+    virtual ~MockResourceSelf() {
+        manager_->unsubscribe(this);
+    }
+    virtual void on_before_dispose(ISrsResource* c) {
+        if (remove_in_before_dispose) {
+            manager_->remove(this);
+        }
+    }
+    virtual void on_disposing(ISrsResource* c) {
+        if (remove_in_disposing) {
+            manager_->remove(this);
+        }
+    }
+    virtual const SrsContextId& get_id() {
+        return _srs_context->get_id();
+    }
+    virtual std::string desc() {
+        return "";
+    }
+};
+
 VOID TEST(KernelRTCTest, ConnectionManagerTest)
 {
     srs_error_t err = srs_success;
 
+    // When hooks disposing, remove itself again.
+    if (true) {
+        SrsResourceManager manager("mgr");
+        HELPER_EXPECT_SUCCESS(manager.start());
+        EXPECT_EQ(0, manager.size()); EXPECT_TRUE(manager.empty());
+
+        MockResourceSelf* resource = new MockResourceSelf(&manager);
+        resource->remove_in_disposing = true;
+        manager.add(resource);
+        EXPECT_EQ(1, manager.size());
+
+        manager.remove(resource);
+        srs_usleep(0);
+        ASSERT_EQ(0, manager.size());
+    }
+
+    // When hooks before-dispose, remove itself again.
+    if (true) {
+        SrsResourceManager manager("mgr");
+        HELPER_EXPECT_SUCCESS(manager.start());
+        EXPECT_EQ(0, manager.size()); EXPECT_TRUE(manager.empty());
+
+        MockResourceSelf* resource = new MockResourceSelf(&manager);
+        resource->remove_in_before_dispose = true;
+        manager.add(resource);
+        EXPECT_EQ(1, manager.size());
+
+        manager.remove(resource);
+        srs_usleep(0);
+        ASSERT_EQ(0, manager.size());
+    }
+
     // Cover all normal scenarios.
     if (true) {
         SrsResourceManager manager("mgr", true);