diff --git a/apps/webapp/app/platform/zodWorker.server.ts b/apps/webapp/app/platform/zodWorker.server.ts index d0a370650..d25d0302c 100644 --- a/apps/webapp/app/platform/zodWorker.server.ts +++ b/apps/webapp/app/platform/zodWorker.server.ts @@ -8,8 +8,9 @@ import type { Task, TaskList, TaskSpec, + WorkerUtils, } from "graphile-worker"; -import { run as graphileRun, parseCronItems } from "graphile-worker"; +import { run as graphileRun, makeWorkerUtils, parseCronItems } from "graphile-worker"; import omit from "lodash.omit"; import { z } from "zod"; @@ -138,6 +139,7 @@ export class ZodWorker { #reporter?: ZodWorkerReporter; #shutdownTimeoutInMs?: number; #shuttingDown = false; + #workerUtils?: WorkerUtils; constructor(options: ZodWorkerOptions) { this.#name = options.name; @@ -166,6 +168,8 @@ export class ZodWorker { const parsedCronItems = parseCronItems(this.#createCronItemsFromRecurringTasks()); + this.#workerUtils = await makeWorkerUtils(this.#runnerOptions); + this.#runner = await graphileRun({ ...this.#runnerOptions, noHandleSignals: true, @@ -330,11 +334,11 @@ export class ZodWorker { >( identifier: K, payload: TPayload extends any[] ? TPayload : never, - options?: ZodWorkerBatchEnqueueOptions + options: ZodWorkerBatchEnqueueOptions ): Promise { const task = this.#tasks[identifier]; - const optionsWithoutTx = removeUndefinedKeys(omit(options ?? {}, ["tx"])); + const optionsWithoutTx = removeUndefinedKeys(omit(options, ["tx"])); // Make sure options passed in to enqueue take precedence over task options const spec = { @@ -364,7 +368,7 @@ export class ZodWorker { spec, }); - const job = await this.#addBatchJob( + const { job, durationInMs } = await this.#addBatchJob( identifier as string, payload, spec as BatchTaskSpec, @@ -376,6 +380,7 @@ export class ZodWorker { payload, spec, job, + durationInMs, }); return job; @@ -444,6 +449,8 @@ export class ZodWorker { spec: BatchTaskSpec, tx: PrismaClientOrTransaction ) { + const now = performance.now(); + const results = await tx.$queryRawUnsafe( `SELECT * FROM add_batch_job( identifier => $1::text, @@ -452,10 +459,10 @@ export class ZodWorker { queue_name => $4::text, run_at => $5::timestamptz, max_attempts => $6::int, - max_payloads => $7::int, - priority => $8::int, - flags => $9::text[], - job_key_mode => $10::text + priority => $7::int, + flags => $8::text[], + job_key_mode => $9::text, + max_payloads => $10::int )`, identifier, spec.jobKey, @@ -463,12 +470,14 @@ export class ZodWorker { spec.queueName || null, spec.runAt || null, spec.maxAttempts || null, - spec.maxPayloads || null, spec.priority || null, spec.flags || null, - spec.jobKeyMode || "preserve_run_at" + spec.jobKeyMode || "preserve_run_at", + spec.maxPayloads || null ); + const durationInMs = performance.now() - now; + const rows = AddJobResultsSchema.safeParse(results); if (!rows.success) { @@ -479,7 +488,7 @@ export class ZodWorker { const job = rows.data[0]; - return job as GraphileJob; + return { job: job as GraphileJob, durationInMs: Math.floor(durationInMs) }; } async #removeJob(jobKey: string, tx: PrismaClientOrTransaction) { @@ -670,6 +679,10 @@ export class ZodWorker { return; } + if (!this.#workerUtils) { + throw new Error("WorkerUtils need to be initialized before running job cleanup."); + } + const job = helpers.job; logger.debug("Received cleanup task", { @@ -696,22 +709,37 @@ export class ZodWorker { }); const rawResults = await this.#prisma.$queryRawUnsafe( - `WITH rows AS (SELECT id FROM ${this.graphileWorkerSchema}.jobs WHERE run_at < $1 AND locked_at IS NULL AND max_attempts = attempts LIMIT $2 FOR UPDATE) DELETE FROM ${this.graphileWorkerSchema}.jobs WHERE id IN (SELECT id FROM rows) RETURNING id`, + `SELECT id + FROM ${this.graphileWorkerSchema}.jobs + WHERE run_at > $1 + AND locked_at IS NULL + AND max_attempts = attempts + LIMIT $2`, expirationDate, this.#cleanup.maxCount ); - const results = Array.isArray(rawResults) ? rawResults : []; + const results = z + .array( + z.object({ + id: z.coerce.string(), + }) + ) + .parse(rawResults); + + const completedJobs = await this.#workerUtils.completeJobs(results.map((job) => job.id)); logger.debug("Cleaned up old jobs", { - count: results.length, + found: results.length, + deleted: completedJobs.length, expirationDate, payload, }); if (this.#reporter) { await this.#reporter("cleanup_stats", { - count: results.length, + found: results.length, + deleted: completedJobs.length, expirationDate, ts: payload._cron.ts, }); diff --git a/apps/webapp/app/routes/admin.api.v1.simulate.batch.ts b/apps/webapp/app/routes/admin.api.v1.simulate.batch.ts index 554e111d7..ffa9f314b 100644 --- a/apps/webapp/app/routes/admin.api.v1.simulate.batch.ts +++ b/apps/webapp/app/routes/admin.api.v1.simulate.batch.ts @@ -1,9 +1,9 @@ -import { ActionArgs, json, redirect } from "@remix-run/server-runtime"; +import { ActionFunctionArgs, json } from "@remix-run/server-runtime"; import { prisma } from "~/db.server"; import { authenticateApiRequest } from "~/services/apiAuth.server"; import { workerQueue } from "~/services/worker.server"; -export async function action({ request }: ActionArgs) { +export async function action({ request }: ActionFunctionArgs) { // Next authenticate the request const authenticationResult = await authenticateApiRequest(request); @@ -36,7 +36,7 @@ export async function action({ request }: ActionArgs) { { jobKey: "simulateBatch", runAt: new Date(Date.now() + (body.deliverAfter ?? 0) * 1000), - maxPayloads: body.maxPayloads ?? null + maxPayloads: body.maxPayloads ?? null, } );