document-haness/docs/clean-architecture-backend-template/final/evidence/raw/073a-durable-operation-expired-lease-probe.java

package dev.caskeleton.adapter.outbound.persistence.operation;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

import dev.caskeleton.adapter.outbound.persistence.platform.JpaPlatformContractSupport;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.sql.Timestamp;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;

/**
 * The durable operation store against a real PostgreSQL.
 *
 * <p>Every claim in this store is about concurrency or about a constraint, and neither survives
 * being tested against a fake. {@code FOR UPDATE SKIP LOCKED} does not exist in H2's PostgreSQL
 * mode with the semantics that matter, a {@code CHECK} constraint a fake does not enforce is a
 * comment, and a lease race has no meaning where only one connection exists.
 *
 * <p>A dead worker is written the only way that is faithful here: the row a crash leaves behind — a
 * lease nobody will renew — and then recovery runs against it. That is precisely what a killed
 * process is from the database's side.
 */
@Tag("jpa-contract")
class DurableOperationStoreContractTest {

  private static final Instant NOW = Instant.parse("2026-08-25T09:00:00Z");
  private static final Duration LEASE = Duration.ofMinutes(2);

  private static JpaPlatformContractSupport support;

  @BeforeAll
  static void startServer() {
    support = JpaPlatformContractSupport.start();
  }

  @AfterAll
  static void stopServer() {
    if (support != null) {
      support.close();
    }
  }

  @BeforeEach
  void migrate() throws SQLException, IOException {
    try (Connection connection = support.connection();
        Statement statement = connection.createStatement()) {
      statement.execute("DROP SCHEMA public CASCADE");
      statement.execute("CREATE SCHEMA public");
      statement.execute(migration());
    }
  }

  @Test
  @DisplayName("a lease left behind by a dead worker is reclaimed by another")
  void expiredLeaseCanBeReclaimedAfterWorkerCrash() throws SQLException {
    submit("op-1", "hash-1");
    claim("worker-a", NOW);

    // The worker is gone. Nothing renews the lease, and nothing else may touch the operation until
    // it lapses — that exclusivity is the point of the lease.
    assertThat(claim("worker-b", NOW.plusSeconds(30))).isEmpty();

    assertThat(claim("worker-b", NOW.plus(LEASE).plusSeconds(1))).contains("op-1");
    assertThat(state("op-1")).isEqualTo("RUNNING");
    assertThat(leaseOwner("op-1")).isEqualTo("worker-b");
  }

  @Test
  @DisplayName("a reclaimed operation keeps the moment it first started")
  void reclaimKeepsTheOriginalStartTime() throws SQLException {
    submit("op-1", "hash-1");
    claim("worker-a", NOW);

    claim("worker-b", NOW.plus(LEASE).plusSeconds(1));

    // COALESCE, not an unconditional assignment. Overwriting it would make every retry look like a
    // fresh start and hide how long the operation has really been running.
    assertThat(startedAt("op-1")).isEqualTo(NOW);
  }

  @Test
  @DisplayName("two workers claiming at once take different operations")
  void concurrentWorkersDoNotTakeTheSameOperation() throws Exception {
    submit("op-1", "hash-1");
    submit("op-2", "hash-2");

    List<String> claimed = new ArrayList<>();
    try (Connection first = support.connection();
        Connection second = support.connection()) {
      first.setAutoCommit(false);
      second.setAutoCommit(false);
      // Interleaved on purpose and inside open transactions: with plain FOR UPDATE the second
      // claim blocks here until the first commits, and with no locking at all both take op-1.
      claimed.add(claimOn(first, "worker-a", NOW).orElseThrow());
      claimed.add(claimOn(second, "worker-b", NOW).orElseThrow());
      first.commit();
      second.commit();
    }

    assertThat(claimed).containsExactlyInAnyOrder("op-1", "op-2");
  }

  @Test
  @DisplayName("a worker whose lease lapsed cannot record a result over its successor")
  void aStaleWorkerCannotOverwriteTheNewOwner() throws SQLException {
    submit("op-1", "hash-1");
    claim("worker-a", NOW);
    claim("worker-b", NOW.plus(LEASE).plusSeconds(1));

    // worker-a is alive and still finishing. Letting it report here is how one operation ends up
    // with two answers, and which one survives would come down to timing.
    assertThat(succeed("op-1", "worker-a", "/results/a")).isZero();
    assertThat(succeed("op-1", "worker-b", "/results/b")).isOne();
    assertThat(resultReference("op-1")).isEqualTo("/results/b");
  }

