Implement opt-in weak reference handling for ElementProxy in automation peers to mitigate memory leaks in virtualized ItemsControls

src/Microsoft.DotNet.Wpf/src/PresentationCore/MS/internal/Automation/ElementProxy.cs

@@ -39,7 +39,8 @@ internal class ElementProxy: IRawElementProviderFragmentRoot, IRawElementProvide
         private ElementProxy(AutomationPeer peer)
         {
             if ((AutomationInteropReferenceType == ReferenceType.Weak) && 
-                (peer is UIElementAutomationPeer || peer is ContentElementAutomationPeer || peer is UIElement3DAutomationPeer))
+                (peer is UIElementAutomationPeer || peer is ContentElementAutomationPeer || peer is UIElement3DAutomationPeer 
+                    || peer.ShouldUseWeakReferenceFromElementProxy))
             {
                 _peer = new WeakReference(peer);
             }

src/Microsoft.DotNet.Wpf/src/PresentationCore/System/Windows/Automation/Peers/AutomationPeer.cs

@@ -523,6 +523,13 @@ internal virtual bool IsDataItemAutomationPeer()
             return false;
         }
 
+        // When true, ElementProxy holds this peer via WeakReference. Opt-in fix for the
+        // UIA-retained-ElementProxy leak in virtualized ItemsControls (see ItemAutomationPeer).
+        internal virtual bool ShouldUseWeakReferenceFromElementProxy
+        {
+            get { return false; }
+        }
+
         // UpdatePeer is called asynchronously.  Between the time the call is
         // posted (InvalidatePeer) and the time the call is executed (UpdatePeer),
         // changes to the visual tree and/or automation tree may have eliminated

src/Microsoft.DotNet.Wpf/src/PresentationFramework/MS/Internal/FrameworkAppContextSwitches.cs

@@ -121,7 +121,23 @@ public static bool ItemAutomationPeerKeepsItsItemAlive
                 return LocalAppContext.GetCachedSwitchValue(ItemAutomationPeerKeepsItsItemAliveSwitchName, ref _ItemAutomationPeerKeepsItsItemAlive);
             }
         }
-
+
+        // Opt-in: when true, ElementProxy holds ItemAutomationPeer (and subclasses) via WeakReference,
+        // letting the GC reclaim peers and their data-item/container chain when UIA Core retains the
+        // ElementProxy aggressively (heap growth in virtualized ItemsControls, e.g. DataGrid).
+        // The parent ItemsControlAutomationPeer's WeakRefElementProxyStorage may reuse still-live
+        // ElementProxy/peer pairs, but it does not guarantee peer re-discovery after GC; callers may
+        // need to handle ElementNotAvailableException and re-walk when the peer has been collected.
+        internal const string UseWeakReferenceFromElementProxyToItemPeerSwitchName = "Switch.System.Windows.Automation.Peers.UseWeakReferenceFromElementProxyToItemPeer";
+        private static int _UseWeakReferenceFromElementProxyToItemPeer;
+        public static bool UseWeakReferenceFromElementProxyToItemPeer
+        {
+            [MethodImpl(MethodImplOptions.AggressiveInlining)]
+            get
+            {
+                return LocalAppContext.GetCachedSwitchValue(UseWeakReferenceFromElementProxyToItemPeerSwitchName, ref _UseWeakReferenceFromElementProxyToItemPeer);
+            }
+        }    
 
         // Switch to opt-out Fluent theme Window Backdrop feature
         internal const string DisableFluentThemeWindowBackdropSwitchName = "Switch.System.Windows.Appearance.DisableFluentThemeWindowBackdrop";

src/Microsoft.DotNet.Wpf/src/PresentationFramework/System/Windows/Automation/Peers/ItemAutomationPeer.cs

@@ -155,6 +155,14 @@ internal override bool IsDataItemAutomationPeer()
             return true;
         }
 
+        // Gated by FrameworkAppContextSwitches.UseWeakReferenceFromElementProxyToItemPeer. When on,
+        // a long-lived UIA ElementProxy no longer pins the data item, container, or parent
+        // ItemsControlAutomationPeer caches; re-discovery goes via WeakRefElementProxyStorage on the parent.
+        internal override bool ShouldUseWeakReferenceFromElementProxy
+        {
+            get { return FrameworkAppContextSwitches.UseWeakReferenceFromElementProxyToItemPeer; }
+        }
+
         internal override void AddToParentProxyWeakRefCache()
         {
             ItemsControlAutomationPeer itemsControlAutomationPeer = ItemsControlAutomationPeer;