fix(migration): counters must describe the run, not the current segment

Ed's completed migration reported `copied: 0` beside
`scanned_count: 2522`. Both numbers were accurate; they were measuring
different things and neither said which.

`scanned_count` was cumulative because `checkpoint` had been persisting
it after every batch. `copied` / `skipped` / `failed` / `source_missing`
were plain locals initialised to zero at the top of the handler, written
to `stats` only via `merge_stats` — which is engine-only and fires on
`Completed`, a state a paused run never reaches. So every pause threw
them away and every resumed segment started counting from nothing.

## The fix has two halves, and only one is the obvious one

Restoring on resume is the obvious half: the four counters now seed from
`stats` exactly as `already_scanned` already did.

The half that actually matters is WHEN they are written. Restoring is
useless if nothing durable exists to restore from, so counters are
persisted per batch through a new handler-callable
`checkpoint_counters`, immediately after the cursor checkpoint.
`merge_stats` stays engine-only; the end-of-run summary write is
unchanged.

Two deliberate choices:

* **Absolute values, not deltas.** The merge is last-write-wins and the
  handler owns the running total. Deltas would double-count on exactly
  the replay path that produced 2522 scanned against 2022 rows.
* **A counter-write failure warns, it does not fail the run.** The
  cursor is the correctness-critical write; these are reporting. Losing
  a migration to a hiccuping stats merge is the wrong trade.

`scanned_count()` is now a default method over the new generic
`stat_u64(key)` rather than a second near-identical query.

## Not fixed, and not claimed to be

