Add suggestion to improve code

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

Author: VietND96

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

@@ -21,16 +21,14 @@
 import static org.openqa.selenium.remote.RemoteTags.SESSION_ID_EVENT;
 
 import java.net.URI;
-import java.util.ArrayList;
 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.ReadWriteLock;
-import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.logging.Logger;
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.events.Event;
@@ -92,17 +90,21 @@ public boolean isReady() {
   @Override
   public boolean add(Session session) {
     Require.nonNull("Session", session);
-    SessionId id = session.getId();
 
-    // IndexedSessionMap now handles internal synchronization
+    SessionId id = session.getId();
     knownSessions.put(id, session);
 
     try (Span span = tracer.getCurrentContext().createSpan("local_sessionmap.add")) {
       AttributeMap attributeMap = tracer.createAttributeMap();
       attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), getClass().getName());
       SESSION_ID.accept(span, id);
       SESSION_ID_EVENT.accept(attributeMap, id);
-      span.addEvent("Added session into local Session Map", attributeMap);
+
+      String sessionAddedMessage =
+          String.format(
+              "Added session to local Session Map, Id: %s, Node: %s", id, session.getUri());
+      span.addEvent(sessionAddedMessage, attributeMap);
+      LOG.info(sessionAddedMessage);
     }
 
     return true;
