feat(queue): support global rate limit (#3468) ref #3019

Author: roggervalf

docs/gitbook/SUMMARY.md

@@ -17,6 +17,7 @@
   * [Auto-removal of jobs](guide/queues/auto-removal-of-jobs.md)
   * [Adding jobs in bulk](guide/queues/adding-bulks.md)
   * [Global Concurrency](guide/queues/global-concurrency.md)
+  * [Global Rate Limit](guide/queues/global-rate-limit.md)
   * [Removing Jobs](guide/queues/removing-jobs.md)
 * [Workers](guide/workers/README.md)
   * [Auto-removal of jobs](guide/workers/auto-removal-of-jobs.md)

docs/gitbook/guide/queues/global-rate-limit.md

@@ -0,0 +1,25 @@
+# Global Rate Limit
+
+The global rate limit config is a queue option that determines how many jobs are allowed to be processed in a specific period of time.
+
+```typescript
+import { Queue } from 'bullmq';
+
+// 1 job per second
+await queue.setGlobalRateLimit(1, 1000);
+```
+
+And in order to get this value:
+
+```typescript
+const globalConcurrency = await queue.getRateLimitTtl();
+```
+
+{% hint style="info" %}
+Note that if you choose a rate limit level in your workers, it won't override the global one.
+{% endhint %}
+
+## Read more:
+
+- 💡 [Set Global Rate Limit API Reference](https://api.docs.bullmq.io/classes/v5.Queue.html#setglobalratelimit)
+- 💡 [Get Rate Limit Ttl API Reference](https://api.docs.bullmq.io/classes/v5.Queue.html#getratelimitttl)

src/classes/queue.ts

@@ -291,6 +291,16 @@ export class Queue<
     return client.hset(this.keys.meta, 'concurrency', concurrency);
   }
 
+  /**
+   * Enable and set rate limit.
+   * @param max - Max number of jobs to process in the time period specified in `duration`
+   * @param duration - Time in milliseconds. During this time, a maximum of `max` jobs will be processed.
+   */
+  async setGlobalRateLimit(max: number, duration: number) {
+    const client = await this.client;
+    return client.hset(this.keys.meta, 'max', max, 'duration', duration);
+  }
+
   /**
    * Remove global concurrency value.
    */

src/classes/scripts.ts

@@ -1223,7 +1223,10 @@ export class Scripts {
   }
 
   getRateLimitTtlArgs(maxJobs?: number): (string | number)[] {
-    const keys: (string | number)[] = [this.queue.keys.limiter];
+    const keys: (string | number)[] = [
+      this.queue.keys.limiter,
+      this.queue.keys.meta,
+    ];
 
     return keys.concat([maxJobs ?? '0']);
   }

src/commands/changePriority-7.lua

@@ -50,7 +50,7 @@ if rcall("EXISTS", jobKey) == 1 then
     local prioritizedKey = KEYS[4]
     local priorityCounterKey = KEYS[6]
     local markerKey = KEYS[7]
-    
+
     -- Re-add with the new priority
     if rcall("ZREM", prioritizedKey, jobId) > 0 then
         reAddJobWithNewPriority( prioritizedKey, markerKey, target,

src/commands/getRateLimitTtl-1.lua renamed to src/commands/getRateLimitTtl-2.lua

@@ -3,6 +3,7 @@
 
     Input:
       KEYS[1] 'limiter'
+      KEYS[2] 'meta'
 
       ARGV[1] maxJobs
 ]]
@@ -16,5 +17,10 @@ local rateLimiterKey = KEYS[1]
 if ARGV[1] ~= "0" then
   return getRateLimitTTL(tonumber(ARGV[1]), rateLimiterKey)
 else
+  local rateLimitMax = rcall("HGET", KEYS[2], "max")
+  if rateLimitMax then
+    return getRateLimitTTL(tonumber(rateLimitMax), rateLimiterKey)
+  end
+
   return rcall("PTTL", rateLimiterKey)
 end

src/commands/includes/getTargetQueueList.lua

@@ -4,19 +4,19 @@
 ]]
 
 local function getTargetQueueList(queueMetaKey, activeKey, waitKey, pausedKey)
-  local queueAttributes = rcall("HMGET", queueMetaKey, "paused", "concurrency")
+  local queueAttributes = rcall("HMGET", queueMetaKey, "paused", "concurrency", "max", "duration")
 
   if queueAttributes[1] then
