fix: optimize LRU cache (#985)

Author: abhinvv1

WebDriverAgentLib/Routing/FBElementCache.m

@@ -73,7 +73,17 @@ - (XCUIElement *)elementForUUID:(NSString *)uuid checkStaleness:(BOOL)checkStale
     @throw [NSException exceptionWithName:FBStaleElementException reason:reason userInfo:@{}];
   }
   if (checkStaleness) {
-    [element fb_takeSnapshot:NO];
+    @try {
+      [element fb_takeSnapshot:NO];
+    } @catch (NSException *exception) {
+      //  if the snapshot method threw FBStaleElementException (implying the element is stale) we need to explicitly remove it from the cache, PR: https://github.com/appium/WebDriverAgent/pull/985
+      if ([exception.name isEqualToString:FBStaleElementException]) {
+        @synchronized (self.elementCache) {
+          [self.elementCache removeObjectForKey:uuid];
+        }
+      }
+      @throw exception;
+    }
   }
   return element;
 }

WebDriverAgentLib/Utilities/LRUCache/LRUCache.h

@@ -54,6 +54,13 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (NSArray *)allObjects;
 
+/**
+ Removes the object associated with the specified key from the cache.
+
+ @param key The key identifying the object to remove.
+ */
+- (void)removeObjectForKey:(id<NSCopying>)key;
+
 @end
 
 NS_ASSUME_NONNULL_END

WebDriverAgentLib/Utilities/LRUCache/LRUCache.m

@@ -135,4 +135,12 @@ - (void)alignSize
   }
 }
 
+- (void)removeObjectForKey:(id<NSCopying>)key
+{
+  LRUCacheNode *node = self.store[key];
+  if (node != nil) {
+    [self removeNode:node];
+  }
+}
+
 @end

WebDriverAgentTests/UnitTests/FBLRUCacheTests.m

@@ -96,4 +96,33 @@ - (void)testInsertionLoop
   XCTAssertEqualObjects(@[@(count)], cache.allObjects);
 }
 
+- (void)testRemoveExistingObjectForKey {
+    LRUCache *cache = [[LRUCache alloc] initWithCapacity:3];
+    [cache setObject:@"foo" forKey:@"bar"];
+    [cache setObject:@"foo2" forKey:@"bar2"];
+    [cache setObject:@"foo3" forKey:@"bar3"];
+    [self assertArray:@[@"foo3", @"foo2", @"foo"] equalsTo:cache.allObjects];
+    [cache removeObjectForKey:@"bar2"];
+    XCTAssertNil([cache objectForKey:@"bar2"]);
+    [self assertArray:@[@"foo3", @"foo"] equalsTo:cache.allObjects];
+}
+
+- (void)testRemoveNonExistingObjectForKey {
+    LRUCache *cache = [[LRUCache alloc] initWithCapacity:2];
+    [cache setObject:@"foo" forKey:@"bar"];
+    [cache removeObjectForKey:@"nonExisting"];
+    XCTAssertNotNil([cache objectForKey:@"bar"]);
+    [self assertArray:@[@"foo"] equalsTo:cache.allObjects];
+}
+
+- (void)testRemoveAndInsertFlow {
+    LRUCache *cache = [[LRUCache alloc] initWithCapacity:2];
+    [cache setObject:@"foo" forKey:@"bar"];
+    [cache setObject:@"foo2" forKey:@"bar2"];
+    [cache removeObjectForKey:@"bar"];
+    XCTAssertNil([cache objectForKey:@"bar"]);
+    [cache setObject:@"foo3" forKey:@"bar3"];
+    [self assertArray:@[@"foo3", @"foo2"] equalsTo:cache.allObjects];
+}
+
 @end