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

@@ -29,8 +29,9 @@
 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.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.events.Event;
@@ -50,18 +51,35 @@
 import org.openqa.selenium.remote.tracing.Span;
 import org.openqa.selenium.remote.tracing.Tracer;
 
-public class LocalSessionMap extends SessionMap {
+public class LocalSessionMap extends SessionMap implements AutoCloseable {
 
   private static final Logger LOG = Logger.getLogger(LocalSessionMap.class.getName());
 
   private final EventBus bus;
   private final IndexedSessionMap knownSessions = new IndexedSessionMap();
+  private final ScheduledExecutorService maintenanceExecutor;
 
   public LocalSessionMap(Tracer tracer, EventBus bus) {
     super(tracer);
 
     this.bus = Require.nonNull("Event bus", bus);
 
+    // Initialize maintenance executor for periodic cleanup
+    this.maintenanceExecutor =
+        Executors.newSingleThreadScheduledExecutor(
+            r -> {
+              Thread t = new Thread(r, "LocalSessionMap-Maintenance");
+              t.setDaemon(true);
+              return t;
+            });
+
+    // Schedule maintenance cleanup every 5 minutes to prevent memory leaks
+    this.maintenanceExecutor.scheduleWithFixedDelay(
+        this::performMaintenanceCleanup,
+        5, // initial delay
+        5, // period
+        TimeUnit.MINUTES);
+
     bus.addListener(SessionClosedEvent.listener(this::remove));
 
     bus.addListener(
@@ -92,17 +110,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;
@@ -132,6 +154,7 @@ public void remove(SessionId id) {
       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",
@@ -157,199 +180,227 @@ private void batchRemoveByUri(URI externalUri, Class<? extends Event> eventClass
       attributeMap.put("node.uri", externalUri.toString());
       attributeMap.put("sessions.count", sessionsToRemove.size());
 
-      LOG.info(
-          String.format(
-              "Event %s triggered batch remove from local Session Map for Node %s",
-              eventClass.getName(), externalUri));
-      String eventMessage = "";
       if (!sessionsToRemove.isEmpty()) {
-        eventMessage =
+        String batchRemoveMessage =
             String.format(
-                "Batch removed %d sessions belonging to Node %s including: %s",
-                sessionsToRemove.size(), externalUri, sessionsToRemove);
+                "Batch removed %d sessions from local Session Map for Node %s (triggered by %s):"
+                    + " %s",
+                sessionsToRemove.size(), externalUri, eventClass.getSimpleName(), sessionsToRemove);
+        span.addEvent(batchRemoveMessage, attributeMap);
+        LOG.info(batchRemoveMessage);
       } else {
-        eventMessage =
+        String noSessionsMessage =
             String.format(
-                "No sessions found to remove from local Session Map for Node %s", externalUri);
+                "No sessions found to remove from local Session Map for Node %s (triggered by %s)",
+                externalUri, eventClass.getSimpleName());
+        span.addEvent(noSessionsMessage, attributeMap);
+      }
+    }
+  }
+
+  private void performMaintenanceCleanup() {
+    knownSessions.performMaintenanceCleanup();
+  }
+
+  /**
+   * Shutdown the maintenance executor and perform final cleanup. This method should be called when
+   * the LocalSessionMap is no longer needed.
+   */
+  @Override
+  public void close() {
+    try {
+      // Perform final maintenance cleanup before shutdown
+      performMaintenanceCleanup();
+
+      // Shutdown the maintenance executor gracefully
+      maintenanceExecutor.shutdown();
+      if (!maintenanceExecutor.awaitTermination(10, TimeUnit.SECONDS)) {
+        LOG.warning("Maintenance executor did not terminate gracefully, forcing shutdown");
+        maintenanceExecutor.shutdownNow();
       }
-      span.addEvent(eventMessage, attributeMap);
-      LOG.info(eventMessage);
+
+      LOG.info("LocalSessionMap maintenance executor shutdown completed");
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      maintenanceExecutor.shutdownNow();
+      LOG.warning("Interrupted during maintenance executor shutdown");
+    } catch (Exception e) {
+      LOG.warning("Error during LocalSessionMap shutdown: " + e.getMessage());
     }
   }
 
   /** 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();
+    // Simplified lock strategy: only for operations requiring coordination between both maps
+    private final Object coordinationLock = new Object();
 
     public Session get(SessionId id) {
-      // Read operations are atomic on ConcurrentHashMap - no lock needed
+      // ConcurrentHashMap is already thread-safe for reads - 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 {
+      // Synchronize only the coordination between maps, not individual map operations
+      synchronized (coordinationLock) {
         Session previous = sessions.put(id, session);
 
+        // Clean up previous session's URI index if it exists
         if (previous != null && previous.getUri() != null) {
           cleanupUriIndex(previous.getUri(), id);
         }
 
+        // Add new session to URI index
         URI sessionUri = session.getUri();
         if (sessionUri != null) {
           sessionsByUri.computeIfAbsent(sessionUri, k -> ConcurrentHashMap.newKeySet()).add(id);
         }
 
         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 {
+      // Synchronize only the coordination between maps
+      synchronized (coordinationLock) {
         Session removed = sessions.remove(id);
 
+        // Clean up URI index if session existed
         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 {
+      // Synchronize the entire batch operation for consistency
+      synchronized (coordinationLock) {
         Map<URI, Set<SessionId>> uriToSessionIds = new HashMap<>();
 
+        // Collect URI mappings for sessions that exist
         for (SessionId id : sessionIds) {
           Session session = sessions.get(id);
           if (session != null && session.getUri() != null) {
             uriToSessionIds.computeIfAbsent(session.getUri(), k -> new HashSet<>()).add(id);
           }
         }
 
+        // Remove all sessions from main map
         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);
         }
-      } 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.
+     * removal with explicit empty set cleanup. Uses atomic operations to prevent race conditions.
+     * Note: This method should only be called from within coordinationLock synchronized blocks.
      */
     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);
-        }
-      }
+      // Use computeIfPresent for atomic operation to prevent race conditions
+      sessionsByUri.computeIfPresent(
+          uri,
+          (key, sessionIds) -> {
+            sessionIds.remove(sessionId);
+            // Return null to remove the entry if empty, otherwise return the modified set
+            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.
+     * removal with explicit empty set cleanup. Uses atomic operations to prevent race conditions.
+     * Note: This method should only be called from within coordinationLock synchronized blocks.
      */
     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);
-        }
-      }
+      // Use computeIfPresent for atomic operation to prevent race conditions
+      sessionsByUri.computeIfPresent(
+          uri,
+          (key, sessionIds) -> {
+            sessionIds.removeAll(sessionIdsToRemove);
+            // Return null to remove the entry if empty, otherwise return the modified set
+            return sessionIds.isEmpty() ? null : sessionIds;
+          });
     }
 
     /**
      * Periodic cleanup to remove any empty sets that may have been missed. Should be called
-     * periodically to prevent memory leaks.
+     * periodically to prevent memory leaks. Enhanced with better error handling and logging.
      */
     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());
