From 38242b6217fa9c91c77fd034ca5b0ec2d8452c8b 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     | 12 +++++++++---
 .../Framework/Utils/ReferenceEqualityComparer.cs | 16 ++++++++++++++++
 2 files changed, 25 insertions(+), 3 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 b66da5e5f4..a3f11f5e61 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;
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);
+    }
+}