Implement queue rate limit

src/actions/create-listener.ts

@@ -14,6 +14,7 @@ import {
   getWorkerShutdownChannel,
   getQueueTaskScheduledChannel,
   getQueueTaskProgressUpdatedChannel,
+  getQueueRateLimitUpdatedChannel,
 } from '../utils/keys';
 import { deSerializeEvent } from '../domain/events/deserialize-event';
 import { Event } from '../domain/events/event';
@@ -72,6 +73,9 @@ export const createListener = ({
       [EventType.WorkerStarted]: getWorkerStartedChannel({ queue }),
       [EventType.WorkerPaused]: getWorkerPausedChannel({ queue }),
       [EventType.WorkerShutdown]: getWorkerShutdownChannel({ queue }),
+      [EventType.QueueRateLimitUpdated]: getQueueRateLimitUpdatedChannel({
+        queue,
+      }),
     };
     client.on('message', (channel, eventString) => {
       const event = deSerializeEvent(eventString);

src/actions/create-manager.ts

@@ -25,13 +25,16 @@ import { resumeQueue } from './resume-queue';
 import { Manager } from '../domain/manager/manager';
 import { TaskResponse } from '../domain/manager/task-response';
 import { OnBeforeEnqueueTask, OnAfterEnqueueTask } from './enqueue-task';
+import { setQueueRateLimit } from './set-queue-rate-limit';
+import { QueueRateLimitConfig } from './get-queue-rate-limit-config';
 
 const debug = debugF('conveyor-mq:manager');
 
 export interface ManagerInput {
   queue: string;
   redisConfig: RedisConfig;
   redisClient?: Redis;
+  queueRateLimitConfig?: QueueRateLimitConfig;
   hooks?: {
     onBeforeEnqueueTask?: OnBeforeEnqueueTask;
     onAfterEnqueueTask?: OnAfterEnqueueTask;
@@ -60,13 +63,15 @@ export interface ManagerInput {
  * - .destroyQueue(): Promise<void> - Destroys the queue by removing all data & data structures.
  * - .pauseQueue(): Promise<void> - Pauses the queue.
  * - .resumeQueue(): Promise<void> - Resumes the queue.
+ * - .setQueueRateLimit({ points, duration }): Promise<void> - Sets the rate limit on the queue.
  * - .onReady(): Promise<void> - Returns a promise which resolves once the manager is ready.
  * - .quit(): Promise<void> - Quits the manager, disconnects the redis clients.
  */
 export const createManager = ({
   queue,
   redisConfig,
   redisClient,
+  queueRateLimitConfig,
   hooks,
 }: ManagerInput): Manager => {
   debug('Starting');
@@ -168,6 +173,14 @@ export const createManager = ({
 
   const ready = async () => {
     await Promise.all([listenerSetupPromise]);
+    if (queueRateLimitConfig) {
+      await setQueueRateLimit({
+        points: queueRateLimitConfig.points,
+        duration: queueRateLimitConfig.duration,
+        queue,
+        client,
+      });
+    }
     debug('Ready');
   };
   const readyPromise = ready();
@@ -238,6 +251,16 @@ export const createManager = ({
       debug(`resumeQueue`);
       return resumeQueue({ queue, client });
     },
+    setQueueRateLimit: async ({
+      points,
+      duration,
+    }: {
+      points: number;
+      duration: number;
+    }) => {
+      await readyPromise;
+      await setQueueRateLimit({ points, duration, queue, client });
+    },
     onReady: async () => {
       await readyPromise;
       debug(`onReady`);

src/actions/create-worker.ts

@@ -44,6 +44,8 @@ import { serializeWorker } from '../domain/worker/serialize-worker';
 import { Task } from '../domain/tasks/task';
 import { Worker } from '../domain/worker/worker';
 import { markTaskProcessing } from './mark-task-processing';
+import { getQueueRateLimitConfig } from './get-queue-rate-limit-config';
+import { createListener } from './create-listener';
 
 const debug = debugF('conveyor-mq:worker');
 
@@ -286,6 +288,42 @@ export const createWorker = ({
     });
   };
 
+  const setQueueRateLimit = ({
+    points,
+    duration,
+  }: {
+    points: number;
+    duration: number;
+  }) => {
+    rateLimiter = new RateLimiterRedis({
+      storeClient: workerClient,
+      points,
+      duration,
+      keyPrefix: queue,
+      inmemoryBlockOnConsumed: points + 1,
+    });
+  };
+
+  const clearQueueRateLimit = () => {
+    rateLimiter = undefined;
+  };
+
+  const setupListener = async () => {
+    const listener = createListener({
+      queue,
+      redisConfig,
+      events: [EventType.QueueRateLimitUpdated],
+    });
+    await listener.onReady();
+    listener.on(EventType.QueueRateLimitUpdated, ({ event }) => {
+      if (event.data?.rateLimitConfig) {
+        setQueueRateLimit(event.data.rateLimitConfig);
+      } else {
+        clearQueueRateLimit();
+      }
+    });
+  };
+
   const pause = async ({
     killProcessingTasks,
   }: {
@@ -341,13 +379,20 @@ export const createWorker = ({
           ensureConnected({ client }),
         ),
       );
-      rateLimiter = new RateLimiterRedis({
-        storeClient: workerClient,
-        points: 1,
-        duration: 5,
-        keyPrefix: queue,
-        inmemoryBlockOnConsumed: 1,
+      const rateLimitConfig = await getQueueRateLimitConfig({
+        queue,
+        client: workerClient,
       });
+      if (
+        rateLimitConfig &&
+        rateLimitConfig.duration &&
+        rateLimitConfig.points
+      ) {
+        setQueueRateLimit({
+          points: rateLimitConfig.points,
+          duration: rateLimitConfig.duration,
+        });
+      }
       forEach([workerQueue, takerQueue], (q) => q.start());
       await upsertWorker();
       upsertInterval = setIntervalAsync(upsertWorker, 15000);
@@ -413,6 +458,7 @@ export const createWorker = ({
   };
 
   const ready = async () => {
+    await setupListener();
     if (autoStart) await start();
     if (onReady) onReady();
     debug('Ready');

src/actions/get-queue-rate-limit-config.ts

@@ -0,0 +1,25 @@
+import { Redis } from 'ioredis';
+import { get } from '../utils/redis';
+import { getQueueRateLimitKey } from '../utils/keys';
+
+export interface QueueRateLimitConfig {
+  points: number;
+  duration: number;
+}
+
+export const getQueueRateLimitConfig = async ({
+  queue,
+  client,
+}: {
+  queue: string;
+  client: Redis;
+}) => {
+  const rateLimitString = await get({
+    key: getQueueRateLimitKey({ queue }),
+    client,
+  });
+  if (!rateLimitString) {
+    return null;
+  }
+  return JSON.parse(rateLimitString) as QueueRateLimitConfig;
+};

src/actions/set-queue-rate-limit.ts

@@ -1,21 +1,35 @@
 import { Redis } from 'ioredis';
-import { set } from '../utils/redis';
-import { getQueueRateLimitKey } from '../utils/keys';
+import { exec } from '../utils/redis';
+import {
+  getQueueRateLimitKey,
+  getQueueRateLimitUpdatedChannel,
+} from '../utils/keys';
+import { EventType } from '../domain/events/event-type';
+import { serializeEvent } from '../domain/events/serialize-event';
 
 export const setQueueRateLimit = async ({
-  amount,
-  interval,
+  points,
+  duration,
   queue,
   client,
 }: {
-  amount: number;
-  interval: number;
+  points: number;
+  duration: number;
   queue: string;
   client: Redis;
 }) => {
-  await set({
-    key: getQueueRateLimitKey({ queue }),
-    value: JSON.stringify({ amount, interval }),
-    client,
-  });
+  const multi = client.multi();
+  multi.set(
+    getQueueRateLimitKey({ queue }),
+    JSON.stringify({ points, duration }),
+  );
+  multi.publish(
+    getQueueRateLimitUpdatedChannel({ queue }),
+    serializeEvent({
+      type: EventType.QueueRateLimitUpdated,
+      createdAt: new Date(),
+      data: { rateLimitConfig: { points, duration } },
+    }),
+  );
+  await exec(multi);
 };

src/domain/events/event-from-json.ts

@@ -12,5 +12,6 @@ export const eventFromJson = (eventJson: any): Event => {
     type: eventJson.type as EventType,
     task: eventJson.task ? taskFromJson(eventJson.task) : undefined,
     worker: eventJson.worker ? workerFromJson(eventJson.worker) : undefined,
+    data: eventJson.data,
   };
 };

src/domain/events/event-to-json.ts

@@ -11,5 +11,6 @@ export const eventToJson = (event: Event) => {
     type: event.type,
     task: event.task ? taskToJson(event.task) : undefined,
     worker: event.worker ? workerToJson(event.worker) : undefined,
+    data: event.data,
   };
 };

src/domain/events/event-type.ts

@@ -12,4 +12,5 @@ export enum EventType {
   WorkerStarted = 'worker_started',
   WorkerPaused = 'worker_paused',
   WorkerShutdown = 'worker_shutdown',
+  QueueRateLimitUpdated = 'queue_rate_limit_updated',
 }

src/domain/events/event.ts

@@ -1,10 +1,14 @@
 import { Task } from '../tasks/task';
 import { EventType } from './event-type';
 import { WorkerInstance } from '../worker/worker-instance';
+import { QueueRateLimitConfig } from '../../actions/get-queue-rate-limit-config';
 
 export interface Event {
   createdAt: Date;
   type: EventType;
   task?: Task;
   worker?: WorkerInstance;
+  data?: {
+    rateLimitConfig?: QueueRateLimitConfig;
+  };
 }

src/domain/manager/manager.ts

@@ -1,3 +1,4 @@
+import { duration } from 'moment';
 import { Task } from '../tasks/task';
 import { TaskResponse } from './task-response';
 import { WorkerInstance } from '../worker/worker-instance';
@@ -22,6 +23,13 @@ export interface Manager {
   destroyQueue: () => Promise<void>;
   pauseQueue: () => Promise<void>;
   resumeQueue: () => Promise<void>;
+  setQueueRateLimit: ({
+    points,
+    duration,
+  }: {
+    points: number;
+    duration: number;
+  }) => Promise<void>;
   onReady: () => Promise<void>;
   quit: () => Promise<void>;
 }

src/utils/keys.ts

@@ -81,6 +81,9 @@ export const getWorkerStartedChannel = ({ queue }: { queue: string }) =>
 export const getWorkerShutdownChannel = ({ queue }: { queue: string }) =>
   `queue:${queue}:worker-shutdown`;
 
+export const getQueueRateLimitUpdatedChannel = ({ queue }: { queue: string }) =>
+  `queue:${queue}:queue-rate-limit-updated`;
+
 export const getQueueTaskQueuedChannel = ({ queue }: { queue: string }) =>
   `queue:${queue}:task-queued`;
 

src/utils/redis.ts

@@ -257,3 +257,9 @@ export const mget = ({
   keys: string[];
   client: Redis;
 }) => client.mget(...mgetKeys);
+
+export const get = ({ key, client }: { key: string; client: Redis }) =>
+  client.get(key);
+
+export const del = ({ key, client }: { key: string; client: Redis }) =>
+  client.del(key);