feat: ghetto durable compute

This commit is contained in:
Mythie
2024-05-16 15:44:39 +10:00
parent 61827ad729
commit 991f808890
20 changed files with 847 additions and 151 deletions

View File

@ -1,3 +1,6 @@
import { JobClient } from './client/client';
import { registerJobs } from './definitions';
export const jobsClient = JobClient.getInstance();
registerJobs(jobsClient);

View File

@ -1,16 +1,20 @@
import { z } from 'zod';
import type { Json } from './json';
export const ZTriggerJobOptionsSchema = z.object({
id: z.string().optional(),
name: z.string(),
payload: z.unknown().refine((x) => x !== undefined, { message: 'payload is required' }),
payload: z.any().refine((x) => x !== undefined, { message: 'payload is required' }),
timestamp: z.number().optional(),
});
// The Omit is a temporary workaround for a "bug" in the zod library
// @see: https://github.com/colinhacks/zod/issues/2966
export type TriggerJobOptions = Omit<z.infer<typeof ZTriggerJobOptionsSchema>, 'payload'> & {
payload: unknown;
// Don't tell the feds
// eslint-disable-next-line @typescript-eslint/no-explicit-any
payload: any;
};
// eslint-disable-next-line @typescript-eslint/no-explicit-any
@ -28,7 +32,7 @@ export type JobDefinition<T = any> = {
export interface JobRunIO {
// stableRun<T extends Json | void>(cacheKey: string, callback: (io: JobRunIO) => T | Promise<T>): Promise<T>;
stableRun<T extends Json | void>(cacheKey: string, callback: () => Promise<T>): Promise<T>;
runTask<T extends Json | void>(cacheKey: string, callback: () => Promise<T>): Promise<T>;
triggerJob(cacheKey: string, options: TriggerJobOptions): Promise<unknown>;
wait(cacheKey: string, ms: number): Promise<void>;
logger: {

View File

@ -2,13 +2,13 @@
* Below type is borrowed from Trigger.dev's SDK, it may be moved elsewhere later.
*/
type JsonPrimitive = string | number | boolean | null | undefined | Date | symbol;
export type JsonPrimitive = string | number | boolean | null | undefined | Date | symbol;
type JsonArray = Json[];
export type JsonArray = Json[];
type JsonRecord<T> = {
export type JsonRecord<T> = {
[Property in keyof T]: Json;
};
// eslint-disable-next-line @typescript-eslint/no-explicit-any
type Json<T = any> = JsonPrimitive | JsonArray | JsonRecord<T>;
export type Json<T = any> = JsonPrimitive | JsonArray | JsonRecord<T>;

View File

@ -1,11 +1,16 @@
import type { NextApiRequest, NextApiResponse } from 'next';
import { sha256 } from '@noble/hashes/sha256';
import { json } from 'micro';
import { prisma } from '@documenso/prisma';
import { BackgroundJobStatus, Prisma } from '@documenso/prisma/client';
import { NEXT_PUBLIC_WEBAPP_URL } from '../../constants/app';
import { sign } from '../../server-only/crypto/sign';
import { verify } from '../../server-only/crypto/verify';
import type { JobDefinition, JobRunIO, TriggerJobOptions } from './_internal/job';
import type { Json } from './_internal/json';
import { BaseJobProvider } from './base';
export class LocalJobProvider extends BaseJobProvider {
@ -33,33 +38,57 @@ export class LocalJobProvider extends BaseJobProvider {
}
public async triggerJob(options: TriggerJobOptions) {
const signature = sign(options);
console.log({ jobDefinitions: this._jobDefinitions });
await Promise.race([
fetch(`${NEXT_PUBLIC_WEBAPP_URL()}/api/jobs/trigger`, {
method: 'POST',
body: JSON.stringify(options),
headers: {
'Content-Type': 'application/json',
'X-Job-Signature': signature,
},
const eligibleJobs = Object.values(this._jobDefinitions).filter(
(job) => job.trigger.name === options.name,
);
console.log({ options });
console.log(
'Eligible jobs:',
eligibleJobs.map((job) => job.name),
);
await Promise.all(
eligibleJobs.map(async (job) => {
// Ideally we will change this to a createMany with returning later once we upgrade Prisma
// @see: https://github.com/prisma/prisma/releases/tag/5.14.0
const pendingJob = await prisma.backgroundJob.create({
data: {
jobId: job.id,
name: job.name,
version: job.version,
payload: options.payload,
},
});
await this.submitJobToEndpoint({
jobId: pendingJob.id,
jobDefinitionId: pendingJob.jobId,
data: options,
});
}),
new Promise((resolve) => {
setTimeout(resolve, 150);
}),
]);
);
}
public getApiHandler() {
return async (req: NextApiRequest, res: NextApiResponse) => {
if (req.method === 'POST') {
const jobId = req.headers['x-job-id'];
const signature = req.headers['x-job-signature'];
const isRetry = req.headers['x-job-retry'] !== undefined;
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
const options = (await json(req)) as TriggerJobOptions;
const definition = this._jobDefinitions[options.name];
if (typeof signature !== 'string' || typeof options !== 'object') {
if (
typeof jobId !== 'string' ||
typeof signature !== 'string' ||
typeof options !== 'object'
) {
res.status(400).send('Bad request');
return;
}
@ -92,10 +121,83 @@ export class LocalJobProvider extends BaseJobProvider {
console.log(`[JOBS]: Triggering job ${options.name} with payload`, options.payload);
await definition.handler({
payload: options.payload,
io: this.createJobRunIO(options.name),
});
let backgroundJob = await prisma.backgroundJob
.update({
where: {
id: jobId,
status: BackgroundJobStatus.PENDING,
},
data: {
status: BackgroundJobStatus.PROCESSING,
retried: {
increment: isRetry ? 1 : 0,
},
lastRetriedAt: isRetry ? new Date() : undefined,
},
})
.catch(() => null);
if (!backgroundJob) {
res.status(404).send('Job not found');
return;
}
try {
await definition.handler({
payload: options.payload,
io: this.createJobRunIO(jobId),
});
backgroundJob = await prisma.backgroundJob.update({
where: {
id: jobId,
status: BackgroundJobStatus.PROCESSING,
},
data: {
status: BackgroundJobStatus.COMPLETED,
completedAt: new Date(),
},
});
} catch (error) {
console.error(`[JOBS]: Job ${options.name} failed`, error);
const taskHasExceededRetries = error instanceof BackgroundTaskExceededRetriesError;
const jobHasExceededRetries =
backgroundJob.retried >= backgroundJob.maxRetries &&
!(error instanceof BackgroundTaskFailedError);
if (taskHasExceededRetries || jobHasExceededRetries) {
backgroundJob = await prisma.backgroundJob.update({
where: {
id: jobId,
status: BackgroundJobStatus.PROCESSING,
},
data: {
status: BackgroundJobStatus.FAILED,
completedAt: new Date(),
},
});
res.status(500).send('Task exceeded retries');
return;
}
backgroundJob = await prisma.backgroundJob.update({
where: {
id: jobId,
status: BackgroundJobStatus.PROCESSING,
},
data: {
status: BackgroundJobStatus.PENDING,
},
});
await this.submitJobToEndpoint({
jobId,
jobDefinitionId: backgroundJob.jobId,
data: options,
});
}
res.status(200).send('OK');
} else {
@ -104,9 +206,105 @@ export class LocalJobProvider extends BaseJobProvider {
};
}
private async submitJobToEndpoint(options: {
jobId: string;
jobDefinitionId: string;
data: TriggerJobOptions;
isRetry?: boolean;
}) {
const { jobId, jobDefinitionId, data, isRetry } = options;
const endpoint = `${NEXT_PUBLIC_WEBAPP_URL()}/api/jobs/${jobDefinitionId}/${jobId}`;
const signature = sign(data);
const headers: Record<string, string> = {
'Content-Type': 'application/json',
'X-Job-Id': jobId,
'X-Job-Signature': signature,
};
if (isRetry) {
headers['X-Job-Retry'] = '1';
}
console.log('Submitting job to endpoint:', endpoint);
await Promise.race([
fetch(endpoint, {
method: 'POST',
body: JSON.stringify(data),
headers,
}).catch(() => null),
new Promise((resolve) => {
setTimeout(resolve, 150);
}),
]);
}
private createJobRunIO(jobId: string): JobRunIO {
return {
stableRun: async (_cacheKey, callback) => await callback(),
runTask: async <T extends void | Json>(cacheKey: string, callback: () => Promise<T>) => {
const hashedKey = Buffer.from(sha256(cacheKey)).toString('hex');
let task = await prisma.backgroundJobTask.findFirst({
where: {
id: `task-${hashedKey}--${jobId}`,
jobId,
},
});
if (!task) {
task = await prisma.backgroundJobTask.create({
data: {
id: `task-${hashedKey}--${jobId}`,
name: cacheKey,
jobId,
status: BackgroundJobStatus.PENDING,
},
});
}
if (task.status === BackgroundJobStatus.COMPLETED) {
// eslint-disable-next-line @typescript-eslint/consistent-type-assertions
return task.result as T;
}
if (task.retried >= 3) {
throw new BackgroundTaskExceededRetriesError('Task exceeded retries');
}
try {
const result = await callback();
task = await prisma.backgroundJobTask.update({
where: {
id: task.id,
jobId,
},
data: {
status: BackgroundJobStatus.COMPLETED,
result: result === null ? Prisma.JsonNull : result,
completedAt: new Date(),
},
});
return result;
} catch {
task = await prisma.backgroundJobTask.update({
where: {
id: task.id,
jobId,
},
data: {
status: BackgroundJobStatus.PENDING,
retried: {
increment: 1,
},
},
});
throw new BackgroundTaskFailedError('Task failed');
}
},
triggerJob: async (_cacheKey, payload) => await this.triggerJob(payload),
logger: {
debug: (...args) => console.debug(`[${jobId}]`, ...args),
@ -122,3 +320,17 @@ export class LocalJobProvider extends BaseJobProvider {
};
}
}
class BackgroundTaskFailedError extends Error {
constructor(message: string) {
super(message);
this.name = 'BackgroundTaskFailedError';
}
}
class BackgroundTaskExceededRetriesError extends Error {
constructor(message: string) {
super(message);
this.name = 'BackgroundTaskExceededRetriesError';
}
}

View File

@ -64,7 +64,7 @@ export class TriggerJobProvider extends BaseJobProvider {
return {
wait: io.wait,
logger: io.logger,
stableRun: async (cacheKey, callback) => io.runTask(cacheKey, callback),
runTask: async (cacheKey, callback) => io.runTask(cacheKey, callback),
triggerJob: async (cacheKey, payload) =>
io.sendEvent(cacheKey, {
...payload,

View File

@ -1 +1,6 @@
export * from './send-confirmation-email';
import type { JobClient } from '../client/client';
import { registerSendConfirmationEmailJob } from './send-confirmation-email';
export const registerJobs = (client: JobClient) => {
registerSendConfirmationEmailJob(client);
};

View File

@ -1,23 +1,47 @@
import { z } from 'zod';
import { sendConfirmationToken } from '../../server-only/user/send-confirmation-token';
import { jobsClient } from '../client';
import type { JobClient } from '../client/client';
jobsClient.defineJob({
id: 'send.confirmation.email',
name: 'Send Confirmation Email',
version: '1.0.0',
trigger: {
name: 'send.confirmation.email',
schema: z.object({
email: z.string().email(),
force: z.boolean().optional(),
}),
},
handler: async ({ payload }) => {
await sendConfirmationToken({
email: payload.email,
force: payload.force,
});
},
});
export const registerSendConfirmationEmailJob = (client: JobClient) => {
client.defineJob({
id: 'send.confirmation.email',
name: 'Send Confirmation Email',
version: '1.0.0',
trigger: {
name: 'send.confirmation.email',
schema: z.object({
email: z.string().email(),
force: z.boolean().optional(),
}),
},
handler: async ({ payload, io }) => {
console.log('---- start job ----');
// eslint-disable-next-line @typescript-eslint/require-await
const result = await io.runTask('console-log-1', async () => {
console.log('Task 1');
return 5;
});
console.log({ result });
console.log('always runs');
// eslint-disable-next-line @typescript-eslint/require-await
await io.runTask('console-log-2', async () => {
await Promise.resolve(null);
throw new Error('dang2');
});
console.log('---- end job ----');
// throw new Error('dang')
await sendConfirmationToken({
email: payload.email,
force: payload.force,
});
},
});
};

View File

@ -45,7 +45,7 @@
"pdf-lib": "^1.17.1",
"pg": "^8.11.3",
"playwright": "1.43.0",
"react": "18.3.1",
"react": "18.2.0",
"remeda": "^1.27.1",
"stripe": "^12.7.0",
"ts-pattern": "^5.0.5",