Fix pooled connection leak from unclosed pipelines and transactions

src/main/java/net/greghaines/jesque/client/AbstractClient.java

@@ -197,6 +197,7 @@ public static void doEnqueue(
    * Helper method that encapsulates the minimum logic for adding jobs to a queue.
    *
    * @param jedis the connection to Redis
+   * @param makePipeline creates the pipeline to send the commands with, which is closed afterwards
    * @param namespace the Resque namespace
    * @param queue the Resque queue name
    * @param jobJsons a list of jobs serialized as JSON
@@ -207,12 +208,13 @@ public static void doBatchEnqueue(
       final String namespace,
       final String queue,
       final List<String> jobJsons) {
-    AbstractPipeline pipelined = makePipeline.get(); // jedis.pipelined();
-    pipelined.sadd(JesqueUtils.createKey(namespace, QUEUES), queue);
-    for (String jobJson : jobJsons) {
-      pipelined.rpush(JesqueUtils.createKey(namespace, QUEUE, queue), jobJson);
+    try (AbstractPipeline pipelined = makePipeline.get()) { // jedis.pipelined();
+      pipelined.sadd(JesqueUtils.createKey(namespace, QUEUES), queue);
+      for (String jobJson : jobJsons) {
+        pipelined.rpush(JesqueUtils.createKey(namespace, QUEUE, queue), jobJson);
+      }
+      pipelined.sync();
     }
-    pipelined.sync();
   }
 
   /**
@@ -372,11 +374,12 @@ public static void doRecurringEnqueue(
     final String hashKey = JesqueUtils.createRecurringHashKey(queueKey);
 
     if (JedisUtils.canUseAsRecurringQueue(jedis, queueKey, hashKey)) {
-      AbstractTransaction transaction = makeTransaction.get(); // jedis.multi();
-      transaction.zadd(queueKey, future.toEpochMilli(), jobJson);
-      transaction.hset(hashKey, jobJson, Long.toString(frequency.toMillis()));
-      if (transaction.exec() == null) {
-        throw new RuntimeException("cannot add " + jobJson + " to recurring queue " + queue);
+      try (AbstractTransaction transaction = makeTransaction.get()) { // jedis.multi();
+        transaction.zadd(queueKey, future.toEpochMilli(), jobJson);
+        transaction.hset(hashKey, jobJson, Long.toString(frequency.toMillis()));
+        if (transaction.exec() == null) {
+          throw new RuntimeException("cannot add " + jobJson + " to recurring queue " + queue);
+        }
       }
     } else {
       throw new IllegalArgumentException(queue + " cannot be used as a recurring queue");
@@ -410,11 +413,12 @@ public static void doRemoveRecurringEnqueue(
     final String hashKey = JesqueUtils.createRecurringHashKey(queueKey);
 
     if (JedisUtils.canUseAsRecurringQueue(jedis, queueKey, hashKey)) {
-      AbstractTransaction transaction = makeTransaction.get(); // jedis.multi();
-      transaction.hdel(hashKey, jobJson);
-      transaction.zrem(queueKey, jobJson);
-      if (transaction.exec() == null) {
-        throw new RuntimeException("cannot remove " + jobJson + " from recurring queue " + queue);
+      try (AbstractTransaction transaction = makeTransaction.get()) { // jedis.multi();
+        transaction.hdel(hashKey, jobJson);
+        transaction.zrem(queueKey, jobJson);
+        if (transaction.exec() == null) {
+          throw new RuntimeException("cannot remove " + jobJson + " from recurring queue " + queue);
+        }
       }
     } else {
       throw new IllegalArgumentException(queue + " cannot be used as a recurring queue");

src/main/java/net/greghaines/jesque/worker/AbstractWorker.java

@@ -567,10 +567,11 @@ protected void failure(final Throwable thrwbl, final Job job, final String curQu
       if (failQueueKey != null) {
         final int failQueueMaxItems = strategy.getFailQueueMaxItems(curQueue);
         if (failQueueMaxItems > 0) {
-          AbstractTransaction tx = createTransaction();
-          tx.rpush(failQueueKey, failMsg(thrwbl, curQueue, job));
-          tx.ltrim(failQueueKey, -failQueueMaxItems, -1);
-          tx.exec();
+          try (AbstractTransaction tx = createTransaction()) {
+            tx.rpush(failQueueKey, failMsg(thrwbl, curQueue, job));
+            tx.ltrim(failQueueKey, -failQueueMaxItems, -1);
+            tx.exec();
+          }
         } else {
           getJedis().rpush(failQueueKey, failMsg(thrwbl, curQueue, job));
         }

src/test/java/net/greghaines/jesque/TestUtils.java

@@ -15,17 +15,21 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
 import java.util.Map;
 import net.greghaines.jesque.client.Client;
 import net.greghaines.jesque.client.ClientImpl;
+import net.greghaines.jesque.utils.PoolUtils;
 import net.greghaines.jesque.worker.JobExecutor;
 import net.greghaines.jesque.worker.JobFactory;
 import net.greghaines.jesque.worker.MapBasedJobFactory;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import redis.clients.jedis.ConnectionPoolConfig;
 import redis.clients.jedis.Jedis;
+import redis.clients.jedis.UnifiedJedis;
 
 /**
  * Test helpers.
@@ -65,6 +69,19 @@ public static Jedis createJedis(final Config config) {
     return jedis;
   }
 
+  /**
+   * @param config the location of the Redis server
+   * @param maxTotal the maximum number of connections the pool hands out at once
+   * @return a new pool that fails to hand out a connection after waiting for one second
+   */
+  public static UnifiedJedis createBoundedJedisPool(final Config config, final int maxTotal) {
+    final ConnectionPoolConfig poolConfig = new ConnectionPoolConfig();
+    poolConfig.setMaxTotal(maxTotal);
+    poolConfig.setBlockWhenExhausted(true);
+    poolConfig.setMaxWait(Duration.ofSeconds(1));
+    return PoolUtils.createJedisPool(config, poolConfig);
+  }
+
   public static void enqueueJobs(final String queue, final List<Job> jobs, final Config config) {
     final Client client = new ClientImpl(config);
     try {

src/test/java/net/greghaines/jesque/client/TestClientPoolImpl.java

@@ -0,0 +1,96 @@
+package net.greghaines.jesque.client;
+
+import static com.google.common.truth.Truth.assertThat;
+import static net.greghaines.jesque.TestUtils.createBoundedJedisPool;
+import static net.greghaines.jesque.TestUtils.createJedis;
+import static net.greghaines.jesque.utils.JesqueUtils.createKey;
+import static net.greghaines.jesque.utils.ResqueConstants.QUEUE;
+
+import java.time.Duration;
+import java.time.Instant;
+import java.util.List;
+import java.util.stream.IntStream;
+import net.greghaines.jesque.Config;
+import net.greghaines.jesque.Job;
+import net.greghaines.jesque.TestAction;
+import net.greghaines.jesque.TestUtils;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import redis.clients.jedis.Jedis;
+import redis.clients.jedis.UnifiedJedis;
+
+public class TestClientPoolImpl {
+
+  private static final Config CONFIG = Config.getDefaultConfig();
+  private static final String TEST_QUEUE = "foo";
+  private static final String TEST_QUEUE_KEY = createKey(CONFIG.getNamespace(), QUEUE, TEST_QUEUE);
+  private static final Duration RECURRING_FREQUENCY = Duration.ofMinutes(1);
+  private static final int POOL_SIZE = 3;
+
+  private UnifiedJedis jedisPool;
+  private Client client;
+
+  @Before
+  public void setUp() {
+    TestUtils.resetRedis(CONFIG);
+    this.jedisPool = createBoundedJedisPool(CONFIG, POOL_SIZE);
+    this.client = new ClientPoolImpl(CONFIG, this.jedisPool);
+  }
+
+  @After
+  public void tearDown() {
+    this.client.end();
+    this.jedisPool.close();
+  }
+
+  @Test
+  public void testBatchEnqueue_ReturnsConnectionToPool() {
+    final int batchCount = POOL_SIZE + 1;
+    final Job job = new Job(TestAction.class.getSimpleName());
+    for (int i = 0; i < batchCount; i++) {
+      this.client.batchEnqueue(TEST_QUEUE, List.of(job));
+    }
+    try (Jedis jedis = createJedis(CONFIG)) {
+      assertThat(jedis.llen(TEST_QUEUE_KEY)).isEqualTo(batchCount);
+    }
+  }
+
+  @Test
+  public void testRecurringEnqueue_ReturnsConnectionToPool() {
+    final List<Job> jobs = createDistinctJobs(POOL_SIZE + 1);
+    final Instant future = Instant.now().plus(RECURRING_FREQUENCY);
+    for (Job job : jobs) {
+      this.client.recurringEnqueue(TEST_QUEUE, job, future, RECURRING_FREQUENCY);
+    }
+    try (Jedis jedis = createJedis(CONFIG)) {
+      assertThat(jedis.zcard(TEST_QUEUE_KEY)).isEqualTo(jobs.size());
+    }
+  }
+
+  @Test
+  public void testRemoveRecurringEnqueue_ReturnsConnectionToPool() {
+    final List<Job> jobs = createDistinctJobs(POOL_SIZE + 1);
+    final Instant future = Instant.now().plus(RECURRING_FREQUENCY);
+    final Client seedingClient = new ClientImpl(CONFIG);
+    try {
+      for (Job job : jobs) {
+        seedingClient.recurringEnqueue(TEST_QUEUE, job, future, RECURRING_FREQUENCY);
+      }
+    } finally {
+      seedingClient.end();
+    }
+    for (Job job : jobs) {
+      this.client.removeRecurringEnqueue(TEST_QUEUE, job);
+    }
+    try (Jedis jedis = createJedis(CONFIG)) {
+      assertThat(jedis.zcard(TEST_QUEUE_KEY)).isEqualTo(0L);
+    }
+  }
+
+  private static List<Job> createDistinctJobs(final int count) {
+    return IntStream.range(0, count)
+        .mapToObj(i -> new Job(TestAction.class.getSimpleName(), i))
+        .toList();
+  }
+}

src/test/java/net/greghaines/jesque/worker/TestWorkerPoolImpl.java

@@ -0,0 +1,63 @@
+package net.greghaines.jesque.worker;
+
+import static com.google.common.truth.Truth.assertThat;
+import static net.greghaines.jesque.TestUtils.createBoundedJedisPool;
+import static net.greghaines.jesque.TestUtils.createJedis;
+import static net.greghaines.jesque.TestUtils.createTestActionJobFactory;
+import static net.greghaines.jesque.utils.JesqueUtils.createKey;
+import static net.greghaines.jesque.utils.ResqueConstants.FAILED;
+
+import java.util.List;
+import net.greghaines.jesque.Config;
+import net.greghaines.jesque.Job;
+import net.greghaines.jesque.TestAction;
+import net.greghaines.jesque.TestUtils;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import redis.clients.jedis.Jedis;
+import redis.clients.jedis.UnifiedJedis;
+
+public class TestWorkerPoolImpl {
+
+  private static final Config CONFIG = Config.getDefaultConfig();
+  private static final String TEST_QUEUE = "foo";
+  private static final String FAIL_QUEUE_KEY = createKey(CONFIG.getNamespace(), FAILED);
+  private static final int POOL_SIZE = 3;
+
+  private UnifiedJedis jedisPool;
+  private WorkerPoolImpl worker;
+
+  @Before
+  public void setUp() {
+    TestUtils.resetRedis(CONFIG);
+    this.jedisPool = createBoundedJedisPool(CONFIG, POOL_SIZE);
+    this.worker =
+        new WorkerPoolImpl(
+            CONFIG, List.of(TEST_QUEUE), createTestActionJobFactory(), this.jedisPool);
+  }
+
+  @After
+  public void tearDown() {
+    this.jedisPool.close();
+  }
+
+  @Test
+  public void testFailure_FailQueueMaxItems_ReturnsConnectionToPool() {
+    final int failureCount = POOL_SIZE + 1;
+    this.worker.setFailQueueStrategy(
+        new DefaultFailQueueStrategy(CONFIG.getNamespace()) {
+          @Override
+          public int getFailQueueMaxItems(final String curQueue) {
+            return failureCount;
+          }
+        });
+    final Job job = new Job(TestAction.class.getSimpleName());
+    for (int i = 0; i < failureCount; i++) {
+      this.worker.failure(new RuntimeException("BOOM!"), job, TEST_QUEUE);
+    }
+    try (Jedis jedis = createJedis(CONFIG)) {
+      assertThat(jedis.llen(FAIL_QUEUE_KEY)).isEqualTo(failureCount);
+    }
+  }
+}