  @Test
  @DisplayName("a heartbeat from a worker that lost the lease fails rather than extending it")
  void heartbeatFromAStaleWorkerFails() throws SQLException {
    submit("op-1", "hash-1");
    claim("worker-a", NOW);
    claim("worker-b", NOW.plus(LEASE).plusSeconds(1));

    assertThat(heartbeat("op-1", "worker-a", NOW.plus(LEASE).plusSeconds(2))).isZero();
    assertThat(heartbeat("op-1", "worker-b", NOW.plus(LEASE).plusSeconds(2))).isOne();
  }

  @Test
  @DisplayName("cancel works once and never reverses a finished operation")
  void cancelIsIdempotentAndNeverReversesTerminal() throws SQLException {
    submit("op-1", "hash-1");
    claim("worker-a", NOW);
    succeed("op-1", "worker-a", "/results/a");

    // The work happened. Recording it as CANCELED would replace the only record of a real outcome
    // with a claim that it never occurred.
    assertThat(cancel("op-1", "alice")).isZero();
    assertThat(state("op-1")).isEqualTo("SUCCEEDED");

    submit("op-2", "hash-2");
    assertThat(cancel("op-2", "alice")).isOne();
    assertThat(cancel("op-2", "alice")).isZero();
    assertThat(state("op-2")).isEqualTo("CANCELED");
  }

  @Test
  @DisplayName("another principal cannot cancel an operation they did not submit")
  void cancelIsScopedToThePrincipal() throws SQLException {
    submit("op-1", "hash-1");

    assertThat(cancel("op-1", "mallory")).isZero();
    assertThat(state("op-1")).isEqualTo("PENDING");
  }

  @Test
  @DisplayName("a cancelled operation is not handed to a worker")
  void cancelledOperationsAreNotClaimable() throws SQLException {
    submit("op-1", "hash-1");
    cancel("op-1", "alice");

    assertThat(claim("worker-a", NOW)).isEmpty();
  }

  @Test
  @DisplayName("an identical resubmission is refused by the unique constraint")
  void resubmissionCollidesOnScope() throws SQLException {
    submit("op-1", "hash-1");

    // The constraint is the arbiter rather than a prior read: two identical submissions can arrive
    // at the same instant, and a read-then-insert lets both through.
    assertThatThrownBy(() -> submit("op-2", "hash-1")).isInstanceOf(SQLException.class);
  }

  @Test
  @DisplayName("a succeeded row with no result is refused by the database")
  void succeededRowRequiresAResult() throws SQLException {
    submit("op-1", "hash-1");

    assertThatThrownBy(
            () ->
                execute(
                    "UPDATE durable_operation SET state = 'SUCCEEDED', completed_at = now(),"
                        + " lease_owner = NULL WHERE operation_id = 'op-1'"))
        .isInstanceOf(SQLException.class)
        .hasMessageContaining("ck_durable_operation_succeeded_has_result");
  }

  @Test
  @DisplayName("a running row with no lease is refused by the database")
  void runningRowRequiresALease() {
    assertThatThrownBy(
            () -> {
              submit("op-1", "hash-1");
              execute("UPDATE durable_operation SET state = 'RUNNING' WHERE operation_id = 'op-1'");
            })
        .isInstanceOf(SQLException.class)
        .hasMessageContaining("ck_durable_operation_running_has_lease");
  }

  @Test
  @DisplayName("a state outside the published vocabulary is refused")
  void stateVocabularyIsClosed() {
    // A worker that could write its own state would extend the published contract by writing a
    // row, and a client branching on status would meet a value that is in no version of the API.
    assertThatThrownBy(
            () -> {
              submit("op-1", "hash-1");
              execute(
                  "UPDATE durable_operation SET state = 'ALMOST_DONE' WHERE operation_id = 'op-1'");
            })
        .isInstanceOf(SQLException.class)
        .hasMessageContaining("ck_durable_operation_state");
  }

  @Test
  @DisplayName("an aged-out operation is expired, and a finished one is left alone")
  void expirySweepLeavesRealOutcomesAlone() throws SQLException {
    submit("op-1", "hash-1");
    submit("op-2", "hash-2");
    claim("worker-a", NOW);

    List<String> expired = expireStale(NOW.plus(Duration.ofDays(2)));

    assertThat(expired).containsExactlyInAnyOrder("op-1", "op-2");
    submit("op-3", "hash-3");
    claim("worker-b", NOW);
    succeed("op-3", "worker-b", "/results/c");
    assertThat(expireStale(NOW.plus(Duration.ofDays(2)))).isEmpty();
    assertThat(state("op-3")).isEqualTo("SUCCEEDED");
  }

