@balena/pinejs 16.1.1-build-renovate-major--balenalint-02ba2563aa4ac82ad90accc30534f744c00a0c16-1 → 16.2.0-build-joshbwlng-tasks-cd1aff5f4bbfacaad189b013578d5cbee51c3932-1

package/src/tasks/index.ts

@@ -0,0 +1,158 @@
+import type { Schema } from 'ajv';
+import * as cronParser from 'cron-parser';
+import { tasks as tasksEnv } from '../config-loader/env';
+import type * as Db from '../database-layer/db';
+import { BadRequestError } from '../sbvr-api/errors';
+import { addPureHook } from '../sbvr-api/hooks';
+import { PinejsClient } from '../sbvr-api/sbvr-utils';
+import type { sbvrUtils } from '../server-glue/module';
+import { ajv, apiRoot, channel } from './common';
+import type { TaskHandler } from './types';
+import { Worker } from './worker';
+
+export * from './types';
+
+// eslint-disable-next-line @typescript-eslint/no-var-requires
+const modelText: string = require('./model.sbvr');
+
+export const config = {
+	models: [
+		{
+			modelName: apiRoot,
+			apiRoot,
+			modelText,
+			customServerCode: exports,
+		},
+	] as sbvrUtils.ExecutableModel[],
+};
+
+// Poll every second and also execute on new task insert
+// Only poll or execute on triggers if worker is not already at max concurrency
+async function createTrigger(tx: Db.Tx): Promise<void> {
+	await tx.executeSql(`
+        CREATE OR REPLACE FUNCTION notify_task_insert()
+        RETURNS TRIGGER AS $$
+        BEGIN
+            PERFORM pg_notify('${channel}', NEW.id::text);
+            RETURN NEW;
+        END;
+        $$ LANGUAGE plpgsql;
+    `);
+
+	// Only trigger if task is pending and not scheduled
+	await tx.executeSql(`
+        CREATE OR REPLACE TRIGGER task_insert_trigger
+        AFTER INSERT ON task
+        FOR EACH ROW WHEN (NEW.status = 'pending' AND NEW."is scheduled to execute on-time" IS NULL)
+        EXECUTE FUNCTION notify_task_insert();
+    `);
+}
+
+let worker: Worker | null = null;
+export async function setup(db: Db.Database, tx: Db.Tx): Promise<void> {
+	// Async task functionality is only supported on Postgres
+	if (db.engine !== 'postgres') {
+		return;
+	}
+
+	// Create trigger function if it doesn't exist
+	await createTrigger(tx);
+
+	const client = new PinejsClient({
+		apiPrefix: `/${apiRoot}/`,
+	});
+	worker = new Worker(client);
+
+	// Add resource hooks
+	addPureHook('POST', apiRoot, 'task', {
+		POSTPARSE: async ({ req, request }) => {
+			// Set the actor
+			request.values.is_created_by__actor =
+				req.user?.actor ?? req.apiKey?.actor;
+			if (request.values.is_created_by__actor == null) {
+				throw new BadRequestError(
+					'Creating tasks with missing actor on req is not allowed',
+				);
+			}
+
+			// Set defaults
+			request.values.status = 'pending';
+			request.values.attempt_count = 0;
+			request.values.priority ??= 1;
+			request.values.attempt_limit ??= 1;
+
+			// Set scheduled start time using cron expression if provided
+			if (
+				request.values.is_scheduled_with__cron_expression != null &&
+				request.values.is_scheduled_to_execute_on__time == null
+			) {
+				try {
+					request.values.is_scheduled_to_execute_on__time = cronParser
+						.parseExpression(request.values.is_scheduled_with__cron_expression)
+						.next()
+						.toDate()
+						.toISOString();
+				} catch (_) {
+					throw new BadRequestError(
+						`Invalid cron expression: ${request.values.is_scheduled_with__cron_expression}`,
+					);
+				}
+			}
+
+			// Assert that the provided start time is far enough in the future
+			if (request.values.is_scheduled_to_execute_on__time != null) {
+				const now = new Date(new Date().getTime() + tasksEnv.queueIntervalMS);
+				const startTime = new Date(
+					request.values.is_scheduled_to_execute_on__time,
+				);
+				if (startTime < now) {
+					throw new BadRequestError(
+						`Task scheduled start time must be greater than ${tasksEnv.queueIntervalMS} milliseconds in the future`,
+					);
+				}
+			}
+
+			// Assert that the requested handler exists
+			const handlerName = request.values.is_executed_by__handler;
+			if (handlerName == null) {
+				throw new BadRequestError(`Must specify a task handler to execute`);
+			}
+			const handler = worker?.handlers[handlerName];
+			if (handler == null) {
+				throw new BadRequestError(
+					`No task handler with name '${handlerName}' registered`,
+				);
+			}
+
+			// Assert that the provided parameter set is valid
+			if (handler.validate != null) {
+				if (!handler.validate(request.values.is_executed_with__parameter_set)) {
+					throw new BadRequestError(
+						`Invalid parameter set: ${ajv.errorsText(handler.validate.errors)}`,
+					);
+				}
+			}
+		},
+	});
+	worker.start();
+}
+
+// Register a task handler
+export function addTaskHandler(
+	name: string,
+	fn: TaskHandler['fn'],
+	schema?: Schema,
+): void {
+	if (worker == null) {
+		return;
+	}
+
+	if (worker.handlers[name] != null) {
+		throw new Error(`Task handler with name '${name}' already registered`);
+	}
+	worker.handlers[name] = {
+		name,
+		fn,
+		validate: schema != null ? ajv.compile(schema) : undefined,
+	};
+}

