diff --git a/schemas/system_db_schema.ts b/schemas/system_db_schema.ts index cff58950f..be7675606 100644 --- a/schemas/system_db_schema.ts +++ b/schemas/system_db_schema.ts @@ -127,6 +127,7 @@ export interface queues { priority_enabled: boolean; partition_queue: boolean; polling_interval_sec: number; + max_dequeues_per_poll: number | null; created_at: number; updated_at: number; } diff --git a/src/adminserver.ts b/src/adminserver.ts index ec3c67b5c..b5f0089d5 100644 --- a/src/adminserver.ts +++ b/src/adminserver.ts @@ -256,6 +256,7 @@ export class DBOSAdminServer { concurrency: q.concurrency, workerConcurrency: q.workerConcurrency, rateLimit: q.rateLimit, + maxDequeuesPerPoll: q.maxDequeuesPerPoll, }); } } catch (e) { @@ -267,6 +268,7 @@ export class DBOSAdminServer { concurrency: q.concurrency, workerConcurrency: q.workerConcurrency, rateLimit: q.rateLimit, + maxDequeuesPerPoll: q.maxDequeuesPerPoll, }); }); sendJson(res, 200, [...merged.values()]); diff --git a/src/conductor/conductor.ts b/src/conductor/conductor.ts index 98c7ed378..380685f85 100644 --- a/src/conductor/conductor.ts +++ b/src/conductor/conductor.ts @@ -845,6 +845,7 @@ export class Conductor { priority_enabled: q.priorityEnabled, partition_queue: q.partitionQueue, polling_interval_sec: q.pollingIntervalSec, + max_dequeues_per_poll: q.maxDequeuesPerPoll, })); } catch (e) { errorMsg = `Exception encountered when listing queues: ${(e as Error).message}`; @@ -868,6 +869,7 @@ export class Conductor { priority_enabled: queue.priorityEnabled, partition_queue: queue.partitionQueue, polling_interval_sec: queue.pollingIntervalSec, + max_dequeues_per_poll: queue.maxDequeuesPerPoll, }; } } catch (e) { diff --git a/src/conductor/protocol.ts b/src/conductor/protocol.ts index 8f1092de5..e98c57f3e 100644 --- a/src/conductor/protocol.ts +++ b/src/conductor/protocol.ts @@ -949,6 +949,7 @@ export interface QueueOutput { priority_enabled: boolean; partition_queue: boolean; polling_interval_sec: number; + max_dequeues_per_poll: number | null; } export class ListQueuesRequest implements BaseMessage { diff --git a/src/debugpoint.ts b/src/debugpoint.ts index e1a4ba554..6dda04e74 100644 --- a/src/debugpoint.ts +++ b/src/debugpoint.ts @@ -72,8 +72,8 @@ export const DEBUG_TRIGGER_STEP_COMMIT = 'DEBUG_TRIGGER_STEP_COMMIT'; export const DEBUG_TRIGGER_INITWF_COMMIT = 'DEBUG_TRIGGER_INITWF_COMMIT'; // Fires inside runQueue between dispatching consecutive partition keys. export const DEBUG_TRIGGER_BETWEEN_PARTITION_DISPATCHES = 'DEBUG_TRIGGER_BETWEEN_PARTITION_DISPATCHES'; -// Fires inside findAndMarkStartableWorkflows after the SELECT FOR UPDATE NOWAIT -// but before COMMIT (i.e. while the row lock is held). Tests can throw a +// Fires inside findAndMarkStartableWorkflows after claiming rows but before +// COMMIT (i.e. while the transaction still holds row locks). Tests can throw a // synthetic 55P03 here to simulate a concurrent executor winning the lock race, // which is the exact condition that triggers the orphan-PENDING bug. export const DEBUG_TRIGGER_FIND_AND_MARK_AFTER_SELECT = 'DEBUG_TRIGGER_FIND_AND_MARK_AFTER_SELECT'; diff --git a/src/sysdb_migrations/internal/migrations.ts b/src/sysdb_migrations/internal/migrations.ts index db37dddac..40185dec5 100644 --- a/src/sysdb_migrations/internal/migrations.ts +++ b/src/sysdb_migrations/internal/migrations.ts @@ -428,11 +428,15 @@ export function allMigrations( "priority_enabled" BOOLEAN NOT NULL DEFAULT FALSE, "partition_queue" BOOLEAN NOT NULL DEFAULT FALSE, "polling_interval_sec" DOUBLE PRECISION NOT NULL DEFAULT 1.0, + "max_dequeues_per_poll" INT4, "created_at" BIGINT NOT NULL DEFAULT (EXTRACT(EPOCH FROM now()) * 1000.0)::bigint, "updated_at" BIGINT NOT NULL DEFAULT (EXTRACT(EPOCH FROM now()) * 1000.0)::bigint )`, ], }, + { + pg: [`ALTER TABLE "${schemaName}"."queues" ADD COLUMN IF NOT EXISTS "max_dequeues_per_poll" INT4`], + }, // Migrations below replace broad indexes on workflow_status with partial // indexes targeted at individual query patterns (recovery, troubleshooting, // dequeue, rate-limit count). Each index DDL runs CONCURRENTLY on Postgres diff --git a/src/system_database.ts b/src/system_database.ts index 3e9ce15a1..ca8bfc715 100644 --- a/src/system_database.ts +++ b/src/system_database.ts @@ -26,7 +26,7 @@ import { } from '../schemas/system_db_schema'; import { globalParams, cancellableSleep, INTERNAL_QUEUE_NAME, Semaphore, sleepConfig, sleepms } from './utils'; import { GlobalLogger } from './telemetry/logs'; -import { WorkflowQueue } from './wfqueue'; +import { DEFAULT_MAX_DEQUEUES_PER_POLL, WorkflowQueue } from './wfqueue'; import { randomUUID } from 'crypto'; import { getClientConfig } from './utils'; import { connectToPGAndReportOutcome, ensurePGDatabase, maskDatabaseUrl } from './database_utils'; @@ -102,6 +102,7 @@ export interface QueueRecord { priorityEnabled: boolean; partitionQueue: boolean; pollingIntervalSec: number; + maxDequeuesPerPoll: number | null; } /** The subset of a queue record that may be changed after creation. */ @@ -115,6 +116,7 @@ const QUEUE_COLUMN_BY_FIELD: Record = { priorityEnabled: 'priority_enabled', partitionQueue: 'partition_queue', pollingIntervalSec: 'polling_interval_sec', + maxDequeuesPerPoll: 'max_dequeues_per_poll', }; function queueRecordFromRow(row: queues): QueueRecord { @@ -127,6 +129,7 @@ function queueRecordFromRow(row: queues): QueueRecord { priorityEnabled: row.priority_enabled, partitionQueue: row.partition_queue, pollingIntervalSec: row.polling_interval_sec, + maxDequeuesPerPoll: row.max_dequeues_per_poll, }; } @@ -2696,6 +2699,7 @@ export class SystemDatabase { executorID: string, appVersion: string, queuePartitionKey?: string, + maxDequeuesThisPoll: number = queue.maxDequeuesPerPoll ?? DEFAULT_MAX_DEQUEUES_PER_POLL, ): Promise { const startTimeMs = Date.now(); const limiterPeriodMS = queue.rateLimit ? queue.rateLimit.periodSec * 1000 : 0; @@ -2782,6 +2786,17 @@ export class SystemDatabase { return claimedIDs; } + if (queue.rateLimit) { + maxTasks = Math.min(maxTasks, queue.rateLimit.limitPerPeriod - numRecentQueries); + } + maxTasks = Math.min(maxTasks, maxDequeuesThisPoll); + + // Return immediately if there are no available tasks due to batch or rate limits. + if (maxTasks <= 0) { + await client.query('COMMIT'); + return claimedIDs; + } + // Retrieve the first max_tasks workflows in the queue. const latestVersionResult = await client.query<{ version_name: string }>( `SELECT version_name @@ -2795,65 +2810,51 @@ export class SystemDatabase { : 'application_version = $3'; const lockMode = queue.concurrency ? 'FOR UPDATE NOWAIT' : 'FOR UPDATE SKIP LOCKED'; - const limitClause = maxTasks !== Infinity ? `LIMIT ${maxTasks}` : ''; + const limitClause = `LIMIT ${maxTasks}`; - const selectParams = [StatusString.ENQUEUED, queue.name, appVersion, ...partitionParams]; - const selectQuery = ` + const selectParams = [StatusString.ENQUEUED, queue.name, appVersion, queuePartitionKey ?? null]; + const claimQuery = ` + WITH candidates AS ( SELECT workflow_uuid FROM "${this.schemaName}".workflow_status WHERE status = $1 AND queue_name = $2 AND ${versionClause} - ${partitionFilter.replace('$PARTITION', '$4')} + AND ($4::text IS NULL OR queue_partition_key = $4) ORDER BY priority ASC, created_at ASC ${limitClause} ${lockMode} + ) + UPDATE "${this.schemaName}".workflow_status wf + SET status = $5, + executor_id = $6, + application_version = $7, + started_at_epoch_ms = $8, + rate_limited = $9, + workflow_deadline_epoch_ms = CASE + WHEN wf.workflow_timeout_ms IS NOT NULL AND wf.workflow_deadline_epoch_ms IS NULL + THEN (EXTRACT(epoch FROM now()) * 1000)::bigint + wf.workflow_timeout_ms + ELSE wf.workflow_deadline_epoch_ms + END + FROM candidates + WHERE wf.workflow_uuid = candidates.workflow_uuid + AND wf.status = $1 + RETURNING wf.workflow_uuid `; - const { rows } = await client.query<{ workflow_uuid: string }>(selectQuery, selectParams); - // Fires while the SELECT FOR UPDATE lock is held — tests can throw a + const { rows } = await client.query<{ workflow_uuid: string }>(claimQuery, [ + ...selectParams, + StatusString.PENDING, + executorID, + appVersion, + startTimeMs, + queue.rateLimit !== undefined, + ]); + // Fires while the claim transaction is still open — tests can throw a // synthetic 55P03 here to simulate a concurrent executor winning the race. await debugTriggerPoint(DEBUG_TRIGGER_FIND_AND_MARK_AFTER_SELECT); - // Start the workflows - const workflowIDs = rows.map((row) => row.workflow_uuid); - for (const id of workflowIDs) { - // If we have a rate limit, stop starting functions when the number - // of functions started this period exceeds the limit. - if (queue.rateLimit && claimedIDs.length + numRecentQueries >= queue.rateLimit.limitPerPeriod) { - break; - } - - // Start the functions by marking them as pending and updating their executor IDs. - // Only claim the workflow if the UPDATE actually transitioned an ENQUEUED row — - // otherwise another worker won the race and we must not re-dispatch it. - const updateRes = await client.query( - `UPDATE "${this.schemaName}".workflow_status - SET status = $1, - executor_id = $2, - application_version = $3, - started_at_epoch_ms = $4, - rate_limited = $5, - workflow_deadline_epoch_ms = CASE - WHEN workflow_timeout_ms IS NOT NULL AND workflow_deadline_epoch_ms IS NULL - THEN (EXTRACT(epoch FROM now()) * 1000)::bigint + workflow_timeout_ms - ELSE workflow_deadline_epoch_ms - END - WHERE workflow_uuid = $6 AND status = $7`, - [ - StatusString.PENDING, - executorID, - appVersion, - startTimeMs, - queue.rateLimit !== undefined, - id, - StatusString.ENQUEUED, - ], - ); - if ((updateRes.rowCount ?? 0) > 0) { - claimedIDs.push(id); - } - } + claimedIDs.push(...rows.map((row) => row.workflow_uuid)); await client.query('COMMIT'); } catch (error) { @@ -3670,7 +3671,7 @@ export class SystemDatabase { async getQueue(name: string): Promise { const { rows } = await this.pool.query( `SELECT name, concurrency, worker_concurrency, rate_limit_max, rate_limit_period_sec, - priority_enabled, partition_queue, polling_interval_sec + priority_enabled, partition_queue, polling_interval_sec, max_dequeues_per_poll FROM "${this.schemaName}".queues WHERE name = $1`, [name], @@ -3681,7 +3682,7 @@ export class SystemDatabase { async listQueues(): Promise { const { rows } = await this.pool.query( `SELECT name, concurrency, worker_concurrency, rate_limit_max, rate_limit_period_sec, - priority_enabled, partition_queue, polling_interval_sec + priority_enabled, partition_queue, polling_interval_sec, max_dequeues_per_poll FROM "${this.schemaName}".queues`, ); return rows.map(queueRecordFromRow); @@ -3724,6 +3725,7 @@ export class SystemDatabase { priority_enabled = EXCLUDED.priority_enabled, partition_queue = EXCLUDED.partition_queue, polling_interval_sec = EXCLUDED.polling_interval_sec, + max_dequeues_per_poll = EXCLUDED.max_dequeues_per_poll, updated_at = EXCLUDED.updated_at` : `ON CONFLICT (name) DO NOTHING`; const client = await this.pool.connect(); @@ -3736,8 +3738,8 @@ export class SystemDatabase { await client.query( `INSERT INTO "${this.schemaName}".queues (name, concurrency, worker_concurrency, rate_limit_max, rate_limit_period_sec, - priority_enabled, partition_queue, polling_interval_sec, updated_at) - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) + priority_enabled, partition_queue, polling_interval_sec, max_dequeues_per_poll, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) ${onConflict}`, [ record.name, @@ -3748,6 +3750,7 @@ export class SystemDatabase { record.priorityEnabled, record.partitionQueue, record.pollingIntervalSec, + record.maxDequeuesPerPoll, now, ], ); diff --git a/src/wfqueue.ts b/src/wfqueue.ts index acefeeb79..768dc87bc 100644 --- a/src/wfqueue.ts +++ b/src/wfqueue.ts @@ -9,6 +9,8 @@ import type { QueueRecord, SystemDatabase } from './system_database'; import type { GlobalLogger } from './telemetry/logs'; import { globalParams, interruptibleSleep, INTERNAL_QUEUE_NAME } from './utils'; +export const DEFAULT_MAX_DEQUEUES_PER_POLL = 100; + /** * Log a single queue's name and its set parameters. Unset parameters are * omitted, matching `Queue: (concurrency=…, worker_concurrency=…, @@ -21,6 +23,7 @@ export function logQueue(logger: GlobalLogger, q: WorkflowQueue): void { if (q.rateLimit !== undefined) opts.push(`limit=${q.rateLimit.limitPerPeriod}/${q.rateLimit.periodSec}s`); if (q.priorityEnabled) opts.push('priority'); if (q.partitionQueue) opts.push('partitioned'); + if (q.maxDequeuesPerPoll !== undefined) opts.push(`max_dequeues_per_poll=${q.maxDequeuesPerPoll}`); const optsStr = opts.length > 0 ? ` (${opts.join(', ')})` : ''; logger.info(`Queue: ${q.name}${optsStr}`); } @@ -55,6 +58,8 @@ export interface QueueParameters { partitionQueue?: boolean; /** Base (minimum) polling interval in ms for this queue's dispatch loop (default 1000) */ minPollingIntervalMs?: number; + /** Maximum workflows this worker claims in one dequeue transaction (default 100) */ + maxDequeuesPerPoll?: number; } /** @@ -122,6 +127,7 @@ async function refreshFromDb(q: WorkflowQueue): Promise { q.priorityEnabled = record.priorityEnabled; q.partitionQueue = record.partitionQueue; q.minPollingIntervalMs = record.pollingIntervalSec * 1000; + q.maxDequeuesPerPoll = record.maxDequeuesPerPoll ?? undefined; } /** @@ -142,6 +148,7 @@ export class WorkflowQueue { priorityEnabled: boolean = false; partitionQueue: boolean = false; minPollingIntervalMs?: number; + maxDequeuesPerPoll?: number; /** * When true, this queue's configuration is persisted in the `queues` system @@ -192,6 +199,7 @@ export class WorkflowQueue { this.priorityEnabled = params.priorityEnabled ?? false; this.partitionQueue = params.partitionQueue ?? false; this.minPollingIntervalMs = params.minPollingIntervalMs; + this.maxDequeuesPerPoll = params.maxDequeuesPerPoll; if (wfQueueRunner.wfQueuesByName.has(name)) { throw Error(`Workflow Queue '${name}' defined multiple times`); @@ -201,13 +209,16 @@ export class WorkflowQueue { /** Throws if any combination of queue parameters is invalid. */ static validateQueueParams(params: QueueParameters): void { - const { concurrency, workerConcurrency, rateLimit, minPollingIntervalMs } = params; + const { concurrency, workerConcurrency, rateLimit, minPollingIntervalMs, maxDequeuesPerPoll } = params; if (workerConcurrency !== undefined && concurrency !== undefined && workerConcurrency > concurrency) { throw new Error('concurrency must be greater than or equal to workerConcurrency'); } if (minPollingIntervalMs !== undefined && minPollingIntervalMs <= 0) { throw new Error('minPollingIntervalMs must be positive'); } + if (maxDequeuesPerPoll !== undefined && (!Number.isInteger(maxDequeuesPerPoll) || maxDequeuesPerPoll <= 0)) { + throw new Error('maxDequeuesPerPoll must be a positive integer'); + } if (rateLimit !== undefined && (rateLimit.limitPerPeriod === undefined || rateLimit.periodSec === undefined)) { throw new Error('rateLimit must specify both limitPerPeriod and periodSec'); } @@ -224,6 +235,7 @@ export class WorkflowQueue { priorityEnabled: params.priorityEnabled ?? false, partitionQueue: params.partitionQueue ?? false, pollingIntervalSec: (params.minPollingIntervalMs ?? 1000) / 1000, + maxDequeuesPerPoll: params.maxDequeuesPerPoll ?? null, }; } @@ -247,6 +259,7 @@ export class WorkflowQueue { q.priorityEnabled = record.priorityEnabled; q.partitionQueue = record.partitionQueue; q.minPollingIntervalMs = record.pollingIntervalSec * 1000; + q.maxDequeuesPerPoll = record.maxDequeuesPerPoll ?? undefined; q.databaseBacked = true; q.clientBound = clientSystemDatabase !== undefined; if (clientSystemDatabase !== undefined) { @@ -306,6 +319,15 @@ export class WorkflowQueue { this.minPollingIntervalMs = value; } + async setMaxDequeuesPerPoll(value: number | undefined): Promise { + requireDatabaseBacked(this); + if (value !== undefined && (!Number.isInteger(value) || value <= 0)) { + throw new Error('maxDequeuesPerPoll must be a positive integer'); + } + await sysDBFor(this).updateQueue(this.name, { maxDequeuesPerPoll: value ?? null }); + this.maxDequeuesPerPoll = value; + } + async getConcurrency(): Promise { await refreshFromDb(this); return this.concurrency; @@ -335,6 +357,11 @@ export class WorkflowQueue { await refreshFromDb(this); return this.minPollingIntervalMs; } + + async getMaxDequeuesPerPoll(): Promise { + await refreshFromDb(this); + return this.maxDequeuesPerPoll; + } } /** Per-queue runtime scheduling state tracked by the shared dispatcher. */ @@ -428,7 +455,11 @@ class WFQueueRunner { private ensureState(queue: WorkflowQueue, now: number): void { if (this.states.has(queue.name)) return; const interval = queue.minPollingIntervalMs ?? WFQueueRunner.defaultMinPollingIntervalMs; - this.states.set(queue.name, { queue, currentPollingMs: interval, nextPollAt: now + interval }); + this.states.set(queue.name, { + queue, + currentPollingMs: interval, + nextPollAt: now + interval, + }); } /** Reconcile DB-backed queues against the queues table in one query: refresh, add, or drop them. */ @@ -538,7 +569,7 @@ class WFQueueRunner { } } - /** Poll one queue once, starting ready workflows; returns true if DB contention was detected. */ + /** Poll one queue once, starting ready workflows. */ private async pollQueue(exec: DBOSExecutor, queue: WorkflowQueue): Promise { let contentionDetected = false; // Helper function that starts dequeued workflows diff --git a/tests/adminserver.test.ts b/tests/adminserver.test.ts index 36f1d2b32..37206d81d 100644 --- a/tests/adminserver.test.ts +++ b/tests/adminserver.test.ts @@ -364,7 +364,11 @@ describe('running-admin-server-tests', () => { } // Database-backed queues should also appear in the metadata endpoint. - await DBOS.registerQueue('admin-db-backed-queue', { concurrency: 7, workerConcurrency: 3 }); + await DBOS.registerQueue('admin-db-backed-queue', { + concurrency: 7, + workerConcurrency: 3, + maxDequeuesPerPoll: 11, + }); try { const dbResponse = await fetch(`http://localhost:3001${WorkflowQueuesMetadataUrl}`, { method: 'GET' }); expect(dbResponse.status).toBe(200); @@ -373,6 +377,7 @@ describe('running-admin-server-tests', () => { expect(dbQueue).toBeDefined(); expect(dbQueue!.concurrency).toBe(7); expect(dbQueue!.workerConcurrency).toBe(3); + expect(dbQueue!.maxDequeuesPerPoll).toBe(11); } finally { await DBOS.deleteQueue('admin-db-backed-queue').catch(() => {}); } diff --git a/tests/client.test.ts b/tests/client.test.ts index 56da514b2..dd8f129f0 100644 --- a/tests/client.test.ts +++ b/tests/client.test.ts @@ -911,6 +911,7 @@ describe('DBOSClient', () => { workerConcurrency: 2, priorityEnabled: true, minPollingIntervalMs: 2500, + maxDequeuesPerPoll: 13, }); expect(registered.name).toBe(queueName); expect(registered.databaseBacked).toBe(true); @@ -924,6 +925,7 @@ describe('DBOSClient', () => { expect(retrieved!.rateLimit).toEqual({ limitPerPeriod: 5, periodSec: 1.5 }); expect(retrieved!.priorityEnabled).toBe(true); expect(retrieved!.minPollingIntervalMs).toBe(2500); + expect(retrieved!.maxDequeuesPerPoll).toBe(13); // Setters write through the client's database; the launched DBOS // executor sees the same row. diff --git a/tests/wfqueue.test.ts b/tests/wfqueue.test.ts index e015699fe..4ce26b4a9 100644 --- a/tests/wfqueue.test.ts +++ b/tests/wfqueue.test.ts @@ -1,6 +1,6 @@ import { StatusString, WorkflowHandle, DBOS, ConfiguredInstance, DBOSClient } from '../src'; import { DBOSConfig, DBOSExecutor, DBOS_QUEUE_MAX_PRIORITY, DBOS_QUEUE_MIN_PRIORITY } from '../src/dbos-executor'; -import { QueueParameters, wfQueueRunner } from '../src/wfqueue'; +import { DEFAULT_MAX_DEQUEUES_PER_POLL, QueueParameters, wfQueueRunner } from '../src/wfqueue'; import { generateDBOSTestConfig, setUpDBOSTestSysDb, @@ -2248,6 +2248,7 @@ describe('database-backed-queue-crud', () => { workerConcurrency: 2, priorityEnabled: true, minPollingIntervalMs: 2500, + maxDequeuesPerPoll: 17, }); expect(registered.name).toBe(queueName); expect(registered.databaseBacked).toBe(true); @@ -2263,6 +2264,7 @@ describe('database-backed-queue-crud', () => { expect(retrieved!.priorityEnabled).toBe(true); expect(retrieved!.partitionQueue).toBe(false); expect(retrieved!.minPollingIntervalMs).toBe(2500); + expect(retrieved!.maxDequeuesPerPoll).toBe(17); expect(retrieved!.databaseBacked).toBe(true); // never_update leaves the existing row alone. @@ -2270,6 +2272,7 @@ describe('database-backed-queue-crud', () => { retrieved = await DBOS.retrieveQueue(queueName); expect(retrieved!.concurrency).toBe(10); expect(retrieved!.rateLimit).toEqual({ limitPerPeriod: 5, periodSec: 1.5 }); + expect(retrieved!.maxDequeuesPerPoll).toBe(17); // always_update overwrites every column, including clearing those that // were previously set but are now omitted. @@ -2280,6 +2283,7 @@ describe('database-backed-queue-crud', () => { expect(retrieved!.rateLimit).toBeUndefined(); expect(retrieved!.priorityEnabled).toBe(false); expect(retrieved!.minPollingIntervalMs).toBe(1000); + expect(retrieved!.maxDequeuesPerPoll).toBeUndefined(); // update_if_latest_version updates when the running version is the latest. await DBOS.registerQueue(queueName, { concurrency: 30, onConflict: 'update_if_latest_version' }); @@ -2310,6 +2314,7 @@ describe('database-backed-queue-crud', () => { workerConcurrency: 2, priorityEnabled: false, minPollingIntervalMs: 1000, + maxDequeuesPerPoll: 10, }); await queue.setConcurrency(8); @@ -2318,6 +2323,7 @@ describe('database-backed-queue-crud', () => { await queue.setPriorityEnabled(true); await queue.setPartitionQueue(true); await queue.setMinPollingIntervalMs(500); + await queue.setMaxDequeuesPerPoll(12); const fresh = await DBOS.retrieveQueue(queueName); for (const q of [queue, fresh]) { @@ -2328,6 +2334,7 @@ describe('database-backed-queue-crud', () => { expect(q!.priorityEnabled).toBe(true); expect(q!.partitionQueue).toBe(true); expect(q!.minPollingIntervalMs).toBe(500); + expect(q!.maxDequeuesPerPoll).toBe(12); } // Cross-validation: workerConcurrency cannot exceed concurrency. @@ -2336,11 +2343,14 @@ describe('database-backed-queue-crud', () => { ); // Polling interval must be positive. await expect(queue.setMinPollingIntervalMs(0)).rejects.toThrow('minPollingIntervalMs must be positive'); + await expect(queue.setMaxDequeuesPerPoll(0)).rejects.toThrow('maxDequeuesPerPoll must be a positive integer'); // Rate limit can be cleared. await queue.setRateLimit(undefined); const cleared = await DBOS.retrieveQueue(queueName); expect(cleared!.rateLimit).toBeUndefined(); + await queue.setMaxDequeuesPerPoll(undefined); + expect((await DBOS.retrieveQueue(queueName))!.maxDequeuesPerPoll).toBeUndefined(); // In-memory queues do not support setters. const legacyName = `legacy_dyn_queue_${randomUUID()}`; @@ -2574,6 +2584,129 @@ describe('database-backed-queue-crud', () => { }); }); +describe('bounded-queue-dequeue-claims', () => { + let config: DBOSConfig; + let client: Client; + + beforeAll(async () => { + config = generateDBOSTestConfig(); + await setUpDBOSTestSysDb(config); + }); + + beforeEach(async () => { + DBOS.setConfig({ ...config, listenQueues: [] }); + await DBOS.launch(); + client = new Client({ connectionString: config.systemDatabaseUrl }); + await client.connect(); + }); + + afterEach(async () => { + await client?.end(); + await DBOS.shutdown(); + }); + + async function enqueueSimple(queueName: string, count: number, partitionKey?: string): Promise { + const handles = await Promise.all( + Array.from({ length: count }, (_, i) => + DBOS.startWorkflow(TestWFs, { + workflowID: `${queueName}-${partitionKey ?? 'none'}-${i}-${randomUUID()}`, + queueName, + enqueueOptions: partitionKey ? { queuePartitionKey: partitionKey } : undefined, + }).testWorkflowSimple('a', String(i)), + ), + ); + return handles.map((h) => h.workflowID); + } + + async function countStatuses(queueName: string): Promise> { + const result = await client.query<{ status: string; count: string }>( + `SELECT status, COUNT(*) AS count + FROM dbos.workflow_status + WHERE queue_name = $1 + GROUP BY status`, + [queueName], + ); + return Object.fromEntries(result.rows.map((row) => [row.status, Number(row.count)])); + } + + test('unbounded queue claims at most the default batch in one dequeue transaction', async () => { + const queueName = `bounded_default_${randomUUID()}`; + const workflowIDs = await enqueueSimple(queueName, DEFAULT_MAX_DEQUEUES_PER_POLL + 5); + const queue = await DBOS.registerQueue(queueName, { onConflict: 'always_update' }); + + try { + const claimed = await DBOSExecutor.globalInstance!.systemDatabase.findAndMarkStartableWorkflows( + queue, + DBOS.executorID, + globalParams.appVersion, + ); + expect(claimed).toHaveLength(DEFAULT_MAX_DEQUEUES_PER_POLL); + const counts = await countStatuses(queueName); + expect(counts[StatusString.PENDING]).toBe(DEFAULT_MAX_DEQUEUES_PER_POLL); + expect(counts[StatusString.ENQUEUED]).toBe(5); + } finally { + await DBOS.cancelWorkflows(workflowIDs); + await DBOS.deleteQueue(queueName); + } + }); + + test('rate-limited queue claims no more than remaining rate-limit capacity', async () => { + const queueName = `bounded_rate_${randomUUID()}`; + const workflowIDs = await enqueueSimple(queueName, 8); + const queue = await DBOS.registerQueue(queueName, { + onConflict: 'always_update', + rateLimit: { limitPerPeriod: 3, periodSec: 60 }, + }); + + try { + const claimed = await DBOSExecutor.globalInstance!.systemDatabase.findAndMarkStartableWorkflows( + queue, + DBOS.executorID, + globalParams.appVersion, + ); + expect(claimed).toHaveLength(3); + const counts = await countStatuses(queueName); + expect(counts[StatusString.PENDING]).toBe(3); + expect(counts[StatusString.ENQUEUED]).toBe(5); + } finally { + await DBOS.cancelWorkflows(workflowIDs); + await DBOS.deleteQueue(queueName); + } + }); + + test('partitioned queue caps each partition dequeue transaction', async () => { + const queueName = `bounded_partition_${randomUUID()}`; + const workflowIDs = [...(await enqueueSimple(queueName, 5, 'a')), ...(await enqueueSimple(queueName, 5, 'b'))]; + const queue = await DBOS.registerQueue(queueName, { + onConflict: 'always_update', + partitionQueue: true, + maxDequeuesPerPoll: 4, + }); + + try { + const partitions = await DBOSExecutor.globalInstance!.systemDatabase.getQueuePartitions(queueName); + const claimed: string[] = []; + for (const partition of partitions) { + const partitionClaimed = await DBOSExecutor.globalInstance!.systemDatabase.findAndMarkStartableWorkflows( + queue, + DBOS.executorID, + globalParams.appVersion, + partition, + ); + claimed.push(...partitionClaimed); + } + + expect(claimed).toHaveLength(8); + const counts = await countStatuses(queueName); + expect(counts[StatusString.PENDING]).toBe(8); + expect(counts[StatusString.ENQUEUED]).toBe(2); + } finally { + await DBOS.cancelWorkflows(workflowIDs); + await DBOS.deleteQueue(queueName); + } + }); +}); + // Regression test for the "orphan-PENDING" bug in partitioned queue dispatch. // // Pre-fix: the dispatch loop iterated partition keys, collected wfids from every