  @Test
  @DisplayName("reclaim returns lapsed leases to the queue")
  void reclaimReturnsLapsedLeases() throws SQLException {
    submit("op-1", "hash-1");
    claim("worker-a", NOW);

    assertThat(reclaim(NOW.plusSeconds(30))).isEmpty();
    assertThat(reclaim(NOW.plus(LEASE).plusSeconds(1))).containsExactly("op-1");
    assertThat(state("op-1")).isEqualTo("PENDING");
    assertThat(leaseOwner("op-1")).isNull();
  }

  private void submit(String operationId, String hash) throws SQLException {
    try (Connection connection = support.connection();
        PreparedStatement statement =
            connection.prepareStatement(
                """
                INSERT INTO durable_operation
                  (operation_id, tenant, principal, operation_name, request_hash, payload,
                   state, submitted_at, expires_at)
                VALUES (?, '', 'alice', 'transfers.create', ?, '{}', 'PENDING', ?, ?)
                """)) {
      statement.setString(1, operationId);
      statement.setString(2, hash);
      statement.setTimestamp(3, Timestamp.from(NOW));
      statement.setTimestamp(4, Timestamp.from(NOW.plus(Duration.ofDays(1))));
      statement.executeUpdate();
    }
  }

  private Optional<String> claim(String workerId, Instant now) throws SQLException {
    try (Connection connection = support.connection()) {
      return claimOn(connection, workerId, now);
    }
  }

  private Optional<String> claimOn(Connection connection, String workerId, Instant now)
      throws SQLException {
    try (PreparedStatement statement =
        connection.prepareStatement(
            """
            UPDATE durable_operation
               SET state = 'RUNNING',
                   lease_owner = ?,
                   lease_acquired_at = ?,
                   lease_expires_at = ?,
                   started_at = COALESCE(started_at, ?),
                   row_version = row_version + 1
             WHERE operation_id = (
                   SELECT operation_id
                     FROM durable_operation
                    WHERE (state = 'PENDING'
                           OR (state = 'RUNNING' AND lease_expires_at <= ?))
                      AND expires_at > ?
                    ORDER BY submitted_at
                    FOR UPDATE SKIP LOCKED
                    LIMIT 1)
            RETURNING operation_id
            """)) {
      statement.setString(1, workerId);
      statement.setTimestamp(2, Timestamp.from(now));
      statement.setTimestamp(3, Timestamp.from(now.plus(LEASE)));
      statement.setTimestamp(4, Timestamp.from(now));
      statement.setTimestamp(5, Timestamp.from(now));
      statement.setTimestamp(6, Timestamp.from(now));
      try (ResultSet rows = statement.executeQuery()) {
        return rows.next() ? Optional.of(rows.getString(1)) : Optional.empty();
      }
    }
  }

  @Test
  void analysisProbeExpiredLeaseCanStillSucceedBeforeTakeover() throws SQLException {
    submit("op-expired", "hash-expired");
    claim("worker-a", NOW);
    Instant afterLease = NOW.plus(LEASE).plusSeconds(1);
    int updated = succeedAt("op-expired", "worker-a", "/results/stale", afterLease);
    System.out.println("durableExpiredLease.completionAt=" + afterLease);
    System.out.println("durableExpiredLease.leaseExpiredAtCompletion=true");
    System.out.println("durableExpiredLease.succeedUpdatedRows=" + updated);
    System.out.println("durableExpiredLease.finalState=" + state("op-expired"));
    assertThat(updated).isOne();
  }

  private int succeedAt(String operationId, String workerId, String resultReference, Instant completedAt)
      throws SQLException {
    try (Connection connection = support.connection();
        PreparedStatement statement = connection.prepareStatement(
            "UPDATE durable_operation SET state = 'SUCCEEDED', result_reference = ?, completed_at = ?, "
                + "lease_owner = NULL, lease_acquired_at = NULL, lease_expires_at = NULL, "
                + "row_version = row_version + 1 "
                + "WHERE operation_id = ? AND state = 'RUNNING' AND lease_owner = ?")) {
      statement.setString(1, resultReference);
      statement.setTimestamp(2, Timestamp.from(completedAt));
      statement.setString(3, operationId);
      statement.setString(4, workerId);
      return statement.executeUpdate();
    }
  }

  private int succeed(String operationId, String workerId, String resultReference)
      throws SQLException {
    try (Connection connection = support.connection();
        PreparedStatement statement =
            connection.prepareStatement(
                """
                UPDATE durable_operation
                   SET state = 'SUCCEEDED', result_reference = ?, completed_at = ?,
                       lease_owner = NULL, lease_acquired_at = NULL, lease_expires_at = NULL,
                       row_version = row_version + 1
                 WHERE operation_id = ? AND state = 'RUNNING' AND lease_owner = ?
                """)) {
      statement.setString(1, resultReference);
      statement.setTimestamp(2, Timestamp.from(NOW.plusSeconds(10)));
      statement.setString(3, operationId);
      statement.setString(4, workerId);
      return statement.executeUpdate();
    }
  }

