fix(engine): a status poll writes only over the status it read, so a cancelled run stays cancelled

cli/cmd/mcp_caveats.go

@@ -67,7 +67,7 @@ var mcpCaveatsCore = []mcpCaveat{
 		ID: mcpCaveatLoadSingleProducer,
 		Text: "A LOAD run is one producer and one consumer, whatever numProducers and numConsumers say, " +
 			"so it cannot show how the cluster behaves under parallel clients. STRESS starts one producer per numProducers.",
-		Refs: []string{mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1640-1649", "case LOAD ->", "case STRESS ->")},
+		Refs: []string{mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1656-1665", "case LOAD ->", "case STRESS ->")},
 	},
 	{
 		ID: mcpCaveatReaperDeadline,
@@ -89,7 +89,7 @@ var mcpCaveatsCore = []mcpCaveat{
 			"kates/src/main/resources/application.properties:336-350",
 			mcpAnchoredRef(mcpJava+"engine/TestTimeoutReaper.java:49-51,67-76,161-164,192-205",
 				"TestResult.TaskStatus.STOPPING", "kates.engine.reaper-grace-ms", "\"-submission\"", "\"Timeout: still \""),
-			mcpJava + "engine/TestOrchestrator.java:132-154,186-189,354,1109-1168,1234-1243,1298-1330",
+			mcpJava + "engine/TestOrchestrator.java:132-154,186-189,354,1125-1184,1250-1259,1314-1346",
 			mcpJava + "persistence/EntityMapper.java:158-167",
 			"kates/src/main/resources/db/migration/V24__run_planned_duration.sql:1-13",
 			mcpJava + "domain/TestRun.java:478-481",
@@ -114,8 +114,8 @@ var mcpCaveatsCore = []mcpCaveat{
 			mcpJava + "domain/TestSpec.java:14-32,111-113",
 			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:165,180-182,225",
 				"applyTypeDefaults(type, request.getSpec())", ".withRequestedSpec(", "buildTasks(type, spec, run.getId())"),
-			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1029-1075", "TestSpec applyTypeDefaults(", "return merged;"),
-			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1619-1623,1714-1718",
+			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1045-1091", "TestSpec applyTypeDefaults(", "return merged;"),
+			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1635-1639,1730-1734",
 				"Only when asked", `"-integrity-0-group")`),
 			mcpAnchoredRef(mcpJava+"engine/NativeKafkaBackend.java:792", `task.getConsumerGroup() + "-integrity"`),
 			mcpAnchoredRef(mcpJava+"persistence/EntityMapper.java:56,84-85,143-145",
@@ -158,7 +158,7 @@ var mcpCaveatsCore = []mcpCaveat{
 		Text: "A TUNE_* run executes one produce task with the spec's single configuration. The tuning report copies " +
 			"that one summary into every step, so all steps show the same numbers and the best step is always step 0.",
 		Refs: []string{
-			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1725-1726", "case TUNE_REPLICATION,", "-tune-0"),
+			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1741-1742", "case TUNE_REPLICATION,", "-tune-0"),
 			mcpJava + "trogdor/SpecFactory.java:38-39",
 			mcpJava + "engine/TuningTestRunner.java:101-139",
 		},
@@ -170,7 +170,7 @@ var mcpCaveatsCore = []mcpCaveat{
 		Refs: []string{
 			mcpJava + "trend/TrendResource.java:31-60",
 			mcpJava + "trend/TrendService.java:53-54,185-189",
-			mcpAnchoredRef(mcpJava+"service/TestRunRepository.java:354-374",
+			mcpAnchoredRef(mcpJava+"service/TestRunRepository.java:335-355",
 				"List<TestRun> findWithResults(", "WHERE r.testType = :type AND r.status = :status",
 				".map(EntityMapper::toDomain)"),
 		},

cli/cmd/mcp_tools_runs.go

@@ -184,7 +184,7 @@ const (
 // 61-63). A stored run can hold others: a schedule saved by a PUT, which runs
 // no bean validation (ScheduleResource.java:137-138), fires its spec as sent,
 // a Kates API without the check of a scenario's specs
-// (TestOrchestrator.java:1110-1114) stored a scenario's base spec as sent, and
+// (TestOrchestrator.java:1126-1130) stored a scenario's base spec as sent, and
 // one without the check of a gRPC request's fields
 // (GrpcTestService.java:116-121) stored a compressionType sent over gRPC
 // unchecked. So a stored value outside these is third-party text.
@@ -659,7 +659,7 @@ func mcpRunDecodeSpec(raw json.RawMessage, w *mcpRunSpecWire) (bool, error) {
 }
 
 // mcpRunSpecFrom reads the stored spec. A run without a topic used one named
-// after its type (TestOrchestrator.java:1598,1925-1926). The topic, acks and
+// after its type (TestOrchestrator.java:1614,1941-1942). The topic, acks and
 // compressionType are shown as identifiers only when they hold values Kafka
 // accepts: a stored spec can hold others (see mcpRunTopicRE), such as a
 // scenario's base spec an older Kates API stored as sent
@@ -785,7 +785,7 @@ func mcpRunRequestedFrom(call *mcpCall, run *client.MCPRun) (*mcpRunRequestedSpe
 func mcpRunTaskFrom(call *mcpCall, t client.MCPRunTask) mcpRunTaskOut {
 	return mcpRunTaskOut{
 		// A scenario run's task ids are the run id and the phase name
-		// (TestOrchestrator.java:1742), which nothing validates.
+		// (TestOrchestrator.java:1758), which nothing validates.
 		TaskID:              call.FenceN(t.TaskID, 128),
 		Phase:               call.FenceN(t.PhaseName, 64),
 		Status:              mcpSanitizeLine(t.Status, 16),
@@ -2230,7 +2230,7 @@ var mcpCaveatsRuns = []mcpCaveat{
 			mcpAnchoredRef(mcpJava+"report/ReportGenerator.java:220-249,283-294",
 				"report.setSummary(MetricUtils.computeSummary(results));",
 				"summaries.put(phase, MetricUtils.computeSummary(rows))"),
-			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:415-424,2073-2089",
+			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:415-424,2089-2105",
 				"TestResult.TaskStatus.PENDING : TestResult.TaskStatus.RUNNING",
 				"private static TestRun withUnfinishedTasksFailed("),
 		},
@@ -2267,7 +2267,7 @@ var mcpCaveatsRuns = []mcpCaveat{
 			"except for a run cancelled before its tasks existed; a cancel through the REST API also leaves a CANCEL " +
 			"audit row.",
 		Refs: []string{
-			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:2024-2102",
+			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:2040-2118",
 				`"Cancelled by user"`, `EventKind.FAILED, "cancelled"`, "return run.withResults(updatedResults);"),
 			mcpAnchoredRef(mcpJava+"api/TestResource.java:497-542",
 				`@Path("/{id}/cancel")`, `auditService.record("CANCEL"`, `"Test cancelled; it is stored as FAILED"`),
@@ -2300,10 +2300,10 @@ var mcpCaveatsRuns = []mcpCaveat{
 				"applyTypeDefaults(type, scenario.getBaseSpec())", ".withRequestedSpec(",
 				"long phaseStartMs = System.currentTimeMillis();", "scenario.resolveSpecForPhase(phase)",
 				"phaseStartMs = saturatedSum(phaseStartMs", "backend.submit(task)", "TestResult.TaskStatus.PENDING"),
-			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1436-1464,1536-1575,1585-1597",
+			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1452-1480,1552-1591,1601-1613",
 				"MAX_RAMP_STEPS = 100", "scenarioInapplicableFields(TestScenario scenario", "is a SPIKE phase",
 				"has none (", "check no record CRCs"),
-			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1722-1791", "List<BenchmarkTask> buildPhaseTask(",
+			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1738-1807", "List<BenchmarkTask> buildPhaseTask(",
 				"int baseTarget = Math.max(1, spec.getThroughput() / steps);", "saturatedSum(startAtMs, s * stepMs)",
 				`taskId + "-spike"`),
 			mcpAnchoredRef(mcpJava+"domain/TestScenario.java:108-183",

cli/cmd/mcp_tools_runs_test.go

@@ -812,7 +812,7 @@ func mcpOutsideFences(h *mcpHarness, raw []byte) string {
 }
 
 // A scenario run's text reaches the store unvalidated: its phase names
-// become task ids (TestOrchestrator.java:1742), and an older Kates API
+// become task ids (TestOrchestrator.java:1758), and an older Kates API
 // stored its base spec's topic, acks and compression as sent
 // (TestOrchestrator.java:347-355,373-377). None of it may reach the model outside a
 // fence.

cli/cmd/mcp_tools_scenario.go

@@ -1244,7 +1244,7 @@ var mcpCaveatsScenario = []mcpCaveat{
 			"plans/mcp-server.md:387-405",
 			mcpJava + "api/TestResource.java:108-135",
 			mcpJava + "domain/TestSpec.java:34-113",
-			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1356-1431",
+			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1372-1447",
 				"Map<String, String> inapplicableFields(", "the trogdor backend cannot run a transactional producer"),
 			mcpAnchoredRef("cli/cmd/apply.go:177-251",
 				"for i, scenario := range sf.Scenarios {", "apiClient.CreateTest(ctx, req)",
@@ -1260,7 +1260,7 @@ var mcpCaveatsScenario = []mcpCaveat{
 			"listed as defaulted may differ on the cluster that runs the scenario.",
 		Refs: []string{
 			mcpJava + "config/TestTypeDefaults.java:22-64,323-464",
-			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1029-1075", "TestSpec applyTypeDefaults(", "return merged;"),
+			mcpAnchoredRef(mcpJava+"engine/TestOrchestrator.java:1045-1091", "TestSpec applyTypeDefaults(", "return merged;"),
 			mcpAnchoredRef("kates/src/main/resources/application.properties:73-114",
 				"# Per-type overrides: STRESS", "kates.tests.roundtrip.throughput=10000"),
 			mcpAnchoredRef("charts/kates/values.yaml:605-673", "defaults:", "tests:", `throughput: "10000"`),

kates/src/main/java/com/bmscomp/kates/engine/TestOrchestrator.java

@@ -623,8 +623,41 @@ public TestRun refreshStatus(String runId) {
         }
         run = run.withResults(updatedResults);
 
-        if (allDone && !updatedResults.isEmpty()) {
+        boolean ended = allDone && !updatedResults.isEmpty();
+        if (ended) {
             run = run.withStatus(anyFailed ? TestResult.TaskStatus.FAILED : TestResult.TaskStatus.DONE);
+        }
+
+        if (!pollSignature(run).equals(signatureBefore)) {
+            try {
+                // Written only over the status this poll read. A cancel, the
+                // timeout reaper or a delete that lands while this poll is under
+                // way has already settled the run, and its write must stay: a
+                // plain write put a cancelled run back to RUNNING, with no
+                // handles and no permit, or stored it DONE.
+                if (!repository.saveIfStatus(run, priorStatus)) {
+                    java.util.Optional<TestRun> stored = repository.findById(runId);
+                    if (stored.isEmpty()) {
+                        settle(runId);
+                        return run;
+                    }
+                    // Not settled: a run that only moved PENDING to RUNNING is
+                    // still going, and its next poll needs its handles.
+                    return stored.get();
+                }
+            } catch (jakarta.persistence.OptimisticLockException e) {
+                // Another writer (the reconciler, the reaper, a concurrent GET)
+                // moved the row first. Their write is as valid as this one, so
+                // return what is actually stored rather than turning a plain
+                // read into a 409 for the client.
+                LOG.debugf("Lost the optimistic-lock race refreshing %s; returning the stored run", runId);
+                return repository.findById(runId).orElse(run);
+            }
+        }
+
+        // Only once this poll's ending is the one stored, so a run that
+        // another writer ended is not announced or counted a second time.
+        if (ended) {
             activeHandles.remove(runId);
             fireEvent(run, anyFailed ? TestLifecycleEvent.EventKind.FAILED : TestLifecycleEvent.EventKind.DONE);
 
@@ -677,23 +710,6 @@ public TestRun refreshStatus(String runId) {
             benchmarkMetrics.endRun(runId);
             releasePermit(runId);
         }
-
-        if (!pollSignature(run).equals(signatureBefore)) {
-            try {
-                // An update only. A delete that lands while this poll is under
-                // way must stay a delete: a plain save inserts the run again.
-                if (!repository.saveIfPresent(run)) {
-                    settle(runId);
-                }
-            } catch (jakarta.persistence.OptimisticLockException e) {
-                // Another writer (the reconciler, the reaper, a concurrent GET)
-                // moved the row first. Their write is as valid as this one, so
-                // return what is actually stored rather than turning a plain
-                // read into a 409 for the client.
-                LOG.debugf("Lost the optimistic-lock race refreshing %s; returning the stored run", runId);
-                return repository.findById(runId).orElse(run);
-            }
-        }
         return run;
     }
 
@@ -910,7 +926,7 @@ private void stopTasks(String runId, List<BenchmarkHandle> handles) {
      * <p>Settled first, so the reconciler stops tracking the run before its
      * row goes. A write already under way cannot bring the row back, because
      * the run's own writes only update a row (see
-     * {@link TestRunRepository#saveIfPresent}). Workers a submission still
+     * {@link TestRunRepository#saveIfStatus}). Workers a submission still
      * under way registers after the settle are stopped by whichever finds the
      * row gone first: that submission's write, or the reconciler.
      *

kates/src/main/java/com/bmscomp/kates/service/TestRunRepository.java

@@ -38,8 +38,11 @@ public class TestRunRepository {
      * re-reads the freshest row and therefore never sees a stale version. This
      * check does, by refusing to write over a state that moved on.
      *
+     * <p>An update, never an insert, so a write that read the run before a
+     * delete cannot bring it back: {@link #delete} takes the same row lock.
+     *
      * @return true if the run was saved, false if another writer changed its
-     *     status first (the caller should leave it alone)
+     *     status first, or its row is gone (the caller should leave it alone)
      */
     @Transactional
     public boolean saveIfStatus(TestRun run, com.bmscomp.kates.domain.TestResult.TaskStatus expectedCurrent) {
@@ -56,28 +59,6 @@ public boolean saveIfStatus(TestRun run, com.bmscomp.kates.domain.TestResult.Tas
         return true;
     }
 
-    /**
-     * Saves a run only if its row still exists: an update, never an insert.
-     *
-     * <p>A plain {@link #save} inserts a run it cannot find, so a write that
-     * read the run before a delete and landed after it brought the deleted run
-     * back. The reconciler's poll and the submission's RUNNING write both could.
-     * The row lock makes the check and the write one step against
-     * {@link #delete}, which takes the same lock.
-     *
-     * @return true if the run was saved, false if its row is gone
-     */
-    @Transactional
-    public boolean saveIfPresent(TestRun run) {
-        TestRunEntity existing =
-                em.find(TestRunEntity.class, run.getId(), jakarta.persistence.LockModeType.PESSIMISTIC_WRITE);
-        if (existing == null) {
-            return false;
-        }
-        save(run);
-        return true;
-    }
-
     @Transactional
     public void save(TestRun run) {
         TestRunEntity existing = em.find(TestRunEntity.class, run.getId());
@@ -154,9 +135,9 @@ public List<TestRun> findByType(TestType type) {
      * over exactly when its row goes. A delete used to queue nothing, so a
      * webhook never heard that a deleted run had ended.
      *
-     * <p>The run's row is locked first, as {@link #saveIfPresent} and
-     * {@link #saveIfStatus} lock it, so a concurrent write either lands before
-     * the delete or finds the row gone. Without the lock a delete could
+     * <p>The run's row is locked first, as {@link #saveIfStatus} locks it, so
+     * a concurrent write either lands before the delete or finds the row
+     * gone. Without the lock a delete could
      * deadlock with a write that holds the run's row and wants its results, or
      * fail on a version that write had just moved. Reading the status under
      * the lock also means a run that ended on its own just before the delete

kates/src/test/java/com/bmscomp/kates/engine/InMemoryEngine.java

@@ -40,8 +40,8 @@
  * and the concurrency permits left.
  *
  * <p>The store keeps the distinction {@link TestRunRepository} draws:
- * {@code save} inserts a run it cannot find, while {@code saveIfPresent} and
- * {@code saveIfStatus} only write over a row that is there.
+ * {@code save} inserts a run it cannot find, while {@code saveIfStatus} only
+ * writes over a row that is there.
  */
 public final class InMemoryEngine {
 
@@ -76,10 +76,6 @@ public InMemoryEngine(BenchmarkBackend... others) {
                 })
                 .when(repository)
                 .save(any());
-        when(repository.saveIfPresent(any())).thenAnswer(invocation -> {
-            TestRun run = invocation.getArgument(0);
-            return rows.computeIfPresent(run.getId(), (id, stored) -> run) != null;
-        });
         when(repository.saveIfStatus(any(), any())).thenAnswer(invocation -> {
             TestRun run = invocation.getArgument(0);
             TestRun stored = rows.get(run.getId());