Revert rebind address to initial one after delay

# Conflicts:
#	src/main/java/redis/clients/jedis/MaintenanceEventListener.java
#	src/main/java/redis/clients/jedis/RebindAware.java
#	src/test/java/redis/clients/jedis/upgrade/UnifiedJedisProactiveRebindTest.java

Author: ggivo

src/main/java/redis/clients/jedis/Connection.java

@@ -11,6 +11,7 @@
 import java.net.SocketException;
 import java.nio.ByteBuffer;
 import java.nio.CharBuffer;
+import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
@@ -839,8 +840,11 @@ public void accept(PushConsumerContext context) {
       }
     }
     private void onMoving(PushMessage message) {
-      HostAndPort rebindTarget = getRebindTarget(message);
-      eventHandler.getListeners().forEach(listener -> listener.onRebind(rebindTarget));
+      RebindEvent rebindEvent = getRebindTarget(message);
+      if (rebindEvent == null) {
+        return;
+      }
+      eventHandler.getListeners().forEach(listener -> listener.onRebind(rebindEvent.target,rebindEvent.rebindTimeout));
     }
 
     private void onMigrating() {
@@ -859,7 +863,7 @@ private void onFailedOver() {
       eventHandler.getListeners().forEach(MaintenanceEventListener::onFailedOver);
     }
 
-    private HostAndPort getRebindTarget(PushMessage message) {
+    private RebindEvent getRebindTarget(PushMessage message) {
       // Extract domain/ip and port from the message
       // MOVING push message format: ["MOVING", slot, "host:port"]
       List<Object> content = message.getContent();
@@ -869,6 +873,14 @@ private HostAndPort getRebindTarget(PushMessage message) {
         return null;
       }
 
+      Object timeObject = content.get(1); // Get the 3rd element (index 2)
+      if (!(timeObject instanceof Long)) {
+        logger.warn("Invalid re-bind message format, expected 2rd element to be a <time> (Long), got {}",
+            timeObject.getClass());
+        return null;
+      }
+
+
       Object addressObject = content.get(2); // Get the 3rd element (index 2)
       if (!(addressObject instanceof byte[])) {
         logger.warn("Invalid re-bind message format, expected 3rd element to be a byte[], got {}",
@@ -887,18 +899,29 @@ private HostAndPort getRebindTarget(PushMessage message) {
 
         String address = parts[0];
         int port = Integer.parseInt(parts[1]);
-        return new HostAndPort(address, port);
+        return new RebindEvent(new HostAndPort(address, port), Duration.ofSeconds((Long) timeObject));
       } catch (Exception e) {
         logger.warn("Error parsing re-bind target from message: {}", message, e);
         return null;
       }
     }
+
+    private static class RebindEvent{
+      Duration rebindTimeout;
+      HostAndPort target;
+
+      public RebindEvent(HostAndPort target, Duration rebindTimeout) {
+        this.target = target;
+        this.rebindTimeout = rebindTimeout;
+      }
+    }
   }
 
   private class ConnectionRebindHandler implements MaintenanceEventListener {
-      public void onRebind(HostAndPort target) {
+    @Override
+    public void onRebind(HostAndPort target, Duration rebindTimeout) {
         rebindRequested = true;
-      }
+    }
   }
 
   private static class AdaptiveTimeoutHandler implements MaintenanceEventListener {
@@ -914,35 +937,40 @@ public AdaptiveTimeoutHandler(Connection connection) {
       this.connectionRef = new WeakReference<>(connection);
     }
 
+    @Override
     public void onMigrating() {
       Connection connection = connectionRef.get();
       if (connection != null) {
         connection.relaxTimeouts();
       }
     }
 
+    @Override
     public void onMigrated() {
       Connection connection = connectionRef.get();
       if (connection != null) {
         connection.disableRelaxedTimeout();
       }
     }
 
+    @Override
     public void onFailOver() {
       Connection connection = connectionRef.get();
       if (connection != null) {
         connection.relaxTimeouts();
       }
     }
 
+    @Override
     public void onFailedOver() {
       Connection connection = connectionRef.get();
       if (connection != null) {
         connection.disableRelaxedTimeout();
       }
     }
 
-    public void onRebind(HostAndPort target) {
+    @Override
+    public void onRebind(HostAndPort target, Duration rebindTimeout) {
       Connection connection = connectionRef.get();
       if (connection != null) {
         connection.relaxTimeouts();

src/main/java/redis/clients/jedis/ConnectionFactory.java

@@ -6,6 +6,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.time.Duration;
 import java.util.function.Supplier;
 
 import redis.clients.jedis.annots.Experimental;
@@ -15,6 +16,7 @@
 import redis.clients.jedis.csc.Cache;
 import redis.clients.jedis.csc.CacheConnection;
 import redis.clients.jedis.exceptions.JedisException;
+import redis.clients.jedis.util.Expirable;
 
 /**
  * PoolableObjectFactory custom impl.
@@ -29,6 +31,7 @@ public class ConnectionFactory implements PooledObjectFactory<Connection> , Rebi
   private final Supplier<Connection> objectMaker;
 
   private final AuthXEventListener authXEventListener;
+  private final RebindAwareHostPortSupplier rebindAwareHostPortSupplier;
 
   public ConnectionFactory(final HostAndPort hostAndPort) {
     this(hostAndPort, DefaultJedisClientConfig.builder().build(), null);
@@ -66,6 +69,22 @@ private ConnectionFactory(final JedisSocketFactory jedisSocketFactory,
       this.authXEventListener = authXManager.getListener();
       authXManager.start();
     }
+
+    if (clientConfig.isProactiveRebindEnabled()) {
+      if (!(jedisSocketFactory instanceof DefaultJedisSocketFactory)) {
+        throw new IllegalStateException("Rebind not supported for custom JedisSocketFactory implementations");
+      }
+      DefaultJedisSocketFactory factory = (DefaultJedisSocketFactory) jedisSocketFactory;
+      this.rebindAwareHostPortSupplier = wrapHostAndPortSupplier(factory);
+    } else {
+      this.rebindAwareHostPortSupplier = null;
+    }
+  }
+
+  private RebindAwareHostPortSupplier wrapHostAndPortSupplier(DefaultJedisSocketFactory factory) {
+    RebindAwareHostPortSupplier hostPortSupplier =  new RebindAwareHostPortSupplier(factory.getHostAndPort(), factory.getHostAndPortSupplier());
+    factory.setHostAndPortSupplier(hostPortSupplier);
+    return  hostPortSupplier;
   }
 
   private Supplier<Connection> connectionSupplier() {
@@ -141,17 +160,39 @@ private void reAuthenticate(Connection jedis) throws Exception {
     }
   }
 
-
   @Override
-  public void rebind(HostAndPort newHostAndPort) {
-    // TODO : extract interface from DefaultJedisSocketFactory so that we can support custom socket factories
-    if (!(jedisSocketFactory instanceof DefaultJedisSocketFactory)) {
-      throw new IllegalStateException("Rebind not supported for custom JedisSocketFactory implementations");
+  public void rebind(HostAndPort newHostAndPort, Duration rebindTimeout) {
+    if (rebindAwareHostPortSupplier != null) {
+      rebindAwareHostPortSupplier.rebind(newHostAndPort, rebindTimeout);
     }
-
-    DefaultJedisSocketFactory factory = (DefaultJedisSocketFactory) jedisSocketFactory;
-    logger.debug("Rebinding to {}", newHostAndPort);
-    factory.updateHostAndPort(newHostAndPort);
   }
 
+  private static class RebindAwareHostPortSupplier implements Supplier<HostAndPort>, RebindAware {
+    private final Supplier<HostAndPort> delegatedSupplier;
+    private final HostAndPort initialHostAndPort;
+    private volatile Expirable<HostAndPort> rebindTarget;
+
+    public RebindAwareHostPortSupplier(HostAndPort initialHostAndPort,
+        Supplier<HostAndPort> hostAndPortSupplier) {
+      this.initialHostAndPort = initialHostAndPort;
+      this.delegatedSupplier = hostAndPortSupplier;
+    }
+
+    public void rebind(HostAndPort rebindTarget, Duration rebindTimeout) {
+      this.rebindTarget = new Expirable<>(rebindTarget, rebindTimeout);
+    }
+
+    public HostAndPort get() {
+      if (rebindTarget != null && rebindTarget.isValid()) {
+        return rebindTarget.getValue();
+      }
+
+      if (delegatedSupplier != null) {
+        return delegatedSupplier.get();
+      }
+
+      return initialHostAndPort;
+    }
+
+  }
 }

src/main/java/redis/clients/jedis/ConnectionPool.java

@@ -3,19 +3,19 @@
 import org.apache.commons.pool2.PooledObjectFactory;
 import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
 
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 import redis.clients.jedis.annots.Experimental;
 import redis.clients.jedis.authentication.AuthXManager;
 import redis.clients.jedis.csc.Cache;
 import redis.clients.jedis.exceptions.JedisException;
 import redis.clients.jedis.util.Pool;
 
+import java.time.Duration;
 import java.util.concurrent.atomic.AtomicReference;
 
 public class ConnectionPool extends Pool<Connection> {
 
   private AuthXManager authXManager;
+  private RebindHandler rebindHandler;
 
   public ConnectionPool(HostAndPort hostAndPort, JedisClientConfig clientConfig) {
     this(new ConnectionFactory(hostAndPort, clientConfig));
@@ -89,7 +89,7 @@ private void attachAuthenticationListener(AuthXManager authXManager) {
 
   private void attachRebindHandler(JedisClientConfig clientConfig, ConnectionFactory factory) {
     if (clientConfig.isProactiveRebindEnabled()) {
-      RebindHandler rebindHandler = new RebindHandler(this, factory);
+      rebindHandler = new RebindHandler(this, factory);
       clientConfig.getMaintenanceEventHandler().addListener(rebindHandler);
     }
   }
@@ -102,13 +102,13 @@ private static class RebindHandler implements MaintenanceEventListener {
     public RebindHandler(ConnectionPool pool, ConnectionFactory factory) {
       this.pool = pool;
       this.factory = factory;
-    }
+     }
 
     @Override
-    public void onRebind(HostAndPort target) {
+    public void onRebind(HostAndPort target, Duration rebindTimeout) {
       HostAndPort previous = rebindTarget.getAndSet(target);
       if (previous != target) {
-        this.factory.rebind(target);
+        this.factory.rebind(target, rebindTimeout);
         this.pool.clear();
       }
     }

src/main/java/redis/clients/jedis/DefaultJedisSocketFactory.java

@@ -8,6 +8,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.function.Supplier;
 import javax.net.ssl.HostnameVerifier;
 import javax.net.ssl.SSLContext;
 import javax.net.ssl.SSLParameters;
@@ -23,6 +24,7 @@ public class DefaultJedisSocketFactory implements JedisSocketFactory {
       Protocol.DEFAULT_PORT);
 
   private volatile HostAndPort hostAndPort = DEFAULT_HOST_AND_PORT;
+  private volatile Supplier<HostAndPort> hostAndPortSupplier;
   private int connectionTimeout = Protocol.DEFAULT_TIMEOUT;
   private int socketTimeout = Protocol.DEFAULT_TIMEOUT;
   private boolean ssl = false;
@@ -33,19 +35,22 @@ public class DefaultJedisSocketFactory implements JedisSocketFactory {
   private HostAndPortMapper hostAndPortMapper = null;
 
   public DefaultJedisSocketFactory() {
+    hostAndPortSupplier = ()->this.hostAndPort;
   }
 
   public DefaultJedisSocketFactory(HostAndPort hostAndPort) {
     this(hostAndPort, null);
   }
 
   public DefaultJedisSocketFactory(JedisClientConfig config) {
-    this(null, config);
+    this((HostAndPort) null, config);
+    hostAndPortSupplier = ()->this.hostAndPort;
   }
 
   public DefaultJedisSocketFactory(HostAndPort hostAndPort, JedisClientConfig config) {
     if (hostAndPort != null) {
       this.hostAndPort = hostAndPort;
+      hostAndPortSupplier = ()->this.hostAndPort;
     }
     if (config != null) {
       this.connectionTimeout = config.getConnectionTimeoutMillis();
@@ -154,17 +159,55 @@ private Socket createSslSocket(HostAndPort _hostAndPort, Socket socket) throws I
     return new SSLSocketWrapper(sslSocket, plainSocket);
   }
 
+  /**
+   *  Updates the HostAndPort used to create the socket.
+   *  <p>
+   *  Note: if Supplier<HostAndPort> is set it has precedence over provided hostAndPort.
+   *  </p>
+   * */
   public void updateHostAndPort(HostAndPort hostAndPort) {
     this.hostAndPort = hostAndPort;
   }
 
+
+  /**
+   * Returns the HostAndPort that is used to create the socket.
+   * <p>
+   * Note: if Supplier<HostAndPort> is set it has precedence over HostAndPort.
+   *</p>
+   * @return the HostAndPort that is used to create the socket.
+   */
   public HostAndPort getHostAndPort() {
     return this.hostAndPort;
   }
 
+
+  /**
+   * Returns the Supplier<HostAndPort> that is used to create the socket.
+   * @return the Supplier<HostAndPort> that is used to create the socket.
+   */
+  public Supplier<HostAndPort> getHostAndPortSupplier() {
+    return this.hostAndPortSupplier;
+  }
+
+  /**
+   * Sets the Supplier<HostAndPort> that is used to create the socket.
+   * @param hostAndPortSupplier the Supplier<HostAndPort> that is used to create the socket.
+   */
+  public void setHostAndPortSupplier(Supplier<HostAndPort> hostAndPortSupplier) {
+    this.hostAndPortSupplier = hostAndPortSupplier;
+  }
+
   protected HostAndPort getSocketHostAndPort() {
     HostAndPortMapper mapper = hostAndPortMapper;
-    HostAndPort hap = this.hostAndPort;
+
+    HostAndPort hap;
+    if (this.hostAndPortSupplier != null) {
+      hap = this.hostAndPortSupplier.get();
+    } else {
+      hap = this.hostAndPort;
+    }
+
     if (mapper != null) {
       HostAndPort mapped = mapper.getHostAndPort(hap);
       if (mapped != null) {

src/main/java/redis/clients/jedis/MaintenanceEventListener.java

@@ -1,5 +1,7 @@
 package redis.clients.jedis;
 
+import java.time.Duration;
+
 public interface MaintenanceEventListener {
 
   default void onMigrating() {
@@ -14,6 +16,6 @@ default void onFailOver() {
   default void onFailedOver() {
   };
 
-  default void onRebind(HostAndPort target) {
+  default void onRebind(HostAndPort target, Duration rebindTimeout) {
   };
-}
+}

src/main/java/redis/clients/jedis/RebindAware.java

@@ -2,6 +2,8 @@
 
 import redis.clients.jedis.annots.Experimental;
 
+import java.time.Duration;
+
 /**
  * Interface for components that support rebinding to a new host and port.
  * <p>
@@ -23,6 +25,7 @@ public interface RebindAware {
    * Components might decide to reject the re-bind request if they are not in a state to support it.
    * </p>
    * @param newHostAndPort The new host and port to use for new connections
+   * @param expiry The duration after which the rebind should be considered expired
    */
-  void rebind(HostAndPort newHostAndPort);
+  void rebind(HostAndPort newHostAndPort, Duration expiry);
 }