From 010c2c934b9e5bf3e4475543cdff5bb0be6e59c8 Mon Sep 17 00:00:00 2001
From: Bananymous <oskari.alaranta@bananymous.com>
Date: Fri, 28 Jun 2024 22:00:29 +0300
Subject: [PATCH] BAN: Write RefPtr and WeakPtr to be thread safe

---
 BAN/include/BAN/RefPtr.h  | 25 +++++++++++++++++++------
 BAN/include/BAN/WeakPtr.h | 33 ++++++++++++++++++++++++++-------
 2 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/BAN/include/BAN/RefPtr.h b/BAN/include/BAN/RefPtr.h
index 8595dac1..6f0e5108 100644
--- a/BAN/include/BAN/RefPtr.h
+++ b/BAN/include/BAN/RefPtr.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <BAN/Atomic.h>
 #include <BAN/Errors.h>
 #include <BAN/Move.h>
 #include <BAN/NoCopyMove.h>
@@ -22,15 +23,27 @@ namespace BAN
 
 		void ref() const
 		{
-			ASSERT(m_ref_count > 0);
-			m_ref_count++;
+			uint32_t old = m_ref_count.fetch_add(1, MemoryOrder::memory_order_relaxed);
+			ASSERT(old > 0);
+		}
+
+		bool try_ref() const
+		{
+			uint32_t expected = m_ref_count.load(MemoryOrder::memory_order_relaxed);
+			for (;;)
+			{
+				if (expected == 0)
+					return false;
+				if (m_ref_count.compare_exchange(expected, expected + 1, MemoryOrder::memory_order_acquire))
+					return true;
+			}
 		}
 
 		void unref() const
 		{
-			ASSERT(m_ref_count > 0);
-			m_ref_count--;
-			if (m_ref_count == 0)
+			uint32_t old = m_ref_count.fetch_sub(1);
+			ASSERT(old > 0);
+			if (old == 1)
 				delete (const T*)this;
 		}
 
@@ -39,7 +52,7 @@ namespace BAN
 		virtual ~RefCounted() { ASSERT(m_ref_count == 0); }
 
 	private:
-		mutable uint32_t m_ref_count = 1;
+		mutable Atomic<uint32_t> m_ref_count = 1;
 	};
 
 	template<typename T>
diff --git a/BAN/include/BAN/WeakPtr.h b/BAN/include/BAN/WeakPtr.h
index b8626875..d68bfd05 100644
--- a/BAN/include/BAN/WeakPtr.h
+++ b/BAN/include/BAN/WeakPtr.h
@@ -2,6 +2,10 @@
 
 #include <BAN/RefPtr.h>
 
+#if __is_kernel
+#include <kernel/Lock/SpinLock.h>
+#endif
+
 namespace BAN
 {
 
@@ -11,22 +15,37 @@ namespace BAN
 	template<typename T>
 	class WeakPtr;
 
+	// FIXME: Write this without using locks...
 	template<typename T>
 	class WeakLink : public RefCounted<WeakLink<T>>
 	{
 	public:
-		RefPtr<T> lock() { ASSERT(m_ptr); return raw_ptr(); }
-		T* raw_ptr() { return m_ptr; }
-
+		RefPtr<T> try_lock()
+		{
+#if __is_kernel
+			Kernel::SpinLockGuard _(m_weak_lock);
+#endif
+			if (m_ptr && m_ptr->try_ref())
+				return RefPtr<T>::adopt(m_ptr);
+			return nullptr;
+		}
 		bool valid() const { return m_ptr; }
-		void invalidate() { m_ptr = nullptr; }
+		void invalidate()
+		{
+#if __is_kernel
+			Kernel::SpinLockGuard _(m_weak_lock);
+#endif
+			m_ptr = nullptr;
+		}
 
 	private:
 		WeakLink(T* ptr) : m_ptr(ptr) {}
 
 	private:
 		T* m_ptr;
-
+#if __is_kernel
+		Kernel::SpinLock m_weak_lock;
+#endif
 		friend class RefPtr<WeakLink<T>>;
 	};
 
@@ -82,8 +101,8 @@ namespace BAN
 
 		RefPtr<T> lock()
 		{
-			if (valid())
-				return m_link->lock();
+			if (m_link)
+				return m_link->try_lock();
 			return nullptr;
 		}