From 949ae2f4034b7eacde612e20c5dc5e9619ca7c45 Mon Sep 17 00:00:00 2001 From: Harry Pierson Date: Wed, 7 Oct 2026 09:39:00 -0700 Subject: [PATCH 1/5] Add workflow rewind, with system-database migration 121 Rewind re-runs a terminal workflow in place, under the same ID, from a start step. Steps before the cut replay from their checkpoints; everything from the cut on is discarded and runs again. A non-terminal workflow is refused ("cancel it first"). In one system-database transaction, from the start step on, a rewind: - rolls each event back to its last value published before the cut, from `workflow_events_history`, then deletes the history rows; - deletes the step checkpoints in `operation_outputs`; - deletes the streams' close sentinels, so the replay can write to them again (stream entries are kept); - deletes the messages consumed by discarded steps, and any not yet consumed; - deletes the `workflow_output` row and clears the legacy output and error columns; - re-enqueues the workflow (internal queue, or the given queue and partition) with recovery attempts, deadline, dedup ID and timestamps reset and `owner_xid` cleared, restamping the application version only when one is given; - re-checks the status it read, failing with "retry the rewind" if it changed. Before the system database, `DBOS.rewindWorkflow` deletes the checkpoints transactional step factories keep in `tx_step_outputs` from the cut on. If that fails, the workflow keeps its terminal status and is not re-enqueued, and the rewind can be retried. Called from a workflow, the rewind is a step. - **API:** `DBOS.rewindWorkflow` and `DBOSClient.rewindWorkflow`, with `RewindOptions` (queue, partition key, application version). A client cannot reach the application databases, so it leaves the step factories' checkpoints in place. - **Migration 121** adds `notifications.consumed_by_function_id`, and `recv` records the consuming step there. - **Conductor** answers `rewind_workflow`. --- .../transact/jdbi/JdbiStepFactoryTest.java | 17 + .../transact/jooq/JooqStepFactoryTest.java | 17 + .../TransactionalStepJdbcIntegrationTest.java | 27 + .../src/main/java/dev/dbos/transact/DBOS.java | 47 ++ .../java/dev/dbos/transact/DBOSClient.java | 39 + .../dbos/transact/conductor/Conductor.java | 20 + .../conductor/protocol/BaseMessage.java | 1 + .../conductor/protocol/MessageType.java | 1 + .../protocol/RewindWorkflowRequest.java | 36 + .../transact/database/SystemDatabase.java | 10 + .../database/dao/NotificationsDAO.java | 13 +- .../transact/database/dao/StreamsDAO.java | 23 + .../transact/database/dao/WorkflowDAO.java | 216 ++++++ .../dbos/transact/execution/DBOSExecutor.java | 28 + .../transact/migrations/MigrationManager.java | 11 +- .../dbos/transact/workflow/RewindOptions.java | 61 ++ .../transact/conductor/ConductorTest.java | 113 +++ .../migrations/MigrationManagerTest.java | 41 +- .../dbos/transact/workflow/RewindTest.java | 703 ++++++++++++++++++ .../transact/workflow/RewindTestService.java | 254 +++++++ 20 files changed, 1670 insertions(+), 8 deletions(-) create mode 100644 transact/src/main/java/dev/dbos/transact/conductor/protocol/RewindWorkflowRequest.java create mode 100644 transact/src/main/java/dev/dbos/transact/workflow/RewindOptions.java create mode 100644 transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java create mode 100644 transact/src/test/java/dev/dbos/transact/workflow/RewindTestService.java diff --git a/transact-jdbi-step-factory/src/test/java/dev/dbos/transact/jdbi/JdbiStepFactoryTest.java b/transact-jdbi-step-factory/src/test/java/dev/dbos/transact/jdbi/JdbiStepFactoryTest.java index 82d81cb7d..c719a058c 100644 --- a/transact-jdbi-step-factory/src/test/java/dev/dbos/transact/jdbi/JdbiStepFactoryTest.java +++ b/transact-jdbi-step-factory/src/test/java/dev/dbos/transact/jdbi/JdbiStepFactoryTest.java @@ -254,6 +254,23 @@ public void testDeleteCheckpointsThroughDbos() throws Exception { assertEquals(0, DBUtils.getTxStepRows(dataSource, "wf-batch-2").size()); } + @Test + public void testRewindRerunsTheTransaction() throws Exception { + var wfid = "wf-rewind"; + var user = "rewindUser"; + try (var _o = new WorkflowOptions(wfid).setContext()) { + assertEquals(1, proxy.insertWorkflow(user).greetCount()); + } + + // The rewind deletes the step's checkpoint, so the transaction runs again rather than + // replaying its recorded result. + WorkflowHandle handle = + dbos.rewindWorkflow(wfid, 0); + assertEquals(new FactoryTestService.TestResult(user, 2), handle.getResult()); + assertEquals(2, getGreetCount(user)); + assertEquals(1, DBUtils.getTxStepRows(dataSource, wfid).size()); + } + @Test public void testInsert() throws Exception { var wfid = "wf1"; diff --git a/transact-jooq-step-factory/src/test/java/dev/dbos/transact/jooq/JooqStepFactoryTest.java b/transact-jooq-step-factory/src/test/java/dev/dbos/transact/jooq/JooqStepFactoryTest.java index c0dc12993..65b99ee97 100644 --- a/transact-jooq-step-factory/src/test/java/dev/dbos/transact/jooq/JooqStepFactoryTest.java +++ b/transact-jooq-step-factory/src/test/java/dev/dbos/transact/jooq/JooqStepFactoryTest.java @@ -244,6 +244,23 @@ public void testDeleteCheckpointsThroughDbos() throws Exception { assertEquals(0, DBUtils.getTxStepRows(dataSource, "wf-batch-2").size()); } + @Test + public void testRewindRerunsTheTransaction() throws Exception { + var wfid = "wf-rewind"; + var user = "rewindUser"; + try (var _o = new WorkflowOptions(wfid).setContext()) { + assertEquals(1, proxy.insertWorkflow(user).greetCount()); + } + + // The rewind deletes the step's checkpoint, so the transaction runs again rather than + // replaying its recorded result. + WorkflowHandle handle = + dbos.rewindWorkflow(wfid, 0); + assertEquals(new FactoryTestService.TestResult(user, 2), handle.getResult()); + assertEquals(2, getGreetCount(user)); + assertEquals(1, DBUtils.getTxStepRows(dataSource, wfid).size()); + } + @Test public void testInsert() throws Exception { var wfid = "wf1"; diff --git a/transact-spring-txstep-starter/src/test/java/dev/dbos/transact/spring/txstep/TransactionalStepJdbcIntegrationTest.java b/transact-spring-txstep-starter/src/test/java/dev/dbos/transact/spring/txstep/TransactionalStepJdbcIntegrationTest.java index 2777fadd0..0cefe8124 100644 --- a/transact-spring-txstep-starter/src/test/java/dev/dbos/transact/spring/txstep/TransactionalStepJdbcIntegrationTest.java +++ b/transact-spring-txstep-starter/src/test/java/dev/dbos/transact/spring/txstep/TransactionalStepJdbcIntegrationTest.java @@ -229,6 +229,33 @@ void deleteCheckpointsThroughDbos() throws SQLException { } } + @Test + void rewindRerunsTheTransactionalStep() throws SQLException { + try (var db = new TransactionalStepTest.TestDatabase()) { + runner(db) + .run( + ctx -> { + assertThat(ctx).hasNotFailed(); + var workflow = ctx.getBean(OrderWorkflowService.class); + var dbos = ctx.getBean(DBOS.class); + var wfid = "wf-jdbc-int-rewind"; + + try (var _o = new WorkflowOptions(wfid).setContext()) { + workflow.processOrder("ord-r", "Widget", 1); + } + assertThat(TransactionalStepTest.getTxRows(db.dataSource, wfid)).hasSize(1); + + // Remove the order the first run placed: a step that runs again places it again, + // while one replayed from a leftover checkpoint does not. + new JdbcTemplate(db.dataSource).update("DELETE FROM orders WHERE id = ?", "ord-r"); + dbos.rewindWorkflow(wfid, 0).getResult(); + + assertThat(orderCount(db.dataSource, "ord-r")).isEqualTo(1); + assertThat(TransactionalStepTest.getTxRows(db.dataSource, wfid)).hasSize(1); + }); + } + } + @Test void isolationLevel() { try (var db = new TransactionalStepTest.TestDatabase()) { diff --git a/transact/src/main/java/dev/dbos/transact/DBOS.java b/transact/src/main/java/dev/dbos/transact/DBOS.java index 2aa6966c6..72bf4b8a4 100644 --- a/transact/src/main/java/dev/dbos/transact/DBOS.java +++ b/transact/src/main/java/dev/dbos/transact/DBOS.java @@ -18,6 +18,7 @@ import dev.dbos.transact.workflow.Queue; import dev.dbos.transact.workflow.QueueConflictResolution; import dev.dbos.transact.workflow.QueueOptions; +import dev.dbos.transact.workflow.RewindOptions; import dev.dbos.transact.workflow.ScheduleStatus; import dev.dbos.transact.workflow.SendMessage; import dev.dbos.transact.workflow.SerializationStrategy; @@ -1044,6 +1045,52 @@ public void updateWorkflowAttributes( return forkWorkflow(workflowId, startStep, new ForkOptions()); } + /** + * Rewind a workflow: re-run it in place, under the same ID, from the step provided. Steps before + * {@code startStep} are replayed from their checkpoints; everything from {@code startStep} on is + * discarded and runs again. Only a workflow in a terminal state can be rewound, so cancel a + * running one first. + * + *

From {@code startStep} on, the rewind deletes the workflow's step checkpoints, including + * those of registered transactional step factories; rolls back the events it published to their + * last value from before the cut; deletes the messages it consumed, and any it has not consumed + * yet; and removes its streams' close markers so the replay can write to them again. Stream + * entries are kept. The workflow is then re-enqueued. + * + *

Called from a workflow, the rewind is a step, so a recovered caller does not rewind its + * target again. + * + * @param Return type of the workflow function + * @param Checked exception thrown by the workflow function, if any + * @param workflowId ID of the workflow to rewind + * @param startStep the first step to discard and run again; 0 re-runs the whole workflow + * @param options {@link RewindOptions} containing the queue, partition key and application + * version to re-enqueue the workflow with + * @return handle to the rewound workflow + * @throws dev.dbos.transact.exceptions.DBOSNonExistentWorkflowException if the workflow does not + * exist + * @throws IllegalStateException if the workflow is not in a terminal state + */ + public @NonNull WorkflowHandle rewindWorkflow( + @NonNull String workflowId, int startStep, @NonNull RewindOptions options) { + return ensureLaunched("rewindWorkflow").rewindWorkflow(workflowId, startStep, options); + } + + /** + * Rewind a workflow: re-run it in place, under the same ID, from the step provided. See {@link + * #rewindWorkflow(String, int, RewindOptions)}. + * + * @param Return type of the workflow function + * @param Checked exception thrown by the workflow function, if any + * @param workflowId ID of the workflow to rewind + * @param startStep the first step to discard and run again; 0 re-runs the whole workflow + * @return handle to the rewound workflow + */ + public @NonNull WorkflowHandle rewindWorkflow( + @NonNull String workflowId, int startStep) { + return rewindWorkflow(workflowId, startStep, new RewindOptions()); + } + /** * List all registered application versions, ordered by timestamp descending. * diff --git a/transact/src/main/java/dev/dbos/transact/DBOSClient.java b/transact/src/main/java/dev/dbos/transact/DBOSClient.java index 39ba2e8c0..5f86dd6fc 100644 --- a/transact/src/main/java/dev/dbos/transact/DBOSClient.java +++ b/transact/src/main/java/dev/dbos/transact/DBOSClient.java @@ -22,6 +22,7 @@ import dev.dbos.transact.workflow.Queue; import dev.dbos.transact.workflow.QueueConflictResolution; import dev.dbos.transact.workflow.QueueOptions; +import dev.dbos.transact.workflow.RewindOptions; import dev.dbos.transact.workflow.ScheduleStatus; import dev.dbos.transact.workflow.SendMessage; import dev.dbos.transact.workflow.SerializationStrategy; @@ -1619,6 +1620,44 @@ public void deleteWorkflows(@NonNull List workflowIds, boolean deleteChi return retrieveWorkflow(forkedWorkflowId); } + /** + * Rewind a workflow: re-run it in place, under the same ID, from the step provided. Only a + * workflow in a terminal state can be rewound. See {@link DBOS#rewindWorkflow(String, int, + * RewindOptions)} for what a rewind discards. + * + *

A client cannot reach the application's databases, so unlike {@link DBOS#rewindWorkflow}, it + * does not delete the checkpoints transactional step factories keep there. A transactional step + * past the cut that still has one replays its recorded result instead of running again. To re-run + * those transactions, rewind from within the application. + * + * @param Type of the workflow's return value + * @param Type of any checked exception thrown by the workflow + * @param workflowId ID of the workflow to rewind + * @param startStep the first step to discard and run again; 0 re-runs the whole workflow + * @param options Options for the rewind + * @return `WorkflowHandle` for the rewound workflow + */ + public @NonNull WorkflowHandle rewindWorkflow( + @NonNull String workflowId, int startStep, @NonNull RewindOptions options) { + systemDatabase.rewindWorkflow(workflowId, startStep, options); + return retrieveWorkflow(workflowId); + } + + /** + * Rewind a workflow: re-run it in place, under the same ID, from the step provided. See {@link + * #rewindWorkflow(String, int, RewindOptions)}. + * + * @param Type of the workflow's return value + * @param Type of any checked exception thrown by the workflow + * @param workflowId ID of the workflow to rewind + * @param startStep the first step to discard and run again; 0 re-runs the whole workflow + * @return `WorkflowHandle` for the rewound workflow + */ + public @NonNull WorkflowHandle rewindWorkflow( + @NonNull String workflowId, int startStep) { + return rewindWorkflow(workflowId, startStep, new RewindOptions()); + } + /** * Get the status of a workflow * diff --git a/transact/src/main/java/dev/dbos/transact/conductor/Conductor.java b/transact/src/main/java/dev/dbos/transact/conductor/Conductor.java index a36b3d22b..4d833450c 100644 --- a/transact/src/main/java/dev/dbos/transact/conductor/Conductor.java +++ b/transact/src/main/java/dev/dbos/transact/conductor/Conductor.java @@ -805,6 +805,7 @@ CompletableFuture getResponseAsync(BaseMessage message, WebSocket case RESUME -> handleResume(this, (ResumeRequest) message); case RESUME_SCHEDULE -> handleResumeSchedule(this, (ResumeScheduleRequest) message); case RETENTION -> handleRetention(this, (RetentionRequest) message); + case REWIND_WORKFLOW -> handleRewind(this, (RewindWorkflowRequest) message); case SET_LATEST_APPLICATION_VERSION -> handleSetLatestApplicationVersion(this, (SetLatestApplicationVersionRequest) message); case TRIGGER_SCHEDULE -> handleTriggerSchedule(this, (TriggerScheduleRequest) message); @@ -928,6 +929,25 @@ static CompletableFuture handleFork( }); } + static CompletableFuture handleRewind( + Conductor conductor, RewindWorkflowRequest request) { + return CompletableFuture.supplyAsync( + () -> { + if (request.body == null || request.body.workflow_id == null) { + return new SuccessResponse( + request, new IllegalArgumentException("Invalid Rewind Workflow Request")); + } + try { + conductor.dbosExecutor.rewindWorkflow( + request.body.workflow_id, request.startStep(), request.toOptions()); + return new SuccessResponse(request, true); + } catch (Exception e) { + logger.error("Exception encountered when rewinding workflow {}", request, e); + return new SuccessResponse(request, e); + } + }); + } + static CompletableFuture handleForkFromFailure( Conductor conductor, ForkFromFailureRequest request) { return CompletableFuture.supplyAsync( diff --git a/transact/src/main/java/dev/dbos/transact/conductor/protocol/BaseMessage.java b/transact/src/main/java/dev/dbos/transact/conductor/protocol/BaseMessage.java index b8613a4f7..d37558b17 100644 --- a/transact/src/main/java/dev/dbos/transact/conductor/protocol/BaseMessage.java +++ b/transact/src/main/java/dev/dbos/transact/conductor/protocol/BaseMessage.java @@ -45,6 +45,7 @@ @JsonSubTypes.Type(value = ResumeRequest.class, name = "resume"), @JsonSubTypes.Type(value = ResumeScheduleRequest.class, name = "resume_schedule"), @JsonSubTypes.Type(value = RetentionRequest.class, name = "retention"), + @JsonSubTypes.Type(value = RewindWorkflowRequest.class, name = "rewind_workflow"), @JsonSubTypes.Type( value = SetLatestApplicationVersionRequest.class, name = "set_latest_application_version"), diff --git a/transact/src/main/java/dev/dbos/transact/conductor/protocol/MessageType.java b/transact/src/main/java/dev/dbos/transact/conductor/protocol/MessageType.java index 5849a8932..0ae44ad97 100644 --- a/transact/src/main/java/dev/dbos/transact/conductor/protocol/MessageType.java +++ b/transact/src/main/java/dev/dbos/transact/conductor/protocol/MessageType.java @@ -32,6 +32,7 @@ public enum MessageType { RESUME("resume"), RESUME_SCHEDULE("resume_schedule"), RETENTION("retention"), + REWIND_WORKFLOW("rewind_workflow"), SET_LATEST_APPLICATION_VERSION("set_latest_application_version"), TRIGGER_SCHEDULE("trigger_schedule"); diff --git a/transact/src/main/java/dev/dbos/transact/conductor/protocol/RewindWorkflowRequest.java b/transact/src/main/java/dev/dbos/transact/conductor/protocol/RewindWorkflowRequest.java new file mode 100644 index 000000000..c3c1a2e57 --- /dev/null +++ b/transact/src/main/java/dev/dbos/transact/conductor/protocol/RewindWorkflowRequest.java @@ -0,0 +1,36 @@ +package dev.dbos.transact.conductor.protocol; + +import dev.dbos.transact.workflow.RewindOptions; + +import com.fasterxml.jackson.annotation.JsonIgnoreProperties; + +public class RewindWorkflowRequest extends BaseMessage { + public RewindWorkflowBody body; + + public RewindWorkflowRequest() {} + + public RewindWorkflowRequest(String requestId, String workflowId, Integer startStep) { + this.type = MessageType.REWIND_WORKFLOW.getValue(); + this.request_id = requestId; + this.body = new RewindWorkflowBody(); + this.body.workflow_id = workflowId; + this.body.start_step = startStep; + } + + @JsonIgnoreProperties(ignoreUnknown = true) + public static class RewindWorkflowBody { + public String workflow_id; + public Integer start_step; // optional: omitted rewinds the whole history + public String application_version; // optional + public String queue_name; // optional + public String queue_partition_key; // optional + } + + public int startStep() { + return body.start_step == null ? 0 : body.start_step; + } + + public RewindOptions toOptions() { + return new RewindOptions(body.application_version, body.queue_name, body.queue_partition_key); + } +} diff --git a/transact/src/main/java/dev/dbos/transact/database/SystemDatabase.java b/transact/src/main/java/dev/dbos/transact/database/SystemDatabase.java index 819d9447d..318834818 100644 --- a/transact/src/main/java/dev/dbos/transact/database/SystemDatabase.java +++ b/transact/src/main/java/dev/dbos/transact/database/SystemDatabase.java @@ -30,6 +30,7 @@ import dev.dbos.transact.workflow.NotificationInfo; import dev.dbos.transact.workflow.Queue; import dev.dbos.transact.workflow.QueueOptions; +import dev.dbos.transact.workflow.RewindOptions; import dev.dbos.transact.workflow.ScheduleStatus; import dev.dbos.transact.workflow.SendMessage; import dev.dbos.transact.workflow.StepAggregateRow; @@ -1245,6 +1246,15 @@ public String forkWorkflow(String originalWorkflowId, int startStep, ForkOptions () -> WorkflowDAO.forkWorkflow(ctx, originalWorkflowId, startStep, options)); } + public void rewindWorkflow(String workflowId, int startStep, RewindOptions options) { + dbRetryIncludingSerializationError( + "rewindWorkflow", () -> WorkflowDAO.rewindWorkflow(ctx, workflowId, startStep, options)); + } + + public void checkRewindable(String workflowId, int startStep) { + dbRetry(() -> WorkflowDAO.checkRewindable(ctx, workflowId, startStep)); + } + public List forkFromFailure(List workflowIds, ForkFromFailureOptions options) { return dbRetryIncludingSerializationError( "forkFromFailure", () -> WorkflowDAO.forkFromFailure(ctx, workflowIds, options)); diff --git a/transact/src/main/java/dev/dbos/transact/database/dao/NotificationsDAO.java b/transact/src/main/java/dev/dbos/transact/database/dao/NotificationsDAO.java index 16f2edaa3..771c34262 100644 --- a/transact/src/main/java/dev/dbos/transact/database/dao/NotificationsDAO.java +++ b/transact/src/main/java/dev/dbos/transact/database/dao/NotificationsDAO.java @@ -291,7 +291,7 @@ public static Object recv( var updateSql = """ UPDATE "%1$s".notifications - SET consumed = TRUE + SET consumed = TRUE, consumed_by_function_id = ? WHERE destination_uuid = ? AND topic = ? AND consumed = FALSE @@ -321,10 +321,13 @@ public static Object recv( String serializedMessage = null; String serialization = null; try (PreparedStatement stmt = conn.prepareStatement(updateSql)) { - stmt.setString(1, workflowId); - stmt.setString(2, recvTopic); - stmt.setString(3, workflowId); - stmt.setString(4, recvTopic); + // The consuming step is recorded so a rewind can delete exactly the + // messages consumed by the steps it discards. + stmt.setInt(1, stepId); + stmt.setString(2, workflowId); + stmt.setString(3, recvTopic); + stmt.setString(4, workflowId); + stmt.setString(5, recvTopic); // Note, if there are two executors running the same workflow waiting on the // same recv, only the first one will return a row here. The second one gets diff --git a/transact/src/main/java/dev/dbos/transact/database/dao/StreamsDAO.java b/transact/src/main/java/dev/dbos/transact/database/dao/StreamsDAO.java index ce9af52b7..526bf89fa 100644 --- a/transact/src/main/java/dev/dbos/transact/database/dao/StreamsDAO.java +++ b/transact/src/main/java/dev/dbos/transact/database/dao/StreamsDAO.java @@ -102,6 +102,29 @@ private static void insertStream( } } + /** + * Deletes the close sentinels a workflow wrote from {@code fromStepId} on, so a rewound workflow + * can append to its streams again. A leftover sentinel would end every reader before the replay's + * entries. + */ + static void deleteCloseSentinels( + Connection conn, String schema, String workflowId, int fromStepId) throws SQLException { + var closed = SerializationUtil.serializeValue(STREAM_CLOSED_SENTINEL, "portable_json", null); + var sql = + """ + DELETE FROM "%s".streams + WHERE workflow_uuid = ? AND function_id >= ? AND value = ? AND serialization = ? + """ + .formatted(schema); + try (var stmt = conn.prepareStatement(sql)) { + stmt.setString(1, workflowId); + stmt.setInt(2, fromStepId); + stmt.setString(3, closed.serializedValue()); + stmt.setString(4, closed.serialization()); + stmt.executeUpdate(); + } + } + private static int getNextOffsetTx(Connection conn, String schema, String workflowId, String key) throws SQLException { String sql = diff --git a/transact/src/main/java/dev/dbos/transact/database/dao/WorkflowDAO.java b/transact/src/main/java/dev/dbos/transact/database/dao/WorkflowDAO.java index 8839eb152..7c417c1ac 100644 --- a/transact/src/main/java/dev/dbos/transact/database/dao/WorkflowDAO.java +++ b/transact/src/main/java/dev/dbos/transact/database/dao/WorkflowDAO.java @@ -28,6 +28,7 @@ import dev.dbos.transact.workflow.GetStepAggregatesInput; import dev.dbos.transact.workflow.GetWorkflowAggregatesInput; import dev.dbos.transact.workflow.ListWorkflowsInput; +import dev.dbos.transact.workflow.RewindOptions; import dev.dbos.transact.workflow.StepAggregateRow; import dev.dbos.transact.workflow.WorkflowAggregateRow; import dev.dbos.transact.workflow.WorkflowEvent; @@ -2121,6 +2122,221 @@ public static String forkWorkflow( return forkedWorkflowId; } + /** + * Drops a terminal workflow's history from {@code startStep} on and re-enqueues it under the same + * ID, so a replay re-executes everything from that step. Unlike a fork, this writes no new + * workflow: peers keep addressing the same ID. + * + *

In one transaction, from {@code startStep} on, it: + * + *

    + *
  • rolls each event published past the cut back to its last value from before the cut, or + * deletes it if it had none; + *
  • deletes the close sentinels of the streams, so the replay can append again. The stream + * entries themselves stay: their offsets are addresses peers read by; + *
  • deletes the step checkpoints and the event history; + *
  • deletes the messages the discarded steps consumed, and any still unconsumed; + *
  • deletes the workflow's outcome; + *
  • re-enqueues the workflow. + *
+ * + * @throws DBOSNonExistentWorkflowException if the workflow does not exist + * @throws IllegalStateException if the workflow is not in a terminal state, or changed status + * while being rewound + */ + public static void rewindWorkflow( + DbContext ctx, String workflowId, int startStep, RewindOptions options) throws SQLException { + // Function IDs start at 0, so 0 is the whole history. + if (startStep < 0) { + throw new IllegalArgumentException("startStep must be >= 0, got " + startStep); + } + Objects.requireNonNull(options, "RewindOptions must not be null"); + var schema = ctx.schema(); + + try (var txConn = ctx.getConnection()) { + SqlTransaction.run( + txConn, + conn -> { + var status = readRewindableStatus(conn, schema, workflowId); + + // Whether a key was published at or past the cut. %2$s is the key column the + // correlated subquery compares against. + var publishedPastCut = + """ + EXISTS ( + SELECT 1 FROM "%1$s".workflow_events_history discarded + WHERE discarded.workflow_uuid = ? AND discarded.key = %2$s + AND discarded.function_id >= ? + ) + """; + + // workflow_events_history is the undo log for workflow_events, so the events are + // rolled back before the history past the cut is deleted. First unpublish every key + // the discarded steps published... + var unpublishSql = + (""" + DELETE FROM "%1$s".workflow_events + WHERE workflow_uuid = ? AND """ + + publishedPastCut) + .formatted(schema, "\"%s\".workflow_events.key".formatted(schema)); + try (var stmt = conn.prepareStatement(unpublishSql)) { + stmt.setString(1, workflowId); + stmt.setString(2, workflowId); + stmt.setInt(3, startStep); + stmt.executeUpdate(); + } + + // ...then restore those keys to the last value published before the cut, if any. + var restoreSql = + (""" + INSERT INTO "%1$s".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 "%1$s".workflow_events_history weh + WHERE weh.workflow_uuid = ? AND weh.function_id < ? AND """ + + publishedPastCut + + """ + ) surviving + WHERE surviving.rn = 1 + """) + .formatted(schema, "weh.key"); + try (var stmt = conn.prepareStatement(restoreSql)) { + stmt.setString(1, workflowId); + stmt.setInt(2, startStep); + stmt.setString(3, workflowId); + stmt.setInt(4, startStep); + stmt.executeUpdate(); + } + + StreamsDAO.deleteCloseSentinels(conn, schema, workflowId, startStep); + + for (var table : List.of("operation_outputs", "workflow_events_history")) { + var deleteSql = + """ + DELETE FROM "%s".%s WHERE workflow_uuid = ? AND function_id >= ? + """ + .formatted(schema, table); + try (var stmt = conn.prepareStatement(deleteSql)) { + stmt.setString(1, workflowId); + stmt.setInt(2, startStep); + stmt.executeUpdate(); + } + } + + var notificationsSql = + """ + DELETE FROM "%s".notifications + WHERE destination_uuid = ? + AND (consumed_by_function_id >= ? OR consumed = FALSE) + """ + .formatted(schema); + try (var stmt = conn.prepareStatement(notificationsSql)) { + stmt.setString(1, workflowId); + stmt.setInt(2, startStep); + stmt.executeUpdate(); + } + + var outputSql = + """ + DELETE FROM "%s".workflow_output WHERE workflow_uuid = ? + """ + .formatted(schema); + try (var stmt = conn.prepareStatement(outputSql)) { + stmt.setString(1, workflowId); + stmt.executeUpdate(); + } + + // Re-enqueue. The legacy output and error columns are cleared too: reads fall back to + // them when there is no workflow_output row, so a value left there would be returned + // as the rewound workflow's result. Re-asserting the status read above keeps a + // workflow that moved on underneath this transaction from being resurrected. + var setVersion = + options.applicationVersion() != null ? ", application_version = ?" : ""; + var enqueueSql = + """ + UPDATE "%1$s".workflow_status + SET status = ?, owner_xid = NULL, queue_name = ?, queue_partition_key = ?, + 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 = %2$s%3$s + WHERE workflow_uuid = ? AND status = ? + """ + .formatted(schema, SystemDatabase.NOW_EPOCH_MS, setVersion); + try (var stmt = conn.prepareStatement(enqueueSql)) { + int i = 1; + stmt.setString(i++, WorkflowState.ENQUEUED.name()); + stmt.setString( + i++, + Objects.requireNonNullElse(options.queueName(), Constants.DBOS_INTERNAL_QUEUE)); + stmt.setString(i++, options.queuePartitionKey()); + if (options.applicationVersion() != null) { + stmt.setString(i++, options.applicationVersion()); + } + stmt.setString(i++, workflowId); + stmt.setString(i++, status); + if (stmt.executeUpdate() != 1) { + throw new IllegalStateException( + "Workflow %s changed status while being rewound; retry the rewind" + .formatted(workflowId)); + } + } + }); + } + } + + /** + * Checks that a workflow can be rewound, without changing anything: it must exist and be in a + * terminal state. A rewind repeats this check in its own transaction; this one lets a caller + * refuse before touching anything outside the system database. + */ + public static void checkRewindable(DbContext ctx, String workflowId, int startStep) + throws SQLException { + if (startStep < 0) { + throw new IllegalArgumentException("startStep must be >= 0, got " + startStep); + } + try (var conn = ctx.getConnection()) { + readRewindableStatus(conn, ctx.schema(), workflowId); + } + } + + private static final Set ACTIVE_STATUSES = + Set.of( + WorkflowState.PENDING.name(), + WorkflowState.ENQUEUED.name(), + WorkflowState.DELAYED.name()); + + private static String readRewindableStatus(Connection conn, String schema, String workflowId) + throws SQLException { + var sql = + """ + SELECT status FROM "%s".workflow_status WHERE workflow_uuid = ? + """ + .formatted(schema); + String status; + try (var stmt = conn.prepareStatement(sql)) { + stmt.setString(1, workflowId); + try (var rs = stmt.executeQuery()) { + if (!rs.next()) { + throw new DBOSNonExistentWorkflowException(workflowId); + } + status = rs.getString("status"); + } + } + // Compared as strings, so a status this SDK does not know (one a newer SDK wrote) counts as + // terminal rather than failing the rewind. + if (ACTIVE_STATUSES.contains(status)) { + throw new IllegalStateException( + ("Cannot rewind %s (%s): only a workflow in a terminal state can be rewound, so cancel it" + + " first") + .formatted(workflowId, status)); + } + return status; + } + public static List forkFromFailure( DbContext ctx, List workflowIds, ForkFromFailureOptions options) throws SQLException { diff --git a/transact/src/main/java/dev/dbos/transact/execution/DBOSExecutor.java b/transact/src/main/java/dev/dbos/transact/execution/DBOSExecutor.java index 96c715eff..c968f0ab3 100644 --- a/transact/src/main/java/dev/dbos/transact/execution/DBOSExecutor.java +++ b/transact/src/main/java/dev/dbos/transact/execution/DBOSExecutor.java @@ -41,6 +41,7 @@ import dev.dbos.transact.workflow.Queue; import dev.dbos.transact.workflow.QueueConflictResolution; import dev.dbos.transact.workflow.QueueOptions; +import dev.dbos.transact.workflow.RewindOptions; import dev.dbos.transact.workflow.ScheduleStatus; import dev.dbos.transact.workflow.SendMessage; import dev.dbos.transact.workflow.SerializationStrategy; @@ -1109,6 +1110,33 @@ public WorkflowHandle forkWorkflow( return retrieveWorkflow(forkedId); } + public WorkflowHandle rewindWorkflow( + String workflowId, int startStep, RewindOptions options) { + Objects.requireNonNull(workflowId); + Objects.requireNonNull(options); + + this.runDbosFunctionAsStep( + () -> { + logger.info("Rewinding workflow: {} to step: {}", workflowId, startStep); + + validateQueue(options.queueName(), options.queuePartitionKey()); + + // The step factories' checkpoints go first and the system database last. If a delete + // fails, the workflow keeps its terminal status and is not re-enqueued, so nothing runs + // against a partly cleared history, and since a delete is idempotent the rewind can + // simply be retried. The checks run first so a running workflow keeps its checkpoints. + if (hasStepCheckpointStores()) { + systemDatabase.checkRewindable(workflowId, startStep); + deleteStepCheckpoints(workflowId, startStep); + } + systemDatabase.rewindWorkflow(workflowId, startStep, options); + return null; // void + }, + "DBOS.rewindWorkflow", + null); + return retrieveWorkflow(workflowId); + } + public List> forkFromFailure( List workflowIds, ForkFromFailureOptions options) { diff --git a/transact/src/main/java/dev/dbos/transact/migrations/MigrationManager.java b/transact/src/main/java/dev/dbos/transact/migrations/MigrationManager.java index 3da1fc517..77b4796c5 100644 --- a/transact/src/main/java/dev/dbos/transact/migrations/MigrationManager.java +++ b/transact/src/main/java/dev/dbos/transact/migrations/MigrationManager.java @@ -577,7 +577,8 @@ public static List getMigrations( migration117(isCockroach), migration118(isCockroach), migration119(isCockroach), - migration120(isCockroach))); + migration120(isCockroach), + MIGRATION_121)); return migrations.stream().map(m -> m.formatted(schema)).toList(); } @@ -1731,4 +1732,12 @@ static String migration120(boolean isCockroach) { + concurrently(isCockroach) + " IF EXISTS \"%1$s\".\"idx_operation_outputs_completed_at_function_name\""; } + + // Migration 121: record which recv consumed a notification, so a rewind can delete the messages + // consumed by the steps it discards. + static final String MIGRATION_121 = + """ + ALTER TABLE "%1$s"."notifications" + ADD COLUMN IF NOT EXISTS "consumed_by_function_id" INTEGER; + """; } diff --git a/transact/src/main/java/dev/dbos/transact/workflow/RewindOptions.java b/transact/src/main/java/dev/dbos/transact/workflow/RewindOptions.java new file mode 100644 index 000000000..7cc451dff --- /dev/null +++ b/transact/src/main/java/dev/dbos/transact/workflow/RewindOptions.java @@ -0,0 +1,61 @@ +package dev.dbos.transact.workflow; + +import org.jspecify.annotations.NonNull; +import org.jspecify.annotations.Nullable; + +/** + * Options for rewinding a workflow. A rewound workflow keeps its ID and is re-enqueued, so these + * options say where it is enqueued and which application version runs it. + * + * @param applicationVersion Application version to run the rewound workflow on; {@code null} keeps + * the version the workflow already has + * @param queueName Queue to re-enqueue the workflow on; {@code null} uses the internal queue + * @param queuePartitionKey Partition key on that queue; {@code null} clears any key the workflow + * had + */ +public record RewindOptions( + @Nullable String applicationVersion, + @Nullable String queueName, + @Nullable String queuePartitionKey) { + + public RewindOptions() { + this(null, null, null); + } + + /** + * Returns a copy of this object with the given applicationVersion. + * + * @param applicationVersion Application version to run the rewound workflow on + */ + public RewindOptions withApplicationVersion(@Nullable String applicationVersion) { + return new RewindOptions(applicationVersion, this.queueName, this.queuePartitionKey); + } + + /** + * Returns a copy of this object with the given queue. + * + * @param queue Queue to re-enqueue the rewound workflow on + * @return a copy with the queue set + */ + public RewindOptions withQueue(@NonNull QueueName queue) { + return withQueue(queue.value()); + } + + /** + * Returns a copy of this object with the given queueName. + * + * @param queueName Queue name to re-enqueue the rewound workflow on + */ + public RewindOptions withQueue(@Nullable String queueName) { + return new RewindOptions(this.applicationVersion, queueName, this.queuePartitionKey); + } + + /** + * Returns a copy of this object with the given queuePartitionKey. + * + * @param queuePartitionKey Queue partition key to re-enqueue the rewound workflow with + */ + public RewindOptions withQueuePartitionKey(@Nullable String queuePartitionKey) { + return new RewindOptions(this.applicationVersion, this.queueName, queuePartitionKey); + } +} diff --git a/transact/src/test/java/dev/dbos/transact/conductor/ConductorTest.java b/transact/src/test/java/dev/dbos/transact/conductor/ConductorTest.java index 5fbb79721..ad4a9b140 100644 --- a/transact/src/test/java/dev/dbos/transact/conductor/ConductorTest.java +++ b/transact/src/test/java/dev/dbos/transact/conductor/ConductorTest.java @@ -38,6 +38,7 @@ import dev.dbos.transact.workflow.GetWorkflowAggregatesInput; import dev.dbos.transact.workflow.ListWorkflowsInput; import dev.dbos.transact.workflow.NotificationInfo; +import dev.dbos.transact.workflow.RewindOptions; import dev.dbos.transact.workflow.StepAggregateRow; import dev.dbos.transact.workflow.StepInfo; import dev.dbos.transact.workflow.VersionInfo; @@ -1607,6 +1608,118 @@ public void canListQueuedWorkflowsWithSingleStringValues() throws Exception { } } + @SuppressWarnings("unchecked") + @RetryingTest(3) + public void canRewind() throws Exception { + MessageListener listener = new MessageListener(); + testServer.setListener(listener); + String workflowId = "sample-wf-id"; + + var mockHandle = (WorkflowHandle) mock(WorkflowHandle.class); + when(mockHandle.workflowId()).thenReturn(workflowId); + when(mockExec.rewindWorkflow(eq(workflowId), anyInt(), any())).thenReturn(mockHandle); + + try (Conductor conductor = builder.build()) { + conductor.start(); + + assertTrue(listener.openLatch.await(5, TimeUnit.SECONDS), "open latch timed out"); + + Map body = + Map.of( + "workflow_id", + workflowId, + "start_step", + 3, + "application_version", + "appver-12345", + "queue_name", + "custom-queue", + "queue_partition_key", + "partition-key", + "unknown-field", + "unknown-field-value"); + listener.send(MessageType.REWIND_WORKFLOW, "12345", Map.of("body", body)); + + assertTrue(listener.messageLatch.await(1, TimeUnit.SECONDS), "message latch timed out"); + ArgumentCaptor optionsCaptor = ArgumentCaptor.forClass(RewindOptions.class); + verify(mockExec).rewindWorkflow(eq(workflowId), eq(3), optionsCaptor.capture()); + RewindOptions options = optionsCaptor.getValue(); + assertEquals("appver-12345", options.applicationVersion()); + assertEquals("custom-queue", options.queueName()); + assertEquals("partition-key", options.queuePartitionKey()); + + JsonNode jsonNode = mapper.readTree(listener.message); + assertEquals("rewind_workflow", jsonNode.get("type").stringValue()); + assertEquals("12345", jsonNode.get("request_id").stringValue()); + assertTrue(jsonNode.get("success").asBoolean()); + assertNull(jsonNode.get("error_message")); + } + } + + @SuppressWarnings("unchecked") + @RetryingTest(3) + public void canRewindWithoutAStartStep() throws Exception { + MessageListener listener = new MessageListener(); + testServer.setListener(listener); + String workflowId = "sample-wf-id"; + + var mockHandle = (WorkflowHandle) mock(WorkflowHandle.class); + when(mockExec.rewindWorkflow(eq(workflowId), anyInt(), any())).thenReturn(mockHandle); + + try (Conductor conductor = builder.build()) { + conductor.start(); + + assertTrue(listener.openLatch.await(5, TimeUnit.SECONDS), "open latch timed out"); + + listener.send( + MessageType.REWIND_WORKFLOW, "12345", Map.of("body", Map.of("workflow_id", workflowId))); + + assertTrue(listener.messageLatch.await(1, TimeUnit.SECONDS), "message latch timed out"); + // An omitted start step rewinds the whole history. + ArgumentCaptor optionsCaptor = ArgumentCaptor.forClass(RewindOptions.class); + verify(mockExec).rewindWorkflow(eq(workflowId), eq(0), optionsCaptor.capture()); + RewindOptions options = optionsCaptor.getValue(); + assertNull(options.applicationVersion()); + assertNull(options.queueName()); + assertNull(options.queuePartitionKey()); + + JsonNode jsonNode = mapper.readTree(listener.message); + assertTrue(jsonNode.get("success").asBoolean()); + } + } + + @RetryingTest(3) + public void canRewindThrow() throws Exception { + MessageListener listener = new MessageListener(); + testServer.setListener(listener); + String workflowId = "sample-wf-id"; + + String errorMessage = "canRewindThrow error"; + doThrow(new IllegalStateException(errorMessage)) + .when(mockExec) + .rewindWorkflow(eq(workflowId), anyInt(), any()); + + try (Conductor conductor = builder.build()) { + conductor.start(); + + assertTrue(listener.openLatch.await(5, TimeUnit.SECONDS), "open latch timed out"); + + listener.send( + MessageType.REWIND_WORKFLOW, + "12345", + Map.of("body", Map.of("workflow_id", workflowId, "start_step", 1))); + + assertTrue(listener.messageLatch.await(1, TimeUnit.SECONDS), "message latch timed out"); + verify(mockExec).rewindWorkflow(eq(workflowId), eq(1), any()); + + JsonNode jsonNode = mapper.readTree(listener.message); + assertEquals("rewind_workflow", jsonNode.get("type").stringValue()); + assertEquals("12345", jsonNode.get("request_id").stringValue()); + assertFalse(jsonNode.get("success").asBoolean()); + assertEquals(errorMessage, jsonNode.get("error_message").stringValue()); + } + } + @RetryingTest(3) public void canListQueuedWorkflowsWithHasParentAndIsFork() throws Exception { MessageListener listener = new MessageListener(); diff --git a/transact/src/test/java/dev/dbos/transact/migrations/MigrationManagerTest.java b/transact/src/test/java/dev/dbos/transact/migrations/MigrationManagerTest.java index 32591b241..43f870537 100644 --- a/transact/src/test/java/dev/dbos/transact/migrations/MigrationManagerTest.java +++ b/transact/src/test/java/dev/dbos/transact/migrations/MigrationManagerTest.java @@ -595,7 +595,7 @@ void testSharedMigrationsUpgradeADatabaseAtTheEndOfThisLanguagesHistory() throws var schema = Constants.DB_SCHEMA; var latest = MigrationManager.getMigrations(schema, true, PgContainer.USE_COCKROACH_DB).size(); - assertEquals(120, latest, "The shared history currently ends at migration 120"); + assertEquals(121, latest, "The shared history currently ends at migration 121"); // A database last migrated by a build that predates the shared base: the runner must walk the // padding between this language's own history and SHARED_MIGRATION_BASE without stalling. @@ -634,6 +634,8 @@ void testSharedMigrationsUpgradeADatabaseAtTheEndOfThisLanguagesHistory() throws assertIndexAbsent(conn, "idx_notifications"); assertIndexExists(conn, "idx_workflow_topic"); + assertColumnExists(conn, "notifications", "consumed_by_function_id"); + for (var index : REBUILT_INDEXES) { assertIndexExists(conn, index[1]); assertIndexAbsent(conn, index[0]); @@ -677,7 +679,7 @@ void testMigrations115To120_RebuildTheHotIndexes() throws Exception { MigrationManager.runMigrations(dbosConfig); try (var conn = dataSource.getConnection()) { - assertEquals(120, getVersion(conn)); + assertEquals(migrations.size(), getVersion(conn)); for (var index : REBUILT_INDEXES) { assertIndexExists(conn, index[1]); assertIndexAbsent(conn, index[0]); @@ -688,7 +690,34 @@ void testMigrations115To120_RebuildTheHotIndexes() throws Exception { // Re-running is a no-op. assertDoesNotThrow(() -> MigrationManager.runMigrations(dbosConfig)); try (var conn = dataSource.getConnection()) { + assertEquals(migrations.size(), getVersion(conn)); + } + } + + @Test + void testMigration121_AddsConsumedByFunctionId() throws Exception { + var schema = Constants.DB_SCHEMA; + var dbosConfig = pgContainer.dbosConfig(); + var useListenNotify = !PgContainer.USE_COCKROACH_DB; + var migrations = + MigrationManager.getMigrations(schema, useListenNotify, PgContainer.USE_COCKROACH_DB); + + // A database last migrated by an SDK that stopped at 120. + MigrationManager.createDatabaseIfNotExists( + pgContainer.jdbcUrl(), pgContainer.username(), pgContainer.password()); + try (var conn = dataSource.getConnection()) { + MigrationManager.ensureDbosSchema(conn, schema); + MigrationManager.ensureMigrationTable(conn, schema); + MigrationManager.runDbosMigrations(conn, schema, migrations.subList(0, 120)); assertEquals(120, getVersion(conn)); + assertColumnAbsent(conn, "notifications", "consumed_by_function_id"); + } + + MigrationManager.runMigrations(dbosConfig); + + try (var conn = dataSource.getConnection()) { + assertEquals(121, getVersion(conn)); + assertColumnExists(conn, "notifications", "consumed_by_function_id"); } } @@ -753,6 +782,14 @@ static void assertColumnExists(Connection conn, String tableName, String columnN } } + static void assertColumnAbsent(Connection conn, String tableName, String columnName) + throws Exception { + try (ResultSet rs = + conn.getMetaData().getColumns(null, Constants.DB_SCHEMA, tableName, columnName)) { + assertFalse(rs.next(), "Column %s.%s should not exist".formatted(tableName, columnName)); + } + } + static void assertTableExists(DatabaseMetaData metaData, String tableName) throws Exception { assertTableExists(metaData, tableName, Constants.DB_SCHEMA); } diff --git a/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java b/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java new file mode 100644 index 000000000..bab0084ef --- /dev/null +++ b/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java @@ -0,0 +1,703 @@ +package dev.dbos.transact.workflow; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import dev.dbos.transact.Constants; +import dev.dbos.transact.DBOS; +import dev.dbos.transact.DBOSTestAccess; +import dev.dbos.transact.StartWorkflowOptions; +import dev.dbos.transact.config.DBOSConfig; +import dev.dbos.transact.context.WorkflowOptions; +import dev.dbos.transact.exceptions.DBOSNonExistentWorkflowException; +import dev.dbos.transact.internal.StepCheckpointStore; +import dev.dbos.transact.json.SerializationUtil; +import dev.dbos.transact.txstep.JdbcStepFactory; +import dev.dbos.transact.utils.DBUtils; +import dev.dbos.transact.utils.DBUtils.EventHistoryRow; +import dev.dbos.transact.utils.PgContainer; +import dev.dbos.transact.utils.TxStepOutputRow; + +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Supplier; + +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; +import org.junit.jupiter.api.AutoClose; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class RewindTest { + + private static final String SCHEMA = Constants.DB_SCHEMA; + + @AutoClose final PgContainer pgContainer = new PgContainer(); + + DBOSConfig dbosConfig; + @AutoClose DBOS dbos; + @AutoClose HikariDataSource dataSource; + // A pool without autocommit: a checkpoint delete that relied on autocommit would be rolled back + // when its connection went back to this pool. + @AutoClose HikariDataSource noAutoCommitDataSource; + + private RewindTestServiceImpl impl; + private RewindTestService proxy; + private JdbcStepFactory first; + private JdbcStepFactory second; + + // When set, the extra checkpoint store registered in beforeEach fails its delete. + private final AtomicBoolean failCheckpointDelete = new AtomicBoolean(false); + + @BeforeEach + void beforeEach() throws SQLException { + dbosConfig = pgContainer.dbosConfig(); + dataSource = pgContainer.dataSource(); + dbos = new DBOS(dbosConfig); + + try (var conn = dataSource.getConnection(); + var stmt = conn.createStatement()) { + stmt.execute("DROP TABLE IF EXISTS rewind_rows"); + stmt.execute("CREATE TABLE rewind_rows (v TEXT NOT NULL)"); + } + var noAutoCommit = new HikariConfig(); + noAutoCommit.setJdbcUrl(pgContainer.jdbcUrl()); + noAutoCommit.setUsername(pgContainer.username()); + noAutoCommit.setPassword(pgContainer.password()); + noAutoCommit.setAutoCommit(false); + noAutoCommitDataSource = new HikariDataSource(noAutoCommit); + + // Two factories with separate checkpoint tables, standing in for two databases. + first = new JdbcStepFactory(dbos, dataSource, "rewind_first"); + second = new JdbcStepFactory(dbos, noAutoCommitDataSource, "rewind_second"); + dbos.integration() + .registerStepCheckpointStore( + new StepCheckpointStore() { + @Override + public void deleteCheckpoints(String workflowId, int fromStepId) throws SQLException { + if (failCheckpointDelete.get()) { + throw new SQLException("datasource down"); + } + } + + @Override + public void deleteCheckpoints(Collection workflowIds) {} + }); + + impl = new RewindTestServiceImpl(dbos); + proxy = dbos.registerProxy(RewindTestService.class, impl); + impl.setProxy(proxy); + impl.setStepFactories(first, second); + + dbos.launch(); + } + + // Runs a workflow to completion under a fresh ID and returns that ID. + private String start(Supplier workflow) { + var workflowId = UUID.randomUUID().toString(); + try (var o = new WorkflowOptions(workflowId).setContext()) { + workflow.get(); + } + return workflowId; + } + + private int stepIdOf(String workflowId, String functionName, int occurrence) { + var matches = + dbos.listWorkflowSteps(workflowId).stream() + .filter( + s -> + s.functionName().equals(functionName) + || s.functionName().endsWith("." + functionName)) + .map(StepInfo::functionId) + .toList(); + assertTrue( + matches.size() > occurrence, + "%s has %d %s steps".formatted(workflowId, matches.size(), functionName)); + return matches.get(occurrence); + } + + /** + * A single-slot queue whose only worker is held until the returned handle is closed. A rewind + * re-enqueues, so without this every assertion about what a rewind leaves behind races the queue + * picking the workflow back up. + */ + private AutoCloseable pausedQueue(String name) throws Exception { + dbos.registerQueue(name, new QueueOptions().withConcurrency(1)); + var started = new java.util.concurrent.CountDownLatch(1); + var released = new java.util.concurrent.CountDownLatch(1); + impl.blockerStarted = started; + impl.blockerReleased = released; + var handle = + dbos.startWorkflow(() -> proxy.blocker(), new StartWorkflowOptions().withQueue(name)); + assertTrue(started.await(10, TimeUnit.SECONDS), name + " blocker never started"); + return () -> { + released.countDown(); + handle.getResult(); + }; + } + + private record MailboxRow(Object message, boolean consumed, Integer consumedBy) {} + + private List mailbox(String workflowId) throws SQLException { + var sql = + """ + SELECT message, serialization, consumed, consumed_by_function_id + FROM "%s".notifications WHERE destination_uuid = ? ORDER BY created_at_epoch_ms + """ + .formatted(SCHEMA); + var rows = new ArrayList(); + try (var conn = dataSource.getConnection(); + var stmt = conn.prepareStatement(sql)) { + stmt.setString(1, workflowId); + try (var rs = stmt.executeQuery()) { + while (rs.next()) { + rows.add( + new MailboxRow( + SerializationUtil.deserializeValue( + rs.getString("message"), rs.getString("serialization"), null), + rs.getBoolean("consumed"), + rs.getObject("consumed_by_function_id", Integer.class))); + } + } + } + return rows; + } + + private List eventHistoryIds(String workflowId) throws SQLException { + return DBUtils.getWorkflowEventHistory(dataSource, workflowId).stream() + .map(EventHistoryRow::stepId) + .sorted() + .toList(); + } + + private List readStream(String workflowId, String key) { + var values = new ArrayList(); + dbos.readStream(workflowId, key).forEachRemaining(values::add); + return values; + } + + private List txCheckpoints(String schema, String workflowId) throws SQLException { + return DBUtils.getTxStepRows(dataSource, workflowId, schema).stream() + .map(TxStepOutputRow::stepId) + .sorted() + .toList(); + } + + private List tableRows() throws SQLException { + var rows = new ArrayList(); + try (var conn = dataSource.getConnection(); + var stmt = conn.createStatement(); + var rs = stmt.executeQuery("SELECT v FROM rewind_rows ORDER BY v")) { + while (rs.next()) { + rows.add(rs.getString(1)); + } + } + return rows; + } + + @Test + public void rewindReplaysTheStepsBeforeTheCut() throws Exception { + var workflowId = start(() -> proxy.fiveSteps("five")); + assertEquals(1, impl.stepRunsOf("three")); + + WorkflowHandle handle = dbos.rewindWorkflow(workflowId, 2); + assertEquals(workflowId, handle.workflowId()); + assertEquals("run2", handle.getResult()); + assertEquals(WorkflowState.SUCCESS, handle.getStatus().status()); + + // Steps 0 and 1 replayed from their checkpoints; everything from step 2 ran again. + assertEquals(1, impl.stepRunsOf("one")); + assertEquals(1, impl.stepRunsOf("two")); + assertEquals(2, impl.stepRunsOf("three")); + assertEquals(2, impl.stepRunsOf("four")); + assertEquals(2, impl.stepRunsOf("five")); + assertEquals(5, dbos.listWorkflowSteps(workflowId).size()); + + // Step 0 is the whole history. + assertEquals("run3", dbos.rewindWorkflow(workflowId, 0).getResult()); + assertEquals(2, impl.stepRunsOf("one")); + assertEquals(3, impl.stepRunsOf("five")); + } + + @Test + public void rewindDeletesNotifications() throws Exception { + var workflowId = UUID.randomUUID().toString(); + WorkflowHandle handle = + dbos.startWorkflow( + () -> proxy.receiver("partial-delete"), new StartWorkflowOptions(workflowId)); + dbos.send(workflowId, "a", "cmd"); + dbos.send(workflowId, "b", "cmd"); + assertEquals("ab:1", handle.getResult()); + + // Each recv stamped the row it took with its own step, which is what lets the rewind delete + // exactly the messages the discarded steps consumed. + var firstRecv = stepIdOf(workflowId, "DBOS.recv", 0); + var secondRecv = stepIdOf(workflowId, "DBOS.recv", 1); + // A message that arrives once the workflow is done sits unconsumed. + dbos.send(workflowId, "stray", "cmd"); + assertEquals( + List.of( + new MailboxRow("a", true, firstRecv), + new MailboxRow("b", true, secondRecv), + new MailboxRow("stray", false, null)), + mailbox(workflowId)); + + try (var q = pausedQueue("rewind_delete_gate")) { + dbos.rewindWorkflow( + workflowId, secondRecv, new RewindOptions().withQueue("rewind_delete_gate")); + // The message the discarded step took is gone, and so is the one still waiting. The first + // recv's message stays consumed: its step survived the cut. + assertEquals(List.of(new MailboxRow("a", true, firstRecv)), mailbox(workflowId)); + // A message that arrives after the cut is what the replayed recv gets. + dbos.send(workflowId, "c", "cmd"); + } + + assertEquals("ac:2", dbos.retrieveWorkflow(workflowId).getResult()); + assertEquals( + List.of(new MailboxRow("a", true, firstRecv), new MailboxRow("c", true, secondRecv)), + mailbox(workflowId)); + } + + @Test + public void rewindUnpublishesEvents() throws Exception { + var workflowId = start(() -> proxy.publisher("events")); + assertEquals( + Map.of("below", "kept", "both", "new", "above", "doomed"), dbos.getAllEvents(workflowId)); + + // Cut at the third setEvent, so "below" and the first "both" survive. + var cut = stepIdOf(workflowId, "DBOS.setEvent", 2); + + try (var q = pausedQueue("rewind_events_gate")) { + dbos.rewindWorkflow(workflowId, cut, new RewindOptions().withQueue("rewind_events_gate")); + // "below" was never touched past the cut; "both" reverts to its last value from below the + // cut; "above" was only ever published past the cut, so it is gone. + assertEquals(Map.of("below", "kept", "both", "old"), dbos.getAllEvents(workflowId)); + assertEquals(List.of(0, 1), eventHistoryIds(workflowId)); + + // And that is what a peer reading by key sees. + assertEquals( + "old", dbos.getEvent(workflowId, "both", java.time.Duration.ofSeconds(1)).orElseThrow()); + assertTrue( + dbos.getEvent(workflowId, "above", java.time.Duration.ofMillis(100)).isEmpty(), + "an event published only past the cut stays unpublished"); + } + + assertEquals("second", dbos.retrieveWorkflow(workflowId).getResult()); + assertEquals(Map.of("below", "kept", "both", "republished"), dbos.getAllEvents(workflowId)); + } + + @Test + public void rewindKeepsStreamEntries() throws Exception { + var workflowId = start(() -> proxy.streamWriter("stream-keep")); + assertEquals(List.of("a1", "b1"), readStream(workflowId, "log")); + + assertEquals("run2", dbos.rewindWorkflow(workflowId, 0).getResult()); + + // Offsets are addresses peers read by, so the discarded run's entries keep theirs and the + // replay appends. Deleting them would hand offset 0 a new value. + assertEquals(List.of("a1", "b1", "a2", "b2"), readStream(workflowId, "log")); + assertEquals( + List.of(0, 1, 2, 3), + DBUtils.getStreamEntries(dataSource, workflowId).stream() + .map(DBUtils.StreamRow::offset) + .toList()); + } + + @Test + public void rewindReopensAClosedStream() throws Exception { + var workflowId = start(() -> proxy.streamCloser("stream-close")); + assertEquals(List.of("v1"), readStream(workflowId, "out")); + + // The close marker ends every reader that reaches it, so one left over from the discarded run + // would hide the replay's entry. + assertEquals("run2", dbos.rewindWorkflow(workflowId, 0).getResult()); + assertEquals(List.of("v1", "v2"), readStream(workflowId, "out")); + + // Cut past the close, the marker is not the discarded run's to undo: its step survives, so + // nothing replays it and it has to stay. + assertEquals("run3", dbos.rewindWorkflow(workflowId, 2).getResult()); + var rows = DBUtils.getStreamEntries(dataSource, workflowId); + assertEquals(3, rows.size()); + assertTrue(rows.get(2).value().contains("__DBOS_STREAM_CLOSED__"), rows.toString()); + assertEquals(List.of("v1", "v2"), readStream(workflowId, "out")); + } + + @Test + public void rewindTheChildThenTheParentToRepairAFailure() throws Exception { + var workflowId = UUID.randomUUID().toString(); + WorkflowHandle handle = + dbos.startWorkflow(() -> proxy.parent("repair"), new StartWorkflowOptions(workflowId)); + var thrown = assertThrows(IllegalStateException.class, handle::getResult); + assertEquals("child is bogus", thrown.getMessage()); + + var childId = + dbos.listWorkflowSteps(workflowId).stream() + .map(StepInfo::childWorkflowId) + .filter(id -> id != null) + .findFirst() + .orElseThrow(); + assertEquals(WorkflowState.ERROR, dbos.retrieveWorkflow(childId).getStatus().status()); + assertEquals(WorkflowState.ERROR, dbos.retrieveWorkflow(workflowId).getStatus().status()); + + // Repair the child on its own first. + assertEquals(42, dbos.rewindWorkflow(childId, 0).getResult()); + + // Then rewind the parent to the getResult that failed. The step that started the child + // survives, and the replay picks up the repaired result. + var getResultStep = stepIdOf(workflowId, "DBOS.getResult", 0); + assertEquals( + 42, dbos.rewindWorkflow(workflowId, getResultStep).getResult()); + assertEquals(2, impl.childRuns.get()); + } + + @Test + public void aRewoundParentAdoptsItsExistingChild() throws Exception { + var workflowId = UUID.randomUUID().toString(); + WorkflowHandle handle = + dbos.startWorkflow( + () -> proxy.adoptingParent("adopt"), new StartWorkflowOptions(workflowId)); + assertEquals(43, handle.getResult()); + assertEquals(1, impl.doublerRuns.get()); + var childId = + dbos.listWorkflowSteps(workflowId).stream() + .map(StepInfo::childWorkflowId) + .filter(id -> id != null) + .findFirst() + .orElseThrow(); + + // The cut is before the step that started the child. The replay starts it again under the + // same ID, finds it already finished, and takes its result without running it. + assertEquals(44, dbos.rewindWorkflow(workflowId, 0).getResult()); + assertEquals(1, impl.doublerRuns.get()); + assertEquals(42, dbos.retrieveWorkflow(childId).getResult()); + } + + @Test + public void rewindOntoAQueueWithAPartitionKey() throws Exception { + dbos.registerQueue("rewind_partitioned", new QueueOptions().withPartitionConcurrency(1)); + + var workflowId = UUID.randomUUID().toString(); + WorkflowHandle handle = + dbos.startWorkflow( + () -> proxy.counter("partition"), + new StartWorkflowOptions(workflowId) + .withQueue("rewind_partitioned") + .withQueuePartitionKey("original")); + assertEquals(1, handle.getResult()); + + dbos.rewindWorkflow( + workflowId, + 0, + new RewindOptions().withQueue("rewind_partitioned").withQueuePartitionKey("repaired")); + assertEquals(2, dbos.retrieveWorkflow(workflowId).getResult()); + var status = dbos.retrieveWorkflow(workflowId).getStatus(); + assertEquals("rewind_partitioned", status.queueName()); + assertEquals("repaired", status.queuePartitionKey()); + + // Omitting the key clears it; omitting the queue falls back to the internal queue. + dbos.rewindWorkflow(workflowId, 0); + assertEquals(3, dbos.retrieveWorkflow(workflowId).getResult()); + status = dbos.retrieveWorkflow(workflowId).getStatus(); + assertEquals(Constants.DBOS_INTERNAL_QUEUE, status.queueName()); + assertNull(status.queuePartitionKey()); + + // A partition key needs a partitioned queue, and nothing is written when it is refused. + assertThrows( + IllegalArgumentException.class, + () -> dbos.rewindWorkflow(workflowId, 0, new RewindOptions().withQueuePartitionKey("pk"))); + assertEquals(WorkflowState.SUCCESS, dbos.retrieveWorkflow(workflowId).getStatus().status()); + } + + @Test + public void rewindOntoADifferentApplicationVersion() throws Exception { + var workflowId = start(() -> proxy.counter("version")); + var runningVersion = DBUtils.getWorkflowRow(dataSource, workflowId).applicationVersion(); + + // Dequeueing matches on application version, so a workflow restamped with a version nothing + // is running stays enqueued instead of replaying. + dbos.rewindWorkflow( + workflowId, 0, new RewindOptions().withApplicationVersion("not-this-deployment")); + var row = DBUtils.getWorkflowRow(dataSource, workflowId); + assertEquals("not-this-deployment", row.applicationVersion()); + Thread.sleep(2500); // several queue polls, any of which would pick it up + assertEquals( + WorkflowState.ENQUEUED.name(), DBUtils.getWorkflowRow(dataSource, workflowId).status()); + assertEquals(1, impl.runsOf("version")); + + // The workflow is ENQUEUED now, and only a terminal workflow can be rewound. + var refused = + assertThrows( + IllegalStateException.class, + () -> + dbos.rewindWorkflow( + workflowId, 0, new RewindOptions().withApplicationVersion(runningVersion))); + assertTrue(refused.getMessage().contains("only a workflow in a terminal state")); + dbos.cancelWorkflow(workflowId); + + // Restamped with the version this executor runs, it replays. + dbos.rewindWorkflow(workflowId, 0, new RewindOptions().withApplicationVersion(runningVersion)); + assertEquals(2, dbos.retrieveWorkflow(workflowId).getResult()); + + // Omitted, the workflow keeps the version it already had. + assertEquals(3, dbos.rewindWorkflow(workflowId, 0).getResult()); + assertEquals( + runningVersion, DBUtils.getWorkflowRow(dataSource, workflowId).applicationVersion()); + } + + @Test + public void databaseStateBetweenRewindAndReplay() throws Exception { + var workflowId = start(() -> proxy.publisher("dbstate")); + var before = DBUtils.getWorkflowRow(dataSource, workflowId); + assertEquals(WorkflowState.SUCCESS.name(), before.status()); + assertNotNull(before.completedAt()); + + // An outcome written by an SDK that predates the workflow_output table sits in the legacy + // columns, and reads fall back to it. Move this one there. + try (var conn = dataSource.getConnection()) { + try (var stmt = + conn.prepareStatement( + """ + UPDATE "%1$s".workflow_status + SET output = (SELECT output FROM "%1$s".workflow_output WHERE workflow_uuid = ?) + WHERE workflow_uuid = ? + """ + .formatted(SCHEMA))) { + stmt.setString(1, workflowId); + stmt.setString(2, workflowId); + assertEquals(1, stmt.executeUpdate()); + } + try (var stmt = + conn.prepareStatement( + "DELETE FROM \"%s\".workflow_output WHERE workflow_uuid = ?".formatted(SCHEMA))) { + stmt.setString(1, workflowId); + assertEquals(1, stmt.executeUpdate()); + } + } + assertEquals("first", dbos.retrieveWorkflow(workflowId).getStatus().output()); + + try (var q = pausedQueue("rewind_gate")) { + dbos.rewindWorkflow( + workflowId, 0, new RewindOptions().withQueue("rewind_gate").withQueuePartitionKey(null)); + + var after = DBUtils.getWorkflowRow(dataSource, workflowId); + assertEquals(WorkflowState.ENQUEUED.name(), after.status()); + assertEquals("rewind_gate", after.queueName()); + assertEquals(0L, after.recoveryAttempts()); + assertNull(after.startedAtEpochMs()); + assertNull(after.completedAt()); + assertNull(after.deadlineEpochMs()); + assertNull(after.deduplicationId()); + // The workflow's identity is untouched. + assertEquals(before.workflowName(), after.workflowName()); + assertEquals(before.createdAt(), after.createdAt()); + assertEquals(before.inputs(), after.inputs()); + + // Both outcome shapes are gone: no workflow_output row, and no legacy columns to fall + // back to. + assertNull(after.output()); + assertNull(after.error()); + assertNull(dbos.retrieveWorkflow(workflowId).getStatus().output()); + + assertEquals(List.of(), DBUtils.getStepRows(dataSource, workflowId)); + assertEquals(List.of(), eventHistoryIds(workflowId)); + assertEquals(Map.of(), dbos.getAllEvents(workflowId)); + } + + assertEquals("second", dbos.retrieveWorkflow(workflowId).getResult()); + } + + @Test + public void rewindRefusals() throws Exception { + assertThrows( + DBOSNonExistentWorkflowException.class, + () -> dbos.rewindWorkflow(UUID.randomUUID().toString(), 0)); + + var workflowId = start(() -> proxy.counter("validation")); + var refused = + assertThrows(IllegalArgumentException.class, () -> dbos.rewindWorkflow(workflowId, -1)); + assertTrue(refused.getMessage().contains("must be >= 0"), refused.getMessage()); + assertThrows( + IllegalArgumentException.class, + () -> dbos.rewindWorkflow(workflowId, 0, new RewindOptions().withQueue("no-such-queue"))); + + // Nothing was written. + assertEquals(WorkflowState.SUCCESS, dbos.retrieveWorkflow(workflowId).getStatus().status()); + assertEquals(1, impl.runsOf("validation")); + } + + @Test + public void theSystemDatabaseRefusesAnActiveWorkflow() throws Exception { + // The client never consults the step factories, so this reaches the system database's own + // status check rather than the pre-check the in-process rewind makes when factories exist. + var started = new java.util.concurrent.CountDownLatch(1); + var released = new java.util.concurrent.CountDownLatch(1); + impl.blockerStarted = started; + impl.blockerReleased = released; + WorkflowHandle handle = dbos.startWorkflow(() -> proxy.blocker()); + assertTrue(started.await(10, TimeUnit.SECONDS)); + try (var client = pgContainer.dbosClient()) { + var refused = + assertThrows( + IllegalStateException.class, () -> client.rewindWorkflow(handle.workflowId(), 0)); + assertTrue(refused.getMessage().contains("only a workflow in a terminal state")); + assertTrue(refused.getMessage().contains("PENDING"), refused.getMessage()); + } finally { + released.countDown(); + } + assertEquals("held", handle.getResult()); + } + + @Test + public void aStatusThisSdkDoesNotKnowCountsAsTerminal() throws Exception { + // A status a newer SDK wrote is not one of the active ones, so the rewind goes ahead. + var workflowId = start(() -> proxy.counter("future-status")); + DBUtils.setWorkflowState(dataSource, workflowId, "SOME_FUTURE_STATUS"); + assertEquals(2, dbos.rewindWorkflow(workflowId, 0).getResult()); + } + + @Test + public void clientRewind() throws Exception { + var workflowId = start(() -> proxy.counter("client")); + try (var client = pgContainer.dbosClient()) { + assertEquals(2, client.rewindWorkflow(workflowId, 0).getResult()); + assertEquals( + 3, + client + .rewindWorkflow( + workflowId, 0, new RewindOptions().withQueue(Constants.DBOS_INTERNAL_QUEUE)) + .getResult()); + assertThrows( + DBOSNonExistentWorkflowException.class, + () -> client.rewindWorkflow(UUID.randomUUID().toString(), 0)); + } + } + + @Test + public void rewindFromInsideAWorkflowIsCheckpointed() throws Exception { + var targetId = start(() -> proxy.counter("repaired")); + assertEquals(1, impl.runsOf("repaired")); + + var repairerId = start(() -> proxy.repairer(targetId)); + assertEquals(2, dbos.retrieveWorkflow(repairerId).getResult()); + assertEquals(2, impl.runsOf("repaired")); + stepIdOf(repairerId, "DBOS.rewindWorkflow", 0); + + // Crash and recover the repairer with its checkpoints intact. A second rewind would either + // run the target again or be refused, because the first left it ENQUEUED. + DBUtils.setWorkflowState(dataSource, repairerId, WorkflowState.PENDING.name()); + var executor = DBOSTestAccess.getDbosExecutor(dbos); + var recovered = executor.recoverPendingWorkflows(List.of(executor.executorId())); + assertTrue(recovered.contains(repairerId), "repairer was not recovered"); + assertEquals(2, dbos.retrieveWorkflow(repairerId).getResult()); + + // The repairer's body ran again, but the rewind did not. + assertEquals(2, impl.runsOf("repairer")); + assertEquals(2, impl.runsOf("repaired")); + } + + @Test + public void rewindFromInsideAStepRunsUnderTheStepsCheckpoint() throws Exception { + var targetId = start(() -> proxy.counter("rewound-in-step")); + var callerId = start(() -> proxy.stepRewinder(targetId)); + assertEquals(2, dbos.retrieveWorkflow(targetId).getResult()); + + // Inside a step the rewind runs directly, with no checkpoint of its own: the caller's only + // step is the one that made the call. + var steps = dbos.listWorkflowSteps(callerId); + assertEquals(1, steps.size(), steps.toString()); + assertTrue(steps.get(0).functionName().endsWith("rewindInStep"), steps.toString()); + + // That step's checkpoint is what keeps a recovered caller from rewinding the target again. + DBUtils.setWorkflowState(dataSource, callerId, WorkflowState.PENDING.name()); + var executor = DBOSTestAccess.getDbosExecutor(dbos); + var recovered = executor.recoverPendingWorkflows(List.of(executor.executorId())); + assertTrue(recovered.contains(callerId), "caller was not recovered"); + assertEquals(targetId, dbos.retrieveWorkflow(callerId).getResult()); + assertEquals(2, impl.runsOf("stepRewinder")); + assertEquals(2, impl.runsOf("rewound-in-step")); + } + + @Test + public void rewindDropsStepFactoryCheckpointsPastTheCut() throws Exception { + var workflowId = start(() -> sneaky(() -> proxy.txWriter("tx"))); + assertEquals(List.of(0, 2), txCheckpoints("rewind_first", workflowId)); + assertEquals(List.of(1, 3), txCheckpoints("rewind_second", workflowId)); + + // A step the system database would reject never reaches the checkpoints. + assertThrows(IllegalArgumentException.class, () -> dbos.rewindWorkflow(workflowId, -1)); + assertEquals(List.of(0, 2), txCheckpoints("rewind_first", workflowId)); + + // Cut at the third step: each factory keeps one checkpoint and loses one. + try (var q = pausedQueue("rewind_tx_gate")) { + dbos.rewindWorkflow(workflowId, 2, new RewindOptions().withQueue("rewind_tx_gate")); + assertEquals(List.of(0), txCheckpoints("rewind_first", workflowId)); + assertEquals(List.of(1), txCheckpoints("rewind_second", workflowId)); + } + + assertEquals(2, dbos.retrieveWorkflow(workflowId).getResult()); + assertEquals(List.of(0, 2), txCheckpoints("rewind_first", workflowId)); + assertEquals(List.of(1, 3), txCheckpoints("rewind_second", workflowId)); + // The transactions before the cut replayed; the ones past it ran again. + assertEquals(List.of("a", "b", "c", "c", "d", "d"), tableRows()); + } + + @Test + public void aFailedCheckpointDeleteLeavesTheWorkflowUntouched() throws Exception { + var workflowId = start(() -> sneaky(() -> proxy.txWriter("delete-failure"))); + + failCheckpointDelete.set(true); + var thrown = assertThrows(RuntimeException.class, () -> dbos.rewindWorkflow(workflowId, 0)); + assertEquals( + "Failed to delete the transactional step checkpoints of workflow " + workflowId, + thrown.getMessage()); + // The system database was not rewound, so the workflow keeps its terminal status and is not + // re-enqueued. + assertEquals(WorkflowState.SUCCESS, dbos.retrieveWorkflow(workflowId).getStatus().status()); + assertEquals(1, impl.runsOf("delete-failure")); + + // A delete is idempotent, so retrying the rewind finishes the job. + failCheckpointDelete.set(false); + assertEquals(2, dbos.rewindWorkflow(workflowId, 0).getResult()); + assertEquals(List.of(0, 2), txCheckpoints("rewind_first", workflowId)); + assertEquals(List.of(1, 3), txCheckpoints("rewind_second", workflowId)); + assertEquals(List.of("a", "a", "b", "b", "c", "c", "d", "d"), tableRows()); + } + + @Test + public void rewindRefusesAnActiveWorkflowBeforeTouchingCheckpoints() throws Exception { + WorkflowHandle handle = dbos.startWorkflow(() -> proxy.txBlocker()); + assertTrue(impl.txCheckpointed.await(10, TimeUnit.SECONDS)); + assertEquals(List.of(0), txCheckpoints("rewind_first", handle.workflowId())); + + var refused = + assertThrows( + IllegalStateException.class, () -> dbos.rewindWorkflow(handle.workflowId(), 0)); + assertTrue(refused.getMessage().contains("only a workflow in a terminal state")); + assertEquals(List.of(0), txCheckpoints("rewind_first", handle.workflowId())); + + impl.txReleased.countDown(); + handle.getResult(); + } + + @FunctionalInterface + private interface SqlSupplier { + T get() throws SQLException; + } + + private static T sneaky(SqlSupplier supplier) { + try { + return supplier.get(); + } catch (SQLException e) { + throw new RuntimeException(e); + } + } +} diff --git a/transact/src/test/java/dev/dbos/transact/workflow/RewindTestService.java b/transact/src/test/java/dev/dbos/transact/workflow/RewindTestService.java new file mode 100644 index 000000000..99f9d0290 --- /dev/null +++ b/transact/src/test/java/dev/dbos/transact/workflow/RewindTestService.java @@ -0,0 +1,254 @@ +package dev.dbos.transact.workflow; + +import dev.dbos.transact.DBOS; +import dev.dbos.transact.txstep.JdbcStepFactory; + +import java.sql.Connection; +import java.sql.SQLException; +import java.time.Duration; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +interface RewindTestService { + + int counter(String name); + + String fiveSteps(String name); + + String step(String label); + + String receiver(String name); + + String publisher(String name); + + String streamWriter(String name); + + String streamCloser(String name); + + int parent(String name) throws Exception; + + int failingChild(int value); + + int adoptingParent(String name) throws Exception; + + int doubler(int value); + + int repairer(String targetId); + + String stepRewinder(String targetId); + + String rewindInStep(String targetId); + + String blocker() throws InterruptedException; + + int txWriter(String name) throws SQLException; + + void txBlocker() throws Exception; +} + +class RewindTestServiceImpl implements RewindTestService { + + private final DBOS dbos; + private RewindTestService proxy; + private JdbcStepFactory first; + private JdbcStepFactory second; + + final Map runs = new ConcurrentHashMap<>(); + final Map stepRuns = new ConcurrentHashMap<>(); + final AtomicInteger childRuns = new AtomicInteger(); + final AtomicInteger doublerRuns = new AtomicInteger(); + + // The blocker holds a queue's only slot until released. + volatile CountDownLatch blockerStarted = new CountDownLatch(1); + volatile CountDownLatch blockerReleased = new CountDownLatch(1); + + // The tx blocker parks after its first transactional step. + final CountDownLatch txCheckpointed = new CountDownLatch(1); + final CountDownLatch txReleased = new CountDownLatch(1); + + RewindTestServiceImpl(DBOS dbos) { + this.dbos = dbos; + } + + void setProxy(RewindTestService proxy) { + this.proxy = proxy; + } + + void setStepFactories(JdbcStepFactory first, JdbcStepFactory second) { + this.first = first; + this.second = second; + } + + int runCount(String name) { + return runs.computeIfAbsent(name, k -> new AtomicInteger()).incrementAndGet(); + } + + int runsOf(String name) { + var count = runs.get(name); + return count == null ? 0 : count.get(); + } + + int stepRunsOf(String label) { + var count = stepRuns.get(label); + return count == null ? 0 : count.get(); + } + + @Override + @Workflow + public int counter(String name) { + return runCount(name); + } + + @Override + @Workflow + public String fiveSteps(String name) { + var run = runCount(name); + proxy.step("one"); + proxy.step("two"); + proxy.step("three"); + proxy.step("four"); + proxy.step("five"); + return "run" + run; + } + + @Override + @Step + public String step(String label) { + stepRuns.computeIfAbsent(label, k -> new AtomicInteger()).incrementAndGet(); + return label; + } + + @Override + @Workflow + public String receiver(String name) { + var run = runCount(name); + String first = dbos.recv("cmd", Duration.ofSeconds(10)).orElse(null); + String second = dbos.recv("cmd", Duration.ofSeconds(10)).orElse(null); + return first + second + ":" + run; + } + + @Override + @Workflow + public String publisher(String name) { + var run = runCount(name); + dbos.setEvent("below", "kept"); + dbos.setEvent("both", "old"); + if (run == 1) { + dbos.setEvent("both", "new"); + dbos.setEvent("above", "doomed"); + return "first"; + } + dbos.setEvent("both", "republished"); + return "second"; + } + + @Override + @Workflow + public String streamWriter(String name) { + var run = runCount(name); + dbos.writeStream("log", "a" + run); + dbos.writeStream("log", "b" + run); + return "run" + run; + } + + @Override + @Workflow + public String streamCloser(String name) { + var run = runCount(name); + dbos.writeStream("out", "v" + run); + dbos.closeStream("out"); + return "run" + run; + } + + @Override + @Workflow + public int parent(String name) throws Exception { + runCount(name); + WorkflowHandle handle = + dbos.startWorkflow(() -> proxy.failingChild(21)); + return handle.getResult(); + } + + @Override + @Workflow + public int failingChild(int value) { + if (childRuns.incrementAndGet() == 1) { + throw new IllegalStateException("child is bogus"); + } + return value * 2; + } + + @Override + @Workflow + public int adoptingParent(String name) throws Exception { + var run = runCount(name); + WorkflowHandle handle = dbos.startWorkflow(() -> proxy.doubler(21)); + return handle.getResult() + run; + } + + @Override + @Workflow + public int doubler(int value) { + doublerRuns.incrementAndGet(); + return value * 2; + } + + @Override + @Workflow + public int repairer(String targetId) { + runCount("repairer"); + WorkflowHandle handle = dbos.rewindWorkflow(targetId, 0); + return handle.getResult(); + } + + @Override + @Workflow + public String stepRewinder(String targetId) { + runCount("stepRewinder"); + return proxy.rewindInStep(targetId); + } + + @Override + @Step + public String rewindInStep(String targetId) { + dbos.rewindWorkflow(targetId, 0); + return targetId; + } + + @Override + @Workflow + public String blocker() throws InterruptedException { + blockerStarted.countDown(); + blockerReleased.await(30, TimeUnit.SECONDS); + return "held"; + } + + @Override + @Workflow + public int txWriter(String name) throws SQLException { + var run = runCount(name); + first.txStep((Connection conn) -> insertRow(conn, "a"), "insertA"); + second.txStep((Connection conn) -> insertRow(conn, "b"), "insertB"); + first.txStep((Connection conn) -> insertRow(conn, "c"), "insertC"); + second.txStep((Connection conn) -> insertRow(conn, "d"), "insertD"); + return run; + } + + @Override + @Workflow + public void txBlocker() throws Exception { + first.txStep((Connection conn) -> insertRow(conn, "a"), "insertA"); + txCheckpointed.countDown(); + txReleased.await(30, TimeUnit.SECONDS); + } + + static void insertRow(Connection conn, String value) throws SQLException { + try (var stmt = conn.prepareStatement("INSERT INTO rewind_rows (v) VALUES (?)")) { + stmt.setString(1, value); + stmt.executeUpdate(); + } + } +} From 2baf7b22be8606b72518a4eee8e4b56153e27866 Mon Sep 17 00:00:00 2001 From: Harry Pierson Date: Wed, 7 Oct 2026 09:49:52 -0700 Subject: [PATCH 2/5] Document rewind's failures, and messages consumed before migration 121 List every exception DBOS.rewindWorkflow and DBOSClient.rewindWorkflow can throw, including the status race that asks the caller to retry and a failed step-factory checkpoint delete. A message consumed before migration 121, or by a DBOS version that does not record the consuming step, has no consumed_by_function_id, so a rewind leaves it consumed and a replayed recv past the cut waits for a new message. Say so in the javadoc. --- transact/src/main/java/dev/dbos/transact/DBOS.java | 11 ++++++++++- .../src/main/java/dev/dbos/transact/DBOSClient.java | 5 +++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/transact/src/main/java/dev/dbos/transact/DBOS.java b/transact/src/main/java/dev/dbos/transact/DBOS.java index 72bf4b8a4..fe2f16909 100644 --- a/transact/src/main/java/dev/dbos/transact/DBOS.java +++ b/transact/src/main/java/dev/dbos/transact/DBOS.java @@ -1057,6 +1057,10 @@ public void updateWorkflowAttributes( * yet; and removes its streams' close markers so the replay can write to them again. Stream * entries are kept. The workflow is then re-enqueued. * + *

A message consumed before the system database reached migration 121, or by a DBOS version + * that does not record which step consumed it, cannot be matched to a step. The rewind leaves it + * consumed, so a replayed {@code recv} past the cut waits for a new message instead. + * *

Called from a workflow, the rewind is a step, so a recovered caller does not rewind its * target again. * @@ -1069,7 +1073,12 @@ public void updateWorkflowAttributes( * @return handle to the rewound workflow * @throws dev.dbos.transact.exceptions.DBOSNonExistentWorkflowException if the workflow does not * exist - * @throws IllegalStateException if the workflow is not in a terminal state + * @throws IllegalArgumentException if {@code startStep} is negative, the queue does not exist, or + * the partition key does not match whether the queue is partitioned + * @throws IllegalStateException if the workflow is not in a terminal state, or its status changed + * while it was being rewound; in the second case, retry the rewind + * @throws RuntimeException if a transactional step factory's checkpoints could not be deleted; + * the workflow is not rewound, and the rewind can be retried */ public @NonNull WorkflowHandle rewindWorkflow( @NonNull String workflowId, int startStep, @NonNull RewindOptions options) { diff --git a/transact/src/main/java/dev/dbos/transact/DBOSClient.java b/transact/src/main/java/dev/dbos/transact/DBOSClient.java index 5f86dd6fc..6670bc1ea 100644 --- a/transact/src/main/java/dev/dbos/transact/DBOSClient.java +++ b/transact/src/main/java/dev/dbos/transact/DBOSClient.java @@ -1636,6 +1636,11 @@ public void deleteWorkflows(@NonNull List workflowIds, boolean deleteChi * @param startStep the first step to discard and run again; 0 re-runs the whole workflow * @param options Options for the rewind * @return `WorkflowHandle` for the rewound workflow + * @throws dev.dbos.transact.exceptions.DBOSNonExistentWorkflowException if the workflow does not + * exist + * @throws IllegalArgumentException if {@code startStep} is negative + * @throws IllegalStateException if the workflow is not in a terminal state, or its status changed + * while it was being rewound; in the second case, retry the rewind */ public @NonNull WorkflowHandle rewindWorkflow( @NonNull String workflowId, int startStep, @NonNull RewindOptions options) { From 12f60bb4957fd178def2d1b3a342f85c6add1672 Mon Sep 17 00:00:00 2001 From: Harry Pierson Date: Wed, 7 Oct 2026 11:45:15 -0700 Subject: [PATCH 3/5] Read the rewind's status through getWorkflowState The rewind's status check ran its own copy of getWorkflowState's query and kept its own set of active statuses. Use getWorkflowState and WorkflowState.isActive instead. The copy compared statuses as strings, so a status this SDK does not know counted as terminal and the rewind went ahead. getWorkflowState rejects an unknown status, as every other status read in the SDK already does, so that test is removed. --- .../transact/database/dao/WorkflowDAO.java | 41 +++++-------------- .../dbos/transact/workflow/RewindTest.java | 8 ---- 2 files changed, 11 insertions(+), 38 deletions(-) diff --git a/transact/src/main/java/dev/dbos/transact/database/dao/WorkflowDAO.java b/transact/src/main/java/dev/dbos/transact/database/dao/WorkflowDAO.java index 7c417c1ac..aedfdb01e 100644 --- a/transact/src/main/java/dev/dbos/transact/database/dao/WorkflowDAO.java +++ b/transact/src/main/java/dev/dbos/transact/database/dao/WorkflowDAO.java @@ -2157,7 +2157,7 @@ public static void rewindWorkflow( SqlTransaction.run( txConn, conn -> { - var status = readRewindableStatus(conn, schema, workflowId); + var state = readRewindableState(conn, schema, workflowId); // Whether a key was published at or past the cut. %2$s is the key column the // correlated subquery compares against. @@ -2277,7 +2277,7 @@ public static void rewindWorkflow( stmt.setString(i++, options.applicationVersion()); } stmt.setString(i++, workflowId); - stmt.setString(i++, status); + stmt.setString(i++, state.name()); if (stmt.executeUpdate() != 1) { throw new IllegalStateException( "Workflow %s changed status while being rewound; retry the rewind" @@ -2299,42 +2299,23 @@ public static void checkRewindable(DbContext ctx, String workflowId, int startSt throw new IllegalArgumentException("startStep must be >= 0, got " + startStep); } try (var conn = ctx.getConnection()) { - readRewindableStatus(conn, ctx.schema(), workflowId); + readRewindableState(conn, ctx.schema(), workflowId); } } - private static final Set ACTIVE_STATUSES = - Set.of( - WorkflowState.PENDING.name(), - WorkflowState.ENQUEUED.name(), - WorkflowState.DELAYED.name()); - - private static String readRewindableStatus(Connection conn, String schema, String workflowId) - throws SQLException { - var sql = - """ - SELECT status FROM "%s".workflow_status WHERE workflow_uuid = ? - """ - .formatted(schema); - String status; - try (var stmt = conn.prepareStatement(sql)) { - stmt.setString(1, workflowId); - try (var rs = stmt.executeQuery()) { - if (!rs.next()) { - throw new DBOSNonExistentWorkflowException(workflowId); - } - status = rs.getString("status"); - } + private static WorkflowState readRewindableState( + Connection conn, String schema, String workflowId) throws SQLException { + var state = getWorkflowState(conn, schema, workflowId); + if (state == null) { + throw new DBOSNonExistentWorkflowException(workflowId); } - // Compared as strings, so a status this SDK does not know (one a newer SDK wrote) counts as - // terminal rather than failing the rewind. - if (ACTIVE_STATUSES.contains(status)) { + if (state.isActive()) { throw new IllegalStateException( ("Cannot rewind %s (%s): only a workflow in a terminal state can be rewound, so cancel it" + " first") - .formatted(workflowId, status)); + .formatted(workflowId, state)); } - return status; + return state; } public static List forkFromFailure( diff --git a/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java b/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java index bab0084ef..d15d31907 100644 --- a/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java +++ b/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java @@ -556,14 +556,6 @@ public void theSystemDatabaseRefusesAnActiveWorkflow() throws Exception { assertEquals("held", handle.getResult()); } - @Test - public void aStatusThisSdkDoesNotKnowCountsAsTerminal() throws Exception { - // A status a newer SDK wrote is not one of the active ones, so the rewind goes ahead. - var workflowId = start(() -> proxy.counter("future-status")); - DBUtils.setWorkflowState(dataSource, workflowId, "SOME_FUTURE_STATUS"); - assertEquals(2, dbos.rewindWorkflow(workflowId, 0).getResult()); - } - @Test public void clientRewind() throws Exception { var workflowId = start(() -> proxy.counter("client")); From 333209606facf263e345e8dba004b6d459042a90 Mon Sep 17 00:00:00 2001 From: Harry Pierson Date: Wed, 7 Oct 2026 13:36:37 -0700 Subject: [PATCH 4/5] Add consumed_by_function_id as INT4 INTEGER is INT4 on Postgres but INT8 on CockroachDB, so migration 121 gave CockroachDB an 8-byte column while every other step-ID column is INT4. JDBC then refuses to read it as an Integer, which failed rewindDeletesNotifications on CockroachDB. --- .../java/dev/dbos/transact/migrations/MigrationManager.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transact/src/main/java/dev/dbos/transact/migrations/MigrationManager.java b/transact/src/main/java/dev/dbos/transact/migrations/MigrationManager.java index 77b4796c5..835af58ac 100644 --- a/transact/src/main/java/dev/dbos/transact/migrations/MigrationManager.java +++ b/transact/src/main/java/dev/dbos/transact/migrations/MigrationManager.java @@ -1738,6 +1738,6 @@ static String migration120(boolean isCockroach) { static final String MIGRATION_121 = """ ALTER TABLE "%1$s"."notifications" - ADD COLUMN IF NOT EXISTS "consumed_by_function_id" INTEGER; + ADD COLUMN IF NOT EXISTS "consumed_by_function_id" INT4; """; } From b3d53a11504426e4dda4ddbb68db3479efcad40a Mon Sep 17 00:00:00 2001 From: Harry Pierson Date: Wed, 7 Oct 2026 13:40:15 -0700 Subject: [PATCH 5/5] Test that a rewind keeps a message consumed without a recorded step A message consumed before migration 121 has no consumed_by_function_id, so the rewind's delete cannot tie it to a discarded step and must leave it consumed. Clear the marker on a consumed message, rewind past its recv, and check that it stays consumed and the replayed recv takes the next message. --- .../dbos/transact/workflow/RewindTest.java | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java b/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java index d15d31907..0821b176a 100644 --- a/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java +++ b/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java @@ -267,6 +267,52 @@ public void rewindDeletesNotifications() throws Exception { mailbox(workflowId)); } + @Test + public void rewindKeepsAMessageConsumedWithoutAConsumingStep() throws Exception { + var workflowId = UUID.randomUUID().toString(); + WorkflowHandle handle = + dbos.startWorkflow( + () -> proxy.receiver("unmarked-consume"), new StartWorkflowOptions(workflowId)); + dbos.send(workflowId, "a", "cmd"); + dbos.send(workflowId, "b", "cmd"); + assertEquals("ab:1", handle.getResult()); + var firstRecv = stepIdOf(workflowId, "DBOS.recv", 0); + var secondRecv = stepIdOf(workflowId, "DBOS.recv", 1); + + // A recv that predates migration 121 consumes a message without recording its step. + var sql = + """ + UPDATE "%s".notifications SET consumed_by_function_id = NULL + WHERE destination_uuid = ? AND consumed_by_function_id = ? + """ + .formatted(SCHEMA); + try (var conn = dataSource.getConnection(); + var stmt = conn.prepareStatement(sql)) { + stmt.setString(1, workflowId); + stmt.setInt(2, secondRecv); + assertEquals(1, stmt.executeUpdate()); + } + + try (var q = pausedQueue("rewind_unmarked_gate")) { + dbos.rewindWorkflow( + workflowId, secondRecv, new RewindOptions().withQueue("rewind_unmarked_gate")); + // Nothing ties the second message to a step past the cut, so it stays consumed. + assertEquals( + List.of(new MailboxRow("a", true, firstRecv), new MailboxRow("b", true, null)), + mailbox(workflowId)); + // The replayed recv cannot take it again and gets the next message instead. + dbos.send(workflowId, "c", "cmd"); + } + + assertEquals("ac:2", dbos.retrieveWorkflow(workflowId).getResult()); + assertEquals( + List.of( + new MailboxRow("a", true, firstRecv), + new MailboxRow("b", true, null), + new MailboxRow("c", true, secondRecv)), + mailbox(workflowId)); + } + @Test public void rewindUnpublishesEvents() throws Exception { var workflowId = start(() -> proxy.publisher("events"));