From 3f7552942c726551a4daf9a9274d2a660b19e726 Mon Sep 17 00:00:00 2001
From: aristotelos <arisvd@gmail.com>
Date: Tue, 14 May 2024 21:17:53 +0200
Subject: [PATCH] Thread safe evicting of elements

Prevent stale enumerators while evicting unavailable elements or windows. Closes #42.
---
 src/FlaUI.WebDriver/Session.cs | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/src/FlaUI.WebDriver/Session.cs b/src/FlaUI.WebDriver/Session.cs
index 9cf9db2..34d69ee 100644
--- a/src/FlaUI.WebDriver/Session.cs
+++ b/src/FlaUI.WebDriver/Session.cs
@@ -1,6 +1,7 @@
 using FlaUI.Core;
 using FlaUI.Core.AutomationElements;
 using FlaUI.UIA3;
+using System.Collections.Concurrent;
 using System.Runtime.InteropServices;
 
 namespace FlaUI.WebDriver
@@ -27,8 +28,8 @@ public Session(Application? app, bool isAppOwnedBySession)
         public UIA3Automation Automation { get; }
         public Application? App { get; }
         public InputState InputState { get; }
-        private Dictionary<string, KnownElement> KnownElementsByElementReference { get; } = new Dictionary<string, KnownElement>();
-        private Dictionary<string, KnownWindow> KnownWindowsByWindowHandle { get; } = new Dictionary<string, KnownWindow>();
+        private ConcurrentDictionary<string, KnownElement> KnownElementsByElementReference { get; } = new ConcurrentDictionary<string, KnownElement>();
+        private ConcurrentDictionary<string, KnownWindow> KnownWindowsByWindowHandle { get; } = new ConcurrentDictionary<string, KnownWindow>();
         public TimeSpan ImplicitWaitTimeout => TimeSpan.FromMilliseconds(TimeoutsConfiguration.ImplicitWaitTimeoutMs);
         public TimeSpan PageLoadTimeout => TimeSpan.FromMilliseconds(TimeoutsConfiguration.PageLoadTimeoutMs);
         public TimeSpan? ScriptTimeout => TimeoutsConfiguration.ScriptTimeoutMs.HasValue ? TimeSpan.FromMilliseconds(TimeoutsConfiguration.ScriptTimeoutMs.Value) : null;
@@ -83,7 +84,7 @@ public KnownElement GetOrAddKnownElement(AutomationElement element)
             if (result == null)
             {
                 result = new KnownElement(element, elementRuntimeId);
-                KnownElementsByElementReference.Add(result.ElementReference, result);
+                KnownElementsByElementReference.TryAdd(result.ElementReference, result);
             }
             return result;
         }
@@ -104,7 +105,7 @@ public KnownWindow GetOrAddKnownWindow(Window window)
             if (result == null)
             {
                 result = new KnownWindow(window, windowRuntimeId);
-                KnownWindowsByWindowHandle.Add(result.WindowHandle, result);
+                KnownWindowsByWindowHandle.TryAdd(result.WindowHandle, result);
             }
             return result;
         }
@@ -123,27 +124,29 @@ public void RemoveKnownWindow(Window window)
             var item = KnownWindowsByWindowHandle.Values.FirstOrDefault(knownElement => knownElement.Window.Equals(window));
             if (item != null)
             {
-                KnownWindowsByWindowHandle.Remove(item.WindowHandle);
+                KnownWindowsByWindowHandle.TryRemove(item.WindowHandle, out _);
             }
         }
 
         public void EvictUnavailableElements()
         {
             // Evict unavailable elements to prevent slowing down
-            var unavailableElements = KnownElementsByElementReference.Where(item => !item.Value.Element.IsAvailable).Select(item => item.Key).ToArray();
+            // (use ToArray to prevent concurrency issues while enumerating)
+            var unavailableElements = KnownElementsByElementReference.ToArray().Where(item => !item.Value.Element.IsAvailable).Select(item => item.Key);
             foreach (var unavailableElementKey in unavailableElements)
             {
-                KnownElementsByElementReference.Remove(unavailableElementKey);
+                KnownElementsByElementReference.TryRemove(unavailableElementKey, out _);
             }
         }
 
         public void EvictUnavailableWindows()
         {
             // Evict unavailable windows to prevent slowing down
-            var unavailableWindows = KnownWindowsByWindowHandle.Where(item => !item.Value.Window.IsAvailable).Select(item => item.Key).ToArray();
+            // (use ToArray to prevent concurrency issues while enumerating)
+            var unavailableWindows = KnownWindowsByWindowHandle.ToArray().Where(item => !item.Value.Window.IsAvailable).Select(item => item.Key).ToArray();
             foreach (var unavailableWindowKey in unavailableWindows)
             {
-                KnownWindowsByWindowHandle.Remove(unavailableWindowKey);
+                KnownWindowsByWindowHandle.TryRemove(unavailableWindowKey, out _);
             }
         }