From d47a7d18c1ce9e24a03323b73fadd239a09a7030 Mon Sep 17 00:00:00 2001
From: reuk <reuk@users.noreply.github.com>
Date: Tue, 22 Aug 2023 14:23:10 +0100
Subject: [PATCH] MessageManager::Lock: Protect Lock from concurrent accesses

Previously, the Lock was not thread-safe when multiple threads were
locking/unlocking a single Lock instance simultaneously. This
isn't normally a problem when using the MessageManagerLock type, because
each MessageManagerLock contains its own private MessageManager::Lock,
and it's not possible for multiple threads to modify that private lock.

This change improves safety when using a MessageManager::Lock in the
manner of a CriticalSection or other JUCE lock type.
---
 .../messages/juce_MessageManager.cpp          | 59 +++++++++++++------
 .../messages/juce_MessageManager.h            | 12 +++-
 2 files changed, 52 insertions(+), 19 deletions(-)

diff --git a/modules/juce_events/messages/juce_MessageManager.cpp b/modules/juce_events/messages/juce_MessageManager.cpp
index da949d2e93..ca9a48cc67 100644
--- a/modules/juce_events/messages/juce_MessageManager.cpp
+++ b/modules/juce_events/messages/juce_MessageManager.cpp
@@ -292,10 +292,7 @@ struct MessageManager::Lock::BlockingMessage   : public MessageManager::MessageB
         std::unique_lock lock { mutex };
 
         if (owner != nullptr)
-        {
-            owner->abort();
-            acquired = true;
-        }
+            owner->setAcquired (true);
 
         condvar.wait (lock, [&] { return owner == nullptr; });
     }
@@ -307,18 +304,11 @@ struct MessageManager::Lock::BlockingMessage   : public MessageManager::MessageB
         owner = nullptr;
     }
 
-    bool wasAcquired()
-    {
-        const std::scoped_lock lock { mutex };
-        return acquired;
-    }
-
 private:
     std::mutex mutex;
     std::condition_variable condvar;
 
     const MessageManager::Lock* owner = nullptr;
-    bool acquired = false;
 
     JUCE_DECLARE_NON_COPYABLE (BlockingMessage)
 };
@@ -326,8 +316,23 @@ private:
 //==============================================================================
 MessageManager::Lock::Lock()                            {}
 MessageManager::Lock::~Lock()                           { exit(); }
-void MessageManager::Lock::enter()    const noexcept    {        tryAcquire (true); }
-bool MessageManager::Lock::tryEnter() const noexcept    { return tryAcquire (false); }
+void MessageManager::Lock::enter()    const noexcept    {        exclusiveTryAcquire (true); }
+bool MessageManager::Lock::tryEnter() const noexcept    { return exclusiveTryAcquire (false); }
+
+bool MessageManager::Lock::exclusiveTryAcquire (bool lockIsMandatory) const noexcept
+{
+    if (lockIsMandatory)
+        entryMutex.enter();
+    else if (! entryMutex.tryEnter())
+        return false;
+
+    const auto result = tryAcquire (lockIsMandatory);
+
+    if (! result)
+        entryMutex.exit();
+
+    return result;
+}
 
 bool MessageManager::Lock::tryAcquire (bool lockIsMandatory) const noexcept
 {
@@ -376,7 +381,7 @@ bool MessageManager::Lock::tryAcquire (bool lockIsMandatory) const noexcept
             condvar.wait (lock, [&] { return std::exchange (abortWait, false); });
         }
 
-        if (blockingMessage->wasAcquired())
+        if (acquired)
         {
             mm->threadWithLock = Thread::getCurrentThreadId();
             return true;
@@ -395,16 +400,28 @@ bool MessageManager::Lock::tryAcquire (bool lockIsMandatory) const noexcept
 
 void MessageManager::Lock::exit() const noexcept
 {
+    const auto wasAcquired = [&]
+    {
+        const std::scoped_lock lock { mutex };
+        return acquired;
+    }();
+
+    if (! wasAcquired)
+        return;
+
+    const ScopeGuard unlocker { [&] { entryMutex.exit(); } };
+
     if (blockingMessage == nullptr)
         return;
 
-    const ScopeGuard scope { [&] { blockingMessage = nullptr; } };
+    const ScopeGuard scope { [&]
+    {
+        blockingMessage = nullptr;
+        acquired = false;
+    } };
 
     blockingMessage->stopWaiting();
 
-    if (! blockingMessage->wasAcquired())
-        return;
-
     if (auto* mm = MessageManager::instance)
     {
         jassert (mm->currentThreadHasLockedMessageManager());
@@ -413,10 +430,16 @@ void MessageManager::Lock::exit() const noexcept
 }
 
 void MessageManager::Lock::abort() const noexcept
+{
+    setAcquired (false);
+}
+
+void MessageManager::Lock::setAcquired (bool x) const noexcept
 {
     const ScopeGuard scope { [&] { condvar.notify_one(); } };
     const std::scoped_lock lock { mutex };
     abortWait = true;
+    acquired = x;
 }
 
 //==============================================================================
diff --git a/modules/juce_events/messages/juce_MessageManager.h b/modules/juce_events/messages/juce_MessageManager.h
index 01daff8e00..9c43083de7 100644
--- a/modules/juce_events/messages/juce_MessageManager.h
+++ b/modules/juce_events/messages/juce_MessageManager.h
@@ -297,11 +297,21 @@ public:
         struct BlockingMessage;
         friend class ReferenceCountedObjectPtr<BlockingMessage>;
 
+        bool exclusiveTryAcquire (bool) const noexcept;
         bool tryAcquire (bool) const noexcept;
 
+        void setAcquired (bool success) const noexcept;
+
         //==============================================================================
-        mutable ReferenceCountedObjectPtr<BlockingMessage> blockingMessage;
+        // This mutex is used to make this lock type behave like a normal mutex.
+        // If multiple threads call enter() simultaneously, only one will succeed in gaining
+        // this mutex. The mutex is released again in exit().
+        mutable CriticalSection entryMutex;
+
+        // This mutex protects the other data members of the lock from concurrent access, which
+        // happens when the BlockingMessage calls setAcquired to indicate that the lock was gained.
         mutable std::mutex mutex;
+        mutable ReferenceCountedObjectPtr<BlockingMessage> blockingMessage;
         mutable std::condition_variable condvar;
         mutable bool abortWait = false, acquired = false;
     };