Merge pull request #45 from FlaUI/thread-safe-evict

Thread safe evicting of elements
FlaUI · May 15, 2024 · bb9eaf0 · bb9eaf0
2 parents e01d093 + b58018f
commit bb9eaf0
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 18 deletions.
diff --git a/src/FlaUI.WebDriver/KnownElement.cs b/src/FlaUI.WebDriver/KnownElement.cs
@@ -4,11 +4,11 @@ namespace FlaUI.WebDriver
 {
     public class KnownElement
     {
-        public KnownElement(AutomationElement element, string? elementRuntimeId)
+        public KnownElement(AutomationElement element, string? elementRuntimeId, string elementReference)
         {
             Element = element;
             ElementRuntimeId = elementRuntimeId;
-            ElementReference = Guid.NewGuid().ToString();
+            ElementReference = elementReference;
         }
 
         public string ElementReference { get; }

diff --git a/src/FlaUI.WebDriver/KnownWindow.cs b/src/FlaUI.WebDriver/KnownWindow.cs
@@ -4,21 +4,22 @@ namespace FlaUI.WebDriver
 {
     public class KnownWindow
     {
-        public KnownWindow(Window window, string? windowRuntimeId)
+        public KnownWindow(Window window, string? windowRuntimeId, string windowHandle)
         {
             Window = window;
             WindowRuntimeId = windowRuntimeId;
-            WindowHandle = Guid.NewGuid().ToString();
+            WindowHandle = windowHandle;
         }
-        public string WindowHandle { get; set; }
+
+        public string WindowHandle { get; }
 
         /// <summary>
         /// A temporarily unique ID, so cannot be used for identity over time, but can be used for improving performance of equality tests.
         /// "The identifier is only guaranteed to be unique to the UI of the desktop on which it was generated. Identifiers can be reused over time."
         /// </summary>
         /// <seealso href="https://learn.microsoft.com/en-us/windows/win32/api/uiautomationclient/nf-uiautomationclient-iuiautomationelement-getruntimeid"/>
-        public string? WindowRuntimeId { get; set; }
+        public string? WindowRuntimeId { get; }
 
-        public Window Window { get; set; }
+        public Window Window { get; }
     }
 }
diff --git 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;
@@ -82,8 +83,11 @@ public KnownElement GetOrAddKnownElement(AutomationElement element)
             var result = KnownElementsByElementReference.Values.FirstOrDefault(knownElement => knownElement.ElementRuntimeId == elementRuntimeId && SafeElementEquals(knownElement.Element, element));
             if (result == null)
             {
-                result = new KnownElement(element, elementRuntimeId);
-                KnownElementsByElementReference.Add(result.ElementReference, result);
+                do
+                {
+                    result = new KnownElement(element, elementRuntimeId, Guid.NewGuid().ToString());
+                }
+                while (!KnownElementsByElementReference.TryAdd(result.ElementReference, result));
             }
             return result;
         }
@@ -103,8 +107,11 @@ public KnownWindow GetOrAddKnownWindow(Window window)
             var result = KnownWindowsByWindowHandle.Values.FirstOrDefault(knownWindow => knownWindow.WindowRuntimeId == windowRuntimeId && SafeElementEquals(knownWindow.Window, window));
             if (result == null)
             {
-                result = new KnownWindow(window, windowRuntimeId);
-                KnownWindowsByWindowHandle.Add(result.WindowHandle, result);
+                do
+                {
+                    result = new KnownWindow(window, windowRuntimeId, Guid.NewGuid().ToString());
+                }
+                while (!KnownWindowsByWindowHandle.TryAdd(result.WindowHandle, result));
             }
             return result;
         }
@@ -123,27 +130,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 _);
             }
         }