Enabled automatic WeakReferenceMessenger trimming

Author: Sergio0694

Microsoft.Toolkit.Mvvm/Messaging/WeakReferenceMessenger.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
+using System.Threading;
 using Microsoft.Collections.Extensions;
 using Microsoft.Toolkit.Mvvm.Messaging.Internals;
 #if NETSTANDARD2_0
@@ -13,14 +14,23 @@
 using RecipientsTable = System.Runtime.CompilerServices.ConditionalWeakTable<object, Microsoft.Collections.Extensions.IDictionarySlim>;
 #endif
 
+#pragma warning disable SA1204
+
 namespace Microsoft.Toolkit.Mvvm.Messaging
 {
     /// <summary>
     /// A class providing a reference implementation for the <see cref="IMessenger"/> interface.
     /// </summary>
     /// <remarks>
+    /// <para>
     /// This <see cref="IMessenger"/> implementation uses weak references to track the registered
     /// recipients, so it is not necessary to manually unregister them when they're no longer needed.
+    /// </para>
+    /// <para>
+    /// The <see cref="WeakReferenceMessenger"/> type will automatically perform internal trimming when
+    /// full GC collections are invoked, so calling <see cref="Cleanup"/> manually is not necessary to
+    /// ensure that on average the internal data structures are as trimmed and compact as possible.
+    /// </para>
     /// </remarks>
     public sealed class WeakReferenceMessenger : IMessenger
     {
@@ -46,6 +56,22 @@ public sealed class WeakReferenceMessenger : IMessenger
         /// </summary>
         private readonly DictionarySlim<Type2, RecipientsTable> recipientsMap = new();
 
+        /// <summary>
+        /// Initializes a new instance of the <see cref="WeakReferenceMessenger"/> class.
+        /// </summary>
+        public WeakReferenceMessenger()
+        {
+            // Register an automatic GC callback to trigger a non-blocking cleanup. This will ensure that the
+            // current messenger instance is trimmed and without leftover recipient maps that are no longer used.
+            // This is necessary (as in, some form of cleanup, either explicit or automatic like in this case)
+            // because the ConditionalWeakTable<TKey, TValue> instances will just remove key-value pairs on their
+            // own as soon as a key (ie. a recipient) is collected, causing their own keys (ie. the Type2 instances
+            // mapping to each conditional table for a pair of message and token types) to potentially remain in the
+            // root mapping structure but without any remaining recipients actually registered there, which just
+            // adds unnecessary overhead when trying to enumerate recipients during broadcasting operations later on.
+            Gen2GcCallback.Register(static obj => ((WeakReferenceMessenger)obj).CleanupWithNonBlockingLock(), this);
+        }
+
         /// <summary>
         /// Gets the default <see cref="WeakReferenceMessenger"/> instance.
         /// </summary>
@@ -224,61 +250,97 @@ public void Cleanup()
         {
             lock (this.recipientsMap)
             {
-                using ArrayPoolBufferWriter<Type2> type2s = ArrayPoolBufferWriter<Type2>.Create();
-                using ArrayPoolBufferWriter<object> emptyRecipients = ArrayPoolBufferWriter<object>.Create();
+                CleanupWithoutLock();
+            }
+        }
 
-                var enumerator = this.recipientsMap.GetEnumerator();
+        /// <inheritdoc/>
+        public void Reset()
+        {
+            lock (this.recipientsMap)
+            {
+                this.recipientsMap.Clear();
+            }
+        }
 
-                // First, we go through all the currently registered pairs of token and message types.
-                // These represents all the combinations of generic arguments with at least one registered
-                // handler, with the exception of those with recipients that have already been collected.
-                while (enumerator.MoveNext())
+        /// <summary>
+        /// Executes a cleanup without locking the current instance. This method has to be
+        /// invoked when a lock on <see cref="recipientsMap"/> has already been acquired.
+        /// </summary>
+        private void CleanupWithNonBlockingLock()
+        {
+            object lockObject = this.recipientsMap;
+            bool lockTaken = false;
+
+            try
+            {
+                Monitor.TryEnter(lockObject, ref lockTaken);
+
+                if (lockTaken)
                 {
-                    emptyRecipients.Reset();
+                    CleanupWithoutLock();
+                }
+            }
+            finally
+            {
+                if (lockTaken)
+                {
+                    Monitor.Exit(lockObject);
+                }
+            }
+        }
 
-                    bool hasAtLeastOneHandler = false;
+        /// <summary>
+        /// Executes a cleanup without locking the current instance. This method has to be
+        /// invoked when a lock on <see cref="recipientsMap"/> has already been acquired.
+        /// </summary>
+        private void CleanupWithoutLock()
+        {
+            using ArrayPoolBufferWriter<Type2> type2s = ArrayPoolBufferWriter<Type2>.Create();
+            using ArrayPoolBufferWriter<object> emptyRecipients = ArrayPoolBufferWriter<object>.Create();
 
-                    // Go through the currently alive recipients to look for those with no handlers left. We track
-                    // the ones we find to remove them outside of the loop (can't modify during enumeration).
-                    foreach (KeyValuePair<object, IDictionarySlim> pair in enumerator.Value)
-                    {
-                        if (pair.Value.Count == 0)
-                        {
-                            emptyRecipients.Add(pair.Key);
-                        }
-                        else
-                        {
-                            hasAtLeastOneHandler = true;
-                        }
-                    }
+            var enumerator = this.recipientsMap.GetEnumerator();
+
+            // First, we go through all the currently registered pairs of token and message types.
+            // These represents all the combinations of generic arguments with at least one registered
+            // handler, with the exception of those with recipients that have already been collected.
+            while (enumerator.MoveNext())
+            {
+                emptyRecipients.Reset();
 
-                    // Remove the handler maps for recipients that are still alive but with no handlers
-                    foreach (object recipient in emptyRecipients.Span)
+                bool hasAtLeastOneHandler = false;
+
+                // Go through the currently alive recipients to look for those with no handlers left. We track
+                // the ones we find to remove them outside of the loop (can't modify during enumeration).
+                foreach (KeyValuePair<object, IDictionarySlim> pair in enumerator.Value)
+                {
+                    if (pair.Value.Count == 0)
                     {
-                        enumerator.Value.Remove(recipient);
+                        emptyRecipients.Add(pair.Key);
                     }
-
-                    // Track the type combinations with no recipients or handlers left
-                    if (!hasAtLeastOneHandler)
+                    else
                     {
-                        type2s.Add(enumerator.Key);
+                        hasAtLeastOneHandler = true;
                     }
                 }
 
-                // Remove all the mappings with no handlers left
-                foreach (Type2 key in type2s.Span)
+                // Remove the handler maps for recipients that are still alive but with no handlers
+                foreach (object recipient in emptyRecipients.Span)
+                {
+                    enumerator.Value.Remove(recipient);
+                }
+
+                // Track the type combinations with no recipients or handlers left
+                if (!hasAtLeastOneHandler)
                 {
-                    this.recipientsMap.TryRemove(key);
+                    type2s.Add(enumerator.Key);
                 }
             }
-        }
 
-        /// <inheritdoc/>
-        public void Reset()
-        {
-            lock (this.recipientsMap)
+            // Remove all the mappings with no handlers left
+            foreach (Type2 key in type2s.Span)
             {
-                this.recipientsMap.Clear();
+                this.recipientsMap.TryRemove(key);
             }
         }