From 581a024ceb3757d0b7dd7e65b22206e99b00f61d Mon Sep 17 00:00:00 2001
From: Christian <6939810+chkr1011@users.noreply.github.com>
Date: Sat, 26 Nov 2022 14:45:34 +0100
Subject: [PATCH] Improve allocations for AsyncSignal.

---
 .github/workflows/ReleaseNotes.md             |  2 +-
 Source/MQTTnet/Internal/AsyncSignal.cs        | 23 ++++++++++++++++++-
 .../Server/InjectedMqttApplicationMessage.cs  |  1 +
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/ReleaseNotes.md b/.github/workflows/ReleaseNotes.md
index 8b137891..c041e0d0 100644
--- a/.github/workflows/ReleaseNotes.md
+++ b/.github/workflows/ReleaseNotes.md
@@ -1 +1 @@
-
+[Core] Fixed a memory leak in _AsyncSignal_ implementation (#1586, thanks to @mario-zelger).
diff --git a/Source/MQTTnet/Internal/AsyncSignal.cs b/Source/MQTTnet/Internal/AsyncSignal.cs
index 2861595a..bc3bcb86 100644
--- a/Source/MQTTnet/Internal/AsyncSignal.cs
+++ b/Source/MQTTnet/Internal/AsyncSignal.cs
@@ -34,7 +34,7 @@ namespace MQTTnet.Internal
                 _isSignaled = true;
 
                 Cleanup();
-                
+
                 // If there is already a waiting task let it run.
                 if (_waiter != null)
                 {
@@ -103,6 +103,8 @@ namespace MQTTnet.Internal
             // ReSharper disable once FieldCanBeMadeReadOnly.Local
             CancellationTokenRegistration _cancellationTokenRegistration;
 
+            volatile bool _isCompleted;
+
             public AsyncSignalWaiter(CancellationToken cancellationToken)
             {
                 if (cancellationToken.CanBeCanceled)
@@ -115,18 +117,37 @@ namespace MQTTnet.Internal
 
             public void Approve()
             {
+                if (_isCompleted)
+                {
+                    return;
+                }
+
+                _isCompleted = true;
                 _promise.TrySetResult(true);
             }
 
             public void Dispose()
             {
                 _cancellationTokenRegistration.Dispose();
+                
+                if (_isCompleted)
+                {
+                    // Avoid allocation of _ObjectDisposedException_ which may not be used.
+                    return;
+                }
 
+                _isCompleted = true;
                 _promise.TrySetException(new ObjectDisposedException(nameof(AsyncSignalWaiter)));
             }
 
             void Cancel()
             {
+                if (_isCompleted)
+                {
+                    return;
+                }
+
+                _isCompleted = true;
                 _promise.TrySetCanceled();
             }
         }
diff --git a/Source/MQTTnet/Server/InjectedMqttApplicationMessage.cs b/Source/MQTTnet/Server/InjectedMqttApplicationMessage.cs
index 0708cee3..90406134 100644
--- a/Source/MQTTnet/Server/InjectedMqttApplicationMessage.cs
+++ b/Source/MQTTnet/Server/InjectedMqttApplicationMessage.cs
@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Threading;
 
 namespace MQTTnet.Server
 {