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 43f61150..2e32ccd4 100644 --- a/transact/src/main/java/dev/dbos/transact/execution/DBOSExecutor.java +++ b/transact/src/main/java/dev/dbos/transact/execution/DBOSExecutor.java @@ -109,6 +109,34 @@ public String fqName() { private record QueueBucket(String queueName, String partitionKey) {} + // Owns this executor's active-ID entry for a workflow. Construction acquires the entry + // (throwing on conflict); release is idempotent, so the try-with-resources backstop never + // removes an entry a resumed execution re-acquired after an explicit release(). + private final class ActiveWorkflowGuard implements AutoCloseable { + private final AtomicBoolean released = new AtomicBoolean(false); + private final String workflowId; + + ActiveWorkflowGuard(String workflowId, QueueBucket bucket) { + if (activeWorkflows.putIfAbsent(workflowId, bucket) != null) { + throw new DBOSWorkflowExecutionConflictException(workflowId); + } + this.workflowId = workflowId; + } + + // Must run before a terminal outcome write becomes durable: once it is visible, a resume + // can re-dispatch this workflow to this executor, and a stale entry would reject it. + void release() { + if (released.compareAndSet(false, true)) { + activeWorkflows.remove(workflowId); + } + } + + @Override + public void close() { + release(); + } + } + private static final ThreadLocal> HOOK_HOLDER = new ThreadLocal<>(); private static final QueueBucket NO_QUEUE = new QueueBucket(null, null); @@ -1743,10 +1771,7 @@ private WorkflowHandle executeWorkflow( finalOptions.isDequeuedRequest() ? new QueueBucket(finalOptions.queueName(), finalOptions.queuePartitionKey()) : NO_QUEUE; - if (activeWorkflows.putIfAbsent(workflowId, bucket) != null) { - throw new DBOSWorkflowExecutionConflictException(workflowId); - } - try { + try (var active = new ActiveWorkflowGuard(workflowId, bucket)) { logger.debug( "executeWorkflow task {}({}) {}", workflow.fullyQualifiedName(), @@ -1778,6 +1803,7 @@ private WorkflowHandle executeWorkflow( return null; } + active.release(); persistWorkflowOutput(workflowId, output, initResult.serialization()); return output; @@ -1809,11 +1835,12 @@ private WorkflowHandle executeWorkflow( throw cancelled; } + // active is already closed here: try-with-resources closes before catch runs, + // so the entry is released before this terminal write becomes durable. persistWorkflowError(workflowId, actual, initResult.serialization()); throw e; } finally { DBOSContextHolder.clear(); - activeWorkflows.remove(workflowId); } }; diff --git a/transact/src/test/java/dev/dbos/transact/workflow/CancelResumeRaceTest.java b/transact/src/test/java/dev/dbos/transact/workflow/CancelResumeRaceTest.java new file mode 100644 index 00000000..3eb7ce6d --- /dev/null +++ b/transact/src/test/java/dev/dbos/transact/workflow/CancelResumeRaceTest.java @@ -0,0 +1,201 @@ +package dev.dbos.transact.workflow; + +import static org.junit.jupiter.api.Assertions.*; + +import dev.dbos.transact.DBOS; +import dev.dbos.transact.DBOSTestAccess; +import dev.dbos.transact.StartWorkflowOptions; +import dev.dbos.transact.utils.PgContainer; + +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import javax.sql.DataSource; + +import com.zaxxer.hikari.HikariDataSource; +import org.junit.jupiter.api.AutoClose; +import org.junit.jupiter.api.Test; + +/** + * Races between a cancelled run's in-flight terminal outcome write and a concurrent resume of the + * same workflow. The test parks run 1's outcome UPDATE (via a JDBC proxy) to hold open the window + * between the workflow function returning and its outcome becoming durable. + */ +public class CancelResumeRaceTest { + + @AutoClose final PgContainer pgContainer = new PgContainer(); + @AutoClose HikariDataSource realDataSource; + @AutoClose DBOS dbos; + + // Latches coordinating the parked stale write. + final CountDownLatch parked = new CountDownLatch(1); + final CountDownLatch releaseStale = new CountDownLatch(1); + final CountDownLatch staleDone = new CountDownLatch(1); + + RaceServiceImpl impl; + RaceService proxy; + + static final String OUTCOME_SQL_MARKER = "SET status = ?, output = ?, error = ?"; + + private static Object invokeUnwrapped(Method m, Object target, Object[] args) throws Throwable { + try { + return m.invoke(target, args); + } catch (InvocationTargetException e) { + throw e.getTargetException(); + } + } + + /** + * Wraps a DataSource so that the FIRST outcome UPDATE writing status SUCCESS for {@code + * targetWorkflowId} parks (signals {@code parked}, awaits {@code releaseStale}) before executing, + * then signals {@code staleDone}. Everything else passes through. + */ + private DataSource parkingDataSource(DataSource delegate, String targetWorkflowId) { + var loader = getClass().getClassLoader(); + var armed = new AtomicBoolean(true); + + return (DataSource) + Proxy.newProxyInstance( + loader, + new Class[] {DataSource.class}, + (dsProxy, dsMethod, dsArgs) -> { + Object result = invokeUnwrapped(dsMethod, delegate, dsArgs); + if (!dsMethod.getName().equals("getConnection") || result == null) { + return result; + } + Connection conn = (Connection) result; + return Proxy.newProxyInstance( + loader, + new Class[] {Connection.class}, + (connProxy, connMethod, connArgs) -> { + Object stmt = invokeUnwrapped(connMethod, conn, connArgs); + if (!connMethod.getName().equals("prepareStatement") + || stmt == null + || connArgs == null + || connArgs.length == 0 + || !(connArgs[0] instanceof String sql) + || !sql.contains(OUTCOME_SQL_MARKER)) { + return stmt; + } + PreparedStatement ps = (PreparedStatement) stmt; + String[] boundStatus = new String[1]; + String[] boundWfId = new String[1]; + return Proxy.newProxyInstance( + loader, + new Class[] {PreparedStatement.class}, + (psProxy, psMethod, psArgs) -> { + if (psMethod.getName().equals("setString") + && psArgs != null + && psArgs.length == 2) { + int idx = (Integer) psArgs[0]; + if (idx == 1) boundStatus[0] = (String) psArgs[1]; + if (idx == 6) boundWfId[0] = (String) psArgs[1]; + } + boolean isTargetStaleWrite = + psMethod.getName().equals("executeUpdate") + && "SUCCESS".equals(boundStatus[0]) + && targetWorkflowId.equals(boundWfId[0]) + && armed.compareAndSet(true, false); + if (!isTargetStaleWrite) { + return invokeUnwrapped(psMethod, ps, psArgs); + } + parked.countDown(); + try { + assertTrue( + releaseStale.await(30, TimeUnit.SECONDS), + "parked stale outcome write was never released"); + return invokeUnwrapped(psMethod, ps, psArgs); + } finally { + staleDone.countDown(); + } + }); + }); + }); + } + + private void setUp(String workflowId) { + realDataSource = pgContainer.dataSource(); + var config = + pgContainer.dbosConfig().withDataSource(parkingDataSource(realDataSource, workflowId)); + dbos = new DBOS(config); + impl = new RaceServiceImpl(); + proxy = dbos.registerProxy(RaceService.class, impl); + dbos.launch(); + } + + // Run the workflow to the point where run 1 has been cancelled, has returned, and its terminal + // outcome write is parked. + private void runCancelAndPark(String workflowId) throws Exception { + dbos.startWorkflow(() -> proxy.raceWorkflow(), new StartWorkflowOptions(workflowId)); + assertTrue(impl.entered.await(15, TimeUnit.SECONDS), "run 1 never entered the workflow"); + + // Cancel while run 1 is still executing; CANCELLED is durably written. + dbos.cancelWorkflow(workflowId); + + // Let run 1's function return. Its terminal outcome write parks before executing. + impl.releaseWorkflow.countDown(); + assertTrue(parked.await(15, TimeUnit.SECONDS), "run 1 outcome write never parked"); + + // Durable status must be CANCELLED (written by cancel; the outcome write is parked). + assertEquals(WorkflowState.CANCELLED, dbos.retrieveWorkflow(workflowId).getStatus().status()); + } + + @Test + public void activeIdReleasedBeforeOutcomeWriteTest() throws Exception { + // The executor's active-workflow-ID entry must be released BEFORE the terminal outcome write + // becomes durable. Otherwise: run 1's stale write is in flight, a client observes CANCELLED + // and resumes, this same executor dequeues the resumed workflow, but the dispatch finds the + // stale active-ID entry and is rejected, leaving the row PENDING with nobody executing it. + String workflowId = "activeIdReleasedBeforeOutcome:%d".formatted(System.currentTimeMillis()); + setUp(workflowId); + DBOSTestAccess.getQueueService(dbos).setSpeedupForTest(); + + runCancelAndPark(workflowId); + + WorkflowHandle resumedHandle = dbos.resumeWorkflow(workflowId); + + // While the stale write is still parked, the resumed workflow must be dequeued and executed + // by this same executor. + assertTrue( + impl.secondRunDone.await(15, TimeUnit.SECONDS), + "resumed dispatch was blocked by a stale active workflow ID"); + assertEquals("completed", resumedHandle.getResult()); + assertEquals(2, impl.runs.get()); + + // Unblock the parked stale write so shutdown is not blocked. + releaseStale.countDown(); + assertTrue(staleDone.await(15, TimeUnit.SECONDS), "stale outcome write never completed"); + } +} + +interface RaceService { + String raceWorkflow() throws InterruptedException; +} + +class RaceServiceImpl implements RaceService { + final AtomicInteger runs = new AtomicInteger(); + final CountDownLatch entered = new CountDownLatch(1); + final CountDownLatch releaseWorkflow = new CountDownLatch(1); + final CountDownLatch secondRunDone = new CountDownLatch(1); + + @Override + @Workflow(name = "raceWorkflow") + public String raceWorkflow() throws InterruptedException { + if (runs.incrementAndGet() > 1) { + secondRunDone.countDown(); + return "completed"; + } + entered.countDown(); + if (!releaseWorkflow.await(30, TimeUnit.SECONDS)) { + throw new IllegalStateException("run 1 was never released"); + } + return ""; + } +}