[+] AuUtility::ADestructionWatcher::DoUnderLock

Include/Aurora/Utility/DestructionWatch.hpp

@@ -22,6 +22,7 @@ namespace Aurora::Utility
         friend struct ADestructionWatcher;
 
         Threading::Waitables::FutexWaitable mutex;
+        Threading::Waitables::FutexWaitable dtorSignal;
         ADestructionWatcher *pFirstWatcher {};
         AuList<ADestructionWatcher *> listpWatches;
 
@@ -35,6 +36,7 @@ namespace Aurora::Utility
         inline ~ADestructionWatcher();
 
         inline void Observe(DestructionWatch *pWatch);
+        inline void Observe(DestructionWatch &pWatch);
         inline bool IsObservedAlive();
         inline bool IsObservedDead();
 
@@ -46,12 +48,49 @@ namespace Aurora::Utility
 
     private:
         friend struct DestructionWatch;
+        friend class Threading::LockGuard<ADestructionWatcher *>;
 
         inline void OnSelfDTOR();
         inline void OnDTOR();
+        inline void Lock();
+        inline void Unlock();
 
+    public:
+        template <typename T, typename ...Args>
+        AuConditional_t<AuIsSame_v<AuResultOf_t<T, Args...>, void>, bool, AuOptional<AuResultOf_t<T, Args...>>>
+            DoUnderLock(const T &callable, Args && ...args)
+        {
+            AU_LOCK_GUARD(this);
+            
+            if (this->pOwnsLock)
+            {
+                if constexpr (AuIsSame_v<AuResultOf_t<T, Args...>, void>)
+                {
+                    callable(AuForward<Args &&>(args)...);
+                    return true;
+                }
+                else
+                {
+                    return callable(AuForward<Args &&>(args)...);
+                }
+            }
+            else
+            {
+                if constexpr (AuIsSame_v<AuResultOf_t<T, Args...>, void>)
+                {
+                    return false;
+                }
+                else
+                {
+                    return AuOptional<AuResultOf_t<T, Args...>> {};
+                }
+            }
+        }
+
+    private:
+        Threading::Waitables::FutexWaitable selfLock;
         DestructionWatch *pParent {};
-        bool bWasCalled {};
+        Threading::Waitables::FutexWaitable *pOwnsLock {};
     };
 
     DestructionWatch::DestructionWatch()
@@ -66,26 +105,35 @@ namespace Aurora::Utility
 
     void DestructionWatch::RemoveAll()
     {
-        AU_LOCK_GUARD(this->mutex);
-
-        if (auto pWatcher = AuExchange(this->pFirstWatcher, nullptr))
         {
-            this->RemoveWatcherCall(pWatcher);
+            AU_LOCK_GUARD(this->mutex);
+
+            if (auto pWatcher = AuExchange(this->pFirstWatcher, nullptr))
+            {
+                this->RemoveWatcherCall(pWatcher);
+            }
+
+            for (const auto pWatcher : AuExchange(this->listpWatches, {}))
+            {
+                this->RemoveWatcherCall(pWatcher);
+            }
         }
 
-        for (const auto pWatcher : AuExchange(this->listpWatches, {}))
         {
-            this->RemoveWatcherCall(pWatcher);
+            AU_LOCK_GUARD(this->dtorSignal);
         }
     }
 
     void DestructionWatch::RemoveWatcher(ADestructionWatcher *pWatcher)
     {
         AU_LOCK_GUARD(this->mutex);
+
         if (this->pFirstWatcher == pWatcher)
         {
             this->pFirstWatcher = nullptr;
+            return;
         }
+
         AuTryRemove(this->listpWatches, pWatcher);
     }
 
@@ -118,18 +166,25 @@ namespace Aurora::Utility
         this->pParent = pWatch;
     }
 
+    void ADestructionWatcher::Observe(DestructionWatch &pWatch)
+    {
+        this->Observe(&pWatch);
+    }
+
     bool ADestructionWatcher::IsObservedAlive()
     {
-        return !this->bWasCalled;
+        return bool(this->pParent);
     }
 
     bool ADestructionWatcher::IsObservedDead()
     {
-        return this->bWasCalled;
+        return !bool(this->pParent);
     }
     
     void ADestructionWatcher::OnSelfDTOR()
     {
+        AU_LOCK_GUARD(this->selfLock);
+
         if (auto pParent = AuExchange(this->pParent, nullptr))
         {
             pParent->RemoveWatcher(this);
@@ -138,13 +193,37 @@ namespace Aurora::Utility
 
     void ADestructionWatcher::OnDTOR()
     {
-        this->pParent = nullptr;
-
-        if (AuExchange(this->bWasCalled, true))
         {
-            return;
+            AU_LOCK_GUARD(this->selfLock);
+            this->pParent = nullptr;
         }
 
         this->OnDestroyedChild();
     }
+
+    void ADestructionWatcher::Lock()
+    {
+        AU_LOCK_GUARD(this->selfLock);
+
+        this->pOwnsLock = nullptr;
+
+        if (!this->pParent)
+        {
+            return;
+        }
+
+        this->pParent->dtorSignal.Lock();
+
+        // noting that once we leave this scope, we could have our parent stolen from us (this->selfLock), but its' destructor will still be blocked until we finish
+        // this is why we end up using a raw pointer here.
+        this->pOwnsLock = &this->pParent->dtorSignal;
+    }
+
+    void ADestructionWatcher::Unlock()
+    {
+        if (this->pOwnsLock)
+        {
+            this->pOwnsLock->Unlock();
+        }
+    }
 }