add timeout

yingzi5920 · Dec 24, 2014 · 8f6bb3b · 8f6bb3b
1 parent 1441266
commit 8f6bb3b
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 7 deletions.
diff --git a/extern/jodis/pom.xml b/extern/jodis/pom.xml
@@ -24,7 +24,7 @@ THE SOFTWARE.
     <modelVersion>4.0.0</modelVersion>
     <groupId>com.wandoulabs.jodis</groupId>
     <artifactId>jodis</artifactId>
-    <version>0.1.1</version>
+    <version>0.1.2</version>
     <packaging>jar</packaging>
     <name>${project.artifactId}</name>
     <description>Round Robin jedis pool for codis</description>

diff --git a/extern/jodis/src/main/java/com/wandoulabs/jodis/RoundRobinJedisPool.java b/extern/jodis/src/main/java/com/wandoulabs/jodis/RoundRobinJedisPool.java
@@ -76,6 +76,8 @@ public class RoundRobinJedisPool implements JedisResourcePool {
 
     private static final int CURATOR_RETRY_MAX_SLEEP_MS = 30 * 1000;
 
+    private static final int JEDIS_POOL_TIMEOUT_UNSET = -1;
+
     private static final ImmutableSet<PathChildrenCacheEvent.Type> RESET_TYPES = Sets
             .immutableEnumSet(PathChildrenCacheEvent.Type.CHILD_ADDED,
                     PathChildrenCacheEvent.Type.CHILD_UPDATED,
@@ -105,6 +107,31 @@ public PooledObject(String addr, JedisPool pool) {
 
     private final JedisPoolConfig poolConfig;
 
+    private final int timeout;
+
+    /**
+     * Create a RoundRobinJedisPool with default timeout.
+     * <p>
+     * We create a CuratorFramework with infinite retry number. If you do not
+     * like the behavior, use the other constructor that allow you pass a
+     * CuratorFramework created by yourself.
+     * 
+     * @param zkAddr
+     *            ZooKeeper connect string. e.g., "zk1:2181"
+     * @param zkSessionTimeoutMs
+     *            ZooKeeper session timeout in ms
+     * @param zkPath
+     *            the codis proxy dir on ZooKeeper. e.g.,
+     *            "/zk/codis/db_xxx/proxy"
+     * @param poolConfig
+     *            same as JedisPool
+     * @see #RoundRobinJedisPool(String, int, String, JedisPoolConfig, int)
+     */
+    public RoundRobinJedisPool(String zkAddr, int zkSessionTimeoutMs, String zkPath,
+            JedisPoolConfig poolConfig) {
+        this(zkAddr, zkSessionTimeoutMs, zkPath, poolConfig, JEDIS_POOL_TIMEOUT_UNSET);
+    }
+
     /**
      * Create a RoundRobinJedisPool.
      * <p>
@@ -121,22 +148,25 @@ public PooledObject(String addr, JedisPool pool) {
      *            "/zk/codis/db_xxx/proxy"
      * @param poolConfig
      *            same as JedisPool
+     * @param timeout
+     *            timeout of JedisPool
      * @see #RoundRobinJedisPool(CuratorFramework, boolean, String,
-     *      JedisPoolConfig)
+     *      JedisPoolConfig, int)
      */
     public RoundRobinJedisPool(String zkAddr, int zkSessionTimeoutMs, String zkPath,
-            JedisPoolConfig poolConfig) {
+            JedisPoolConfig poolConfig, int timeout) {
         this(CuratorFrameworkFactory
                 .builder()
                 .connectString(zkAddr)
                 .sessionTimeoutMs(zkSessionTimeoutMs)
                 .retryPolicy(
                         new BoundedExponentialBackoffRetryUntilElapsed(CURATOR_RETRY_BASE_SLEEP_MS,
-                                CURATOR_RETRY_MAX_SLEEP_MS, -1L)).build(), true, zkPath, poolConfig);
+                                CURATOR_RETRY_MAX_SLEEP_MS, -1L)).build(), true, zkPath,
+                poolConfig, timeout);
     }
 
     /**
-     * Create a RoundRobinJedisPool.
+     * Create a RoundRobinJedisPool with default timeout.
      * 
      * @param curatorClient
      *            We will start it if it has not started yet.
@@ -150,7 +180,28 @@ public RoundRobinJedisPool(String zkAddr, int zkSessionTimeoutMs, String zkPath,
      */
     public RoundRobinJedisPool(CuratorFramework curatorClient, boolean closeCurator, String zkPath,
             JedisPoolConfig poolConfig) {
+        this(curatorClient, closeCurator, zkPath, poolConfig, JEDIS_POOL_TIMEOUT_UNSET);
+    }
+
+    /**
+     * Create a RoundRobinJedisPool.
+     * 
+     * @param curatorClient
+     *            We will start it if it has not started yet.
+     * @param closeCurator
+     *            Whether to close the curatorClient passed in when close.
+     * @param zkPath
+     *            the codis proxy dir on ZooKeeper. e.g.
+     *            "/zk/codis/db_xxx/proxy"
+     * @param poolConfig
+     *            same as JedisPool
+     * @param timeout
+     *            timeout of JedisPool
+     */
+    public RoundRobinJedisPool(CuratorFramework curatorClient, boolean closeCurator, String zkPath,
+            JedisPoolConfig poolConfig, int timeout) {
         this.poolConfig = poolConfig;
+        this.timeout = timeout;
         this.curatorClient = curatorClient;
         this.closeCurator = closeCurator;
         watcher = new PathChildrenCache(curatorClient, zkPath, true);
@@ -203,8 +254,14 @@ private void resetPools() {
                 if (pool == null) {
                     LOG.info("Add new proxy: " + addr);
                     String[] hostAndPort = addr.split(":");
-                    pool = new PooledObject(addr, new JedisPool(poolConfig, hostAndPort[0],
-                            Integer.parseInt(hostAndPort[1])));
+                    String host = hostAndPort[0];
+                    int port = Integer.parseInt(hostAndPort[1]);
+                    if (timeout == JEDIS_POOL_TIMEOUT_UNSET) {
+                        pool = new PooledObject(addr, new JedisPool(poolConfig, host, port));
+                    } else {
+                        pool = new PooledObject(addr,
+                                new JedisPool(poolConfig, host, port, timeout));
+                    }
                 }
                 builder.add(pool);
             } catch (Exception e) {