Add a metric for number of event repo subscribers and project subscribers

Author: ericallam

apps/webapp/app/v3/eventRepository.server.ts

@@ -26,6 +26,9 @@ import { env } from "~/env.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
 import { DynamicFlushScheduler } from "./dynamicFlushScheduler.server";
+import { singleton } from "~/utils/singleton";
+import { Gauge } from "prom-client";
+import { metricsRegister } from "~/metrics.server";
 
 export type CreatableEvent = Omit<
   Prisma.TaskEventCreateInput,
@@ -148,6 +151,11 @@ export class EventRepository {
   private readonly _flushScheduler: DynamicFlushScheduler<CreatableEvent>;
   private _randomIdGenerator = new RandomIdGenerator();
   private _redisPublishClient: Redis;
+  private _subscriberCount = 0;
+
+  get subscriberCount() {
+    return this._subscriberCount;
+  }
 
   constructor(private db: PrismaClient = prisma, private readonly _config: EventRepoConfig) {
     this._flushScheduler = new DynamicFlushScheduler({
@@ -721,6 +729,9 @@ export class EventRepository {
     // Subscribe to the channel.
     await redis.psubscribe(channel);
 
+    // Increment the subscriber count.
+    this._subscriberCount++;
+
     const eventEmitter = new EventEmitter();
 
     // Define the message handler.
@@ -733,6 +744,8 @@ export class EventRepository {
     // Return a function that can be used to unsubscribe.
     const unsubscribe = async () => {
       await redis.punsubscribe(channel);
+      redis.quit();
+      this._subscriberCount--;
     };
 
     return {
@@ -796,19 +809,34 @@ export class EventRepository {
   }
 }
 
-export const eventRepository = new EventRepository(prisma, {
-  batchSize: env.EVENTS_BATCH_SIZE,
-  batchInterval: env.EVENTS_BATCH_INTERVAL,
-  retentionInDays: env.EVENTS_DEFAULT_LOG_RETENTION,
-  redis: {
-    port: env.REDIS_PORT,
-    host: env.REDIS_HOST,
-    username: env.REDIS_USERNAME,
-    password: env.REDIS_PASSWORD,
-    enableAutoPipelining: true,
-    ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
-  },
-});
+export const eventRepository = singleton("eventRepo", initializeEventRepo);
+
+function initializeEventRepo() {
+  const repo = new EventRepository(prisma, {
+    batchSize: env.EVENTS_BATCH_SIZE,
+    batchInterval: env.EVENTS_BATCH_INTERVAL,
+    retentionInDays: env.EVENTS_DEFAULT_LOG_RETENTION,
+    redis: {
+      port: env.REDIS_PORT,
+      host: env.REDIS_HOST,
+      username: env.REDIS_USERNAME,
+      password: env.REDIS_PASSWORD,
+      enableAutoPipelining: true,
+      ...(env.REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
+    },
+  });
+
+  new Gauge({
+    name: "event_repository_subscriber_count",
+    help: "Number of event repository subscribers",
+    collect() {
+      this.set(repo.subscriberCount);
+    },
+    registers: [metricsRegister],
+  });
+
+  return repo;
+}
 
 export function stripAttributePrefix(attributes: Attributes, prefix: string) {
   const result: Attributes = {};

apps/webapp/app/v3/services/projectPubSub.server.ts

@@ -2,6 +2,8 @@ import { z } from "zod";
 import { singleton } from "~/utils/singleton";
 import { ZodPubSub, ZodSubscriber } from "../utils/zodPubSub.server";
 import { env } from "~/env.server";
+import { Gauge } from "prom-client";
+import { metricsRegister } from "~/metrics.server";
 
 const messageCatalog = {
   WORKER_CREATED: z.object({
@@ -18,7 +20,7 @@ export type ProjectSubscriber = ZodSubscriber<typeof messageCatalog>;
 export const projectPubSub = singleton("projectPubSub", initializeProjectPubSub);
 
 function initializeProjectPubSub() {
-  return new ZodPubSub({
+  const pubSub = new ZodPubSub({
     redis: {
       port: env.REDIS_PORT,
       host: env.REDIS_HOST,
@@ -29,4 +31,15 @@ function initializeProjectPubSub() {
     },
     schema: messageCatalog,
   });
+
+  new Gauge({
+    name: "project_pub_sub_subscribers",
+    help: "Number of project pub sub subscribers",
+    collect() {
+      this.set(pubSub.subscriberCount);
+    },
+    registers: [metricsRegister],
+  });
+
+  return pubSub;
 }

apps/webapp/app/v3/utils/zodPubSub.server.ts

@@ -59,6 +59,8 @@ class RedisZodSubscriber<TMessageCatalog extends ZodMessageCatalogSchema>
     await this._subscriber.punsubscribe();
 
     this.onUnsubscribed.post({ pattern: this._pattern });
+
+    this._subscriber.quit();
   }
 
   async #onMessage(pattern: string, channel: string, serializedMessage: string) {