Re-stamp workflow executor_id when successfully checkpointing a step

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();

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);
     }
   }
 

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();

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);

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();
 

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, List<Strin
     }
   }
 
+  /**
+   * Claims {@code workflowId} for {@code executorId}, skipping the write when the workflow is
+   * already stamped with it. No-ops when {@code executorId} is null, as it is for contexts that
+   * have no executor identity of their own (such as {@link dev.dbos.transact.DBOSClient}).
+   *
+   * <p>Runs 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,

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);

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

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");
+      }
+    }
+  }
+}