refactor: Aggregate worker metrics in background

runner/src/metrics.ts

@@ -1,25 +1,25 @@
 import express from 'express';
-import promClient from 'prom-client';
+import { Gauge, Histogram, Counter, AggregatorRegistry } from 'prom-client';
 
-const UNPROCESSED_STREAM_MESSAGES = new promClient.Gauge({
+const UNPROCESSED_STREAM_MESSAGES = new Gauge({
   name: 'queryapi_runner_unprocessed_stream_messages',
   help: 'Number of Redis Stream messages not yet processed',
   labelNames: ['indexer', 'type'],
 });
 
-const EXECUTION_DURATION = new promClient.Histogram({
+const EXECUTION_DURATION = new Histogram({
   name: 'queryapi_runner_execution_duration_milliseconds',
   help: 'Time taken to execute an indexer function',
   labelNames: ['indexer', 'type'],
 });
 
-const CACHE_HIT = new promClient.Counter({
+const CACHE_HIT = new Counter({
   name: 'queryapi_runner_cache_hit',
   help: 'The number of times cache was hit successfully',
   labelNames: ['type', 'key']
 });
 
-const CACHE_MISS = new promClient.Counter({
+const CACHE_MISS = new Counter({
   name: 'queryapi_runner_cache_miss',
   help: 'The number of times cache was missed',
   labelNames: ['type', 'key']
@@ -32,15 +32,22 @@ export const METRICS = {
   CACHE_MISS
 };
 
+const aggregatorRegistry = new AggregatorRegistry();
+const workerMetrics: Record<number, string> = {};
+
+export const registerWorkerMetrics = (workerId: number, metrics: string): void => {
+  workerMetrics[workerId] = metrics;
+};
+
 export const startServer = async (): Promise<void> => {
   const app = express();
 
   // https://github.com/DefinitelyTyped/DefinitelyTyped/issues/50871
   // eslint-disable-next-line @typescript-eslint/no-misused-promises
   app.get('/metrics', async (_req, res) => {
-    res.set('Content-Type', promClient.register.contentType);
+    res.set('Content-Type', aggregatorRegistry.contentType);
 
-    const metrics = await promClient.register.metrics();
+    const metrics = await AggregatorRegistry.aggregate(Object.values(workerMetrics)).metrics();
     res.send(metrics);
   });
 

runner/src/stream-handler/stream-handler.ts

@@ -1,12 +1,10 @@
 import path from 'path';
 import { Worker, isMainThread } from 'worker_threads';
 
-import { type Message } from './types';
-import { METRICS } from '../metrics';
-import { Gauge, Histogram } from 'prom-client';
+import { registerWorkerMetrics } from '../metrics';
 
 export default class StreamHandler {
-  private readonly worker?: Worker;
+  private readonly worker: Worker;
 
   constructor (
     public readonly streamKey: string
@@ -18,27 +16,21 @@ export default class StreamHandler {
         },
       });
 
-      this.worker.on('message', this.handleMessage);
-      this.worker.on('error', this.handleError);
+      this.worker.on('message', this.handleMessage.bind(this));
+      this.worker.on('error', this.handleError.bind(this));
     } else {
       throw new Error('StreamHandler should not be instantiated in a worker thread');
     }
   }
 
   private handleError (error: Error): void {
     console.log(`Encountered error processing stream: ${this.streamKey}, terminating thread`, error);
-    this.worker?.terminate().catch(() => {
+    this.worker.terminate().catch(() => {
       console.log(`Failed to terminate thread for stream: ${this.streamKey}`);
     });
   }
 
-  private handleMessage (message: Message): void {
-    if (METRICS[message.type] instanceof Gauge) {
-      (METRICS[message.type] as Gauge).labels(message.labels).set(message.value);
-    }
-
-    if (METRICS[message.type] instanceof Histogram) {
-      (METRICS[message.type] as Histogram).labels(message.labels).observe(message.value);
-    }
+  private handleMessage (message: string): void {
+    registerWorkerMetrics(this.worker.threadId, message);
   }
 }

runner/src/stream-handler/worker.ts

@@ -1,8 +1,9 @@
 import { isMainThread, parentPort, workerData } from 'worker_threads';
+import promClient from 'prom-client';
 
 import Indexer from '../indexer';
 import RedisClient from '../redis-client';
-import { type Message } from './types';
+import { METRICS } from '../metrics';
 
 if (isMainThread) {
   throw new Error('Worker should not be run on main thread');
@@ -53,24 +54,17 @@ void (async function main () {
 
       await redisClient.deleteStreamMessage(streamKey, id);
 
-      parentPort?.postMessage({
-        type: 'EXECUTION_DURATION',
-        labels: { indexer: indexerName, type: streamType },
-        value: performance.now() - startTime,
-      } satisfies Message);
+      METRICS.EXECUTION_DURATION.labels({ indexer: indexerName, type: streamType }).observe(performance.now() - startTime);
 
       console.log(`Success: ${indexerName}`);
     } catch (err) {
       await sleep(10000);
       console.log(`Failed: ${indexerName}`, err);
     } finally {
       const unprocessedMessages = await redisClient.getUnprocessedStreamMessages(streamKey);
+      METRICS.UNPROCESSED_STREAM_MESSAGES.labels({ indexer: indexerName, type: streamType }).set(unprocessedMessages?.length ?? 0);
 
-      parentPort?.postMessage({
-        type: 'UNPROCESSED_STREAM_MESSAGES',
-        labels: { indexer: indexerName, type: streamType },
-        value: unprocessedMessages?.length ?? 0,
-      } satisfies Message);
+      parentPort?.postMessage(await promClient.register.getMetricsAsJSON());
     }
   }
 })();