Feature: Cross-instance notification system (#689)

* Add notification system

* Graceful worker shutdown on migration notification

* Add notification catalog

* Rename pgListen call to on

apps/webapp/app/platform/zodWorker.server.ts

@@ -15,6 +15,8 @@ import omit from "lodash.omit";
 import { z } from "zod";
 import { PrismaClient, PrismaClientOrTransaction } from "~/db.server";
 import { workerLogger as logger } from "~/services/logger.server";
+import { PgListenService } from "~/services/db/pgListen.server";
+import { safeJsonParse } from "~/utils/json";
 
 export interface MessageCatalogSchema {
   [key: string]: z.ZodFirstPartySchemaTypes | z.ZodDiscriminatedUnion<any, any>;
@@ -164,8 +166,23 @@ export class ZodWorker<TMessageCatalog extends MessageCatalogSchema> {
       this.#logDebug("pool:create", { attempts });
     });
 
-    this.#runner?.events.on("pool:listen:success", ({ workerPool, client }) => {
+    this.#runner?.events.on("pool:listen:success", async ({ workerPool, client }) => {
       this.#logDebug("pool:listen:success");
+
+      // hijack client instance to listen and react to incoming NOTIFY events
+      const pgListen = new PgListenService(client, this.#name, logger);
+
+      await pgListen.call("trigger:graphile:migrate", async ({ latestMigration }) => {
+        this.#logDebug("Detected incoming migration", { latestMigration });
+
+        if (latestMigration > 10) {
+          // already migrated past v0.14 - nothing to do
+          return;
+        }
+
+        // simulate SIGTERM to trigger graceful shutdown
+        this._handleSignal("SIGTERM");
+      });
     });
 
     this.#runner?.events.on("pool:listen:error", ({ error }) => {

apps/webapp/app/services/db/pgListen.server.ts

@@ -0,0 +1,66 @@
+import type { PoolClient } from "pg";
+import { z } from "zod";
+import { Logger } from "@trigger.dev/core";
+import { logger } from "~/services/logger.server";
+import { NotificationCatalog, NotificationChannel, notificationCatalog } from "./types";
+import { safeJsonParse } from "~/utils/json";
+
+export class PgListenService {
+  #poolClient: PoolClient;
+  #logger: Logger;
+  #loggerNamespace: string;
+
+  constructor(poolClient: PoolClient, loggerNamespace?: string, loggerInstance?: Logger) {
+    this.#poolClient = poolClient;
+    this.#logger = loggerInstance ?? logger;
+    this.#loggerNamespace = loggerNamespace ?? "";
+  }
+
+  public async on<TChannel extends NotificationChannel>(
+    channelName: TChannel,
+    callback: (payload: z.infer<NotificationCatalog[TChannel]>) => Promise<void>
+  ) {
+    this.#logDebug("Registering notification handler", { channelName });
+
+    const isValidChannel = channelName.match(/^[a-zA-Z0-9:-_]+$/);
+
+    if (!isValidChannel) {
+      throw new Error(`Invalid channel name: ${channelName}`);
+    }
+
+    this.#poolClient.query(`LISTEN "${channelName}"`).then(null, (error) => {
+      this.#logDebug("LISTEN error", error);
+    });
+
+    this.#poolClient.on("notification", async (notification) => {
+      if (notification.channel !== channelName) {
+        return;
+      }
+
+      this.#logDebug("Notification received", { notification });
+
+      if (!notification.payload) {
+        return;
+      }
+
+      const payload = safeJsonParse(notification.payload);
+
+      const parsedPayload = notificationCatalog[channelName].safeParse(payload);
+
+      if (!parsedPayload.success) {
+        throw new Error(
+          `Failed to parse notification payload: ${channelName} - ${JSON.stringify(
+            parsedPayload.error
+          )}`
+        );
+      }
+
+      await callback(parsedPayload.data);
+    });
+  }
+
+  #logDebug(message: string, args?: any) {
+    const namespace = this.#loggerNamespace ? `[${this.#loggerNamespace}]` : "";
+    this.#logger.debug(`[pgListen]${namespace} ${message}`, args);
+  }
+}

