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

Author: pujagani

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) {

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) {