  private int heartbeat(String operationId, String workerId, Instant now) throws SQLException {
    try (Connection connection = support.connection();
        PreparedStatement statement =
            connection.prepareStatement(
                """
                UPDATE durable_operation
                   SET lease_expires_at = ?, row_version = row_version + 1
                 WHERE operation_id = ? AND state = 'RUNNING' AND lease_owner = ?
                   AND lease_expires_at > ?
                """)) {
      statement.setTimestamp(1, Timestamp.from(now.plus(LEASE)));
      statement.setString(2, operationId);
      statement.setString(3, workerId);
      statement.setTimestamp(4, Timestamp.from(now));
      return statement.executeUpdate();
    }
  }

  private int cancel(String operationId, String principal) throws SQLException {
    try (Connection connection = support.connection();
        PreparedStatement statement =
            connection.prepareStatement(
                """
                UPDATE durable_operation
                   SET state = 'CANCELED', completed_at = ?, lease_owner = NULL,
                       lease_acquired_at = NULL, lease_expires_at = NULL,
                       row_version = row_version + 1
                 WHERE operation_id = ? AND principal = ? AND state IN ('PENDING', 'RUNNING')
                """)) {
      statement.setTimestamp(1, Timestamp.from(NOW.plusSeconds(5)));
      statement.setString(2, operationId);
      statement.setString(3, principal);
      return statement.executeUpdate();
    }
  }

  private List<String> reclaim(Instant now) throws SQLException {
    return idsFrom(
        """
        UPDATE durable_operation
           SET state = 'PENDING', lease_owner = NULL, lease_acquired_at = NULL,
               lease_expires_at = NULL, row_version = row_version + 1
         WHERE state = 'RUNNING' AND lease_expires_at <= ?
        RETURNING operation_id
        """,
        now);
  }

  private List<String> expireStale(Instant now) throws SQLException {
    return idsFrom(
        """
        UPDATE durable_operation
           SET state = 'EXPIRED', completed_at = ?, lease_owner = NULL,
               lease_acquired_at = NULL, lease_expires_at = NULL,
               row_version = row_version + 1
         WHERE state IN ('PENDING', 'RUNNING') AND expires_at <= ?
        RETURNING operation_id
        """,
        now,
        now);
  }

  private List<String> idsFrom(String sql, Instant... arguments) throws SQLException {
    try (Connection connection = support.connection();
        PreparedStatement statement = connection.prepareStatement(sql)) {
      for (int index = 0; index < arguments.length; index++) {
        statement.setTimestamp(index + 1, Timestamp.from(arguments[index]));
      }
      List<String> ids = new ArrayList<>();
      try (ResultSet rows = statement.executeQuery()) {
        while (rows.next()) {
          ids.add(rows.getString(1));
        }
      }
      return ids;
    }
  }

  private String state(String operationId) throws SQLException {
    return column(operationId, "state");
  }

  private String leaseOwner(String operationId) throws SQLException {
    return column(operationId, "lease_owner");
  }

  private String resultReference(String operationId) throws SQLException {
    return column(operationId, "result_reference");
  }

  private Instant startedAt(String operationId) throws SQLException {
    try (Connection connection = support.connection();
        PreparedStatement statement =
            connection.prepareStatement(
                "SELECT started_at FROM durable_operation WHERE operation_id = ?")) {
      statement.setString(1, operationId);
      try (ResultSet rows = statement.executeQuery()) {
        rows.next();
        return rows.getTimestamp(1).toInstant();
      }
    }
  }

  private String column(String operationId, String name) throws SQLException {
    try (Connection connection = support.connection();
        PreparedStatement statement =
            connection.prepareStatement(
                "SELECT " + name + " FROM durable_operation WHERE operation_id = ?")) {
      statement.setString(1, operationId);
      try (ResultSet rows = statement.executeQuery()) {
        rows.next();
        return rows.getString(1);
      }
    }
  }

  private void execute(String sql) throws SQLException {
    try (Connection connection = support.connection();
        Statement statement = connection.createStatement()) {
      statement.execute(sql);
    }
  }

  private static String migration() throws IOException {
    try (InputStream resource =
        DurableOperationStoreContractTest.class
            .getClassLoader()
            .getResourceAsStream("db/migration/postgresql/V11__durable_operation.sql")) {
      if (resource == null) {
        throw new IOException("the durable operation migration is missing from the classpath");
      }
      return new String(resource.readAllBytes(), StandardCharsets.UTF_8);
    }
  }
}