Add ownership checkings for FlagLock under debug (#432)

* add ownership checkings for FlagLock under debug

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* fix owner reset ordering

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* addresss CR

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* fix include and constexpr problem

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* use thread_local variable as thread identity

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* change default form for initialization

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* address CR

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* fix typo and format

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* add more assertions

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* adjust flag lock and comments

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

* address CR

Signed-off-by: SchrodingerZhu <i@zhuyi.fan>

src/backend/address_space.h

@@ -28,7 +28,7 @@ namespace snmalloc
      * This is infrequently used code, a spin lock simplifies the code
      * considerably, and should never be on the fast path.
      */
-    std::atomic_flag spin_lock = ATOMIC_FLAG_INIT;
+    FlagWord spin_lock{};
 
   public:
     /**

src/backend/globalconfig.h

@@ -46,7 +46,7 @@ namespace snmalloc
     inline static std::atomic<bool> initialised{false};
 
     SNMALLOC_REQUIRE_CONSTINIT
-    inline static std::atomic_flag initialisation_lock{};
+    inline static FlagWord initialisation_lock{};
 
   public:
     static ChunkAllocatorState&

src/ds/flaglock.h

@@ -6,21 +6,137 @@
 
 namespace snmalloc
 {
+  /**
+   * @brief The DebugFlagWord struct
+   * Wrapper for std::atomic_flag so that we can examine
+   * the re-entrancy problem at debug mode.
+   */
+  struct DebugFlagWord
+  {
+    /**
+     * @brief flag
+     * The underlying atomic field.
+     */
+    std::atomic_flag flag = ATOMIC_FLAG_INIT;
+
+    constexpr DebugFlagWord() = default;
+
+    template<typename... Args>
+    constexpr DebugFlagWord(Args&&... args) : flag(std::forward<Args>(args)...)
+    {}
+
+    /**
+     * @brief set_owner
+     * Record the identity of the locker.
+     */
+    void set_owner()
+    {
+      SNMALLOC_ASSERT(nullptr == owner);
+      owner = get_thread_identity();
+    }
+
+    /**
+     * @brief clear_owner
+     * Set the identity to null.
+     */
+    void clear_owner()
+    {
+      SNMALLOC_ASSERT(get_thread_identity() == owner);
+      owner = nullptr;
+    }
+
+    /**
+     * @brief assert_not_owned_by_current_thread
+     * Assert the lock should not be held already by current thread.
+     */
+    void assert_not_owned_by_current_thread()
+    {
+      SNMALLOC_ASSERT(get_thread_identity() != owner);
+    }
+
+  private:
+    using ThreadIdentity = int const*;
+
+    /**
+     * @brief owner
+     * We use a pointer to TLS field as the thread identity.
+     * std::thread::id can be another solution but it does not
+     * support `constexpr` initialisation on some platforms.
+     */
+    ThreadIdentity owner = nullptr;
+
+    /**
+     * @brief get_thread_identity
+     * @return The identity of current thread.
+     */
+    inline ThreadIdentity get_thread_identity()
+    {
+      static thread_local int SNMALLOC_THREAD_IDENTITY = 0;
+      return &SNMALLOC_THREAD_IDENTITY;
+    }
+  };
+
+  /**
+   * @brief The ReleaseFlagWord struct
+   * The shares the same structure with DebugFlagWord but
+   * all member functions associated with ownership checkings
+   * are empty so that they can be optimised out at Release mode.
+   */
+  struct ReleaseFlagWord
+  {
+    std::atomic_flag flag = ATOMIC_FLAG_INIT;
+
+    constexpr ReleaseFlagWord() = default;
+
+    template<typename... Args>
+    constexpr ReleaseFlagWord(Args&&... args)
+    : flag(std::forward<Args>(args)...)
+    {}
+
+    void set_owner() {}
+    void clear_owner() {}
+    void assert_not_owned_by_current_thread() {}
+  };
+
+#ifdef NDEBUG
+  using FlagWord = ReleaseFlagWord;
+#else
+  using FlagWord = DebugFlagWord;
+#endif
+
   class FlagLock
   {
   private:
-    std::atomic_flag& lock;
+    FlagWord& lock;
 
   public:
-    FlagLock(std::atomic_flag& lock) : lock(lock)
+    FlagLock(FlagWord& lock) : lock(lock)
     {
-      while (lock.test_and_set(std::memory_order_acquire))
+      while (lock.flag.test_and_set(std::memory_order_acquire))
+      {
+        // assert_not_owned_by_current_thread is only called when the first
+        // acquiring is failed; which means the lock is already held somewhere
+        // else.
+        lock.assert_not_owned_by_current_thread();
+#ifdef __cpp_lib_atomic_flag_test
+        // acquire ordering because we need other thread's release to be
+        // visible. This loop is better for spin-waiting because it won't issue
+        // expensive write operation (xchg for example).
+        while (lock.flag.test(std::memory_order_acquire))
+        {
+          Aal::pause();
+        }
+#else
         Aal::pause();
+#endif
+      }
+      lock.set_owner();
     }
 
     ~FlagLock()
     {
-      lock.clear(std::memory_order_release);
+      lock.clear_owner();
+      lock.flag.clear(std::memory_order_release);
     }
   };
 } // namespace snmalloc

src/ds/helpers.h

@@ -17,7 +17,7 @@ namespace snmalloc
   template<class Object, void init(Object*) noexcept>
   class Singleton
   {
-    inline static std::atomic_flag flag;
+    inline static FlagWord flag;
     inline static std::atomic<bool> initialised{false};
     inline static Object obj;
 

src/mem/pool.h

@@ -29,7 +29,7 @@ namespace snmalloc
 
   private:
     MPMCStack<T, PreZeroed> stack;
-    std::atomic_flag lock = ATOMIC_FLAG_INIT;
+    FlagWord lock{};
     T* list{nullptr};
 
   public: