[grid] Add node removed listener to Redis and Jdbc session map implem…

…entations
SeleniumHQ · Aug 31, 2021 · 95c78af · 95c78af
1 parent 97b6c1b
commit 95c78af
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/java/src/org/openqa/selenium/grid/sessionmap/jdbc/JdbcBackedSessionMap.java b/java/src/org/openqa/selenium/grid/sessionmap/jdbc/JdbcBackedSessionMap.java
@@ -23,6 +23,7 @@
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.grid.config.Config;
 import org.openqa.selenium.grid.config.ConfigException;
+import org.openqa.selenium.grid.data.NodeRemovedEvent;
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.grid.data.SessionClosedEvent;
 import org.openqa.selenium.grid.log.LoggingOptions;
@@ -83,6 +84,11 @@ public JdbcBackedSessionMap(Tracer tracer, Connection jdbcConnection, EventBus b
 
     this.connection = jdbcConnection;
     this.bus.addListener(SessionClosedEvent.listener(this::remove));
+
+    this.bus.addListener(NodeRemovedEvent.listener(nodeStatus -> nodeStatus.getSlots().stream()
+      .filter(slot -> slot.getSession() != null)
+      .map(slot -> slot.getSession().getId())
+      .forEach(this::remove)));
   }
 
   public static SessionMap create(Config config) {

diff --git a/java/src/org/openqa/selenium/grid/sessionmap/redis/RedisBackedSessionMap.java b/java/src/org/openqa/selenium/grid/sessionmap/redis/RedisBackedSessionMap.java
@@ -24,6 +24,7 @@
 import org.openqa.selenium.NoSuchSessionException;
 import org.openqa.selenium.events.EventBus;
 import org.openqa.selenium.grid.config.Config;
+import org.openqa.selenium.grid.data.NodeRemovedEvent;
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.grid.data.SessionClosedEvent;
 import org.openqa.selenium.grid.log.LoggingOptions;
@@ -79,6 +80,11 @@ public RedisBackedSessionMap(Tracer tracer, URI serverUri, EventBus bus) {
     this.connection = new GridRedisClient(serverUri);
     this.serverUri = serverUri;
     this.bus.addListener(SessionClosedEvent.listener(this::remove));
+
+    this.bus.addListener(NodeRemovedEvent.listener(nodeStatus -> nodeStatus.getSlots().stream()
+      .filter(slot -> slot.getSession() != null)
+      .map(slot -> slot.getSession().getId())
+      .forEach(this::remove)));
   }
 
   public static SessionMap create(Config config) {