@balena/pinejs 16.2.0-build-joshbwlng-tasks-a83c83b4c78803915d0cb6297cc8cc1862622d08-1 → 17.0.0-build-wip-large-file-uploads-0c8ef752deac19fa0d6a7dfa9f7173813cab7867-1

package/src/tasks/worker.ts

@@ -1,246 +0,0 @@
-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;
-	}
-
-	// Check if instance can execute more tasks
-	private canExecute(): boolean {
-		return (
-			this.executing < this.concurrency && Object.keys(this.handlers).length > 0
-		);
-	}
-
-	// Execute a task
-	private async execute(tx: Db.Tx, task: PartialTask): Promise<void> {
-		this.executing++;
-		try {
-			// Get specified handler
-			const handler = this.handlers[task.is_executed_by__handler];
-			const startedOnTime = new Date();
-			if (handler == null) {
-				await this.finalize(
-					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.finalize(
-					tx,
-					task,
-					startedOnTime,
-					'failed',
-					`Invalid parameter set: ${ajv.errorsText(handler.validate.errors)}`,
-				);
-				return;
-			}
-
-			// Execute handler
-			const result = await handler.fn({
-				api: new PinejsClient({
-					passthrough: {
-						tx,
-					},
-				}),
-				params: task.is_executed_with__parameter_set ?? {},
-				tx,
-			});
-
-			// Update task with results
-			await this.finalize(tx, task, startedOnTime, result.status, result.error);
-		} catch (err) {
-			// This shouldn't happen, but if it does we want to log and kill the process
-			console.error('Task execution failed:', err);
-			process.exit(1);
-		} finally {
-			this.executing--;
-		}
-	}
-
-	// Update task and schedule next attempt if needed
-	private async finalize(
-		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,
-				},
-			});
-		}
-	}
-
-	// Calculate next attempt time using exponential backoff
-	private getNextAttemptTime(attempt: number): Date | null {
-		const delay = Math.ceil(Math.exp(Math.min(10, attempt))) * 1000;
-		return new Date(Date.now() + delay);
-	}
-
-	// Poll for tasks to execute
-	private poll(): void {
-		let executed = false;
-		const handlerNames = Object.keys(this.handlers);
-		const binds = handlerNames.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`,
-					handlerNames,
-				);
-				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(): void {
-		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();
-	}
-}