-    return pausedKey, true
+    return pausedKey, true, queueAttributes[3], queueAttributes[4]
   else
     if queueAttributes[2] then
       local activeCount = rcall("LLEN", activeKey)
       if activeCount >= tonumber(queueAttributes[2]) then
-        return waitKey, true
+        return waitKey, true, queueAttributes[3], queueAttributes[4]
       else
-        return waitKey, false
+        return waitKey, false, queueAttributes[3], queueAttributes[4]
       end
     end
   end
-  return waitKey, false
+  return waitKey, false, queueAttributes[3], queueAttributes[4]
 end

src/commands/includes/prepareJobForProcessing.lua

@@ -10,24 +10,22 @@
 --- @include "addBaseMarkerIfNeeded"
 
 local function prepareJobForProcessing(keyPrefix, rateLimiterKey, eventStreamKey,
-    jobId, processedOn, maxJobs, markerKey, opts)
+    jobId, processedOn, maxJobs, limiterDuration, markerKey, opts)
   local jobKey = keyPrefix .. jobId
 
   -- Check if we need to perform rate limiting.
   if maxJobs then
     local jobCounter = tonumber(rcall("INCR", rateLimiterKey))
 
     if jobCounter == 1 then
-      local limiterDuration = opts['limiter'] and opts['limiter']['duration']
       local integerDuration = math.floor(math.abs(limiterDuration))
       rcall("PEXPIRE", rateLimiterKey, integerDuration)
     end
   end
 
-  local lockKey = jobKey .. ':lock'
-
   -- get a lock
   if opts['token'] ~= "0" then
+    local lockKey = jobKey .. ':lock'
     rcall("SET", lockKey, opts['token'], "PX", opts['lockDuration'])
   end
 

src/commands/moveStalledJobsToWait-8.lua

@@ -24,7 +24,6 @@ local rcall = redis.call
 -- Includes
 --- @include "includes/addJobInTargetList"
 --- @include "includes/batches"
---- @include "includes/getTargetQueueList"
 --- @include "includes/moveJobToWait"
 --- @include "includes/trimEvents"
 

src/commands/moveToActive-11.lua

@@ -51,14 +51,15 @@ local opts = cmsgpack.unpack(ARGV[3])
 --- @include "includes/prepareJobForProcessing"
 --- @include "includes/promoteDelayedJobs"
 
-local target, isPausedOrMaxed = getTargetQueueList(KEYS[9], activeKey, waitKey, KEYS[8])
+local target, isPausedOrMaxed, rateLimitMax, rateLimitDuration = getTargetQueueList(KEYS[9],
+    activeKey, waitKey, KEYS[8])
 
 -- Check if there are delayed jobs that we can move to wait.
 local markerKey = KEYS[11]
 promoteDelayedJobs(delayedKey, markerKey, target, KEYS[3], eventStreamKey, ARGV[1],
                    ARGV[2], KEYS[10], isPausedOrMaxed)
 
-local maxJobs = tonumber(opts['limiter'] and opts['limiter']['max'])
+local maxJobs = tonumber(rateLimitMax or (opts['limiter'] and opts['limiter']['max']))
 local expireTime = getRateLimitTTL(maxJobs, rateLimiterKey)
 
 -- Check if we are rate limited first.
@@ -67,6 +68,8 @@ if expireTime > 0 then return {0, 0, expireTime, 0} end
 -- paused or maxed queue
 if isPausedOrMaxed then return {0, 0, 0, 0} end
 
+local limiterDuration = (opts['limiter'] and opts['limiter']['duration']) or rateLimitDuration
+
 -- no job ID, try non-blocking move from wait to active
 local jobId = rcall("RPOPLPUSH", waitKey, activeKey)
 
@@ -78,12 +81,12 @@ end
 
 if jobId then
     return prepareJobForProcessing(ARGV[1], rateLimiterKey, eventStreamKey, jobId, ARGV[2],
-                                   maxJobs, markerKey, opts)
+                                   maxJobs, limiterDuration, markerKey, opts)
 else
     jobId = moveJobFromPrioritizedToActive(KEYS[3], activeKey, KEYS[10])
     if jobId then
         return prepareJobForProcessing(ARGV[1], rateLimiterKey, eventStreamKey, jobId, ARGV[2],
-                                       maxJobs, markerKey, opts)
+                                       maxJobs, limiterDuration, markerKey, opts)
     end
 end