Merge pull request #37 from launchdarkly/jko/fix-jedis-race

Use a thread-safe Jedis pool

Author: jkodumal

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

@@ -1,5 +1,6 @@
 package com.launchdarkly.client;
 
+import java.io.Closeable;
 import java.util.Map;
 
 /**
@@ -13,7 +14,7 @@
  * of features based on update messages that may be received out-of-order.
  *
  */
-public interface FeatureStore {
+public interface FeatureStore extends Closeable {
   /**
    *
    * Returns the {@link com.launchdarkly.client.FeatureRep} to which the specified key is mapped, or
@@ -71,4 +72,5 @@ public interface FeatureStore {
    * @return true if this store has been initialized
    */
   boolean initialized();
+
 }

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

@@ -1,5 +1,6 @@
 package com.launchdarkly.client;
 
+import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
@@ -141,4 +142,14 @@ public void upsert(String key, FeatureRep<?> feature) {
   public boolean initialized() {
     return initialized;
   }
+
+  /**
+   * Does nothing; this class does not have any resources to release
+   * @throws IOException
+   */
+  @Override
+  public void close() throws IOException
+  {
+    return;
+  }
 }

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

@@ -5,10 +5,12 @@
 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;
 import redis.clients.jedis.Transaction;
 
+import java.io.IOException;
 import java.lang.reflect.Type;
 import java.net.URI;
 import java.util.HashMap;
@@ -22,7 +24,7 @@
  */
 public class RedisFeatureStore implements FeatureStore {
   private static final String DEFAULT_PREFIX = "launchdarkly";
-  private final Jedis jedis;
+  private final JedisPool pool;
   private LoadingCache<String, FeatureRep<?>> cache;
   private LoadingCache<String, Boolean> initCache;
   private String prefix;
@@ -36,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) {
-    jedis = new Jedis(host, port);
+    pool = new JedisPool(getPoolConfig(), host, port);
     setPrefix(prefix);
     createCache(cacheTimeSecs);
     createInitCache(cacheTimeSecs);
@@ -50,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) {
-    jedis = new Jedis(uri);
+    pool = new JedisPool(getPoolConfig(), uri);
     setPrefix(prefix);
     createCache(cacheTimeSecs);
     createInitCache(cacheTimeSecs);
@@ -61,7 +63,7 @@ public RedisFeatureStore(URI uri, String prefix, long cacheTimeSecs) {
    *
    */
   public RedisFeatureStore() {
-    jedis = new Jedis("localhost");
+    pool = new JedisPool(getPoolConfig(), "localhost");
     this.prefix = DEFAULT_PREFIX;
   }
 
@@ -128,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
@@ -149,16 +159,25 @@ public Map<String, FeatureRep<?>> all() {
    */
   @Override
   public void init(Map<String, FeatureRep<?>> features) {
-    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();
+      }
+    }
   }
 
 
@@ -172,7 +191,9 @@ public void init(Map<String, FeatureRep<?>> features) {
    */
   @Override
   public void delete(String key, int version) {
+    Jedis jedis = null;
     try {
+      jedis = getJedis();
       Gson gson = new Gson();
       jedis.watch(featuresKey());
 
@@ -192,7 +213,10 @@ public void delete(String key, int version) {
       }
     }
     finally {
-      jedis.unwatch();
+      if (jedis != null) {
+        jedis.unwatch();
+        jedis.close();
+      }
     }
 
   }
@@ -206,7 +230,9 @@ public void delete(String key, int version) {
    */
   @Override
   public void upsert(String key, FeatureRep<?> feature) {
+    Jedis jedis = null;
     try {
+      jedis = getJedis();
       Gson gson = new Gson();
       jedis.watch(featuresKey());
 
@@ -223,7 +249,10 @@ public void upsert(String key, FeatureRep<?> feature) {
       }
     }
     finally {
-      jedis.unwatch();
+      if (jedis != null) {
+        jedis.unwatch();
+        jedis.close();
+      }
     }
   }
 
@@ -245,26 +274,64 @@ public boolean initialized() {
     return getInit();
   }
 
+  /**
+   * Releases all resources associated with the store. The store must no longer be used once closed.
+   * @throws IOException
+   */
+  public void close() throws IOException
+  {
+    pool.destroy();
+  }
+
+
 
   private String featuresKey() {
     return prefix + ":features";
   }
 
   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;
+      }
+
+      Type type = new TypeToken<FeatureRep<?>>() {}.getType();
+      FeatureRep<?> f = gson.fromJson(featureJson, type);
 
-    if (featureJson == null) {
-      return null;
+      return f.deleted ? null : f;
+    } finally {
+      if (jedis != null) {
+        jedis.close();
+      }
     }
+  }
 
-    Type type = new TypeToken<FeatureRep<?>>() {}.getType();
-    FeatureRep<?> f = gson.fromJson(featureJson, type);
+  private final Jedis getJedis() {
+    return pool.getResource();
+  }
 
-    return f.deleted ? null : f;
+  private final JedisPoolConfig getPoolConfig() {
+    JedisPoolConfig config = new JedisPoolConfig();
+    config.setMaxTotal(256);
+    config.setBlockWhenExhausted(false);
+    return config;
   }
 }

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

@@ -101,6 +101,9 @@ public void close() throws IOException {
     if (es != null) {
       es.close();
     }
+    if (store != null) {
+      store.close();
+    }
   }
 
   boolean initialized() {