@@ -112,7 +114,6 @@ public boolean add(Session session) {
   public Session get(SessionId id) {
     Require.nonNull("Session ID", id);
 
-    // IndexedSessionMap now handles internal synchronization
     Session session = knownSessions.get(id);
     if (session == null) {
       throw new NoSuchSessionException("Unable to find session with ID: " + id);
@@ -124,14 +125,14 @@ public Session get(SessionId id) {
   public void remove(SessionId id) {
     Require.nonNull("Session ID", id);
 
-    // IndexedSessionMap now handles internal synchronization
     Session removedSession = knownSessions.remove(id);
 
     try (Span span = tracer.getCurrentContext().createSpan("local_sessionmap.remove")) {
       AttributeMap attributeMap = tracer.createAttributeMap();
       attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), getClass().getName());
       SESSION_ID.accept(span, id);
       SESSION_ID_EVENT.accept(attributeMap, id);
+
       String sessionDeletedMessage =
           String.format(
               "Deleted session from local Session Map, Id: %s, Node: %s",
@@ -142,57 +143,42 @@ public void remove(SessionId id) {
     }
   }
 
-  /** Batch remove sessions by URI with proper locking to prevent race conditions */
   private void batchRemoveByUri(URI externalUri, Class<? extends Event> eventClass) {
-    // IndexedSessionMap now handles internal synchronization
     Set<SessionId> sessionsToRemove = knownSessions.getSessionsByUri(externalUri);
-    if (!sessionsToRemove.isEmpty()) {
-      knownSessions.batchRemove(sessionsToRemove);
+
+    if (sessionsToRemove.isEmpty()) {
+      return; // Early return for empty operations - no tracing overhead
     }
 
+    knownSessions.batchRemove(sessionsToRemove);
+
     try (Span span = tracer.getCurrentContext().createSpan("local_sessionmap.batch_remove")) {
       AttributeMap attributeMap = tracer.createAttributeMap();
       attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), getClass().getName());
       attributeMap.put("event.class", eventClass.getName());
       attributeMap.put("node.uri", externalUri.toString());
       attributeMap.put("sessions.count", sessionsToRemove.size());
 
-      LOG.info(
+      String batchRemoveMessage =
           String.format(
-              "Event %s triggered batch remove from local Session Map for Node %s",
-              eventClass.getName(), externalUri));
-      String eventMessage = "";
-      if (!sessionsToRemove.isEmpty()) {
-        eventMessage =
-            String.format(
-                "Batch removed %d sessions belonging to Node %s including: %s",
-                sessionsToRemove.size(), externalUri, sessionsToRemove);
-      } else {
-        eventMessage =
-            String.format(
-                "No sessions found to remove from local Session Map for Node %s", externalUri);
-      }
-      span.addEvent(eventMessage, attributeMap);
-      LOG.info(eventMessage);
+              "Batch removed %d sessions from local Session Map for Node %s (triggered by %s)",
+              sessionsToRemove.size(), externalUri, eventClass.getSimpleName());
+      span.addEvent(batchRemoveMessage, attributeMap);
+      LOG.info(batchRemoveMessage);
     }
   }
 
-  /** Custom ConcurrentMap implementation that automatically maintains a URI-to-SessionId index */
   private static class IndexedSessionMap {
     private final ConcurrentMap<SessionId, Session> sessions = new ConcurrentHashMap<>();
     private final ConcurrentMap<URI, Set<SessionId>> sessionsByUri = new ConcurrentHashMap<>();
-    // Internal lock to ensure atomicity of multi-step operations across both maps
-    private final ReadWriteLock internalLock = new ReentrantReadWriteLock();
+    private final Object coordinationLock = new Object();
 
     public Session get(SessionId id) {
-      // Read operations are atomic on ConcurrentHashMap - no lock needed
       return sessions.get(id);
     }
 
     public Session put(SessionId id, Session session) {
-      // Write lock needed: multiple operations across both maps must be atomic
-      internalLock.writeLock().lock();
-      try {
+      synchronized (coordinationLock) {
         Session previous = sessions.put(id, session);
 
         if (previous != null && previous.getUri() != null) {
@@ -205,133 +191,69 @@ public Session put(SessionId id, Session session) {
         }
 
         return previous;
-      } finally {
-        internalLock.writeLock().unlock();
       }
     }
 
     public Session remove(SessionId id) {
-      // Write lock needed: multiple operations across both maps must be atomic
-      internalLock.writeLock().lock();
-      try {
+      synchronized (coordinationLock) {
         Session removed = sessions.remove(id);
 
         if (removed != null && removed.getUri() != null) {
           cleanupUriIndex(removed.getUri(), id);
         }
 
         return removed;
-      } finally {
-        internalLock.writeLock().unlock();
       }
     }
 
     public void batchRemove(Set<SessionId> sessionIds) {
-      // Write lock needed: multiple operations across both maps must be atomic
-      internalLock.writeLock().lock();
-      try {
+      synchronized (coordinationLock) {
         Map<URI, Set<SessionId>> uriToSessionIds = new HashMap<>();
 
+        // Single loop: remove sessions and collect URI mappings in one pass
         for (SessionId id : sessionIds) {
-          Session session = sessions.get(id);
+          Session session = sessions.remove(id);
           if (session != null && session.getUri() != null) {
             uriToSessionIds.computeIfAbsent(session.getUri(), k -> new HashSet<>()).add(id);
           }
         }
 
-        for (SessionId id : sessionIds) {
-          sessions.remove(id);
-        }
-
-        // Robust cleanup for each URI
+        // Clean up URI index for all affected URIs
         for (Map.Entry<URI, Set<SessionId>> entry : uriToSessionIds.entrySet()) {
-          URI uri = entry.getKey();
-          Set<SessionId> idsToRemove = entry.getValue();
-          cleanupUriIndex(uri, idsToRemove);
+          cleanupUriIndex(entry.getKey(), entry.getValue());
         }
-      } finally {
-        internalLock.writeLock().unlock();
       }
     }
 
-    /**
-     * Robust cleanup of URI index to prevent memory leaks from empty sets. Handles single session
-     * removal with explicit empty set cleanup.
-     */
     private void cleanupUriIndex(URI uri, SessionId sessionId) {
-      Set<SessionId> sessionIds = sessionsByUri.get(uri);
-      if (sessionIds != null) {
-        sessionIds.remove(sessionId);
-        // Explicit check and removal to prevent memory leaks
-        if (sessionIds.isEmpty()) {
-          sessionsByUri.remove(uri, sessionIds);
-        }
-      }
+      sessionsByUri.computeIfPresent(
+          uri,
+          (key, sessionIds) -> {
+            sessionIds.remove(sessionId);
+            return sessionIds.isEmpty() ? null : sessionIds;
+          });
     }
 
-    /**
-     * Robust cleanup of URI index to prevent memory leaks from empty sets. Handles batch session
-     * removal with explicit empty set cleanup.
-     */
     private void cleanupUriIndex(URI uri, Set<SessionId> sessionIdsToRemove) {
-      Set<SessionId> sessionIds = sessionsByUri.get(uri);
-      if (sessionIds != null) {
-        sessionIds.removeAll(sessionIdsToRemove);
-        // Explicit check and removal to prevent memory leaks
-        if (sessionIds.isEmpty()) {
-          sessionsByUri.remove(uri, sessionIds);
-        }
-      }
-    }
-
-    /**
-     * Periodic cleanup to remove any empty sets that may have been missed. Should be called
-     * periodically to prevent memory leaks.
-     */
-    public void performMaintenanceCleanup() {
-      internalLock.writeLock().lock();
-      try {
-        // Find and remove empty URI sets
-        Set<URI> emptyUris = new HashSet<>();
-        for (Map.Entry<URI, Set<SessionId>> entry : sessionsByUri.entrySet()) {
-          if (entry.getValue().isEmpty()) {
-            emptyUris.add(entry.getKey());
-          }
-        }
-
-        for (URI emptyUri : emptyUris) {
-          sessionsByUri.remove(emptyUri);
-        }
-      } finally {
-        internalLock.writeLock().unlock();
-      }
+      sessionsByUri.computeIfPresent(
+          uri,
+          (key, sessionIds) -> {
+            sessionIds.removeAll(sessionIdsToRemove);
+            return sessionIds.isEmpty() ? null : sessionIds;
+          });
     }
 
     public Set<SessionId> getSessionsByUri(URI uri) {
-      // Read operations are atomic on ConcurrentHashMap - no lock needed
       Set<SessionId> result = sessionsByUri.get(uri);
-      // Return empty set instead of null, and ensure we don't return empty sets
       return (result != null && !result.isEmpty()) ? result : Set.of();
     }
 
     public Set<Map.Entry<SessionId, Session>> entrySet() {
-      // Read lock to ensure consistent view during iteration
-      internalLock.readLock().lock();
-      try {
-        return new HashSet<>(sessions.entrySet());
-      } finally {
-        internalLock.readLock().unlock();
-      }
+      return Collections.unmodifiableSet(sessions.entrySet());
     }
 
     public Collection<Session> values() {
-      // Read lock to ensure consistent view during iteration
-      internalLock.readLock().lock();
-      try {
-        return new ArrayList<>(sessions.values());
-      } finally {
-        internalLock.readLock().unlock();
-      }
+      return Collections.unmodifiableCollection(sessions.values());
     }
 
     public int size() {
@@ -343,13 +265,9 @@ public boolean isEmpty() {
     }
 
     public void clear() {
-      // Write lock needed: multiple operations across both maps must be atomic
-      internalLock.writeLock().lock();
-      try {
+      synchronized (coordinationLock) {
         sessions.clear();
         sessionsByUri.clear();
-      } finally {
-        internalLock.writeLock().unlock();
       }
     }
   }