feat: add jobs memory usage avg field (#17)

Co-authored-by: Boris Dorofeev <[email protected]>

Author: intpro

src/helpers/index.ts

@@ -7,3 +7,4 @@ export * from './wrapQueueArgs';
 export * from './composeFC';
 export * from './deleteKeys';
 export * from './fixDelayStream';
+export * from './memoryUsage';

src/helpers/memoryUsage.ts

@@ -0,0 +1,41 @@
+import { getBullConnection } from './getBullConnection';
+import { Options } from '../definitions';
+import Redis from 'ioredis';
+import path from 'path';
+import util from 'util';
+import fs from 'fs';
+
+let commandDefined = false;
+
+async function defineCommand(redis: Redis.Redis) {
+  const readFile = util.promisify(fs.readFile);
+  const lua = await readFile(path.join(__dirname, 'zsetKeysMemoryUsage.lua'));
+
+  redis.defineCommand('zsetKeysMemoryUsage', { numberOfKeys: 0, lua: lua.toString() });
+
+  commandDefined = true;
+}
+
+export async function getZsetKeysMemoryUsageAvg(
+  prefix: string,
+  queueName: string,
+  keySetName: string,
+  limit: number,
+  opts: Options
+): Promise<number> {
+  const redis = getBullConnection(opts);
+
+  if (!commandDefined) {
+    await defineCommand(redis);
+  }
+
+  const fullPrefix = [prefix, queueName].join(':');
+
+  const [amountBytes, jobsCount] = await (redis as any).zsetKeysMemoryUsage(
+    fullPrefix + ':',
+    keySetName,
+    limit
+  );
+
+  return amountBytes / (jobsCount || 1);
+}

src/helpers/zsetKeysMemoryUsage.lua

@@ -0,0 +1,21 @@
+--[[
+  Input:
+    ARGV[1]  keys prefix
+    ARGV[2]  set name
+    ARGV[3]  limit
+
+  Output:
+    (integer) bytesAmount
+]]
+
+local jobs = redis.call("ZREVRANGE", ARGV[1]..ARGV[2], 0, ARGV[3])
+
+local bytesAmount = 0
+
+if (#jobs > 0) then
+  for _, jobId in ipairs(jobs) do
+    bytesAmount = bytesAmount + redis.call('memory', 'usage', ARGV[1]..jobId)
+  end  
+end
+
+return {bytesAmount, #jobs}

src/types/queue/Queue.jobsMemoryUsageAvg.ts

@@ -0,0 +1,41 @@
+import { SchemaComposer, ObjectTypeComposerFieldConfigDefinition } from 'graphql-compose';
+import { Queue } from 'bullmq';
+import { Options } from '../../definitions';
+import { getZsetKeysMemoryUsageAvg } from '../../helpers';
+
+export function createJobsMemoryUsageAvgFC(
+  sc: SchemaComposer<any>,
+  opts: Options
+): ObjectTypeComposerFieldConfigDefinition<any, any> {
+  return {
+    type: 'Int!',
+    args: {
+      keySetName: {
+        type: sc.createEnumTC({
+          name: 'keySetNamesEnum',
+          values: {
+            COMPLETED: { value: 'completed' },
+            FAILED: { value: 'failed' },
+          },
+        }),
+        defaultValue: 'completed',
+      },
+      limit: {
+        type: 'Int',
+        defaultValue: 100,
+      },
+    },
+    resolve: async (queue: Queue, { limit, keySetName }) => {
+      const avgBytes =
+        (await getZsetKeysMemoryUsageAvg(
+          queue.opts?.prefix || 'bull',
+          queue.name,
+          keySetName,
+          limit,
+          opts
+        )) || 0;
+
+      return avgBytes.toFixed(0);
+    },
+  };
+}

src/types/queue/Queue.ts

@@ -9,6 +9,7 @@ import { createDelayedJobsFC } from './Queue.delayedJobs';
 import { createFailedJobsFC } from './Queue.failedJobs';
 import { createWorkersTC } from './Queue.workers';
 import { createDurationAvgFC } from './Queue.durationAvg';
+import { createJobsMemoryUsageAvgFC } from './Queue.jobsMemoryUsageAvg';
 import { SchemaComposer } from 'graphql-compose';
 import { Options } from '../../definitions';
 
@@ -29,6 +30,7 @@ export function getQueueTC(sc: SchemaComposer<any>, opts: Options) {
       jobsFailed: createFailedJobsFC(sc, opts),
       activeWorkers: createWorkersTC(sc, opts),
       durationAvg: createDurationAvgFC(sc, opts),
+      jobsMemoryUsageAvg: createJobsMemoryUsageAvgFC(sc, opts),
     });
   });
 }