Feat: Customize retry error type and processing handler (#8)

Signed-off-by: bodong.ybd <[email protected]>

Author: yangbodong22011

src/main/java/io/valkey/ExceptionHandler.java

@@ -0,0 +1,37 @@
+package io.valkey;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Predicate;
+
+public class ExceptionHandler {
+    private Map<Predicate<String>, ErrorCallback> patternCallbacks;
+
+    public interface ErrorCallback {
+        void onError(String errorMessage);
+    }
+
+    public ExceptionHandler() {
+        this.patternCallbacks = new HashMap<>();
+    }
+
+    public void register(Predicate<String> pattern, ErrorCallback callback) {
+        patternCallbacks.put(pattern, callback);
+    }
+
+    // This method allows the registration of a new error pattern and its corresponding callback.
+    // It takes a Predicate<String> that defines the error pattern to match,
+    // and an ErrorCallback that will be executed when the pattern is matched.
+    // This enables dynamic handling of different types of errors based on their messages.
+    public void handleException(Exception e) {
+        Optional.ofNullable(e.getMessage()).ifPresent(errorMessage -> {
+            patternCallbacks.forEach((pattern, callback) -> {
+                if (pattern.test(errorMessage)) {
+                    callback.onError(errorMessage);
+                }
+            });
+        });
+    }
+}
+

src/main/java/io/valkey/UnifiedJedis.java

@@ -291,6 +291,11 @@ public UnifiedJedis(ConnectionProvider provider, int maxAttempts, Duration maxTo
     this(new RetryableCommandExecutor(provider, maxAttempts, maxTotalRetriesDuration), provider);
   }
 
+  public UnifiedJedis(ConnectionProvider provider, int maxAttempts, Duration maxTotalRetriesDuration,
+      ExceptionHandler handler) {
+    this(new RetryableCommandExecutor(provider, maxAttempts, maxTotalRetriesDuration, handler), provider);
+  }
+
   /**
    * Constructor which supports multiple cluster/database endpoints each with their own isolated connection pool.
    * <p>

src/main/java/io/valkey/executors/RetryableCommandExecutor.java

@@ -6,7 +6,9 @@
 
 import io.valkey.CommandObject;
 import io.valkey.Connection;
+import io.valkey.ExceptionHandler;
 import io.valkey.annots.VisibleForTesting;
+import io.valkey.exceptions.JedisDataException;
 import io.valkey.util.IOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -22,12 +24,19 @@ public class RetryableCommandExecutor implements CommandExecutor {
   protected final ConnectionProvider provider;
   protected final int maxAttempts;
   protected final Duration maxTotalRetriesDuration;
+  protected final ExceptionHandler handler;
 
   public RetryableCommandExecutor(ConnectionProvider provider, int maxAttempts,
       Duration maxTotalRetriesDuration) {
+    this(provider, maxAttempts, maxTotalRetriesDuration, null);
+  }
+
+  public RetryableCommandExecutor(ConnectionProvider provider, int maxAttempts, Duration maxTotalRetriesDuration,
+      ExceptionHandler handler) {
     this.provider = provider;
     this.maxAttempts = maxAttempts;
     this.maxTotalRetriesDuration = maxTotalRetriesDuration;
+    this.handler = handler;
   }
 
   @Override
@@ -58,6 +67,11 @@ public final <T> T executeCommand(CommandObject<T> commandObject) {
         if (reset) {
           consecutiveConnectionFailures = 0;
         }
+      } catch (JedisException e) {
+        lastException = e;
+        if (handler != null) {
+          handler.handleException(e);
+        }
       } finally {
         if (connection != null) {
           connection.close();

src/test/java/io/valkey/util/ExceptionHandlerTest.java

@@ -0,0 +1,108 @@
+package io.valkey.util;
+
+import java.time.Duration;
+import java.util.Map;
+import java.util.UUID;
+
+import io.valkey.DefaultJedisClientConfig;
+import io.valkey.ExceptionHandler;
+import io.valkey.HostAndPort;
+import io.valkey.HostAndPorts;
+import io.valkey.Jedis;
+import io.valkey.UnifiedJedis;
+import io.valkey.providers.PooledConnectionProvider;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+
+public class ExceptionHandlerTest {
+    private final HostAndPort hostAndPort = HostAndPorts.getRedisServers().get(0);
+    private Jedis jedis;
+
+    @Before
+    public void setUp() {
+        jedis = new Jedis(hostAndPort);
+        jedis.auth("foobared");
+    }
+
+    @Test
+    public void retryWithExceptionHandler() {
+        int maxAttempts = 10;
+        Duration maxTotalRetriesDuration = Duration.ofSeconds(10);
+        PooledConnectionProvider provider = new PooledConnectionProvider(hostAndPort,
+            DefaultJedisClientConfig.builder().password("foobared").build());
+        ExceptionHandler handler = new ExceptionHandler();
+        handler.register(
+            message -> message.contains("WRONGTYPE Operation against a key holding the wrong kind of value"),
+            errorMessage -> {
+                try {
+                    System.out.println("Retrying...");
+                    Thread.sleep(1000);
+                } catch (InterruptedException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        );
+
+        String key = UUID.randomUUID().toString();
+        Assert.assertEquals("OK", jedis.set(key, key));
+        new Thread(new Runnable() {
+            @Override
+            public void run() {
+                try {
+                    // after 5s, del this key.
+                    Thread.sleep(5000);
+                    Assert.assertEquals(1, jedis.del(key));
+                } catch (Exception e) {
+                    e.printStackTrace();
+                }
+            }
+        }).start();
+
+        UnifiedJedis unifiedJedis = new UnifiedJedis(provider, maxAttempts, maxTotalRetriesDuration, handler);
+        Map<String, String> map = unifiedJedis.hgetAll(key);
+        Assert.assertTrue(map.isEmpty());
+    }
+
+    @Test
+    public void retryWithExponentialBackoffCallback() {
+        int maxAttempts = 4;
+        Duration maxTotalRetriesDuration = Duration.ofSeconds(20);
+        PooledConnectionProvider provider = new PooledConnectionProvider(hostAndPort,
+            DefaultJedisClientConfig.builder().password("foobared").build());
+        String key = UUID.randomUUID().toString();
+        Assert.assertEquals("OK", jedis.set(key, key));
+        ExceptionHandler handler = new ExceptionHandler();
+        class ExponentialBackoffCallback implements ExceptionHandler.ErrorCallback {
+            private int attempt = 0; // 计数器
+
+            @Override
+            public void onError(String errorMessage) {
+                // 计算 sleep 时间（2^attempt 秒）
+                int sleepTime = (int) Math.pow(2, attempt);
+                try {
+                    System.out.println("Sleeping for " + sleepTime + " seconds before handling: " + errorMessage);
+                    Thread.sleep(sleepTime * 1000); // 转换为毫秒
+                    if (attempt == maxAttempts - 2) {
+                        Assert.assertEquals(1, jedis.del(key));
+                    }
+                } catch (InterruptedException ie) {
+                    Thread.currentThread().interrupt(); // 恢复中断状态
+                }
+                attempt++; // 增加计数
+            }
+        }
+
+        handler.register(
+            message -> message.contains("WRONGTYPE Operation against a key holding the wrong kind of value"),
+            new ExponentialBackoffCallback()
+        );
+
+        UnifiedJedis unifiedJedis = new UnifiedJedis(provider, maxAttempts, maxTotalRetriesDuration, handler);
+        long start = System.currentTimeMillis();
+        Map<String, String> map = unifiedJedis.hgetAll(key);
+        long end = System.currentTimeMillis();
+        Assert.assertTrue(map.isEmpty());
+        Assert.assertTrue(end - start >= 7000); // 2^0 + 2^1 + 2^2 = 7 seconds
+    }
+}