package/src/tasks/model.sbvr

@@ -0,0 +1,60 @@
+Vocabulary: tasks
+
+Term: id
+    Concept Type: Big Serial (Type)
+Term: actor
+	Concept Type: Integer (Type)
+Term: attempt count
+	Concept Type: Integer (Type)
+Term: attempt limit
+	Concept Type: Integer (Type)
+Term: cron expression
+	Concept Type: Short Text (Type)
+Term: error message
+	Concept Type: Short Text (Type)
+Term: handler
+	Concept Type: Short Text (Type)
+Term: key
+	Concept Type: Short Text (Type)
+Term: parameter set
+	Concept Type: JSON (Type)
+Term: priority
+	Concept Type: Integer (Type)
+Term: status
+	Concept Type: Short Text (Type)
+Term: time
+	Concept Type: Date Time (Type)
+
+Term: task
+Fact type: task has id
+	Necessity: each task has exactly one id
+Fact type: task has key
+	Necessity: each task has at most one key
+Fact type: task is created by actor
+	Necessity: each task is created by exactly one actor
+Fact type: task is executed by handler
+	Necessity: each task is executed by exactly one handler
+Fact type: task is executed with parameter set
+	Necessity: each task is executed with at most one parameter set
+Fact type: task has priority
+	Necessity: each task has exactly one priority
+	Necessity: each task has a priority that is greater than or equal to 0
+Fact type: task is scheduled with cron expression
+	Necessity: each task is scheduled with at most one cron expression
+Fact type: task is scheduled to execute on time
+	Necessity: each task is scheduled to execute on at most one time
+Fact type: task has status
+	Necessity: each task has exactly one status
+	Definition: "pending" or "cancelled" or "success" or "failed"
+Fact type: task started on time
+	Necessity: each task started on at most one time
+Fact type: task ended on time
+	Necessity: each task ended on at most one time
+Fact type: task has error message
+	Necessity: each task has at most one error message
+Fact type: task has attempt count
+	Necessity: each task has exactly one attempt count
+Fact type: task has attempt limit
+	Necessity: each task has exactly one attempt limit
+	Necessity: each task has an attempt limit that is greater than or equal to 1
+

package/src/tasks/types.ts

