From 69a2f33edc60e987cd144046b520f6b8a4d059f9 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Wed, 4 Jan 2023 10:08:59 +0800
Subject: [PATCH] fix(core): force delete thread_local data on android: at
 dlopen/dlclose case, android do not call cb(exit) register from
 pthread_key_create(&m_key, exit)

GitOrigin-RevId: c78352e7114ad362a9c72d15dc9a84331d3bb00a
---
 src/core/impl/utils/thread.cpp                | 14 +++++++++++
 .../include/megbrain/utils/thread_local.h     | 24 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/src/core/impl/utils/thread.cpp b/src/core/impl/utils/thread.cpp
index 5e5b249cf..02fe6f1bb 100644
--- a/src/core/impl/utils/thread.cpp
+++ b/src/core/impl/utils/thread.cpp
@@ -257,6 +257,20 @@ void SyncableCounter::wait_zero() {
     }
 }
 
+/* =============== ThreadLocalForceFree ===============  */
+#if defined(__ANDROID__) && !USE_STL_THREAD_LOCAL
+void ThreadLocalForceFree::push(void* d) {
+    MGB_LOCK_GUARD(m_mutex);
+    td.push_back(d);
+}
+
+//! make ff init as soon as possible
+static ThreadLocalForceFree ff;
+ThreadLocalForceFree& get_thread_local_force_free_instance() {
+    return ff;
+}
+#endif
+
 #else  // MGB_HAVE_THREAD
 #pragma message "threading support is disabled"
 #if MGB_CUDA
diff --git a/src/core/include/megbrain/utils/thread_local.h b/src/core/include/megbrain/utils/thread_local.h
index 3b968e02b..b453850cf 100644
--- a/src/core/include/megbrain/utils/thread_local.h
+++ b/src/core/include/megbrain/utils/thread_local.h
@@ -22,6 +22,24 @@
 #pragma message("force disable USE_STL_THREAD_LOCAL for thread_local mem leak at dlopen/dlclose")
 #undef USE_STL_THREAD_LOCAL
 #define USE_STL_THREAD_LOCAL 0
+class ThreadData;
+class ThreadLocalForceFree {
+public:
+    ThreadLocalForceFree() = default;
+    ~ThreadLocalForceFree() {
+        MGB_LOCK_GUARD(m_mutex);
+        for (auto& d : td) {
+            delete (ThreadData*)d;
+        }
+    }
+
+    void push(void* d);
+
+private:
+    std::vector<void*> td;
+    MGB_MUTEX m_mutex;
+};
+ThreadLocalForceFree& get_thread_local_force_free_instance();
 #endif
 
 #if USE_STL_THREAD_LOCAL
@@ -57,6 +75,10 @@ class ThreadLocalPtr {
             return static_cast<ThreadData*>(d)->data;
         }
         ThreadData* t_data = new ThreadData();
+#if defined(__ANDROID__)
+        get_thread_local_force_free_instance().push((void*)t_data);
+#endif
+
         t_data->data = m_constructor();
         t_data->self = this;
         pthread_setspecific(m_key, t_data);
@@ -67,7 +89,9 @@ class ThreadLocalPtr {
         ThreadData* td = static_cast<ThreadData*>(d);
         if (td && td->self->m_destructor)
             td->self->m_destructor(td->data);
+#if !defined(__ANDROID__)
         delete td;
+#endif
     }
 
 public:
-- 
GitLab