Add MariaDB queue table. Fix Bugs

2026-10-18 15:58:02 +01:00
parent d872c3661a
commit f9ba5f8657
28 changed files with 785 additions and 49 deletions
@@ -0,0 +1,17 @@
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+type Params<Table extends string = string> = {
+    table: Table;
+    id: number | string;
+    /** `queue_lock_token` from the claimed row */
+    lock_token: string;
+    config?: BunMariaDBConfig;
+};
+/**
+ * Mark a claimed job as done: deletes it, or sets it to `done` when the
+ * table has `keep_completed`.
+ *
+ * Fails when the lock token no longer matches, i.e. the lease expired and
+ * another worker reclaimed the job.
+ */
+export default function QueueAck<Table extends string = string>({ table, id, lock_token, config, }: Params<Table>): Promise<DBResponseObject>;
+export {};
@@ -0,0 +1,61 @@
+import dbHandler from "../db-handler";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+/**
+ * Mark a claimed job as done: deletes it, or sets it to `done` when the
+ * table has `keep_completed`.
+ *
+ * Fails when the lock token no longer matches, i.e. the lease expired and
+ * another worker reclaimed the job.
+ */
+export default async function QueueAck({ table, id, lock_token, config, }) {
+    try {
+        const options = requireQueueTable(table);
+        const tableSql = quoteIdentifier(table);
+        const claimedWhere = `WHERE id = ? AND queue_status = 'processing' AND queue_lock_token = ?`;
+        if (!options.keep_completed) {
+            const res = await dbHandler({
+                query: `DELETE FROM ${tableSql} ${claimedWhere} RETURNING *`,
+                values: [id, lock_token],
+                config,
+            });
+            if (!res.success) {
+                return { ...res, success: false };
+            }
+            if (!res.payload?.length) {
+                return {
+                    ...res,
+                    success: false,
+                    msg: "Job is not claimed with this lock token",
+                };
+            }
+            return res;
+        }
+        const res = await dbHandler({
+            query: `UPDATE ${tableSql} SET queue_status = 'done', queue_lease_expires_at = NULL, queue_lock_token = NULL, updated_at = ${QUEUE_NOW_SQL} ${claimedWhere}`,
+            values: [id, lock_token],
+            config,
+        });
+        if (!res.success) {
+            return { ...res, success: false };
+        }
+        if (!res.db_res?.affectedRows) {
+            return {
+                ...res,
+                success: false,
+                msg: "Job is not claimed with this lock token",
+            };
+        }
+        return await dbHandler({
+            query: `SELECT * FROM ${tableSql} WHERE id = ?`,
+            values: [id],
+            config,
+        });
+    }
+    catch (error) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}
@@ -0,0 +1,12 @@
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+type Params<Table extends string = string> = {
+    table: Table;
+    id: number | string;
+    config?: BunMariaDBConfig;
+};
+/**
+ * Take a `pending` job out of the queue without deleting it. Jobs a worker
+ * has already claimed can't be cancelled.
+ */
+export default function QueueCancel<Table extends string = string>({ table, id, config, }: Params<Table>): Promise<DBResponseObject>;
+export {};
@@ -0,0 +1,39 @@
+import dbHandler from "../db-handler";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+/**
+ * Take a `pending` job out of the queue without deleting it. Jobs a worker
+ * has already claimed can't be cancelled.
+ */
+export default async function QueueCancel({ table, id, config, }) {
+    try {
+        requireQueueTable(table);
+        const tableSql = quoteIdentifier(table);
+        const res = await dbHandler({
+            query: `UPDATE ${tableSql} SET queue_status = 'cancelled', updated_at = ${QUEUE_NOW_SQL} WHERE id = ? AND queue_status = 'pending'`,
+            values: [id],
+            config,
+        });
+        if (!res.success) {
+            return { ...res, success: false };
+        }
+        if (!res.db_res?.affectedRows) {
+            return {
+                ...res,
+                success: false,
+                msg: "Job is not pending",
+            };
+        }
+        return await dbHandler({
+            query: `SELECT * FROM ${tableSql} WHERE id = ?`,
+            values: [id],
+            config,
+        });
+    }
+    catch (error) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}
@@ -0,0 +1,25 @@
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+type Params<Table extends string = string> = {
+    table: Table;
+    /** Max jobs to claim. Defaults to 1 */
+    limit?: number;
+    /** How long the claim lasts before another worker may take the job. Defaults to the table's `lease_ms` */
+    lease_ms?: number;
+    /** Jobs whose lease expired after this many attempts are marked `failed`. Defaults to the table's `max_attempts` */
+    max_attempts?: number;
+    config?: BunMariaDBConfig;
+};
+/**
+ * Claim the next available jobs, oldest `queue_available_at` first. Each
+ * returned row carries a `queue_lock_token` that must be passed to `ack`,
+ * `fail` or `extend`.
+ *
+ * `FOR UPDATE SKIP LOCKED` lets concurrent workers claim different jobs
+ * without blocking on each other or claiming the same job twice.
+ */
+export default function QueueClaim<Schema extends {
+    [k: string]: any;
+} = {
+    [k: string]: any;
+}, Table extends string = string>({ table, limit, lease_ms, max_attempts, config, }: Params<Table>): Promise<DBResponseObject<Schema>>;
+export {};
@@ -0,0 +1,79 @@
+import dbHandler, { holdEventLoop, resolveClient } from "../db-handler";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+/**
+ * Claim the next available jobs, oldest `queue_available_at` first. Each
+ * returned row carries a `queue_lock_token` that must be passed to `ack`,
+ * `fail` or `extend`.
+ *
+ * `FOR UPDATE SKIP LOCKED` lets concurrent workers claim different jobs
+ * without blocking on each other or claiming the same job twice.
+ */
+export default async function QueueClaim({ table, limit = 1, lease_ms, max_attempts, config, }) {
+    try {
+        const options = requireQueueTable(table);
+        const tableSql = quoteIdentifier(table);
+        const finalLimit = Math.max(1, Math.floor(Number(limit) || 1));
+        const leaseMs = Math.max(0, Math.floor(Number(lease_ms ?? options.lease_ms) || 0));
+        const maxAttempts = Math.max(1, Math.floor(Number(max_attempts ?? options.max_attempts) || 1));
+        /**
+         * Release jobs whose worker died or stalled past its lease: back to
+         * `pending`, or `failed` if that was the last attempt.
+         */
+        const expired_res = await dbHandler({
+            query: `UPDATE ${tableSql} SET queue_status = IF(queue_attempts >= ?, 'failed', 'pending'), queue_lease_expires_at = NULL, queue_lock_token = NULL, queue_last_error = 'Lease expired', updated_at = ${QUEUE_NOW_SQL} WHERE queue_status = 'processing' AND queue_lease_expires_at <= ${QUEUE_NOW_SQL}`,
+            values: [maxAttempts],
+            config,
+        });
+        if (!expired_res.success) {
+            return { ...expired_res, success: false };
+        }
+        const CLIENT = resolveClient(config);
+        if (!CLIENT) {
+            throw new Error(`Couldn't grab MariaDB Client.`);
+        }
+        const lock_token = crypto.randomUUID();
+        return await holdEventLoop(() => CLIENT.begin(async (tx) => {
+            /**
+             * The ORDER BY must match `idx_queue_claim`. A filesort would
+             * lock every pending row before LIMIT applies, and SKIP LOCKED
+             * would then hide the whole queue from other workers.
+             */
+            const matched_res = await dbHandler({
+                query: `SELECT id FROM ${tableSql} FORCE INDEX (idx_queue_claim) WHERE queue_status = 'pending' AND queue_available_at <= ${QUEUE_NOW_SQL} ORDER BY queue_status, queue_available_at, id LIMIT ? FOR UPDATE SKIP LOCKED`,
+                values: [finalLimit],
+                client: tx,
+            });
+            if (!matched_res.success) {
+                return {
+                    ...matched_res,
+                    success: false,
+                };
+            }
+            const ids = (matched_res.payload || []).map((row) => row.id);
+            if (ids.length === 0) {
+                return { success: true, payload: [] };
+            }
+            const idPlaceholders = ids.map(() => "?").join(",");
+            const update_res = await dbHandler({
+                query: `UPDATE ${tableSql} SET queue_status = 'processing', queue_attempts = queue_attempts + 1, queue_lease_expires_at = ${QUEUE_NOW_SQL} + ?, queue_lock_token = ?, updated_at = ${QUEUE_NOW_SQL} WHERE id IN (${idPlaceholders})`,
+                values: [leaseMs, lock_token, ...ids],
+                client: tx,
+            });
+            if (!update_res.success) {
+                throw new Error(update_res.error || "Queue claim failed");
+            }
+            return await dbHandler({
+                query: `SELECT * FROM ${tableSql} WHERE id IN (${idPlaceholders}) ORDER BY queue_available_at, id`,
+                values: ids,
+                client: tx,
+            });
+        }));
+    }
+    catch (error) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}
@@ -0,0 +1,16 @@
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+type Params<Table extends string = string> = {
+    table: Table;
+    id: number | string;
+    /** `queue_lock_token` from the claimed row */
+    lock_token: string;
+    /** New lease length, counted from now. Defaults to the table's `lease_ms` */
+    lease_ms?: number;
+    config?: BunMariaDBConfig;
+};
+/**
+ * Extend the lease on a claimed job, for jobs that run longer than the
+ * lease given to `claim`.
+ */
+export default function QueueExtend<Table extends string = string>({ table, id, lock_token, lease_ms, config, }: Params<Table>): Promise<DBResponseObject>;
+export {};
@@ -0,0 +1,39 @@
+import dbHandler from "../db-handler";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+/**
+ * Extend the lease on a claimed job, for jobs that run longer than the
+ * lease given to `claim`.
+ */
+export default async function QueueExtend({ table, id, lock_token, lease_ms, config, }) {
+    try {
+        const options = requireQueueTable(table);
+        const leaseMs = Math.max(0, Math.floor(Number(lease_ms ?? options.lease_ms) || 0));
+        const res = await dbHandler({
+            query: `UPDATE ${quoteIdentifier(table)} SET queue_lease_expires_at = ${QUEUE_NOW_SQL} + ?, updated_at = ${QUEUE_NOW_SQL} WHERE id = ? AND queue_status = 'processing' AND queue_lock_token = ?`,
+            values: [leaseMs, id, lock_token],
+            config,
+        });
+        if (!res.success) {
+            return { ...res, success: false };
+        }
+        if (!res.db_res?.affectedRows) {
+            return {
+                ...res,
+                success: false,
+                msg: "Job is not claimed with this lock token",
+            };
+        }
+        return await dbHandler({
+            query: `SELECT * FROM ${quoteIdentifier(table)} WHERE id = ?`,
+            values: [id],
+            config,
+        });
+    }
+    catch (error) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}
@@ -0,0 +1,19 @@
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+type Params<Table extends string = string> = {
+    table: Table;
+    id: number | string;
+    /** `queue_lock_token` from the claimed row */
+    lock_token: string;
+    error?: string;
+    /** Wait this long before the job can be retried. Defaults to 0 */
+    retry_delay_ms?: number;
+    /** Mark the job `failed` instead of retrying once it reaches this many attempts. Defaults to the table's `max_attempts` */
+    max_attempts?: number;
+    config?: BunMariaDBConfig;
+};
+/**
+ * Release a claimed job after a failed attempt. It goes back to `pending`
+ * (after `retry_delay_ms`) or to `failed` once `max_attempts` is reached.
+ */
+export default function QueueFail<Table extends string = string>({ table, id, lock_token, error, retry_delay_ms, max_attempts, config, }: Params<Table>): Promise<DBResponseObject>;
+export {};
@@ -0,0 +1,41 @@
+import dbHandler from "../db-handler";
+import { quoteIdentifier } from "../../utils/sql-identifier";
+import { QUEUE_NOW_SQL, requireQueueTable } from "./queue-utils";
+/**
+ * Release a claimed job after a failed attempt. It goes back to `pending`
+ * (after `retry_delay_ms`) or to `failed` once `max_attempts` is reached.
+ */
+export default async function QueueFail({ table, id, lock_token, error, retry_delay_ms = 0, max_attempts, config, }) {
+    try {
+        const options = requireQueueTable(table);
+        const tableSql = quoteIdentifier(table);
+        const retryDelayMs = Math.max(0, Math.floor(Number(retry_delay_ms) || 0));
+        const maxAttempts = Math.max(1, Math.floor(Number(max_attempts ?? options.max_attempts) || 1));
+        const res = await dbHandler({
+            query: `UPDATE ${tableSql} SET queue_status = IF(queue_attempts >= ?, 'failed', 'pending'), queue_available_at = ${QUEUE_NOW_SQL} + ?, queue_lease_expires_at = NULL, queue_lock_token = NULL, queue_last_error = ?, updated_at = ${QUEUE_NOW_SQL} WHERE id = ? AND queue_status = 'processing' AND queue_lock_token = ?`,
+            values: [maxAttempts, retryDelayMs, error ?? null, id, lock_token],
+            config,
+        });
+        if (!res.success) {
+            return { ...res, success: false };
+        }
+        if (!res.db_res?.affectedRows) {
+            return {
+                ...res,
+                success: false,
+                msg: "Job is not claimed with this lock token",
+            };
+        }
+        return await dbHandler({
+            query: `SELECT * FROM ${quoteIdentifier(table)} WHERE id = ?`,
+            values: [id],
+            config,
+        });
+    }
+    catch (error) {
+        return {
+            success: false,
+            error: error.message,
+        };
+    }
+}
@@ -0,0 +1,21 @@
+import type { BunMariaDBConfig, DBResponseObject } from "../../types";
+type Params<Schema extends {
+    [k: string]: any;
+} = {
+    [k: string]: any;
+}, Table extends string = string> = {
+    table: Table;
+    data: Schema[];
+    /** Wait this long before the jobs can be claimed */
+    delay_ms?: number;
+    config?: BunMariaDBConfig;
+};
+/**
+ * Add jobs to a queue table.
+ */
+export default function QueuePush<Schema extends {
+    [k: string]: any;
+} = {
+    [k: string]: any;
+}, Table extends string = string>({ table, data, delay_ms, config, }: Params<Schema, Table>): Promise<DBResponseObject>;
+export {};
@@ -0,0 +1,11 @@
+import DbInsert from "../mariadb/db-insert";
+/**
+ * Add jobs to a queue table.
+ */
+export default async function QueuePush({ table, data, delay_ms, config, }) {
+    const jobs = data.map((d) => ({
+        ...(delay_ms ? { queue_available_at: Date.now() + delay_ms } : {}),
+        ...d,
+    }));
+    return DbInsert({ table, data: jobs, config });
+}
@@ -0,0 +1,34 @@
+import { type BUN_MARIADB_QueueTableOptions } from "../../types";
+/**
+ * Current time in ms from the database clock. Leases are set and checked
+ * with this so workers on hosts with drifting clocks agree on expiry.
+ */
+export declare const QUEUE_NOW_SQL = "CAST(UNIX_TIMESTAMP(NOW(3)) * 1000 AS SIGNED)";
+/**
+ * Extra WHERE condition for generic update/delete on queue tables: rows a
+ * worker has claimed are only changed through the queue API.
+ */
+export declare const QUEUE_NOT_CLAIMED_SQL = "`queue_status` <> 'processing'";
+export type ResolvedQueueTableOptions = Required<BUN_MARIADB_QueueTableOptions>;
+/**
+ * A table's `queue_table` options with defaults filled in (looked up
+ * directly or through its parent table), or `undefined` if it is not a
+ * queue table.
+ */
+export declare function grabQueueTableOptions(tableName: string): ResolvedQueueTableOptions | undefined;
+/**
+ * Whether a table is declared with `queue_table` (directly or through its
+ * parent table).
+ */
+export declare function isQueueTable(tableName: string): boolean;
+/**
+ * Queue options for a table the queue API was called on. Throws if the table
+ * is not a queue table.
+ */
+export declare function requireQueueTable(tableName: string): ResolvedQueueTableOptions;
+/**
+ * Throw if `data` sets queue state the queue API owns.
+ */
+export declare function assertNoProtectedQueueFields(table: string, data: {
+    [k: string]: any;
+}[]): void;
@@ -0,0 +1,69 @@
+import _ from "lodash";
+import { readLiveSchema } from "../../functions/live-schema";
+import { QueueProtectedFields, } from "../../types";
+/**
+ * Current time in ms from the database clock. Leases are set and checked
+ * with this so workers on hosts with drifting clocks agree on expiry.
+ */
+export const QUEUE_NOW_SQL = "CAST(UNIX_TIMESTAMP(NOW(3)) * 1000 AS SIGNED)";
+/**
+ * Extra WHERE condition for generic update/delete on queue tables: rows a
+ * worker has claimed are only changed through the queue API.
+ */
+export const QUEUE_NOT_CLAIMED_SQL = "`queue_status` <> 'processing'";
+const QUEUE_TABLE_DEFAULTS = {
+    max_attempts: 5,
+    lease_ms: 30_000,
+    keep_completed: false,
+};
+/**
+ * A table's `queue_table` options with defaults filled in (looked up
+ * directly or through its parent table), or `undefined` if it is not a
+ * queue table.
+ */
+export function grabQueueTableOptions(tableName) {
+    const dbSchema = global.DB_SCHEMA || readLiveSchema();
+    const table = dbSchema?.tables?.find((t) => t.tableName === tableName);
+    if (!table)
+        return undefined;
+    if (table.queue_table === undefined) {
+        return table.parentTableName
+            ? grabQueueTableOptions(table.parentTableName)
+            : undefined;
+    }
+    if (!table.queue_table)
+        return undefined;
+    return {
+        ...QUEUE_TABLE_DEFAULTS,
+        ...(typeof table.queue_table === "object"
+            ? _.omitBy(table.queue_table, _.isUndefined)
+            : {}),
+    };
+}
+/**
+ * Whether a table is declared with `queue_table` (directly or through its
+ * parent table).
+ */
+export function isQueueTable(tableName) {
+    return Boolean(grabQueueTableOptions(tableName));
+}
+/**
+ * Queue options for a table the queue API was called on. Throws if the table
+ * is not a queue table.
+ */
+export function requireQueueTable(tableName) {
+    const options = grabQueueTableOptions(tableName);
+    if (!options) {
+        throw new Error(`\`${tableName}\` is not a queue table`);
+    }
+    return options;
+}
+/**
+ * Throw if `data` sets queue state the queue API owns.
+ */
+export function assertNoProtectedQueueFields(table, data) {
+    const protectedFields = QueueProtectedFields.filter((field) => data.some((d) => d && d[field] !== undefined));
+    if (protectedFields.length > 0) {
+        throw new Error(`\`${table}\` is a queue table: ${protectedFields.join(", ")} can only be changed through BunMariaDB.queue`);
+    }
+}