+      synchronized (coordinationLock) {
+        try {
+          int initialSize = sessionsByUri.size();
+          int emptyUrisRemoved = 0;
+
+          // Use iterator to safely remove empty entries during iteration
+          var iterator = sessionsByUri.entrySet().iterator();
+          while (iterator.hasNext()) {
+            var entry = iterator.next();
+            if (entry.getValue().isEmpty()) {
+              iterator.remove();
+              emptyUrisRemoved++;
+            }
           }
-        }
 
-        for (URI emptyUri : emptyUris) {
-          sessionsByUri.remove(emptyUri);
+          if (emptyUrisRemoved > 0) {
+            LOG.info(
+                String.format(
+                    "Maintenance cleanup removed %d empty URI entries from sessionsByUri map (size:"
+                        + " %d -> %d)",
+                    emptyUrisRemoved, initialSize, sessionsByUri.size()));
+          }
+        } catch (Exception e) {
+          LOG.warning("Error during maintenance cleanup: " + e.getMessage());
         }
-      } finally {
-        internalLock.writeLock().unlock();
       }
     }
 
     public Set<SessionId> getSessionsByUri(URI uri) {
-      // Read operations are atomic on ConcurrentHashMap - no lock needed
+      // ConcurrentHashMap is already thread-safe for reads - 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();
-      }
+      // ConcurrentHashMap provides thread-safe iteration - no lock needed
+      // Create defensive copy to prevent external modification
+      return new HashSet<>(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();
-      }
+      // ConcurrentHashMap provides thread-safe iteration - no lock needed
+      // Create defensive copy to prevent external modification
+      return new ArrayList<>(sessions.values());
     }
 
     public int size() {
+      // ConcurrentHashMap size() is already thread-safe
       return sessions.size();
     }
 
     public boolean isEmpty() {
+      // ConcurrentHashMap isEmpty() is already thread-safe
       return sessions.isEmpty();
     }
 
     public void clear() {
-      // Write lock needed: multiple operations across both maps must be atomic
-      internalLock.writeLock().lock();
-      try {
+      // Synchronize clearing both maps for consistency
+      synchronized (coordinationLock) {
         sessions.clear();
         sessionsByUri.clear();
-      } finally {
-        internalLock.writeLock().unlock();
       }
     }
   }