@@ -0,0 +1,58 @@
+import type { ValidateFunction } from 'ajv';
+import type { AnyObject } from 'pinejs-client-core';
+import type * as Db from '../database-layer/db';
+import type { PinejsClient } from '../sbvr-api/sbvr-utils';
+
+export const taskStatuses = [
+	'pending',
+	'cancelled',
+	'success',
+	'failed',
+] as const;
+export type TaskStatus = (typeof taskStatuses)[number];
+export interface Task {
+	id: number;
+	created_at: Date;
+	modified_at: Date;
+	is_created_by__actor: number;
+	is_executed_by__handler: string;
+	is_executed_with__parameter_set: object | null;
+	is_scheduled_with__cron_expression: string | null;
+	is_scheduled_to_execute_on__time: Date | null;
+	priority: number;
+	status: TaskStatus;
+	started_on__time: Date | null;
+	ended_on__time: Date | null;
+	error_message: string | null;
+	attempt_count: number;
+	attempt_limit: number;
+}
+
+export type PartialTask = Pick<
+	Task,
+	| 'id'
+	| 'is_created_by__actor'
+	| 'is_executed_by__handler'
+	| 'is_executed_with__parameter_set'
+	| 'is_scheduled_with__cron_expression'
+	| 'priority'
+	| 'attempt_count'
+	| 'attempt_limit'
+>;
+
+export interface TaskArgs {
+	api: PinejsClient;
+	params: AnyObject;
+	tx: Db.Tx;
+}
+
+export type TaskResponse = Promise<{
+	status: TaskStatus;
+	error?: string;
+}>;
+
+export interface TaskHandler {
+	name: string;
+	fn: (options: TaskArgs) => TaskResponse;
+	validate?: ValidateFunction;
+}

package/src/tasks/worker.ts