The 2522-vs-2022 overshoot itself. This makes it legible — cumulative
and per-segment values now both land on the row — but whether the final
segment re-walked rows it had already counted is a cursor question that
needs reproducing, not inferring. The counters should let it be observed
next time rather than reconstructed afterwards.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Edouard Vanbelle
2026-09-08 00:26:12 +02:00
parent bea9e51128
commit d99b718d43
3 changed files with 135 additions and 15 deletions
+11 -8
View File
@@ -229,19 +229,22 @@ impl JobStore for PgJobStore {
Ok(row.and_then(|(v,)| v))
}
async fn scanned_count(&self) -> Result<u64, DomainError> {
// `(stats->>'scanned_count')::BIGINT` — text cast rather than
// `->` numeric extraction because the stored value has been
// written via `((...)::text)::jsonb` in `checkpoint`, which
// may present as either a JSON number or a JSON string
// depending on prior versions. `::BIGINT` handles both.
// `(stats ->> $2)::BIGINT` — text extraction then cast, rather
// than `->` numeric extraction, because the stored value has been
// written via `((...)::text)::jsonb` in `checkpoint` and may
// present as either a JSON number or a JSON string depending on
// prior versions. `::BIGINT` handles both.
//
// `scanned_count()` is the trait's default wrapper around this.
async fn stat_u64(&self, key: &str) -> Result<u64, DomainError> {
let row: Option<(Option<i64>,)> = sqlx::query_as(
"SELECT (stats ->> 'scanned_count')::BIGINT FROM jobs.recoverable_runs WHERE id = $1",
"SELECT (stats ->> $2)::BIGINT FROM jobs.recoverable_runs WHERE id = $1",
)
.bind(self.run_id)
.bind(key)
.fetch_optional(self.pool.as_ref())
.await
.map_err(|e| map_sqlx_err("scanned_count", e))?;
.map_err(|e| map_sqlx_err("stat_u64", e))?;
Ok(row.and_then(|(v,)| v).unwrap_or(0).max(0) as u64)
}
+80 -3
View File
@@ -560,7 +560,44 @@ pub trait JobStore: Send + Sync {
/// Returns `0` if the key is absent (fresh row) or not a
/// number. Callers on a Fresh run can safely skip this — the
/// answer is trivially 0 and the write path starts fresh.
async fn scanned_count(&self) -> Result<u64, DomainError>;
async fn scanned_count(&self) -> Result<u64, DomainError> {
self.stat_u64("scanned_count").await
}
/// Read any numeric key out of the run's `stats` JSONB.
///
/// The generalisation of [`Self::scanned_count`], which is now
/// one caller of it. Handlers use this on a Resume path to
/// restore their own cumulative counters — see
/// [`Self::checkpoint_counters`].
///
/// Returns `0` when the key is absent or not a number, so a
/// fresh run and a run that never wrote the key are the same
/// answer.
async fn stat_u64(&self, key: &str) -> Result<u64, DomainError>;
/// Handler-callable. Merge the handler's OWN cumulative counters
/// into `stats` mid-run.
///
/// Distinct from [`Self::merge_stats`], which stays engine-only and
/// runs once at `Completed`. That timing is the problem this
/// solves: counters written only at the end are lost by a pause,
/// so every resumed segment restarts them at zero and the final
/// row reports the LAST segment rather than the run. `backend_
/// migration` showed this as `copied: 0` on a migration that had
/// copied plenty, next to a `scanned_count` that was cumulative
/// because `checkpoint` had been persisting it all along.
///
/// Pass ABSOLUTE values, not deltas — the merge is
/// `stats = stats || $1`, so each write displaces the last. Keys
/// are the handler's own; do not write engine-owned
/// `scanned_count` / `finding_count` through here.
async fn checkpoint_counters(
&self,
counters: &serde_json::Map<String, serde_json::Value>,
) -> Result<(), DomainError> {
self.merge_stats(counters).await
}
/// Persist one finding to `jobs.run_findings` and bump
/// `stats.finding_count` on the parent run. Consistency handlers
@@ -1453,8 +1490,14 @@ mod tests {
async fn get_string_param(&self, key: &str) -> Result<Option<String>, DomainError> {
Ok(self.state.lock().unwrap().string_params.get(key).cloned())
}
async fn scanned_count(&self) -> Result<u64, DomainError> {
Ok(self.state.lock().unwrap().scanned_count)
/// Mirrors the PG row: `scanned_count` is its own column-like
/// field, every other counter lives in the merged stats map.
async fn stat_u64(&self, key: &str) -> Result<u64, DomainError> {
let s = self.state.lock().unwrap();
if key == "scanned_count" {
return Ok(s.scanned_count);
}
Ok(s.extra_stats.get(key).and_then(|v| v.as_u64()).unwrap_or(0))
}
async fn merge_stats(
&self,
@@ -2126,6 +2169,40 @@ mod tests {
assert_eq!(*seen.lock().unwrap(), Some(b"halfway".to_vec()));
}
/// Counters written mid-run must survive to be read back, because
/// that round-trip is the whole mechanism by which a resumed
/// segment continues its totals instead of restarting them at zero.
/// `backend_migration` reported `copied: 0` on a migration that had
/// copied thousands precisely because nothing persisted them until
/// `Completed`, which a paused run never reaches.
#[tokio::test]
async fn checkpoint_counters_round_trip_through_stats() {
let provider = Arc::new(MemProvider::new());
let store = provider.open_or_start("counter_job").await.unwrap();
let store: Arc<dyn JobStore> = match store {
OpenedRun::Fresh { store: s } | OpenedRun::Resumed { store: s, .. } => s,
OpenedRun::AlreadyActive { .. } => panic!("fresh provider cannot be active"),
};
// Absent keys read as 0, so a fresh run needs no special case.
assert_eq!(store.stat_u64("copied").await.unwrap(), 0);
let mut counters = serde_json::Map::new();
counters.insert("copied".into(), serde_json::json!(120u64));
counters.insert("skipped".into(), serde_json::json!(7u64));
store.checkpoint_counters(&counters).await.unwrap();
assert_eq!(store.stat_u64("copied").await.unwrap(), 120);
assert_eq!(store.stat_u64("skipped").await.unwrap(), 7);
// Absolute, not additive: a later batch's write displaces the
// earlier one rather than summing with it. The handler owns the
// running total; the store only records it.
counters.insert("copied".into(), serde_json::json!(300u64));
store.checkpoint_counters(&counters).await.unwrap();
assert_eq!(store.stat_u64("copied").await.unwrap(), 300);
}
#[tokio::test]
async fn concurrent_trigger_hits_already_active() {
let provider = Arc::new(MemProvider::new());
@@ -557,16 +557,33 @@ impl RecoverableJobHandler for BackendMigrationService {
},
};
let mut copied_count = 0u64;
// Restored on Resume, exactly like `already_scanned` above.
//
// These used to start at zero on every segment while
// `scanned_count` was restored, so one counter described the
// migration and the other four described the current segment.
// A run that paused and resumed then reported `copied: 0`
// beside a `scanned_count` in the thousands — the numbers were
// measuring different things and only one of them said so.
// `checkpoint_counters` below persists them per batch so a
// pause cannot discard them.
let restore = |key: &'static str| async move {
if is_fresh {
0
} else {
store.stat_u64(key).await.unwrap_or(0)
}
};
let mut copied_count = restore("copied").await;
// Populated by the smart-skip probe below: target blob
// already exists at the current head format+key, so a
// rewrite would be identical bytes. Cheap (15-byte range
// read via `is_at_head_format`), massive latency win on
// resume + on backends where the source was rotated to the
// same key as the target already had.
let mut skipped_count: u64 = 0;
let mut failed_count = 0u64;
let mut source_missing_count = 0u64;
let mut skipped_count: u64 = restore("skipped").await;
let mut failed_count = restore("failed").await;
let mut source_missing_count = restore("source_missing").await;
loop {
// Cooperative cancel poll between batches.
@@ -923,6 +940,29 @@ impl RecoverableJobHandler for BackendMigrationService {
message: format!("checkpoint: {e}"),
};
}
// Persist the counters alongside the cursor. Absolute
// values, not deltas — the merge is last-write-wins, and
// the checkpoint above already made this batch's work part
// of the durable position. A failure here is logged but
// does NOT fail the run: the cursor is the correctness-
// critical write, these are reporting.
let counters: serde_json::Map<String, serde_json::Value> = serde_json::json!({
"copied": copied_count,
"skipped": skipped_count,
"failed": failed_count,
"source_missing": source_missing_count,
})
.as_object()
.cloned()
.unwrap_or_default();
if let Err(e) = store.checkpoint_counters(&counters).await {
tracing::warn!(
target: "oxicloud::migration",
event = "backend_migration.counter_persist_failed",
error = %e,
"could not persist per-batch counters; totals may under-report after a resume"
);
}
// Bump the shared progress snapshot so the server-status
// header middleware surfaces fresh numbers on every
// user's next API call. Guard is held only for a struct