feat(queue): add rateLimit method (#2896)

taskforcesh · Nov 19, 2024 · db84ad5 · db84ad5
1 parent ff70613
commit db84ad5
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 6 deletions.
diff --git a/src/classes/queue.ts b/src/classes/queue.ts
@@ -470,6 +470,34 @@ export class Queue<
       await super.close();
     });
   }
+
+  /**
+   * Overrides the rate limit to be active for the next jobs.
+   *
+   * @param expireTimeMs - expire time in ms of this rate limit.
+   */
+  async rateLimit(expireTimeMs: number): Promise<void> {
+    await this.trace<void>(
+      SpanKind.INTERNAL,
+      'rateLimit',
+      this.name,
+      async span => {
+        span?.setAttributes({
+          [TelemetryAttributes.QueueRateLimit]: expireTimeMs,
+        });
+
+        await this.client.then(client =>
+          client.set(
+            this.keys.limiter,
+            Number.MAX_SAFE_INTEGER,
+            'PX',
+            expireTimeMs,
+          ),
+        );
+      },
+    );
+  }
+
   /**
    * Resumes the processing of this queue globally.
    *

diff --git a/src/classes/worker.ts b/src/classes/worker.ts
@@ -626,7 +626,7 @@ export class Worker<
 
   /**
    * Overrides the rate limit to be active for the next jobs.
-   *
+   * @deprecated This method is deprecated and will be removed in v6. Use queue.rateLimit method instead.
    * @param expireTimeMs - expire time in ms of this rate limit.
    */
   async rateLimit(expireTimeMs: number): Promise<void> {

diff --git a/src/enums/telemetry-attributes.ts b/src/enums/telemetry-attributes.ts
@@ -13,6 +13,7 @@ export enum TelemetryAttributes {
   QueueDrainDelay = 'bullmq.queue.drain.delay',
   QueueGrace = 'bullmq.queue.grace',
   QueueCleanLimit = 'bullmq.queue.clean.limit',
+  QueueRateLimit = 'bullmq.queue.rate.limit',
   JobType = 'bullmq.job.type',
   QueueOptions = 'bullmq.queue.options',
   QueueEventMaxLength = 'bullmq.queue.event.max.length',

diff --git a/tests/test_concurrency.ts b/tests/test_concurrency.ts
@@ -1,5 +1,11 @@
 import { default as IORedis } from 'ioredis';
-import { FlowProducer, QueueEvents, Queue, Worker } from '../src/classes';
+import {
+  FlowProducer,
+  QueueEvents,
+  Queue,
+  Worker,
+  RateLimitError,
+} from '../src/classes';
 import { delay, removeAllQueueData } from '../src/utils';
 import { beforeEach, describe, it, after as afterAll } from 'mocha';
 import { v4 } from 'uuid';
@@ -156,8 +162,8 @@ describe('Concurrency', () => {
         queueName,
         async job => {
           if (job.attemptsStarted === 1) {
-            await worker.rateLimit(dynamicLimit);
-            throw Worker.RateLimitError();
+            await queue.rateLimit(dynamicLimit);
+            throw new RateLimitError();
           }
         },
         {
@@ -240,8 +246,8 @@ describe('Concurrency', () => {
           queueName,
           async job => {
             if (job.attemptsStarted === 1) {
-              await worker.rateLimit(dynamicLimit);
-              throw Worker.RateLimitError();
+              await queue.rateLimit(dynamicLimit);
+              throw new RateLimitError();
             }
           },
           {