@dbos-inc/dbos-sdk 5.1.6-preview → 5.1.8-preview

package/dist/src/system_database.js

@@ -15,6 +15,7 @@ const pg_1 = require("pg");
 const error_1 = require("./error");
 const workflow_1 = require("./workflow");
 const utils_1 = require("./utils");
+const async_hooks_1 = require("async_hooks");
 const crypto_1 = require("crypto");
 const utils_2 = require("./utils");
 const database_utils_1 = require("./database_utils");
@@ -564,6 +565,7 @@ class SystemDatabase {
     dbPollingIntervalResultMs = 1000;
     dbPollingIntervalEventMs = 10000;
     dbPollingIntervalStreamMs = 1000;
+    dbPollingIntervalCancelMs = 1000;
     shouldUseDBNotifications = true;
     notificationsMap = new NotificationMap();
     workflowEventsMap = new NotificationMap();
@@ -587,6 +589,11 @@ class SystemDatabase {
      */
     pollLimiter;
     runningWorkflowMap = new Map(); // Map from workflowID to workflow promise, queue name and partition key
+    #cancelWatchers = new Map();
+    #cancelPollerLoop = undefined;
+    #cancelPollerWake = null;
+    // Runs the lazily started cancel poller outside the async context of the step that happens to start it
+    #cancelPollerScope = new async_hooks_1.AsyncResource('DBOSCancelPoller');
     // Per-partition-key created_at cursors: keep per-key queue order monotonic across batches
     #batchCreatedAtCursors = new Map();
     // Set by destroy(), so polling waits end instead of running on against a pool that outlives this handle.
@@ -773,6 +780,10 @@ class SystemDatabase {
             await this.#notifierLoop;
             this.#notifierLoop = undefined;
         }
+        this.#cancelPollerWake?.();
+        if (this.#cancelPollerLoop) {
+            await this.#cancelPollerLoop;
+        }
         if (this.notificationsClient) {
             this.#retireNotificationsClient(this.notificationsClient);
         }
@@ -1493,6 +1504,106 @@ class SystemDatabase {
             this.runningWorkflowMap.delete(wfid);
         }
     }
+    /**
+     * Drop a workflow's history from `startStep` onwards and re-enqueue it under the same
+     * workflow ID, so a replay re-executes everything from that step.
+     *
+     * Unlike forkWorkflow this writes no new workflow: peers keep addressing the same ID,
+     * and the workflow's mailbox, events, and streams are not copied anywhere. Unlike
+     * resumeWorkflows it applies only to workflows in a terminal state.
+     *
+     * Events published at or past the cut are rolled back to the last value published
+     * below it, using workflow_events_history as an undo log; a key that only the
+     * discarded run ever published is unpublished outright.
+     *
+     * Messages the discarded run consumed or received after the rewind point are deleted.
+     *
+     * Stream entries written by the discarded run remain in place, with the exception of
+     * the close sentinel, which has to go so new entries can be appended.
+     */
+    async rewindWorkflow(workflowID, startStep, options = {}) {
+        // Function IDs start at 0 in this SDK, so 0 is the whole history, not an error.
+        if (startStep < 0) {
+            throw new error_1.DBOSError(`startStep must be >= 0, got ${startStep}`);
+        }
+        const schema = this.schemaName;
+        const client = await this.#connect();
+        try {
+            await client.query('BEGIN ISOLATION LEVEL READ COMMITTED');
+            const { rows: statusRows } = await client.query(`SELECT status FROM "${schema}".workflow_status WHERE workflow_uuid = $1`, [workflowID]);
+            if (statusRows.length === 0) {
+                throw new error_1.DBOSNonExistentWorkflowError(`Workflow ${workflowID} does not exist`);
+            }
+            const status = statusRows[0].status;
+            if (status === workflow_1.StatusString.PENDING || status === workflow_1.StatusString.ENQUEUED || status === workflow_1.StatusString.DELAYED) {
+                throw new error_1.DBOSError(`Cannot rewind ${workflowID} (${status}): only a workflow in a terminal state can be rewound, so cancel it first`);
+            }
+            // Roll workflow_events back to the last value published before the cut, using
+            // workflow_events_history as an undo log, then rewind that history too.
+            // `publishedPastCut` is the correlated EXISTS both halves key off.
+            const publishedPastCut = `EXISTS (
+        SELECT 1 FROM "${schema}".workflow_events_history discarded
+        WHERE discarded.workflow_uuid = $1 AND discarded.key = %KEY% AND discarded.function_id >= $2
+      )`;
+            // First unpublish everything the discarded run touched.
+            await client.query(`DELETE FROM "${schema}".workflow_events
+         WHERE workflow_uuid = $1 AND ${publishedPastCut.replace('%KEY%', '"' + schema + '".workflow_events.key')}`, [workflowID, startStep]);
+            // Then restore those keys, if any, to the value published below the cut.
+            await client.query(`INSERT INTO "${schema}".workflow_events (workflow_uuid, key, value, serialization)
+         SELECT surviving.workflow_uuid, surviving.key, surviving.value, surviving.serialization
+         FROM (
+           SELECT weh.workflow_uuid, weh.key, weh.value, weh.serialization,
+                  ROW_NUMBER() OVER (PARTITION BY weh.key ORDER BY weh.function_id DESC) AS rn
+           FROM "${schema}".workflow_events_history weh
+           WHERE weh.workflow_uuid = $1 AND weh.function_id < $2
+             AND ${publishedPastCut.replace('%KEY%', 'weh.key')}
+         ) surviving
+         WHERE surviving.rn = 1`, [workflowID, startStep]);
+            // Clear the streams' close sentinels so the replay can append again.
+            await client.query(`DELETE FROM "${schema}".streams
+         WHERE workflow_uuid = $1 AND function_id >= $2 AND value = $3 AND serialization = $4`, [workflowID, startStep, exports.DBOS_STREAM_CLOSED_SENTINEL_SERIALIZED, serialization_1.DBOSPortableJSON.name()]);
+            // Discard the steps and the events history past the cut.
+            for (const table of ['operation_outputs', 'workflow_events_history']) {
+                await client.query(`DELETE FROM "${schema}".${table} WHERE workflow_uuid = $1 AND function_id >= $2`, [
+                    workflowID,
+                    startStep,
+                ]);
+            }
+            // Delete messages consumed or received after the rewind point
+            await client.query(`DELETE FROM "${schema}".notifications WHERE destination_uuid = $1 AND (consumed_by_function_id >= $2 OR consumed = false)`, [workflowID, startStep]);
+            await client.query(`DELETE FROM "${schema}".workflow_output WHERE workflow_uuid = $1`, [workflowID]);
+            // Re-enqueue the workflow. Re-asserting the status we read keeps a workflow that
+            // moved on underneath us from being resurrected.
+            const setVersion = options.applicationVersion !== undefined ? ', application_version = $6' : '';
+            const params = [
+                workflow_1.StatusString.ENQUEUED,
+                options.queueName ?? utils_1.INTERNAL_QUEUE_NAME,
+                options.queuePartitionKey ?? null,
+                workflowID,
+                status,
+            ];
+            if (options.applicationVersion !== undefined) {
+                params.push(options.applicationVersion);
+            }
+            const { rowCount } = await client.query(`UPDATE "${schema}".workflow_status
+         SET status = $1, queue_name = $2, queue_partition_key = $3, recovery_attempts = 0,
+             workflow_deadline_epoch_ms = NULL, deduplication_id = NULL, started_at_epoch_ms = NULL,
+             completed_at = NULL, output = NULL, error = NULL,
+             updated_at = (EXTRACT(EPOCH FROM now()) * 1000)::bigint${setVersion}
+         WHERE workflow_uuid = $4 AND status = $5`, params);
+            if (rowCount !== 1) {
+                throw new error_1.DBOSError(`Workflow ${workflowID} changed status while being rewound; retry the rewind`);
+            }
+            await client.query('COMMIT');
+        }
+        catch (e) {
+            await client.query('ROLLBACK').catch(() => undefined);
+            throw e;
+        }
+        finally {
+            client.release();
+        }
+    }
     async forkWorkflow(workflowID, startStep, options = {}) {
         const newWorkflowID = options.newWorkflowID ?? (0, crypto_1.randomUUID)();
         const result = await this.bulkForkWorkflows([workflowID], [newWorkflowID], [startStep], options);
@@ -1990,6 +2101,74 @@ class SystemDatabase {
     async checkIfCanceledLimited(workflowID) {
         await this.#pollWithLimiter(() => this.#checkIfCanceled(this.pool, workflowID));
     }
+    watchForCancellation(workflowID) {
+        const controller = new AbortController();
+        let state = 'idle';
+        const start = () => {
+            state = 'watching';
+            if (this.#destroyed)
+                return;
+            let watchers = this.#cancelWatchers.get(workflowID);
+            if (watchers === undefined) {
+                watchers = new Set();
+                this.#cancelWatchers.set(workflowID, watchers);
+            }
+            watchers.add(controller);
+            this.#cancelPollerLoop ??= this.#cancelPollerScope.runInAsyncScope(() => this.#runCancelPoller());
+        };
+        const stop = () => {
+            const watchers = this.#cancelWatchers.get(workflowID);
+            if (watchers?.delete(controller) && watchers.size === 0) {
+                this.#cancelWatchers.delete(workflowID);
+                if (this.#cancelWatchers.size === 0)
+                    this.#cancelPollerWake?.();
+            }
+        };
+        return {
+            get signal() {
+                if (state === 'idle')
+                    start();
+                return controller.signal;
+            },
+            get cancelled() {
+                return controller.signal.aborted;
+            },
+            release: () => {
+                if (state === 'watching')
+                    stop();
+                state = 'released';
+            },
+        };
+    }
+    async #runCancelPoller() {
+        while (this.#cancelWatchers.size > 0 && !this.#destroyed) {
+            const { promise, cancel } = (0, utils_1.cancellableSleep)(this.dbPollingIntervalCancelMs);
+            this.#cancelPollerWake = cancel;
+            await promise;
+            this.#cancelPollerWake = null;
+            if (this.#cancelWatchers.size === 0 || this.#destroyed)
+                break;
+            try {
+                await this.#abortCancelledWatchers();
+            }
+            catch (e) {
+                this.logger.debug(`Workflow cancellation poll failed: ${String(e)}`);
+            }
+        }
+        this.#cancelPollerLoop = undefined;
+    }
+    async #abortCancelledWatchers() {
+        const { rows } = await this.#pollWithLimiter(() => this.pool.query(`SELECT workflow_uuid FROM "${this.schemaName}".workflow_status WHERE workflow_uuid = ANY($1) AND status = $2`, [Array.from(this.#cancelWatchers.keys()), workflow_1.StatusString.CANCELLED]));
+        for (const { workflow_uuid: workflowID } of rows) {
+            const watchers = this.#cancelWatchers.get(workflowID);
+            if (watchers === undefined)
+                continue;
+            this.#cancelWatchers.delete(workflowID);
+            const reason = new error_1.DBOSWorkflowCancelledError(workflowID);
+            for (const controller of watchers)
+                controller.abort(reason);
+        }
+    }
     // A missing row normally means the workflow has not been inserted yet, so
     // polling for it is correct. Callers that know the row must already exist
     // (e.g. a run parking on an outcome it just failed to write) pass
@@ -2208,7 +2387,7 @@ class SystemDatabase {
         try {
             await client.query(`BEGIN ISOLATION LEVEL READ COMMITTED`);
             const finalRecvRows = (await client.query(`UPDATE "${this.schemaName}".notifications
-        SET consumed = true
+        SET consumed = true, consumed_by_function_id = $3
         WHERE destination_uuid = $1
           AND topic = $2
           AND consumed = false
@@ -2221,7 +2400,7 @@ class SystemDatabase {
             ORDER BY created_at_epoch_ms ASC
             LIMIT 1
           )
-        RETURNING notifications.message, notifications.serialization;`, [workflowID, topic])).rows;
+        RETURNING notifications.message, notifications.serialization;`, [workflowID, topic, functionID])).rows;
             if (finalRecvRows.length > 0) {
                 message = finalRecvRows[0].message;
                 serialization = finalRecvRows[0].serialization;
@@ -4691,6 +4870,12 @@ __decorate([
     __metadata("design:paramtypes", [Object]),
     __metadata("design:returntype", Promise)
 ], SystemDatabase.prototype, "debounceDelayedWorkflowStandalone", null);
+__decorate([
+    dbRetry(),
+    __metadata("design:type", Function),
+    __metadata("design:paramtypes", [String, Number, Object]),
+    __metadata("design:returntype", Promise)
+], SystemDatabase.prototype, "rewindWorkflow", null);
 __decorate([
     dbRetry()
     // A missing row normally means the workflow has not been inserted yet, so