@dbos-inc/dbos-sdk 4.20.7-preview → 4.20.9-preview

package/dist/src/system_database.js

@@ -386,8 +386,14 @@ class SystemDatabase {
     workflowEventsMap = new NotificationMap();
     streamsMap = new NotificationMap();
     customPool = false;
+    /**
+     * Caps how many DB-backed polling reads (from wait operations) may run
+     * concurrently against the pool, so a polling storm cannot check out every
+     * client and starve control-plane operations. See {@link #pollWithLimiter}.
+     */
+    pollLimiter;
     runningWorkflowMap = new Map(); // Map from workflowID to workflow promise, queue name and partition key
-    constructor(systemDatabaseUrl, logger, serializer, sysDbPoolSize = exports.DEFAULT_POOL_SIZE, systemDatabasePool, schemaName = 'dbos', useListenNotify = true) {
+    constructor(systemDatabaseUrl, logger, serializer, sysDbPoolSize = exports.DEFAULT_POOL_SIZE, systemDatabasePool, schemaName = 'dbos', useListenNotify = true, pollingConcurrency) {
         this.systemDatabaseUrl = systemDatabaseUrl;
         this.logger = logger;
         this.serializer = serializer;
@@ -406,6 +412,11 @@ class SystemDatabase {
             };
             this.pool = new pg_1.Pool(systemPoolConfig);
         }
+        // Default the polling limit to half the pool (minimum 1), reserving the rest
+        // of the pool for control-plane operations.
+        const effectivePoolSize = this.pool.options.max ?? sysDbPoolSize;
+        const pollingLimit = pollingConcurrency ?? Math.max(1, Math.floor(effectivePoolSize / 2));
+        this.pollLimiter = new utils_1.Semaphore(pollingLimit);
         this.pool.on('error', (err) => {
             this.logger.warn(`Unexpected error in pool: ${err}`);
         });
@@ -1169,16 +1180,38 @@ class SystemDatabase {
             //throw new Error('Message notification map is not empty - shutdown is not clean.');
         }
     }
+    /**
+     * Run a DB-backed polling read under the polling concurrency limiter, so that
+     * high-fan-out wait loops cannot check out every pool client and starve
+     * control-plane operations. Only polling reads should go through here;
+     * control-plane work hits the pool directly and bypasses the limiter.
+     */
+    #pollWithLimiter(query) {
+        return this.pollLimiter.runExclusive(query);
+    }
+    /**
+     * Cancellation check for use inside polling wait loops: the status read runs
+     * under the polling limiter so it counts against the same concurrency budget
+     * as the rest of the loop's reads.
+     */
+    async #checkIfCanceledLimited(workflowID) {
+        await this.#pollWithLimiter(() => this.#checkIfCanceled(this.pool, workflowID));
+    }
     async awaitWorkflowResult(workflowID, timeoutSeconds, callerID, timerFuncID, pollingIntervalMs) {
         const timeoutms = timeoutSeconds !== undefined ? timeoutSeconds * 1000 : undefined;
         let finishTime = timeoutms !== undefined ? Date.now() + timeoutms : undefined;
         const pollIntervalMs = pollingIntervalMs ?? this.dbPollingIntervalResultMs;
+        // Record the durable timeout deadline once before polling. #durableSleep persists it on the
+        // first call and reads back the same value on recovery, so it never changes across iterations.
+        if (timerFuncID !== undefined && callerID !== undefined && timeoutms !== undefined) {
+            finishTime = await this.#durableSleep(callerID, timerFuncID, timeoutms);
+        }
         while (true) {
             if (callerID)
-                await this.checkIfCanceled(callerID);
+                await this.#checkIfCanceledLimited(callerID);
             try {
-                const { rows } = await this.pool.query(`SELECT status, output, error, serialization FROM "${this.schemaName}".workflow_status
-           WHERE workflow_uuid=$1`, [workflowID]);
+                const { rows } = await this.#pollWithLimiter(() => this.pool.query(`SELECT status, output, error, serialization FROM "${this.schemaName}".workflow_status
+             WHERE workflow_uuid=$1`, [workflowID]));
                 if (rows.length > 0) {
                     const status = rows[0].status;
                     if (status === workflow_1.StatusString.SUCCESS) {
@@ -1206,9 +1239,6 @@ class SystemDatabase {
             const ct = Date.now();
             if (finishTime && ct > finishTime)
                 return undefined; // Time's up
-            if (timerFuncID !== undefined && callerID !== undefined && timeoutms !== undefined) {
-                finishTime = await this.#durableSleep(callerID, timerFuncID, timeoutms);
-            }
             let poll = finishTime ? finishTime - Date.now() : pollIntervalMs;
             poll = Math.min(pollIntervalMs, poll);
             await (0, utils_1.sleepms)(poll);
@@ -1219,11 +1249,11 @@ class SystemDatabase {
         const pollIntervalMs = pollingIntervalMs ?? this.dbPollingIntervalResultMs;
         while (true) {
             if (callerID)
-                await this.checkIfCanceled(callerID);
-            const { rows } = await this.pool.query(`SELECT workflow_uuid FROM "${this.schemaName}".workflow_status
-         WHERE workflow_uuid IN (${placeholders})
-           AND status NOT IN ('${workflow_1.StatusString.PENDING}', '${workflow_1.StatusString.ENQUEUED}', '${workflow_1.StatusString.DELAYED}')
-         LIMIT 1`, workflowIds);
+                await this.#checkIfCanceledLimited(callerID);
+            const { rows } = await this.#pollWithLimiter(() => this.pool.query(`SELECT workflow_uuid FROM "${this.schemaName}".workflow_status
+           WHERE workflow_uuid IN (${placeholders})
+             AND status NOT IN ('${workflow_1.StatusString.PENDING}', '${workflow_1.StatusString.ENQUEUED}', '${workflow_1.StatusString.DELAYED}')
+           LIMIT 1`, workflowIds));
             if (rows.length > 0) {
                 return rows[0].workflow_uuid;
             }
@@ -1236,10 +1266,10 @@ class SystemDatabase {
         while (remainingWorkflowIds.size > 0) {
             const currentWorkflowIds = [...remainingWorkflowIds];
             if (callerID)
-                await this.checkIfCanceled(callerID);
-            const { rows } = await this.pool.query(`SELECT workflow_uuid FROM "${this.schemaName}".workflow_status
-         WHERE workflow_uuid = ANY($1::text[])
-           AND status NOT IN ('${workflow_1.StatusString.PENDING}', '${workflow_1.StatusString.ENQUEUED}', '${workflow_1.StatusString.DELAYED}')`, [currentWorkflowIds]);
+                await this.#checkIfCanceledLimited(callerID);
+            const { rows } = await this.#pollWithLimiter(() => this.pool.query(`SELECT workflow_uuid FROM "${this.schemaName}".workflow_status
+           WHERE workflow_uuid = ANY($1::text[])
+             AND status NOT IN ('${workflow_1.StatusString.PENDING}', '${workflow_1.StatusString.ENQUEUED}', '${workflow_1.StatusString.DELAYED}')`, [currentWorkflowIds]));
             for (const row of rows) {
                 remainingWorkflowIds.delete(row.workflow_uuid);
             }
@@ -1318,6 +1348,11 @@ class SystemDatabase {
         const timeoutms = timeoutSeconds !== undefined ? timeoutSeconds * 1000 : undefined;
         let finishTime = timeoutms !== undefined ? Date.now() + timeoutms : undefined;
         const pollIntervalMs = pollingIntervalMs ?? this.dbPollingIntervalEventMs;
+        // Record the durable timeout deadline once before polling. #durableSleep persists it on the
+        // first call and reads back the same value on recovery, so it never changes across iterations.
+        if (timeoutms) {
+            finishTime = await this.#durableSleep(workflowID, timeoutFunctionID, timeoutms);
+        }
         while (true) {
             // register the key with the global notifications listener.
             let resolveNotification;
@@ -1327,17 +1362,14 @@ class SystemDatabase {
             const payload = `${workflowID}::${topic}`;
             const cbr = this.notificationsMap.registerCallback(payload, resolveNotification);
             try {
-                await this.checkIfCanceled(workflowID);
+                await this.#checkIfCanceledLimited(workflowID);
                 // Check if the key is already in the DB, then wait for the notification if it isn't.
-                const initRecvRows = (await this.pool.query(`SELECT topic FROM "${this.schemaName}".notifications WHERE destination_uuid=$1 AND topic=$2 AND consumed = false;`, [workflowID, topic])).rows;
+                const initRecvRows = (await this.#pollWithLimiter(() => this.pool.query(`SELECT topic FROM "${this.schemaName}".notifications WHERE destination_uuid=$1 AND topic=$2 AND consumed = false;`, [workflowID, topic]))).rows;
                 if (initRecvRows.length !== 0)
                     break;
                 const ct = Date.now();
                 if (finishTime && ct > finishTime)
                     break; // Time's up
-                if (timeoutms) {
-                    finishTime = await this.#durableSleep(workflowID, timeoutFunctionID, timeoutms);
-                }
                 let poll = finishTime ? finishTime - Date.now() : pollIntervalMs;
                 poll = Math.min(pollIntervalMs, poll);
                 const { promise, cancel } = (0, utils_1.cancellableSleep)(poll);
@@ -1442,6 +1474,12 @@ class SystemDatabase {
         const timeoutms = timeoutSeconds !== undefined ? timeoutSeconds * 1000 : undefined;
         let finishTime = timeoutms !== undefined ? Date.now() + timeoutms : undefined;
         const pollIntervalMs = pollingIntervalMs ?? this.dbPollingIntervalEventMs;
+        // If we have a callerWorkflow, we want a durable sleep, otherwise, not. Record the durable
+        // timeout deadline once before polling. #durableSleep persists it on the first call and reads
+        // back the same value on recovery, so it never changes across iterations.
+        if (callerWorkflow && timeoutms) {
+            finishTime = await this.#durableSleep(callerWorkflow.workflowID, callerWorkflow.timeoutFunctionID ?? -1, timeoutms);
+        }
         // Register the key with the global notifications listener first... we do not want to look in the DB first
         //  or that would cause a timing hole.
         while (true) {
@@ -1452,11 +1490,11 @@ class SystemDatabase {
             const cbr = this.workflowEventsMap.registerCallback(payloadKey, resolveNotification);
             try {
                 if (callerWorkflow?.workflowID)
-                    await this.checkIfCanceled(callerWorkflow?.workflowID);
+                    await this.#checkIfCanceledLimited(callerWorkflow?.workflowID);
                 // Check if the key is already in the DB, then wait for the notification if it isn't.
-                const initRecvRows = (await this.pool.query(`SELECT key, value, serialization
+                const initRecvRows = (await this.#pollWithLimiter(() => this.pool.query(`SELECT key, value, serialization
              FROM "${this.schemaName}".workflow_events
-             WHERE workflow_uuid=$1 AND key=$2;`, [workflowID, key])).rows;
+             WHERE workflow_uuid=$1 AND key=$2;`, [workflowID, key]))).rows;
                 if (initRecvRows.length > 0) {
                     value = initRecvRows[0].value;
                     valueSer = initRecvRows[0].serialization;
@@ -1465,10 +1503,6 @@ class SystemDatabase {
                 const ct = Date.now();
                 if (finishTime && ct > finishTime)
                     break; // Time's up
-                // If we have a callerWorkflow, we want a durable sleep, otherwise, not
-                if (callerWorkflow && timeoutms) {
-                    finishTime = await this.#durableSleep(callerWorkflow.workflowID, callerWorkflow.timeoutFunctionID ?? -1, timeoutms);
-                }
                 let poll = finishTime ? finishTime - Date.now() : pollIntervalMs;
                 poll = Math.min(pollIntervalMs, poll);
                 const { promise, cancel } = (0, utils_1.cancellableSleep)(poll);