diff --git a/transact/src/main/java/dev/dbos/transact/database/DbContext.java b/transact/src/main/java/dev/dbos/transact/database/DbContext.java index d92dbfce..a5b29834 100644 --- a/transact/src/main/java/dev/dbos/transact/database/DbContext.java +++ b/transact/src/main/java/dev/dbos/transact/database/DbContext.java @@ -9,7 +9,11 @@ import javax.sql.DataSource; public record DbContext( - DataSource dataSource, String schema, DBOSSerializer serializer, BooleanSupplier closed) { + DataSource dataSource, + String schema, + DBOSSerializer serializer, + BooleanSupplier closed, + String executorId) { public Connection getConnection() throws SQLException { return dataSource.getConnection(); 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 95c17cbf..ac2cd6a7 100644 --- a/transact/src/main/java/dev/dbos/transact/database/SystemDatabase.java +++ b/transact/src/main/java/dev/dbos/transact/database/SystemDatabase.java @@ -107,14 +107,15 @@ private SystemDatabase( String schema, boolean created, DBOSSerializer serializer, - boolean useListenNotify) { + boolean useListenNotify, + String executorId) { validatePostgresDataSource(dataSource); schema = sanitizeSchema(schema); if (schema.contains("\"")) { throw new IllegalArgumentException("Schema name must not contain double quotes"); } - this.ctx = new DbContext(dataSource, schema, serializer, this.closed::get); + this.ctx = new DbContext(dataSource, schema, serializer, this.closed::get, executorId); this.created = created; try { useListenNotify = isCockroach(dataSource) ? false : useListenNotify; @@ -136,32 +137,42 @@ public SystemDatabase( String schema, DBOSSerializer serializer, boolean useListenNotify) { - this(createDataSource(url, user, password), schema, true, serializer, useListenNotify); + this(createDataSource(url, user, password), schema, true, serializer, useListenNotify, null); } public SystemDatabase(String url, String user, String password, String schema) { - this(createDataSource(url, user, password), schema, true, null, true); + this(createDataSource(url, user, password), schema, true, null, true, null); } public SystemDatabase(DataSource dataSource, String schema) { - this(dataSource, schema, false, null, true); + this(dataSource, schema, false, null, true, null); } public SystemDatabase(DataSource dataSource, String schema, DBOSSerializer serializer) { - this(dataSource, schema, false, serializer, true); + this(dataSource, schema, false, serializer, true, null); } public static SystemDatabase create(DBOSConfig config) { + return create(config, null); + } + + public static SystemDatabase create(DBOSConfig config, String executorId) { if (config.dataSource() == null) { return new SystemDatabase( - config.databaseUrl(), - config.dbUser(), - config.dbPassword(), + createDataSource(config.databaseUrl(), config.dbUser(), config.dbPassword()), config.databaseSchema(), + true, config.serializer(), - config.useListenNotify()); + config.useListenNotify(), + executorId); } else { - return new SystemDatabase(config.dataSource(), config.databaseSchema(), config.serializer()); + return new SystemDatabase( + config.dataSource(), + config.databaseSchema(), + false, + config.serializer(), + true, + executorId); } } 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 3871e823..5d365a38 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 @@ -204,8 +204,7 @@ ON CONFLICT (message_uuid) DO NOTHING if (workflowId != null) { var output = new StepResult(workflowId, stepId, functionName, null, null, null, null); - StepsDAO.recordStepResult( - conn, ctx.schema(), output, startTime, System.currentTimeMillis()); + StepsDAO.recordStepResult(ctx, conn, output, startTime, System.currentTimeMillis()); } conn.commit(); @@ -341,7 +340,7 @@ public static Object recv( var output = new StepResult( workflowId, stepId, stepName, serializedMessage, null, null, serialization); - StepsDAO.recordStepResult(conn, ctx.schema(), output, startTime); + StepsDAO.recordStepResult(ctx, conn, output, startTime); conn.commit(); return deserializedMessage; @@ -443,7 +442,7 @@ public static void setEvent( if (asStep) { StepResult output = new StepResult(workflowId, functionId, functionName, null, null, null, null); - StepsDAO.recordStepResult(conn, ctx.schema(), output, startTime); + StepsDAO.recordStepResult(ctx, conn, output, startTime); } conn.commit(); diff --git a/transact/src/main/java/dev/dbos/transact/database/dao/StepsDAO.java b/transact/src/main/java/dev/dbos/transact/database/dao/StepsDAO.java index 25bab0e4..fc9d1eda 100644 --- a/transact/src/main/java/dev/dbos/transact/database/dao/StepsDAO.java +++ b/transact/src/main/java/dev/dbos/transact/database/dao/StepsDAO.java @@ -34,22 +34,22 @@ public static void recordStepResult( DbContext ctx, StepResult result, long startTimeEpochMs, long endTimeEpochMs) throws SQLException { try (var conn = ctx.getConnection()) { - recordStepResult(conn, ctx.schema(), result, startTimeEpochMs, endTimeEpochMs); + recordStepResult(ctx, conn, result, startTimeEpochMs, endTimeEpochMs); } DebugTriggers.debugTriggerPoint(DebugTriggers.DEBUG_TRIGGER_STEP_COMMIT); } static void recordStepResult( - Connection conn, String schema, StepResult result, long startTimeEpochMs) + DbContext ctx, Connection conn, StepResult result, long startTimeEpochMs) throws SQLException { - recordStepResult(conn, schema, result, startTimeEpochMs, System.currentTimeMillis()); + recordStepResult(ctx, conn, result, startTimeEpochMs, System.currentTimeMillis()); } static void recordStepResult( - Connection conn, String schema, StepResult result, Long startTimeEpochMs, Long endTimeEpochMs) + DbContext ctx, Connection conn, StepResult result, Long startTimeEpochMs, Long endTimeEpochMs) throws SQLException { - Objects.requireNonNull(schema); + String schema = Objects.requireNonNull(ctx.schema()); String sql = """ INSERT INTO "%s".operation_outputs @@ -59,6 +59,7 @@ static void recordStepResult( """ .formatted(schema); + boolean won = false; try (var stmt = conn.prepareStatement(sql)) { stmt.setString(1, result.workflowId()); stmt.setInt(2, result.stepId()); @@ -86,14 +87,17 @@ static void recordStepResult( stmt.setObject(8, endTimeEpochMs); try (ResultSet rs = stmt.executeQuery()) { - if (rs.next() && endTimeEpochMs != null) { - long completedAt = rs.getLong("completed_at_epoch_ms"); - if (completedAt != endTimeEpochMs) { - logger.warn( - String.format( - "Step output for %s:%d-%s was already recorded", - result.workflowId(), result.stepId(), result.stepName())); - throw new DBOSWorkflowExecutionConflictException(result.workflowId()); + if (rs.next()) { + won = true; + if (endTimeEpochMs != null) { + long completedAt = rs.getLong("completed_at_epoch_ms"); + if (completedAt != endTimeEpochMs) { + logger.warn( + String.format( + "Step output for %s:%d-%s was already recorded", + result.workflowId(), result.stepId(), result.stepName())); + throw new DBOSWorkflowExecutionConflictException(result.workflowId()); + } } } } @@ -105,6 +109,10 @@ static void recordStepResult( throw e; } } + + if (won) { + WorkflowDAO.restampExecutorId(conn, schema, result.workflowId(), ctx.executorId()); + } } static StepResult checkStepResult( @@ -297,7 +305,7 @@ public static boolean patch(DbContext ctx, String workflowId, int functionId, St var checkpointName = getCheckpointName(conn, ctx.schema(), workflowId, functionId); if (checkpointName == null) { var output = new StepResult(workflowId, functionId, patchName, null, null, null, null); - recordStepResult(conn, ctx.schema(), output, System.currentTimeMillis(), null); + recordStepResult(ctx, conn, output, System.currentTimeMillis(), null); return true; } else { return patchName.equals(checkpointName); 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 39dd4d24..9366d552 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 @@ -64,8 +64,7 @@ public static void writeStreamFromWorkflow( insertStream(conn, ctx.schema(), workflowId, functionId, key, value, serializationFormat); var output = new StepResult(workflowId, functionId, functionName, null, null, null, null); - StepsDAO.recordStepResult( - conn, ctx.schema(), output, startTime, System.currentTimeMillis()); + StepsDAO.recordStepResult(ctx, conn, output, startTime, System.currentTimeMillis()); conn.commit(); 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 eba95ad2..dd8dc140 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 @@ -1276,7 +1276,7 @@ public static void recordChildWorkflow( new StepResult(parentId, functionId, functionName, null, null, null, null) .withChildWorkflowId(childId); try (var conn = ctx.getConnection()) { - StepsDAO.recordStepResult(conn, ctx.schema(), result, null, null); + StepsDAO.recordStepResult(ctx, conn, result, null, null); } } @@ -1802,6 +1802,32 @@ private static void markWasForkedFrom(Connection conn, String schema, ListRuns on the caller's connection so it joins the caller's transaction. + */ + static void restampExecutorId( + Connection conn, String schema, String workflowId, @Nullable String executorId) + throws SQLException { + if (executorId == null) { + return; + } + String sql = + """ + UPDATE "%s".workflow_status SET executor_id = ? WHERE workflow_uuid = ? AND executor_id IS DISTINCT FROM ? + """ + .formatted(schema); + try (var stmt = conn.prepareStatement(sql)) { + stmt.setString(1, executorId); + stmt.setString(2, workflowId); + stmt.setString(3, executorId); + stmt.executeUpdate(); + } + } + private static void batchCopyWorkflowData( Connection conn, String schema, 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 2e32ccd4..5e19fb43 100644 --- a/transact/src/main/java/dev/dbos/transact/execution/DBOSExecutor.java +++ b/transact/src/main/java/dev/dbos/transact/execution/DBOSExecutor.java @@ -261,7 +261,7 @@ public void start( executorService = executorServiceSupplier.get(); timeoutScheduler = Executors.newScheduledThreadPool(2); - systemDatabase = SystemDatabase.create(config); + systemDatabase = SystemDatabase.create(config, this.executorId); systemDatabase.start(); systemDatabase.createApplicationVersion(this.appVersion); diff --git a/transact/src/test/java/dev/dbos/transact/database/SystemDatabaseTest.java b/transact/src/test/java/dev/dbos/transact/database/SystemDatabaseTest.java index 71ad43a4..6f6b4457 100644 --- a/transact/src/test/java/dev/dbos/transact/database/SystemDatabaseTest.java +++ b/transact/src/test/java/dev/dbos/transact/database/SystemDatabaseTest.java @@ -2645,7 +2645,7 @@ public boolean isWrapperFor(Class i) throws SQLException { private DbContext recordingCtx(IsolationRecordingDataSource ds) { String schema = SystemDatabase.sanitizeSchema(dbosConfig.databaseSchema()); - return new DbContext(ds, schema, null, () -> false); + return new DbContext(ds, schema, null, () -> false, null); } @Test diff --git a/transact/src/test/java/dev/dbos/transact/execution/ExecutorIdRestampTest.java b/transact/src/test/java/dev/dbos/transact/execution/ExecutorIdRestampTest.java new file mode 100644 index 00000000..1d26bbda --- /dev/null +++ b/transact/src/test/java/dev/dbos/transact/execution/ExecutorIdRestampTest.java @@ -0,0 +1,130 @@ +package dev.dbos.transact.execution; + +import static org.junit.jupiter.api.Assertions.*; + +import dev.dbos.transact.DBOS; +import dev.dbos.transact.StartWorkflowOptions; +import dev.dbos.transact.config.DBOSConfig; +import dev.dbos.transact.utils.PgContainer; +import dev.dbos.transact.workflow.Step; +import dev.dbos.transact.workflow.Workflow; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +import javax.sql.DataSource; + +import com.zaxxer.hikari.HikariDataSource; +import org.junit.jupiter.api.AutoClose; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class ExecutorIdRestampTest { + + static final String EXECUTOR_ID = "restamp-test-executor"; + + @AutoClose final PgContainer pgContainer = new PgContainer(); + + DBOSConfig dbosConfig; + @AutoClose HikariDataSource dataSource; + + interface RestampService { + String twoStepWorkflow(); + + String stepOne(); + + String stepTwo(); + } + + static class RestampServiceImpl implements RestampService { + static CountDownLatch midpoint = new CountDownLatch(1); + static CountDownLatch proceed = new CountDownLatch(1); + + private RestampService self; + + void setSelf(RestampService self) { + this.self = self; + } + + @Override + @Workflow(name = "twoStepWorkflow") + public String twoStepWorkflow() { + var one = self.stepOne(); + midpoint.countDown(); + try { + proceed.await(30, TimeUnit.SECONDS); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + return one + self.stepTwo(); + } + + @Override + @Step(name = "stepOne") + public String stepOne() { + return "one"; + } + + @Override + @Step(name = "stepTwo") + public String stepTwo() { + return "two"; + } + } + + @BeforeEach + void setUp() { + dbosConfig = pgContainer.dbosConfig().withExecutorId(EXECUTOR_ID); + dataSource = pgContainer.dataSource(); + RestampServiceImpl.midpoint = new CountDownLatch(1); + RestampServiceImpl.proceed = new CountDownLatch(1); + } + + @Test + void restampExecutorIdOnStepCheckpoint() throws Exception { + try (var dbos = new DBOS(dbosConfig)) { + var impl = new RestampServiceImpl(); + var service = dbos.registerProxy(RestampService.class, impl); + impl.setSelf(service); + dbos.launch(); + + String wfid = "restamp-wf-1"; + var handle = dbos.startWorkflow(service::twoStepWorkflow, new StartWorkflowOptions(wfid)); + + assertTrue(RestampServiceImpl.midpoint.await(30, TimeUnit.SECONDS)); + setExecutorId(dataSource, wfid, "stale-executor"); + assertEquals("stale-executor", getExecutorId(dataSource, wfid)); + RestampServiceImpl.proceed.countDown(); + + assertEquals("onetwo", handle.getResult()); + assertEquals(EXECUTOR_ID, getExecutorId(dataSource, wfid)); + } + } + + private static void setExecutorId(DataSource ds, String workflowId, String executorId) + throws SQLException { + String sql = "UPDATE dbos.workflow_status SET executor_id = ? WHERE workflow_uuid = ?"; + try (Connection conn = ds.getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql)) { + pstmt.setString(1, executorId); + pstmt.setString(2, workflowId); + assertEquals(1, pstmt.executeUpdate()); + } + } + + private static String getExecutorId(DataSource ds, String workflowId) throws SQLException { + String sql = "SELECT executor_id FROM dbos.workflow_status WHERE workflow_uuid = ?"; + try (Connection conn = ds.getConnection(); + PreparedStatement pstmt = conn.prepareStatement(sql)) { + pstmt.setString(1, workflowId); + try (var rs = pstmt.executeQuery()) { + assertTrue(rs.next()); + return rs.getString("executor_id"); + } + } + } +}