Added configurable timeout to connection string #152

redis · Nov 17, 2015 · d937e87 · d937e87
1 parent 24b04cc
commit d937e87
Show file tree

Hide file tree

Showing 2 changed files with 103 additions and 0 deletions.
diff --git a/src/main/java/com/lambdaworks/redis/RedisURI.java b/src/main/java/com/lambdaworks/redis/RedisURI.java
@@ -8,8 +8,12 @@
 import java.net.SocketAddress;
 import java.net.URI;
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
+import java.util.StringTokenizer;
 import java.util.concurrent.TimeUnit;
 
 import com.google.common.collect.ImmutableSet;
@@ -57,6 +61,20 @@ public class RedisURI implements Serializable, ConnectionPoint {
     public static final String URI_SCHEME_REDIS = "redis";
     public static final String URI_SCHEME_REDIS_SECURE = "rediss";
     public static final String URI_SCHEME_REDIS_SOCKET = "redis-socket";
+    public static final String TIMEOUT_PARAMETER_NAME = "timeout";
+    public static final Map<String, TimeUnit> TIME_UNIT_MAP;
+
+    static {
+        Map<String, TimeUnit> unitMap = new HashMap<String, TimeUnit>();
+        unitMap.put("ns", TimeUnit.NANOSECONDS);
+        unitMap.put("us", TimeUnit.MICROSECONDS);
+        unitMap.put("ms", TimeUnit.MILLISECONDS);
+        unitMap.put("s", TimeUnit.SECONDS);
+        unitMap.put("m", TimeUnit.MINUTES);
+        unitMap.put("h", TimeUnit.HOURS);
+        unitMap.put("d", TimeUnit.DAYS);
+        TIME_UNIT_MAP = Collections.unmodifiableMap(unitMap);
+    }
 
     /**
      * The default sentinel port.
@@ -169,10 +187,54 @@ public static RedisURI create(URI uri) {
             }
         }
 
+        if (isNotEmpty(uri.getQuery())) {
+            StringTokenizer st = new StringTokenizer(uri.getQuery(), "&;");
+            while (st.hasMoreTokens()) {
+                String queryParam = st.nextToken().toLowerCase();
+                if (queryParam.startsWith(TIMEOUT_PARAMETER_NAME)) {
+                    parseTimeout(builder, queryParam);
+                }
+            }
+        }
+
         return builder.build();
 
     }
 
+    private static void parseTimeout(Builder builder, String queryParam) {
+        int index = queryParam.indexOf('=');
+        if (index < 0) {
+            return;
+        }
+
+        String timeoutString = queryParam.substring(index + 1);
+
+        int numbersEnd = 0;
+        while (numbersEnd < timeoutString.length() && Character.isDigit(timeoutString.charAt(numbersEnd))) {
+            numbersEnd++;
+        }
+
+        if (numbersEnd == 0) {
+            if (timeoutString.startsWith("-")) {
+                builder.withTimeout(0, TimeUnit.MILLISECONDS);
+            } else {
+                // no-op, leave defaults
+            }
+        } else {
+            String timeoutValueString = timeoutString.substring(0, numbersEnd);
+            long timeoutValue = Long.parseLong(timeoutValueString);
+            builder.withTimeout(timeoutValue, TimeUnit.MILLISECONDS);
+
+            String suffix = timeoutString.substring(numbersEnd);
+            TimeUnit timeoutUnit = TIME_UNIT_MAP.get(suffix);
+            if (timeoutUnit == null) {
+                timeoutUnit = TimeUnit.MILLISECONDS;
+            }
+
+            builder.withTimeout(timeoutValue, timeoutUnit);
+        }
+    }
+
     private static Builder configureStandalone(URI uri) {
         Builder builder;
         Set<String> allowedSchemes = ImmutableSet.of(URI_SCHEME_REDIS, URI_SCHEME_REDIS_SECURE, URI_SCHEME_REDIS_SOCKET);

diff --git a/src/test/java/com/lambdaworks/redis/RedisURITest.java b/src/test/java/com/lambdaworks/redis/RedisURITest.java
@@ -4,6 +4,7 @@
 
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 
 import org.junit.Test;
 
@@ -85,4 +86,44 @@ public void socketEqualsTest() throws Exception {
         assertThat(redisURI3).isNotEqualTo(redisURI2);
         assertThat(redisURI3.hashCode()).isNotEqualTo(redisURI2.hashCode());
     }
+
+    @Test
+    public void timeoutParsingTest() throws Exception {
+        checkUriTimeout("redis://auth@localhost:1234/5?timeout=5000", 5000, TimeUnit.MILLISECONDS);
+        checkUriTimeout("redis://auth@localhost:1234/5?timeout=5000ms", 5000, TimeUnit.MILLISECONDS);
+        checkUriTimeout("redis://auth@localhost:1234/5?timeout=5s", 5, TimeUnit.SECONDS);
+        checkUriTimeout("redis://auth@localhost:1234/5?timeout=100us", 100, TimeUnit.MICROSECONDS);
+        checkUriTimeout("redis://auth@localhost:1234/5?TIMEOUT=1000000NS", 1000000, TimeUnit.NANOSECONDS);
+        checkUriTimeout("redis://auth@localhost:1234/5?timeout=60m", 60, TimeUnit.MINUTES);
+        checkUriTimeout("redis://auth@localhost:1234/5?timeout=24h", 24, TimeUnit.HOURS);
+        checkUriTimeout("redis://auth@localhost:1234/5?timeout=1d", 1, TimeUnit.DAYS);
+
+        checkUriTimeout("redis://auth@localhost:1234/5?timeout=-1", 0, TimeUnit.MILLISECONDS);
+
+        RedisURI defaultUri = new RedisURI();
+        checkUriTimeout("redis://auth@localhost:1234/5?timeout=junk", defaultUri.getTimeout(), defaultUri.getUnit());
+    }
+
+    @Test
+    public void queryStringDecodingTest() throws Exception {
+        String timeout = "%74%69%6D%65%6F%75%74";
+        String eq = "%3d";
+        String s = "%73";
+        checkUriTimeout("redis://auth@localhost:1234/5?" + timeout + eq + "5" + s, 5, TimeUnit.SECONDS);
+    }
+
+    @Test
+    public void timeoutParsingWithJunkParamTest() throws Exception {
+        RedisURI redisURI1 = RedisURI.create("redis-sentinel://auth@localhost:1234/5?timeout=5s;junkparam=#master-instance");
+        assertThat(redisURI1.getTimeout()).isEqualTo(5);
+        assertThat(redisURI1.getUnit()).isEqualTo(TimeUnit.SECONDS);
+        assertThat(redisURI1.getSentinelMasterId()).isEqualTo("master-instance");
+    }
+
+    private RedisURI checkUriTimeout(String uri, long expectedTimeout, TimeUnit expectedUnit) {
+        RedisURI redisURI1 = RedisURI.create(uri);
+        assertThat(redisURI1.getTimeout()).isEqualTo(expectedTimeout);
+        assertThat(redisURI1.getUnit()).isEqualTo(expectedUnit);
+        return redisURI1;
+    }
 }