@@ -0,0 +1,241 @@
+import type { AnyObject } from 'pinejs-client-core';
+import { tasks as tasksEnv } from '../config-loader/env';
+import type * as Db from '../database-layer/db';
+import * as permissions from '../sbvr-api/permissions';
+import { PinejsClient } from '../sbvr-api/sbvr-utils';
+import { sbvrUtils } from '../server-glue/module';
+import { ajv, channel } from './common';
+import type { PartialTask, TaskHandler, TaskStatus } from './types';
+
+// Map of column names with SBVR names used in SELECT queries
+const selectColumns = Object.entries({
+	id: 'id',
+	'is executed by-handler': 'is_executed_by__handler',
+	'is executed with-parameter set': 'is_executed_with__parameter_set',
+	'is scheduled with-cron expression': 'is_scheduled_with__cron_expression',
+	'attempt count': 'attempt_count',
+	'attempt limit': 'attempt_limit',
+	priority: 'priority',
+	'is created by-actor': 'is_created_by__actor',
+})
+	.map(([key, value]) => `t."${key}" AS "${value}"`)
+	.join(', ');
+
+// The worker is responsible for executing tasks in the queue. It listens for
+// notifications and polls the database for tasks to execute. It will execute
+// tasks in parallel up to a certain concurrency limit.
+export class Worker {
+	public handlers: Record<string, TaskHandler> = {};
+	private readonly concurrency: number;
+	private readonly interval: number;
+	private client: PinejsClient;
+	private executing = 0;
+
+	constructor(client: PinejsClient) {
+		this.client = client;
+		this.concurrency = tasksEnv.queueConcurrency;
+		this.interval = tasksEnv.queueIntervalMS;
+	}
+
+	private canExecute(): boolean {
+		return (
+			this.executing < this.concurrency && Object.keys(this.handlers).length > 0
+		);
+	}
+
+	private poll(): void {
+		let executed = false;
+		const names = Object.keys(this.handlers);
+		const binds = names.map((_, index) => `$${index + 1}`).join(', ');
+		sbvrUtils.db
+			.transaction(async (tx) => {
+				if (!this.canExecute()) {
+					return;
+				}
+
+				const result = await tx.executeSql(
+					`SELECT ${selectColumns}
+				FROM task AS t
+				WHERE
+					t."is executed by-handler" IN (${binds}) AND
+					t."status" = 'pending' AND
+					t."attempt count" <= t."attempt limit" AND
+					(
+						t."is scheduled to execute on-time" IS NULL OR
+						t."is scheduled to execute on-time" <= CURRENT_TIMESTAMP + INTERVAL '${Math.ceil(this.interval / 1000)} second'
+					)
+				ORDER BY
+					t."is scheduled to execute on-time" ASC,
+					t."priority" DESC,
+					t."id" ASC
+				LIMIT ${Math.max(this.concurrency - this.executing, 0)}
+				FOR UPDATE SKIP LOCKED`,
+					names,
+				);
+				if (result.rows.length === 0) {
+					return;
+				}
+
+				// Tasks found, execute them in parallel
+				await Promise.all(
+					result.rows.map(async (row) => {
+						await this.execute(tx, row as PartialTask);
+					}),
+				);
+				executed = true;
+			})
+			.catch((err) => {
+				console.error('Failed polling for tasks:', err);
+			})
+			.finally(() => {
+				setTimeout(() => this.poll(), executed ? 0 : this.interval);
+			});
+	}
+
+	// Start listening and polling for tasks
+	public start() {
+		sbvrUtils.db.on?.(
+			'notification',
+			async (msg) => {
+				if (this.canExecute()) {
+					await sbvrUtils.db.transaction(async (tx) => {
+						const result = await tx.executeSql(
+							`SELECT ${selectColumns} FROM task AS t WHERE id = $1 FOR UPDATE SKIP LOCKED`,
+							[msg.payload],
+						);
+						if (result.rows.length > 0) {
+							await this.execute(tx, result.rows[0] as PartialTask);
+						}
+					});
+				}
+			},
+			{
+				channel,
+			},
+		);
+		this.poll();
+	}
+
+	private async execute(tx: Db.Tx, task: PartialTask): Promise<void> {
+		this.executing++;
+		try {
+			// Get the handler
+			const handler = this.handlers[task.is_executed_by__handler];
+			const startedOnTime = new Date();
+			if (handler == null) {
+				await this.update(
+					tx,
+					task,
+					startedOnTime,
+					'failed',
+					'Matching task handler not found',
+				);
+				return;
+			}
+
+			// Validate parameters before execution so we can fail early if
+			// the parameter set is invalid. This can happen if the handler
+			// definition changes after a task is added to the queue.
+			if (
+				handler.validate != null &&
+				!handler.validate(task.is_executed_with__parameter_set)
+			) {
+				await this.update(
+					tx,
+					task,
+					startedOnTime,
+					'failed',
+					`Invalid parameter set: ${ajv.errorsText(handler.validate.errors)}`,
+				);
+				return;
+			}
+
+			// Execute the handler
+			const result = await handler.fn({
+				api: new PinejsClient({
+					passthrough: {
+						tx,
+					},
+				}),
+				params: task.is_executed_with__parameter_set ?? {},
+				tx,
+			});
+
+			// Update the task with the results
+			await this.update(tx, task, startedOnTime, result.status, result.error);
+		} catch (err) {
+			// This shouldn't normally happen, but if it does, we want to log it and kill the process
+			console.error('Task execution failed:', err);
+			process.exit(1);
+		} finally {
+			this.executing--;
+		}
+	}
+
+	private getNextAttemptTime(attempt: number): Date | null {
+		const delay = Math.ceil(Math.exp(Math.min(10, attempt))) * 1000;
+		return new Date(Date.now() + delay);
+	}
+
+	private async update(
+		tx: Db.Tx,
+		task: PartialTask,
+		startedOnTime: Date,
+		status: TaskStatus,
+		errorMessage?: string,
+	): Promise<void> {
+		const attemptCount = task.attempt_count + 1;
+		const body: AnyObject = {
+			started_on__time: startedOnTime,
+			ended_on__time: new Date(),
+			status,
+			attempt_count: attemptCount,
+			...(errorMessage != null && { error_message: errorMessage }),
+		};
+
+		// Re-enqueue if the task failed but has retries left, remember that
+		// attemptCount includes the initial attempt while attempt_limit does not
+		if (status === 'failed' && attemptCount < task.attempt_limit) {
+			body.status = 'pending';
+
+			// Schedule next attempt using exponential backoff
+			body.is_scheduled_to_execute_on__time =
+				this.getNextAttemptTime(attemptCount);
+		}
+
+		// Patch current task
+		await this.client.patch({
+			resource: 'task',
+			passthrough: {
+				tx,
+				req: permissions.root,
+			},
+			id: task.id,
+			body,
+		});
+
+		// Create new task with same configuration if previous
+		// iteration completed and has a cron expression
+		if (
+			['failed', 'success'].includes(body.status) &&
+			task.is_scheduled_with__cron_expression != null
+		) {
+			await this.client.post({
+				resource: 'task',
+				passthrough: {
+					tx,
+					req: permissions.root,
+				},
+				body: {
+					attempt_limit: task.attempt_limit,
+					is_created_by__actor: task.is_created_by__actor,
+					is_executed_by__handler: task.is_executed_by__handler,
+					is_executed_with__parameter_set: task.is_executed_with__parameter_set,
+					is_scheduled_with__cron_expression:
+						task.is_scheduled_with__cron_expression,
+					priority: task.priority,
+				},
+			});
+		}
+	}
+}