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 82d81cb7..c719a058 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 c0dc1299..65b99ee9 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 2777fadd..0cefe812 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 2aa6966c..fe2f1690 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,61 @@ 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. + * + *

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. + * + * @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 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) { + 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 39ba2e8c..6670bc1e 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,49 @@ 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 + * @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) { + 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 a36b3d22..4d833450 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 b8613a4f..d37558b1 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 5849a893..0ae44ad9 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 00000000..c3c1a2e5 --- /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 819d9447..31883481 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 16f2edaa..771c3426 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 ce9af52b..526bf89f 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 8839eb15..aedfdb01 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,202 @@ 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: + * + *

+ * + * @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 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. + 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++, state.name()); + 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()) { + readRewindableState(conn, ctx.schema(), workflowId); + } + } + + 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); + } + 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, state)); + } + return state; + } + 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 96c715ef..c968f0ab 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 3da1fc51..835af58a 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" INT4; + """; } 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 00000000..7cc451df --- /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 5fbb7972..ad4a9b14 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 32591b24..43f87053 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 00000000..0821b176 --- /dev/null +++ b/transact/src/test/java/dev/dbos/transact/workflow/RewindTest.java @@ -0,0 +1,741 @@ +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 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")); + 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 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 00000000..99f9d029 --- /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(); + } + } +}