apps/webapp/app/services/db/pgNotify.server.ts

@@ -0,0 +1,27 @@
+import { z } from "zod";
+import { PrismaClient, prisma } from "~/db.server";
+import { logger } from "~/services/logger.server";
+import { NotificationCatalog, NotificationChannel } from "./types";
+
+export class PgNotifyService {
+  #prismaClient: PrismaClient;
+
+  constructor(prismaClient: PrismaClient = prisma) {
+    this.#prismaClient = prismaClient;
+  }
+
+  public async call<TChannel extends NotificationChannel>(
+    channelName: TChannel,
+    payload: z.infer<NotificationCatalog[TChannel]>
+  ) {
+    this.#logDebug("Sending notification", { channelName, notifyPayload: payload });
+
+    await this.#prismaClient.$executeRaw`
+      SELECT pg_notify(${channelName}, ${JSON.stringify(payload)})
+    `;
+  }
+
+  #logDebug(message: string, args?: any) {
+    logger.debug(`[pgNotify] ${message}`, args);
+  }
+}

apps/webapp/app/services/db/types.ts

@@ -0,0 +1,11 @@
+import { z } from "zod";
+
+export const notificationCatalog = {
+  "trigger:graphile:migrate": z.object({
+    latestMigration: z.number(),
+  }),
+};
+
+export type NotificationCatalog = typeof notificationCatalog;
+
+export type NotificationChannel = keyof NotificationCatalog;

apps/webapp/app/services/worker.server.ts

@@ -21,6 +21,7 @@ import { DeliverHttpSourceRequestService } from "./sources/deliverHttpSourceRequ
 import { PerformTaskOperationService } from "./tasks/performTaskOperation.server";
 import { ProcessCallbackTimeoutService } from "./tasks/processCallbackTimeout";
 import { ProbeEndpointService } from "./endpoints/probeEndpoint.server";
+import { PgNotifyService } from "./db/pgNotify.server";
 
 const workerCatalog = {
   indexEndpoint: z.object({
@@ -121,6 +122,10 @@ if (env.NODE_ENV === "production") {
 }
 
 export async function init() {
+  // const pgNotify = new PgNotifyService();
+  // await pgNotify.call("trigger:graphile:migrate", { latestMigration: 10 });
+  // await new Promise((resolve) => setTimeout(resolve, 10000))
+
   if (env.WORKER_ENABLED === "true") {
     await workerQueue.initialize();
   }

apps/webapp/package.json

@@ -65,6 +65,7 @@
     "@trigger.dev/core": "workspace:*",
     "@trigger.dev/database": "workspace:*",
     "@trigger.dev/sdk": "workspace:*",
+    "@types/pg": "8.6.6",
     "@uiw/react-codemirror": "^4.19.5",
     "class-variance-authority": "^0.5.2",
     "clsx": "^1.2.1",

pnpm-lock.yaml

@@ -121,6 +121,7 @@ importers:
       '@types/morgan': ^1.9.3
       '@types/node': ^18.11.15
       '@types/node-fetch': ^2.6.2
+      '@types/pg': 8.6.6
       '@types/prismjs': ^1.26.0
       '@types/qs': ^6.9.7
       '@types/react': 18.2.17
@@ -236,6 +237,7 @@ importers:
       '@trigger.dev/core': link:../../packages/core
       '@trigger.dev/database': link:../../packages/database
       '@trigger.dev/sdk': link:../../packages/trigger-sdk
+      '@types/pg': 8.6.6
       '@uiw/react-codemirror': 4.19.5_th22fcplkuhrqjnlojwclcaim4
       class-variance-authority: 0.5.2_typescript@5.2.2
       clsx: 1.2.1