From 33c69b5398a3bffb3094af6d20e9697b96c72260 Mon Sep 17 00:00:00 2001
From: joerg1985 <16140691+joerg1985@users.noreply.github.com>
Date: Thu, 8 Aug 2024 11:16:43 +0200
Subject: [PATCH] [grid] ensure the local_sessionmap.remove event is raised
 (#14337)

Co-authored-by: Diego Molina <diemol@users.noreply.github.com>
---
 .../sessionmap/local/LocalSessionMap.java     | 72 ++++++++-----------
 1 file changed, 28 insertions(+), 44 deletions(-)

diff --git a/java/src/org/openqa/selenium/grid/sessionmap/local/LocalSessionMap.java b/java/src/org/openqa/selenium/grid/sessionmap/local/LocalSessionMap.java
index 2dae2e5779e4f..6e2ff03859170 100644
--- a/java/src/org/openqa/selenium/grid/sessionmap/local/LocalSessionMap.java
+++ b/java/src/org/openqa/selenium/grid/sessionmap/local/LocalSessionMap.java
@@ -20,12 +20,12 @@
 import static org.openqa.selenium.remote.RemoteTags.SESSION_ID;
 import static org.openqa.selenium.remote.RemoteTags.SESSION_ID_EVENT;
 
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReadWriteLock;
-import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.concurrent.ConcurrentMap;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.grid.config.Config;
@@ -48,28 +48,14 @@ public class LocalSessionMap extends SessionMap {
   private static final Logger LOG = Logger.getLogger(LocalSessionMap.class.getName());
 
   private final EventBus bus;
-  private final Map<SessionId, Session> knownSessions = new ConcurrentHashMap<>();
-  private final ReadWriteLock lock = new ReentrantReadWriteLock(/* be fair */ true);
+  private final ConcurrentMap<SessionId, Session> knownSessions = new ConcurrentHashMap<>();
 
   public LocalSessionMap(Tracer tracer, EventBus bus) {
     super(tracer);
 
     this.bus = Require.nonNull("Event bus", bus);
 
-    bus.addListener(
-        SessionClosedEvent.listener(
-            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);
-                knownSessions.remove(id);
-                String sessionDeletedMessage = "Deleted session from local Session Map";
-                span.addEvent(sessionDeletedMessage, attributeMap);
-                LOG.info(String.format("%s, Id: %s", sessionDeletedMessage, id));
-              }
-            }));
+    bus.addListener(SessionClosedEvent.listener(this::remove));
 
     bus.addListener(
         NodeRemovedEvent.listener(
@@ -77,14 +63,19 @@ public LocalSessionMap(Tracer tracer, EventBus bus) {
                 nodeStatus.getSlots().stream()
                     .filter(slot -> slot.getSession() != null)
                     .map(slot -> slot.getSession().getId())
-                    .forEach(knownSessions::remove)));
+                    .forEach(this::remove)));
 
     bus.addListener(
         NodeRestartedEvent.listener(
-            nodeStatus ->
-                knownSessions
-                    .values()
-                    .removeIf(value -> value.getUri().equals(nodeStatus.getExternalUri()))));
+            nodeStatus -> {
+              List<SessionId> toRemove =
+                  knownSessions.entrySet().stream()
+                      .filter((e) -> e.getValue().getUri().equals(nodeStatus.getExternalUri()))
+                      .map(Map.Entry::getKey)
+                      .collect(Collectors.toList());
+
+              toRemove.forEach(this::remove);
+            }));
   }
 
   public static SessionMap create(Config config) {
@@ -103,8 +94,6 @@ public boolean isReady() {
   public boolean add(Session session) {
     Require.nonNull("Session", session);
 
-    Lock writeLock = lock.writeLock();
-    writeLock.lock();
     try (Span span = tracer.getCurrentContext().createSpan("local_sessionmap.add")) {
       AttributeMap attributeMap = tracer.createAttributeMap();
       attributeMap.put(AttributeKey.LOGGER_CLASS.getKey(), getClass().getName());
@@ -115,8 +104,6 @@ public boolean add(Session session) {
       span.addEvent("Added session into local session map", attributeMap);
 
       return true;
-    } finally {
-      writeLock.unlock();
     }
   }
 
@@ -124,30 +111,27 @@ public boolean add(Session session) {
   public Session get(SessionId id) {
     Require.nonNull("Session ID", id);
 
-    Lock readLock = lock.readLock();
-    readLock.lock();
-    try {
-      Session session = knownSessions.get(id);
-      if (session == null) {
-        throw new NoSuchSessionException("Unable to find session with ID: " + id);
-      }
-
-      return session;
-    } finally {
-      readLock.unlock();
+    Session session = knownSessions.get(id);
+    if (session == null) {
+      throw new NoSuchSessionException("Unable to find session with ID: " + id);
     }
+
+    return session;
   }
 
   @Override
   public void remove(SessionId id) {
     Require.nonNull("Session ID", id);
 
-    Lock writeLock = lock.writeLock();
-    writeLock.lock();
-    try {
+    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);
       knownSessions.remove(id);
-    } finally {
-      writeLock.unlock();
+      String sessionDeletedMessage = "Deleted session from local Session Map";
+      span.addEvent(sessionDeletedMessage, attributeMap);
+      LOG.info(String.format("%s, Id: %s", sessionDeletedMessage, id));
     }
   }
 }