Change RedisBackedJobService to use a connection pool

serving/pom.xml

@@ -250,6 +250,12 @@
       <groupId>com.fasterxml.jackson.dataformat</groupId>
       <artifactId>jackson-dataformat-yaml</artifactId>
     </dependency>
+
+    <dependency>
+      <groupId>com.github.kstyrc</groupId>
+      <artifactId>embedded-redis</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <profiles>

serving/src/main/java/feast/serving/service/RedisBackedJobService.java

@@ -43,34 +43,40 @@ public RedisBackedJobService(JedisPool jedisPool) {
 
   @Override
   public Optional<Job> get(String id) {
-    Jedis jedis = jedisPool.getResource();
-    String json = jedis.get(id);
-    if (json == null) {
-      return Optional.empty();
-    }
+    Jedis jedis = null;
     Job job = null;
-    Builder builder = Job.newBuilder();
     try {
+      jedis = jedisPool.getResource();
+      String json = jedis.get(id);
+      if (json == null) {
+        return Optional.empty();
+      }
+      Builder builder = Job.newBuilder();
       JsonFormat.parser().merge(json, builder);
       job = builder.build();
     } catch (Exception e) {
       log.error(String.format("Failed to parse JSON for Feast job: %s", e.getMessage()));
     } finally {
-      jedis.close();
+      if (jedis != null) {
+        jedis.close();
+      }
     }
     return Optional.ofNullable(job);
   }
 
   @Override
   public void upsert(Job job) {
-    Jedis jedis = jedisPool.getResource();
+    Jedis jedis = null;
     try {
+      jedis = jedisPool.getResource();
       jedis.set(job.getId(), JsonFormat.printer().omittingInsignificantWhitespace().print(job));
       jedis.expire(job.getId(), defaultExpirySeconds);
     } catch (Exception e) {
       log.error(String.format("Failed to upsert job: %s", e.getMessage()));
     } finally {
-      jedis.close();
+      if (jedis != null) {
+        jedis.close();
+      }
     }
   }
 }

serving/src/test/java/feast/serving/service/RedisBackedJobServiceTest.java

@@ -0,0 +1,44 @@
+package feast.serving.service;
+
+import java.io.IOException;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import redis.clients.jedis.JedisPool;
+import redis.clients.jedis.JedisPoolConfig;
+import redis.embedded.RedisServer;
+
+public class RedisBackedJobServiceTest {
+  private static String REDIS_HOST = "localhost";
+  private static int REDIS_PORT = 51235;
+  private RedisServer redis;
+
+  @Before
+  public void setUp() throws IOException {
+    redis = new RedisServer(REDIS_PORT);
+    redis.start();
+  }
+
+  @After
+  public void teardown() {
+    redis.stop();
+  }
+
+  @Test
+  public void shouldRecoverIfRedisConnectionIsLost() {
+    JedisPoolConfig jedisPoolConfig = new JedisPoolConfig();
+    jedisPoolConfig.setMaxTotal(1);
+    jedisPoolConfig.setMaxWaitMillis(10);
+    JedisPool jedisPool = new JedisPool(jedisPoolConfig, REDIS_HOST, REDIS_PORT);
+    RedisBackedJobService jobService = new RedisBackedJobService(jedisPool);
+    jobService.get("does not exist");
+    redis.stop();
+    try {
+      jobService.get("does not exist");
+    } catch (Exception e) {
+      // pass, this should fail, and return a broken connection to the pool
+    }
+    redis.start();
+    jobService.get("does not exist");
+  }
+}