Fix review comment as suggestion

Co-authored-by: Copilot <[email protected]>
Signed-off-by: Viet Nguyen Duc <[email protected]>

Author: VietND96

java/src/org/openqa/selenium/grid/sessionmap/local/LocalSessionMap.java

@@ -22,12 +22,14 @@
 
 import java.net.URI;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.logging.Logger;
@@ -115,14 +117,9 @@ public boolean add(Session session) {
   public Session get(SessionId id) {
     Require.nonNull("Session ID", id);
 
-    Session session = knownSessions.get(id);
-    if (session != null) {
-      return session;
-    }
-
     sessionMapLock.readLock().lock();
     try {
-      session = knownSessions.get(id);
+      Session session = knownSessions.get(id);
       if (session == null) {
         throw new NoSuchSessionException("Unable to find session with ID: " + id);
       }
@@ -201,6 +198,8 @@ private void batchRemoveByUri(URI externalUri, Class<? extends Event> eventClass
   private static class IndexedSessionMap {
     private final ConcurrentMap<SessionId, Session> sessions = new ConcurrentHashMap<>();
     private final ConcurrentMap<URI, Set<SessionId>> sessionsByUri = new ConcurrentHashMap<>();
+    private final ReadWriteLock lock = new ReentrantReadWriteLock();
+    private final Lock writeLock = lock.writeLock();
 
     public Session get(SessionId id) {
       return sessions.get(id);
@@ -242,29 +241,34 @@ public Session remove(SessionId id) {
     }
 
     public void batchRemove(Set<SessionId> sessionIds) {
-      Map<URI, Set<SessionId>> uriToSessionIds = new HashMap<>();
-
-      for (SessionId id : sessionIds) {
-        Session session = sessions.get(id);
-        if (session != null && session.getUri() != null) {
-          uriToSessionIds.computeIfAbsent(session.getUri(), k -> new HashSet<>()).add(id);
+      writeLock.lock();
+      try {
+        Map<URI, Set<SessionId>> uriToSessionIds = new HashMap<>();
+
+        for (SessionId id : sessionIds) {
+          Session session = sessions.get(id);
+          if (session != null && session.getUri() != null) {
+            uriToSessionIds.computeIfAbsent(session.getUri(), k -> new HashSet<>()).add(id);
+          }
         }
-      }
 
-      for (SessionId id : sessionIds) {
-        sessions.remove(id);
-      }
+        for (SessionId id : sessionIds) {
+          sessions.remove(id);
+        }
 
-      for (Map.Entry<URI, Set<SessionId>> entry : uriToSessionIds.entrySet()) {
-        URI uri = entry.getKey();
-        Set<SessionId> idsToRemove = entry.getValue();
+        for (Map.Entry<URI, Set<SessionId>> entry : uriToSessionIds.entrySet()) {
+          URI uri = entry.getKey();
+          Set<SessionId> idsToRemove = entry.getValue();
 
-        sessionsByUri.computeIfPresent(
-            uri,
-            (k, existingIds) -> {
-              existingIds.removeAll(idsToRemove);
-              return existingIds.isEmpty() ? null : existingIds;
-            });
+          sessionsByUri.computeIfPresent(
+              uri,
+              (k, existingIds) -> {
+                existingIds.removeAll(idsToRemove);
+                return existingIds.isEmpty() ? null : existingIds;
+              });
+        }
+      } finally {
+        writeLock.unlock();
       }
     }