Shut down existing workers prior to new migrations

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,28 @@ 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 (payload) => {
+        const parsedPayload = safeJsonParse(payload);
+
+        const MigrationNotificationPayloadSchema = z.object({
+          latestMigration: z.number(),
+        });
+
+        const migrationPayload = MigrationNotificationPayloadSchema.parse(parsedPayload);
+
+        this.#logDebug("Detected incoming migration", {
+          latestMigration: migrationPayload.latestMigration,
+        });
+
+        // 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,48 @@
+import { logger } from "~/services/logger.server";
+import { Logger } from "@trigger.dev/core";
+import type { PoolClient } from "pg";
+
+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 call(channelName: string, callback: (payload: string) => 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;
+      }
+
+      await callback(notification.payload);
+    });
+  }
+
+  #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,23 @@
+import { SerializableJson } from "@trigger.dev/core";
+import { PrismaClient, prisma } from "~/db.server";
+import { logger } from "~/services/logger.server";
+
+export class PgNotifyService {
+  #prismaClient: PrismaClient;
+
+  constructor(prismaClient: PrismaClient = prisma) {
+    this.#prismaClient = prismaClient;
+  }
+
+  public async call(channelName: string, payload: SerializableJson) {
+    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/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,7 +122,7 @@ if (env.NODE_ENV === "production") {
 }
 
 export async function init() {
-  await addMigrationDelay();
+  await addMigrationDelayAndNotify();
 
   if (env.WORKER_ENABLED === "true") {
     await workerQueue.initialize();
@@ -133,10 +134,11 @@ export async function init() {
 }
 
 /** Helper for graphile-worker v0.14.0 migration. No-op if already migrated. */
-async function addMigrationDelay() {
-  const migrationQueryResult = await prisma.$queryRawUnsafe(`
+async function addMigrationDelayAndNotify() {
+  const migrationQueryResult = await prisma.$queryRaw`
     SELECT id FROM graphile_worker.migrations
-    ORDER BY id DESC LIMIT 1`);
+    ORDER BY id DESC LIMIT 1
+  `;
 
   const MigrationQueryResultSchema = z.array(z.object({ id: z.number() }));
 
@@ -155,9 +157,15 @@ async function addMigrationDelay() {
     return;
   }
 
-  console.log(`⚠️  delaying worker startup for migration: ${env.WORKER_MIGRATION_DELAY}ms`);
+  console.log(`⚠️  detected pending graphile migration`);
+  console.log(`⚠️  delaying worker startup by ${env.WORKER_MIGRATION_DELAY}ms`);
 
   await new Promise((resolve) => setTimeout(resolve, env.WORKER_MIGRATION_DELAY));
+
+  console.log(`⚠️  notifying running workers about incoming migration`);
+
+  const pgNotify = new PgNotifyService();
+  await pgNotify.call("trigger:graphile:migrate", { latestMigration });
 }
 
 function getWorkerQueue() {

apps/webapp/package.json

@@ -64,6 +64,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",