refactor: use new object to mutex holder

metaopt · Sep 23, 2024 · 56a0c47 · 56a0c47
1 parent e35b643
commit 56a0c47
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/include/treespec.h b/include/treespec.h
@@ -393,7 +393,14 @@ class PyTreeIter {
           m_none_is_leaf{none_is_leaf},
           m_namespace{registry_namespace},
           m_is_dict_insertion_ordered{PyTreeSpec::IsDictInsertionOrdered(registry_namespace)},
-          m_mutex{} {};
+#ifdef Py_GIL_DISABLED
+          // NOLINTNEXTLINE[whitespace/braces]
+          m_mutex_holder_object{py::handle{reinterpret_cast<PyObject *>(&PyBaseObject_Type)}()}
+#else
+          m_mutex_holder_object{}  // NOLINT[whitespace/braces]
+#endif
+    {
+    }
 
     PyTreeIter() = delete;
     ~PyTreeIter() = default;
@@ -417,7 +424,7 @@ class PyTreeIter {
     const bool m_none_is_leaf;
     const std::string m_namespace;
     const bool m_is_dict_insertion_ordered;
-    mutable mutex m_mutex;
+    const py::object m_mutex_holder_object;
 
     template <bool NoneIsLeaf>
     [[nodiscard]] py::object NextImpl();

diff --git a/src/treespec/traversal.cpp b/src/treespec/traversal.cpp
@@ -23,7 +23,6 @@ limitations under the License.
 
 #include "include/critical_section.h"
 #include "include/exceptions.h"
-#include "include/mutex.h"
 #include "include/registry.h"
 #include "include/treespec.h"
 #include "include/utils.h"
@@ -167,7 +166,7 @@ py::object PyTreeIter::NextImpl() {
 }
 
 py::object PyTreeIter::Next() {
-    const scoped_lock_guard lock{m_mutex};
+    const scoped_critical_section cs{m_mutex_holder_object};
 
     if (m_none_is_leaf) [[unlikely]] {
         return NextImpl<NONE_IS_LEAF>();

diff --git a/src/treespec/treespec.cpp b/src/treespec/treespec.cpp
@@ -1600,6 +1600,7 @@ size_t PyTreeSpec::ThreadIndentTypeHash::operator()(
         Py_VISIT(pair.first.ptr());
     }
     Py_VISIT(self.m_root.ptr());
+    Py_VISIT(self.m_mutex_holder_object.ptr());
     return 0;
 }