Merge pull request #7 from valkey-io/feature-for-throttle

Feature: support retry when throttled
valkey-io · Feb 27, 2025 · e6791d5 · e6791d5
2 parents 52bb5a1 + 0ccc1f9
commit e6791d5
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 0 deletions.
diff --git a/src/main/java/io/valkey/Protocol.java b/src/main/java/io/valkey/Protocol.java
@@ -66,6 +66,7 @@ public final class Protocol {
   private static final String WRONGPASS_PREFIX = "WRONGPASS";
   private static final String NOPERM_PREFIX = "NOPERM";
   private static final String REDIRECT_PREFIX = "REDIRECT ";
+  private static final String THROTTLED_PREFIX = "THROTTLED ";
 
   private Protocol() {
     throw new InstantiationError("Must not instantiate this class");
@@ -114,6 +115,8 @@ private static void processError(final RedisInputStream is) {
         || message.startsWith(WRONGPASS_PREFIX)
         || message.startsWith(NOPERM_PREFIX)) {
       throw new JedisAccessControlException(message);
+    } else if (message.startsWith(THROTTLED_PREFIX)) {
+      throw new JedisThrottledDataException(message);
     }
     throw new JedisDataException(message);
   }

diff --git a/src/main/java/io/valkey/UnifiedJedis.java b/src/main/java/io/valkey/UnifiedJedis.java
@@ -239,12 +239,24 @@ public UnifiedJedis(ClusterConnectionProvider provider, int maxAttempts, Duratio
         new ClusterCommandObjects());
   }
 
+  public UnifiedJedis(ClusterConnectionProvider provider, int maxAttempts, Duration maxTotalRetriesDuration,
+      boolean retryOnThrottled) {
+    this(new ClusterCommandExecutor(provider, maxAttempts, maxTotalRetriesDuration, retryOnThrottled), provider,
+        new ClusterCommandObjects());
+  }
+
   protected UnifiedJedis(ClusterConnectionProvider provider, int maxAttempts, Duration maxTotalRetriesDuration,
       RedisProtocol protocol) {
     this(new ClusterCommandExecutor(provider, maxAttempts, maxTotalRetriesDuration), provider,
         new ClusterCommandObjects(), protocol);
   }
 
+  public UnifiedJedis(ClusterConnectionProvider provider, int maxAttempts, Duration maxTotalRetriesDuration,
+      boolean retryOnThrottled, RedisProtocol protocol) {
+    this(new ClusterCommandExecutor(provider, maxAttempts, maxTotalRetriesDuration, retryOnThrottled), provider,
+        new ClusterCommandObjects(), protocol);
+  }
+
   /**
    * Constructs a {@code UnifiedJedis} instance with the given parameters.
    *
@@ -291,6 +303,11 @@ public UnifiedJedis(ConnectionProvider provider, int maxAttempts, Duration maxTo
     this(new RetryableCommandExecutor(provider, maxAttempts, maxTotalRetriesDuration), provider);
   }
 
+  public UnifiedJedis(ConnectionProvider provider, int maxAttempts, Duration maxTotalRetriesDuration,
+      boolean retryOnThrottled) {
+    this(new RetryableCommandExecutor(provider, maxAttempts, maxTotalRetriesDuration, retryOnThrottled), provider);
+  }
+
   /**
    * Constructor which supports multiple cluster/database endpoints each with their own isolated connection pool.
    * <p>

diff --git a/src/main/java/io/valkey/exceptions/JedisThrottledDataException.java b/src/main/java/io/valkey/exceptions/JedisThrottledDataException.java
@@ -0,0 +1,18 @@
+package io.valkey.exceptions;
+
+public class JedisThrottledDataException extends JedisDataException {
+
+    private static final long serialVersionUID = 3900401231195549148L;
+
+    public JedisThrottledDataException(String message) {
+        super(message);
+    }
+
+    public JedisThrottledDataException(Throwable cause) {
+        super(cause);
+    }
+
+    public JedisThrottledDataException(String message, Throwable cause) {
+        super(message, cause);
+    }
+}
diff --git a/src/main/java/io/valkey/executors/ClusterCommandExecutor.java b/src/main/java/io/valkey/executors/ClusterCommandExecutor.java
@@ -18,6 +18,7 @@
 import io.valkey.exceptions.JedisConnectionException;
 import io.valkey.exceptions.JedisMovedDataException;
 import io.valkey.exceptions.JedisRedirectionException;
+import io.valkey.exceptions.JedisThrottledDataException;
 import io.valkey.util.IOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -31,12 +32,19 @@ public class ClusterCommandExecutor implements CommandExecutor {
   public final ClusterConnectionProvider provider;
   protected final int maxAttempts;
   protected final Duration maxTotalRetriesDuration;
+  protected final boolean retryOnThrottled;
 
   public ClusterCommandExecutor(ClusterConnectionProvider provider, int maxAttempts,
       Duration maxTotalRetriesDuration) {
+    this(provider, maxAttempts, maxTotalRetriesDuration, false);
+  }
+
+  public ClusterCommandExecutor(ClusterConnectionProvider provider, int maxAttempts,
+      Duration maxTotalRetriesDuration, boolean retryOnThrottled) {
     this.provider = provider;
     this.maxAttempts = maxAttempts;
     this.maxTotalRetriesDuration = maxTotalRetriesDuration;
+    this.retryOnThrottled = retryOnThrottled;
   }
 
   @Override
@@ -124,6 +132,11 @@ public final <T> T executeCommand(CommandObject<T> commandObject) {
           // it rebuilds cluster's slot cache recommended by Redis cluster specification
           provider.renewSlotCache(connection);
         }
+      } catch (JedisThrottledDataException ex) {
+        if (retryOnThrottled) {
+          lastException = ex;
+          sleep(getBackoffSleepMillis(attemptsLeft, deadline));
+        }
       } finally {
         IOUtils.closeQuietly(connection);
       }

diff --git a/src/main/java/io/valkey/executors/RetryableCommandExecutor.java b/src/main/java/io/valkey/executors/RetryableCommandExecutor.java
@@ -7,6 +7,7 @@
 import io.valkey.CommandObject;
 import io.valkey.Connection;
 import io.valkey.annots.VisibleForTesting;
+import io.valkey.exceptions.JedisThrottledDataException;
 import io.valkey.util.IOUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -22,12 +23,19 @@ public class RetryableCommandExecutor implements CommandExecutor {
   protected final ConnectionProvider provider;
   protected final int maxAttempts;
   protected final Duration maxTotalRetriesDuration;
+  protected final boolean retryOnThrottled;
 
   public RetryableCommandExecutor(ConnectionProvider provider, int maxAttempts,
       Duration maxTotalRetriesDuration) {
+    this(provider, maxAttempts, maxTotalRetriesDuration, false);
+  }
+
+  public RetryableCommandExecutor(ConnectionProvider provider, int maxAttempts,
+      Duration maxTotalRetriesDuration, boolean retryOnThrottled) {
     this.provider = provider;
     this.maxAttempts = maxAttempts;
     this.maxTotalRetriesDuration = maxTotalRetriesDuration;
+    this.retryOnThrottled = retryOnThrottled;
   }
 
   @Override
@@ -58,6 +66,11 @@ public final <T> T executeCommand(CommandObject<T> commandObject) {
         if (reset) {
           consecutiveConnectionFailures = 0;
         }
+      } catch (JedisThrottledDataException ex) {
+        if (retryOnThrottled) {
+          lastException = ex;
+          sleep(getBackoffSleepMillis(attemptsLeft, deadline));
+        }
       } finally {
         if (connection != null) {
           connection.close();