diff --git a/apps/webapp/app/platform/zodWorker.server.ts b/apps/webapp/app/platform/zodWorker.server.ts index 780454642..ff24992f7 100644 --- a/apps/webapp/app/platform/zodWorker.server.ts +++ b/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; @@ -164,8 +166,28 @@ export class ZodWorker { 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 }) => { diff --git a/apps/webapp/app/services/db/pgListen.server.ts b/apps/webapp/app/services/db/pgListen.server.ts new file mode 100644 index 000000000..74ecc26e7 --- /dev/null +++ b/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) { + 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); + } +} diff --git a/apps/webapp/app/services/db/pgNotify.server.ts b/apps/webapp/app/services/db/pgNotify.server.ts new file mode 100644 index 000000000..2b41e09ad --- /dev/null +++ b/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); + } +} diff --git a/apps/webapp/app/services/worker.server.ts b/apps/webapp/app/services/worker.server.ts index 307abcf7b..2ac9efa7b 100644 --- a/apps/webapp/app/services/worker.server.ts +++ b/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() { diff --git a/apps/webapp/package.json b/apps/webapp/package.json index e5d963ec6..16797db82 100644 --- a/apps/webapp/package.json +++ b/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",