[grid] Remove sessions from Redis backed session map on restart

SeleniumHQ · Sep 2, 2021 · 8e6f4fe · 8e6f4fe
1 parent a98b38f
commit 8e6f4fe
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/java/src/org/openqa/selenium/grid/sessionmap/redis/RedisBackedSessionMap.java b/java/src/org/openqa/selenium/grid/sessionmap/redis/RedisBackedSessionMap.java
@@ -25,6 +25,7 @@
 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.NodeRestartedEvent;
 import org.openqa.selenium.grid.data.Session;
 import org.openqa.selenium.grid.data.SessionClosedEvent;
 import org.openqa.selenium.grid.log.LoggingOptions;
@@ -85,6 +86,8 @@ public RedisBackedSessionMap(Tracer tracer, URI serverUri, EventBus bus) {
       .filter(slot -> slot.getSession() != null)
       .map(slot -> slot.getSession().getId())
       .forEach(this::remove)));
+
+    bus.addListener(NodeRestartedEvent.listener(nodeStatus -> this.removeByUri(nodeStatus.getExternalUri())));
   }
 
   public static SessionMap create(Config config) {
@@ -275,6 +278,26 @@ public void remove(SessionId id) {
     }
   }
 
+  public void removeByUri(URI uri) {
+    List<String> uriKeys = connection.getKeysByPattern("session:*:uri");
+
+    if (uriKeys.isEmpty()) {
+      return;
+    }
+
+    String[] keys = new String[uriKeys.size()];
+    keys = uriKeys.toArray(keys);
+
+    List<KeyValue<String, String>> keyValues = connection.mget(keys);
+    keyValues.stream()
+      .filter(entry -> entry.getValue().equals(uri.toString()))
+      .map(KeyValue::getKey)
+      .map(key -> {
+        String[] sessionKey = key.split(":");
+        return new SessionId(sessionKey[1]);
+      }).forEach(this::remove);
+  }
+
   @Override
   public boolean isReady() {
     return connection.isOpen();

diff --git a/java/src/org/openqa/selenium/redis/GridRedisClient.java b/java/src/org/openqa/selenium/redis/GridRedisClient.java
@@ -76,6 +76,10 @@ public String get(String key) {
     return connection.sync().get(key);
   }
 
+  public List<String> getKeysByPattern(String pattern) {
+    return connection.sync().keys(pattern);
+  }
+
   public boolean isOpen() {
     return connection.isOpen();
   }