Fix review comment as suggestion

Co-authored-by: Copilot <[email protected]>

Author: VietND96

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

@@ -242,29 +242,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();
       }
     }