From 80e85bffa84d80336fc535162f2070a624f8a928 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Standa=20Luke=C5=A1?= <exyi@outlook.cz>
Date: Tue, 19 Nov 2024 20:43:50 +0100
Subject: [PATCH] Fix Repeater memoization memory leak
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

If the ConditionalWeakTable references itself, it will get leaked by
.NET Core GC: https://github.com/dotnet/runtime/issues/12255

Fortunately, we don't really need the ConditinalWeakTable here:
* Repeater should generally be a short-lived object, so
   strong references will also be collected reasonably soon
* Between Load and PreRender, the viewmodels were pinned
   in the DataContext of all children and the children are in Children
   collection. Only after PreRender could those references be useful,
   but we can as well Clear the dictionary at that point.

Co-authored-by: Adam Štěpánek <13067679+cafour@users.noreply.github.com>
---
 src/Framework/Framework/Controls/Repeater.cs     | 15 +++++++++++----
 .../Framework/Utils/ReferenceEqualityComparer.cs | 16 ++++++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)
 create mode 100644 src/Framework/Framework/Utils/ReferenceEqualityComparer.cs

diff --git a/src/Framework/Framework/Controls/Repeater.cs b/src/Framework/Framework/Controls/Repeater.cs
index c6a3073be1..ccc4cf5f87 100644
--- a/src/Framework/Framework/Controls/Repeater.cs
+++ b/src/Framework/Framework/Controls/Repeater.cs
@@ -10,6 +10,7 @@
 using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Collections.Generic;
+using DotVVM.Framework.Utils;
 
 namespace DotVVM.Framework.Controls
 {
@@ -139,6 +140,7 @@ protected internal override void OnLoad(IDotvvmRequestContext context)
         protected internal override void OnPreRender(IDotvvmRequestContext context)
         {
             SetChildren(context, renderClientTemplate: !RenderOnServer, memoizeReferences: false);
+            childrenCache.Clear(); // not going to need the unreferenced children anymore
             base.OnPreRender(context);
         }
 
@@ -248,7 +250,7 @@ private DotvvmControl CreateEmptyItem(IDotvvmRequestContext context)
             return emptyDataContainer;
         }
 
-        private ConditionalWeakTable<object, DataItemContainer> childrenCache = new ConditionalWeakTable<object, DataItemContainer>();
+        private readonly Dictionary<object, DataItemContainer> childrenCache = new(ReferenceEqualityComparer<object>.Instance);
         private DotvvmControl AddItem(IList<DotvvmControl> c, IDotvvmRequestContext context, object? item = null, int? index = null, bool allowMemoizationRetrieve = false, bool allowMemoizationStore = false)
         {
             if (allowMemoizationRetrieve && item != null && childrenCache.TryGetValue(item, out var container2) && container2.Parent == null)
@@ -276,8 +278,12 @@ private DotvvmControl AddItem(IList<DotvvmControl> c, IDotvvmRequestContext cont
             // write it to the cache after the content is build. If it would be before that, exception could be suppressed
             if (allowMemoizationStore && item != null)
             {
-                // this GetValue call adds the value without raising exception when the value is already added...
-                childrenCache.GetValue(item, _ => container);
+#if DotNetCore
+                childrenCache.TryAdd(item, container);
+#else
+                if (!childrenCache.ContainsKey(item))
+                    childrenCache.Add(item, container);
+#endif
             }
 
             return container;
@@ -306,6 +312,7 @@ private void SetChildren(IDotvvmRequestContext context, bool renderClientTemplat
             if (DataSource != null)
             {
                 var index = 0;
+                var isCommand = context.RequestType == DotvvmRequestType.Command; // on GET request we are not initializing the Repeater twice
                 foreach (var item in GetIEnumerableFromDataSource()!)
                 {
                     if (SeparatorTemplate != null && index > 0)
@@ -313,7 +320,7 @@ private void SetChildren(IDotvvmRequestContext context, bool renderClientTemplat
                         AddSeparator(Children, context);
                     }
                     AddItem(Children, context, item, index,
-                        allowMemoizationRetrieve: context.RequestType == DotvvmRequestType.Command && !memoizeReferences, // on GET request we are not initializing the Repeater twice
+                        allowMemoizationRetrieve: isCommand && !memoizeReferences,
                         allowMemoizationStore: memoizeReferences
                     );
                     index++;
diff --git a/src/Framework/Framework/Utils/ReferenceEqualityComparer.cs b/src/Framework/Framework/Utils/ReferenceEqualityComparer.cs
new file mode 100644
index 0000000000..4c34d41263
--- /dev/null
+++ b/src/Framework/Framework/Utils/ReferenceEqualityComparer.cs
@@ -0,0 +1,16 @@
+using System.Collections.Generic;
+using System.Runtime.CompilerServices;
+
+namespace DotVVM.Framework.Utils
+{
+    // TODO next version: Replace with System.Collections.Generic.ReferenceEqualityComparer
+    internal class ReferenceEqualityComparer<T> : IEqualityComparer<T>
+        where T : class
+    {
+        public static ReferenceEqualityComparer<T> Instance { get; } = new ReferenceEqualityComparer<T>();
+
+
+        public bool Equals(T? x, T? y) => ReferenceEquals(x, y);
+        public int GetHashCode(T obj) => obj == null ? 0 : RuntimeHelpers.GetHashCode(obj);
+    }
+}