Fix possible deadlock due to blocking on pool waits

Author: jkodumal

src/main/java/com/launchdarkly/client/RedisFeatureStore.java

@@ -5,7 +5,6 @@
 import com.google.common.cache.LoadingCache;
 import com.google.gson.Gson;
 import com.google.gson.reflect.TypeToken;
-import org.apache.http.util.EntityUtils;
 import redis.clients.jedis.Jedis;
 import redis.clients.jedis.JedisPool;
 import redis.clients.jedis.JedisPoolConfig;
@@ -39,7 +38,7 @@ public class RedisFeatureStore implements FeatureStore {
    * @param cacheTimeSecs an optional timeout for the in-memory cache. If set to 0, no in-memory caching will be performed
    */
   public RedisFeatureStore(String host, int port, String prefix, long cacheTimeSecs) {
-    pool = new JedisPool(new JedisPoolConfig(), host, port);
+    pool = new JedisPool(getPoolConfig(), host, port);
     setPrefix(prefix);
     createCache(cacheTimeSecs);
     createInitCache(cacheTimeSecs);
@@ -53,7 +52,7 @@ public RedisFeatureStore(String host, int port, String prefix, long cacheTimeSec
    * @param cacheTimeSecs an optional timeout for the in-memory cache. If set to 0, no in-memory caching will be performed
    */
   public RedisFeatureStore(URI uri, String prefix, long cacheTimeSecs) {
-    pool = new JedisPool(new JedisPoolConfig(), uri);
+    pool = new JedisPool(getPoolConfig(), uri);
     setPrefix(prefix);
     createCache(cacheTimeSecs);
     createInitCache(cacheTimeSecs);
@@ -64,7 +63,7 @@ public RedisFeatureStore(URI uri, String prefix, long cacheTimeSecs) {
    *
    */
   public RedisFeatureStore() {
-    pool = new JedisPool(new JedisPoolConfig(), "localhost");
+    pool = new JedisPool(getPoolConfig(), "localhost");
     this.prefix = DEFAULT_PREFIX;
   }
 
@@ -131,18 +130,26 @@ public FeatureRep<?> get(String key) {
    */
   @Override
   public Map<String, FeatureRep<?>> all() {
-    Map<String,String> featuresJson = jedis().hgetAll(featuresKey());
-    Map<String, FeatureRep<?>> result = new HashMap<String, FeatureRep<?>>();
-    Gson gson = new Gson();
+    Jedis jedis = null;
+    try {
+      jedis = getJedis();
+      Map<String,String> featuresJson = getJedis().hgetAll(featuresKey());
+      Map<String, FeatureRep<?>> result = new HashMap<String, FeatureRep<?>>();
+      Gson gson = new Gson();
 
-    Type type = new TypeToken<FeatureRep<?>>() {}.getType();
+      Type type = new TypeToken<FeatureRep<?>>() {}.getType();
 
-    for (Map.Entry<String, String> entry : featuresJson.entrySet()) {
-      FeatureRep<?> rep =  gson.fromJson(entry.getValue(), type);
-      result.put(entry.getKey(), rep);
+      for (Map.Entry<String, String> entry : featuresJson.entrySet()) {
+        FeatureRep<?> rep =  gson.fromJson(entry.getValue(), type);
+        result.put(entry.getKey(), rep);
+      }
+      return result;
+    } finally {
+      if (jedis != null) {
+        jedis.close();
+      }
     }
 
-    return result;
   }
   /**
    * Initializes (or re-initializes) the store with the specified set of features. Any existing entries
@@ -152,17 +159,25 @@ public Map<String, FeatureRep<?>> all() {
    */
   @Override
   public void init(Map<String, FeatureRep<?>> features) {
-    Jedis jedis = jedis();
-    Gson gson = new Gson();
-    Transaction t = jedis.multi();
+    Jedis jedis = null;
 
-    t.del(featuresKey());
+    try {
+      jedis = getJedis();
+      Gson gson = new Gson();
+      Transaction t = jedis.multi();
 
-    for (FeatureRep<?> f: features.values()) {
-      t.hset(featuresKey(), f.key, gson.toJson(f));
-    }
+      t.del(featuresKey());
 
-    t.exec();
+      for (FeatureRep<?> f: features.values()) {
+        t.hset(featuresKey(), f.key, gson.toJson(f));
+      }
+
+      t.exec();
+    } finally {
+      if (jedis != null) {
+        jedis.close();
+      }
+    }
   }
 
 
@@ -176,8 +191,9 @@ public void init(Map<String, FeatureRep<?>> features) {
    */
   @Override
   public void delete(String key, int version) {
-    Jedis jedis = jedis();
+    Jedis jedis = null;
     try {
+      jedis = getJedis();
       Gson gson = new Gson();
       jedis.watch(featuresKey());
 
@@ -197,7 +213,10 @@ public void delete(String key, int version) {
       }
     }
     finally {
-      jedis.unwatch();
+      if (jedis != null) {
+        jedis.unwatch();
+        jedis.close();
+      }
     }
 
   }
@@ -211,8 +230,9 @@ public void delete(String key, int version) {
    */
   @Override
   public void upsert(String key, FeatureRep<?> feature) {
-    Jedis jedis = jedis();
+    Jedis jedis = null;
     try {
+      jedis = getJedis();
       Gson gson = new Gson();
       jedis.watch(featuresKey());
 
@@ -229,7 +249,10 @@ public void upsert(String key, FeatureRep<?> feature) {
       }
     }
     finally {
-      jedis.unwatch();
+      if (jedis != null) {
+        jedis.unwatch();
+        jedis.close();
+      }
     }
   }
 
@@ -267,24 +290,48 @@ private String featuresKey() {
   }
 
   private Boolean getInit() {
-    return jedis().exists(featuresKey());
+    Jedis jedis = null;
+
+    try {
+      jedis = getJedis();
+      return jedis.exists(featuresKey());
+    } finally {
+      if (jedis != null) {
+        jedis.close();
+      }
+    }
   }
 
   private FeatureRep<?> getRedis(String key) {
-    Gson gson = new Gson();
-    String featureJson = jedis().hget(featuresKey(), key);
+    Jedis jedis = null;
+    try {
+      jedis = getJedis();
+      Gson gson = new Gson();
+      String featureJson = getJedis().hget(featuresKey(), key);
 
-    if (featureJson == null) {
-      return null;
-    }
+      if (featureJson == null) {
+        return null;
+      }
 
-    Type type = new TypeToken<FeatureRep<?>>() {}.getType();
-    FeatureRep<?> f = gson.fromJson(featureJson, type);
+      Type type = new TypeToken<FeatureRep<?>>() {}.getType();
+      FeatureRep<?> f = gson.fromJson(featureJson, type);
 
-    return f.deleted ? null : f;
+      return f.deleted ? null : f;
+    } finally {
+      if (jedis != null) {
+        jedis.close();
+      }
+    }
   }
 
-  private final Jedis jedis() {
+  private final Jedis getJedis() {
     return pool.getResource();
   }
+
+  private final JedisPoolConfig getPoolConfig() {
+    JedisPoolConfig config = new JedisPoolConfig();
+    config.setMaxTotal(256);
+    config.setBlockWhenExhausted(false);
+    return config;
+  }
 }