[automatic-failover] Support for dynamic add/remove endpoints (#3517)

* - Add/Remove databases safely

* - secure switchToDatabase

* - guard listeners and db switch against race conditions.

* feedbacks from @ggivo
- add close to both MultiDbConnection and CircuitBreaker
- skip switchToDatabase when source and destination is same db

* - add test around attempt to switch to same db

Author: atakavci

src/main/java/io/lettuce/core/failover/CircuitBreaker.java

@@ -1,5 +1,6 @@
 package io.lettuce.core.failover;
 
+import java.io.Closeable;
 import java.io.IOException;
 import java.net.ConnectException;
 import java.util.Arrays;
@@ -25,7 +26,7 @@
  * @author Ali Takavci
  * @since 7.1
  */
-public class CircuitBreaker {
+public class CircuitBreaker implements Closeable {
 
     private static final Logger log = LoggerFactory.getLogger(CircuitBreaker.class);
 
@@ -134,6 +135,11 @@ private void fireStateChanged(State previousState, State newState) {
         }
     }
 
+    @Override
+    public void close() {
+        listeners.clear();
+    }
+
     public static enum State {
         CLOSED, OPEN
     }

src/main/java/io/lettuce/core/failover/RedisDatabase.java

@@ -1,5 +1,7 @@
 package io.lettuce.core.failover;
 
+import java.io.Closeable;
+
 import io.lettuce.core.RedisURI;
 import io.lettuce.core.api.StatefulRedisConnection;
 import io.lettuce.core.failover.CircuitBreaker.CircuitBreakerConfig;
@@ -12,7 +14,7 @@
  * @author Ali Takavci
  * @since 7.1
  */
-public class RedisDatabase<C extends StatefulRedisConnection<?, ?>> {
+public class RedisDatabase<C extends StatefulRedisConnection<?, ?>> implements Closeable {
 
     private final float weight;
 
@@ -53,6 +55,12 @@ public CircuitBreaker getCircuitBreaker() {
         return circuitBreaker;
     }
 
+    @Override
+    public void close() {
+        connection.close();
+        circuitBreaker.close();
+    }
+
     public static class RedisDatabaseConfig {
 
         private RedisURI redisURI;

src/main/java/io/lettuce/core/failover/StatefulRedisMultiDbConnectionImpl.java

@@ -8,6 +8,9 @@
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 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.function.Supplier;
 
 import io.lettuce.core.AbstractRedisClient;
@@ -61,6 +64,12 @@ public class StatefulRedisMultiDbConnectionImpl<C extends StatefulRedisConnectio
 
     protected final DatabaseConnectionFactory<C, K, V> connectionFactory;
 
+    private final ReadWriteLock multiDbLock = new ReentrantReadWriteLock();
+
+    private final Lock readLock = multiDbLock.readLock();
+
+    private final Lock writeLock = multiDbLock.writeLock();
+
     public StatefulRedisMultiDbConnectionImpl(Map<RedisURI, RedisDatabase<C>> connections, ClientResources resources,
             RedisCodec<K, V> codec, Supplier<JsonParser> parser, DatabaseConnectionFactory<C, K, V> connectionFactory) {
         if (connections == null || connections.isEmpty()) {
@@ -151,14 +160,18 @@ public RedisCommands<K, V> sync() {
 
     @Override
     public void addListener(RedisConnectionStateListener listener) {
-        connectionStateListeners.add(listener);
-        current.getConnection().addListener(listener);
+        doBySharedLock(() -> {
+            connectionStateListeners.add(listener);
+            current.getConnection().addListener(listener);
+        });
     }
 
     @Override
     public void removeListener(RedisConnectionStateListener listener) {
-        connectionStateListeners.remove(listener);
-        current.getConnection().removeListener(listener);
+        doBySharedLock(() -> {
+            connectionStateListeners.remove(listener);
+            current.getConnection().removeListener(listener);
+        });
     }
 
     @Override
@@ -224,14 +237,18 @@ public boolean isMulti() {
 
     @Override
     public void addListener(PushListener listener) {
-        pushListeners.add(listener);
-        current.getConnection().addListener(listener);
+        doBySharedLock(() -> {
+            pushListeners.add(listener);
+            current.getConnection().addListener(listener);
+        });
     }
 
     @Override
     public void removeListener(PushListener listener) {
-        pushListeners.remove(listener);
-        current.getConnection().removeListener(listener);
+        doBySharedLock(() -> {
+            pushListeners.remove(listener);
+            current.getConnection().removeListener(listener);
+        });
     }
 
     @Override
@@ -251,21 +268,45 @@ public Iterable<RedisURI> getEndpoints() {
 
     @Override
     public void switchToDatabase(RedisURI redisURI) {
-        RedisDatabase<C> fromDb = current;
-        RedisDatabase<C> toDb = databases.get(redisURI);
-        if (fromDb == null || toDb == null) {
-            throw new UnsupportedOperationException("Cannot initiate switch without a current and target database!");
-        }
-        current = toDb;
-        connectionStateListeners.forEach(listener -> {
-            toDb.getConnection().addListener(listener);
-            fromDb.getConnection().removeListener(listener);
-        });
-        pushListeners.forEach(listener -> {
-            toDb.getConnection().addListener(listener);
-            fromDb.getConnection().removeListener(listener);
+        doByExclusiveLock(() -> {
+            RedisDatabase<C> fromDb = current;
+            RedisDatabase<C> toDb = databases.get(redisURI);
+            if (fromDb == null || toDb == null) {
+                throw new UnsupportedOperationException(
+                        "Unable to switch between endpoints - the driver was not able to locate the source or destination endpoint.");
+            }
+            if (fromDb.equals(toDb)) {
+                return;
+            }
+            current = toDb;
+            connectionStateListeners.forEach(listener -> {
+                toDb.getConnection().addListener(listener);
+                fromDb.getConnection().removeListener(listener);
+            });
+            pushListeners.forEach(listener -> {
+                toDb.getConnection().addListener(listener);
+                fromDb.getConnection().removeListener(listener);
+            });
+            fromDb.getDatabaseEndpoint().handOverCommandQueue(toDb.getDatabaseEndpoint());
         });
-        fromDb.getDatabaseEndpoint().handOverCommandQueue(toDb.getDatabaseEndpoint());
+    }
+
+    protected void doBySharedLock(Runnable operation) {
+        readLock.lock();
+        try {
+            operation.run();
+        } finally {
+            readLock.unlock();
+        }
+    }
+
+    protected void doByExclusiveLock(Runnable operation) {
+        writeLock.lock();
+        try {
+            operation.run();
+        } finally {
+            writeLock.unlock();
+        }
     }
 
     @Override
@@ -294,39 +335,44 @@ public void addDatabase(DatabaseConfig databaseConfig) {
         }
 
         RedisURI redisURI = databaseConfig.getRedisURI();
-        if (databases.containsKey(redisURI)) {
-            throw new IllegalArgumentException("Database already exists: " + redisURI);
-        }
 
-        // Create new database connection using the factory
-        RedisDatabase<C> database = connectionFactory.createDatabase(databaseConfig, codec);
+        doByExclusiveLock(() -> {
+            if (databases.containsKey(redisURI)) {
+                throw new IllegalArgumentException("Database already exists: " + redisURI);
+            }
+
+            // Create new database connection using the factory
+            RedisDatabase<C> database = connectionFactory.createDatabase(databaseConfig, codec);
 
-        // Add listeners to the new connection if it's the current one
-        // (though it won't be current initially since we're just adding it)
-        databases.put(redisURI, database);
+            // Add listeners to the new connection if it's the current one
+            // (though it won't be current initially since we're just adding it)
+            databases.put(redisURI, database);
+
+            database.getCircuitBreaker().addListener(this::onCircuitBreakerStateChange);
+        });
 
-        database.getCircuitBreaker().addListener(this::onCircuitBreakerStateChange);
     }
 
     @Override
     public void removeDatabase(RedisURI redisURI) {
         if (redisURI == null) {
             throw new IllegalArgumentException("RedisURI must not be null");
         }
-
-        RedisDatabase<C> database = databases.get(redisURI);
-        if (database == null) {
-            throw new IllegalArgumentException("Database not found: " + redisURI);
-        }
-
-        if (current.getRedisURI().equals(redisURI)) {
-            throw new UnsupportedOperationException("Cannot remove the currently active database: " + redisURI);
-        }
-
-        // Remove the database and close its connection
-        databases.remove(redisURI);
-        database.getConnection().close();
-        database.getCircuitBreaker().removeListener(this::onCircuitBreakerStateChange);
+        doByExclusiveLock(() -> {
+            RedisDatabase<C> database = null;
+            database = databases.get(redisURI);
+            if (database == null) {
+                throw new IllegalArgumentException("Database not found: " + redisURI);
+            }
+
+            if (current.getRedisURI().equals(redisURI)) {
+                throw new UnsupportedOperationException("Cannot remove the currently active database: " + redisURI);
+            }
+
+            // Remove the database and close its connection
+            databases.remove(redisURI);
+            database.close();
+        });
     }
 
 }

src/main/java/io/lettuce/core/failover/StatefulRedisMultiDbPubSubConnectionImpl.java

@@ -43,14 +43,18 @@ public StatefulRedisMultiDbPubSubConnectionImpl(
 
     @Override
     public void addListener(RedisPubSubListener<K, V> listener) {
-        pubSubListeners.add(listener);
-        current.getConnection().addListener(listener);
+        doBySharedLock(() -> {
+            pubSubListeners.add(listener);
+            current.getConnection().addListener(listener);
+        });
     }
 
     @Override
     public void removeListener(RedisPubSubListener<K, V> listener) {
-        pubSubListeners.remove(listener);
-        current.getConnection().removeListener(listener);
+        doBySharedLock(() -> {
+            pubSubListeners.remove(listener);
+            current.getConnection().removeListener(listener);
+        });
     }
 
     @Override
@@ -87,13 +91,15 @@ protected RedisPubSubReactiveCommandsImpl<K, V> newRedisReactiveCommandsImpl() {
     public void switchToDatabase(RedisURI redisURI) {
 
         RedisDatabase<StatefulRedisPubSubConnection<K, V>> fromDb = current;
-        super.switchToDatabase(redisURI);
-        pubSubListeners.forEach(listener -> {
-            current.getConnection().addListener(listener);
-            fromDb.getConnection().removeListener(listener);
+        doByExclusiveLock(() -> {
+            super.switchToDatabase(redisURI);
+            pubSubListeners.forEach(listener -> {
+                current.getConnection().addListener(listener);
+                fromDb.getConnection().removeListener(listener);
+            });
+
+            moveSubscriptions(fromDb, current);
         });
-
-        moveSubscriptions(fromDb, current);
     }
 
     @SuppressWarnings("unchecked")