From 02780d3bff6b3ae48cdfa7edc09bad7ccf32a3ca Mon Sep 17 00:00:00 2001
From: calebboyd <caleb.boyd@gmail.com>
Date: Fri, 13 Jul 2018 15:35:00 -0500
Subject: [PATCH] feat: support custom strategies on a queue

---
 README.md          |  2 +-
 index.d.ts         |  3 ++-
 lib/job.js         |  4 +---
 lib/queue.js       |  5 ++++-
 test/queue-test.js | 41 +++++++++++++++++++++++++++++++++++++++++
 5 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 127fdae8..706c1c9c 100644
--- a/README.md
+++ b/README.md
@@ -580,7 +580,7 @@ const counts = await queue.checkHealth();
 console.log('job state counts:', counts);
 ```
 
-#### Queue#close([cb])
+#### Queue#close([timeout], [cb])
 
 Closes the queue's connections to Redis. Idempotent.
 
diff --git a/index.d.ts b/index.d.ts
index 4a9c9533..9b0c35d6 100644
--- a/index.d.ts
+++ b/index.d.ts
@@ -10,6 +10,7 @@ declare class BeeQueue<T = any> extends EventEmitter {
   jobs: any;
   paused: boolean;
   settings: any;
+  backoffStrategies: Map<string, (job: BeeQueue.Job) => number>;
 
   constructor(name: string, settings?: BeeQueue.QueueSettings);
 
@@ -116,7 +117,7 @@ declare namespace BeeQueue {
     setId(id: string): this;
     retries(n: number): this;
     backoff(
-      strategy: 'immediate' | 'fixed' | 'exponential',
+      strategy: keyof typeof this.backoffStrategies,
       delayFactor?: number
     ): this;
     delayUntil(dateOrTimestamp: Date | number): this;
diff --git a/lib/job.js b/lib/job.js
index fc694c31..71bbbcde 100644
--- a/lib/job.js
+++ b/lib/job.js
@@ -1,9 +1,7 @@
 'use strict';
 
 const Emitter = require('events').EventEmitter;
-
 const helpers = require('./helpers');
-const strategies = require('./backoff');
 
 class Job extends Emitter {
   constructor(queue, jobId, data, options) {
@@ -146,7 +144,7 @@ class Job extends Emitter {
   }
 
   backoff(strategy, delay) {
-    if (!strategies.has(strategy)) {
+    if (!this.queue.backoffStrategies.has(strategy)) {
       throw new Error('unknown strategy');
     }
 
diff --git a/lib/queue.js b/lib/queue.js
index 38eb0b04..a931b47c 100644
--- a/lib/queue.js
+++ b/lib/queue.js
@@ -20,6 +20,7 @@ class Queue extends Emitter {
     this.jobs = new Map();
     this.activeJobs = new Set();
     this.checkTimer = null;
+    this.backoffStrategies = new Map(backoff);
 
     this._closed = null;
     this._isClosed = false;
@@ -581,7 +582,9 @@ class Queue extends Emitter {
         ? job.options.backoff.strategy
         : 'immediate';
       const strategy =
-        job.options.retries > 0 ? backoff.get(strategyName) : null;
+        job.options.retries > 0
+          ? this.backoffStrategies.get(strategyName)
+          : null;
       const delay = strategy ? strategy(job) : -1;
       if (delay < 0) {
         job.status = 'failed';
diff --git a/test/queue-test.js b/test/queue-test.js
index 52c8ac08..075ae2f4 100644
--- a/test/queue-test.js
+++ b/test/queue-test.js
@@ -1710,6 +1710,47 @@ describe('Queue', (it) => {
       t.throws(() => job.backoff('fixed', 44.5), /positive integer/i);
     });
 
+    it('should support custom backoff strategies', async (t) => {
+      const queue = t.context.makeQueue({
+        activateDelayedJobs: true,
+      });
+
+      queue.backoffStrategies.set(
+        'my-custom-backoff',
+        (job) => job.options.backoff.delay
+      );
+
+      const calls = [];
+
+      queue.process(async (job) => {
+        t.deepEqual(job.options.backoff, {
+          strategy: 'my-custom-backoff',
+          delay: 100,
+        });
+        t.deepEqual(job.data, {is: 'my-custom-backoff'});
+        calls.push(Date.now());
+        if (calls.length === 1) {
+          throw new Error('forced retry');
+        }
+        t.is(calls.length, 2);
+      });
+
+      const succeed = helpers.waitOn(queue, 'succeeded', true);
+
+      await queue
+        .createJob({is: 'my-custom-backoff'})
+        .retries(2)
+        .backoff('my-custom-backoff', 100)
+        .save();
+
+      await succeed;
+
+      t.is(calls.length, 2);
+
+      // Ensure there was a delay.
+      t.true(calls[1] - calls[0] >= 100);
+    });
+
     it('should handle fixed backoff', async (t) => {
       const queue = t.context.makeQueue({
         activateDelayedJobs: true,