osm-katana/src/cancel.rs
//! Cooperative cancellation for a convert.
//!
//! A Sweden cut measures **2 665 s** end to end (`pass1` 1 042 s, `sort_and_stree`
//! 0.9 s, `pass2` 1 622 s — measured 2026-08-02, `scratch/clipleak/recut.jsonl`),
//! against a collector `refresh_secs` of 3 600. Before this module
//! [`crate::convert`] had no way to stop, so a daemon that owned one could not be
//! shut down for well over half its refresh interval.
//!
//! # Shape
//!
//! The convert is already phased (`pass1` → `sort_and_stree` → `pass2`, see
//! [`crate::phase_log`]) and already fanned over gatling, where the collector
//! thread calls [`znippy_zoomies::gatling::TypedSink::process`] **once per work unit, in strict
//! stream order**. That existing call is the checkpoint — there is no second
//! control path, no extra thread and no polling loop:
//!
//! * [`Checkpointed`] wraps any [`znippy_zoomies::gatling::TypedSink`] and tests the token once
//! per work unit before delegating. One relaxed atomic load per unit.
//! * [`CancelToken::check`] is called at each phase boundary, so the one
//! genuinely uninterruptible step — the k-way merge inside
//! `CoordFileWriter::finalize`, i.e. the `sort_and_stree` phase — is bracketed
//! rather than entered blind. Measured at 0.862 s of a 2 664.653 s Sweden cut
//! (0.03 %), that is the whole structural floor under a cancel.
//!
//! # What it costs, and how fast it stops
//!
//! MEASURED (2026-08-03, oden, `--release`):
//!
//! * **The checkpoint is free.** `Checkpointed::process` vs the same sink called
//! directly: 1.716 vs 1.719 ns/call (100 M calls × 7 alternating rounds,
//! minimum of rounds) — a −0.004 ns "difference", i.e. under the measurement's
//! own resolution. An end-to-end A/B on the liechtenstein fixture has no power
//! to see it either: median 0.796 s with the seam vs 0.845 s without, n = 9
//! each, ranges 0.697–0.886 and 0.701–0.960 — the arm *with* the seam came out
//! nominally faster, which is noise, not a speed-up.
//! * **It stops in milliseconds.** Cancel request → `convert` returning:
//! 1–15 ms across 25 fixture runs (`tests/convert_cancel.rs`).
//!
//! INFERRED for a planet-scale run (not measured — no planet convert was run for
//! this): `BZ2_CHUNK_SIZE` is 200 MB and `bz2_split` cuts a chunk into about
//! `n_workers` segments, so a 173 GB `.osm.bz2` pass yields ≈ 865 chunks × 31
//! workers ≈ 27 000 work units. At ~1.7 ns per checkpoint that is ~45 µs added to
//! a 1 042 s pass — 4 × 10⁻⁶ %. The stop latency is bounded not by the checkpoint
//! interval but by the in-flight segment each worker must finish, ~one chunk of
//! decode.
//!
//! # What a cancelled convert leaves behind
//!
//! Nothing that could be mistaken for a finished extract. A parquet writer that
//! is dropped mid-run leaves a **footerless** file — unreadable, but still a file
//! where a complete one is expected. [`crate::convert`] therefore deletes the
//! parquet outputs it opened (and the node-coord intermediates) on any failed or
//! cancelled run, so the output directory is left with no half-written product at
//! all. See `crate::remove_partial_outputs`.
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
/// The error a cancelled convert fails with.
///
/// `phase` is the phase whose checkpoint observed the flag (`"pass1"`,
/// `"sort_and_stree"`, `"pass2"`, or `"convert"` for the pre-flight check);
/// `units_done` is how many work units that phase had already completed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Cancelled {
pub phase: &'static str,
pub units_done: u64,
}
impl std::fmt::Display for Cancelled {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"convert cancelled during {} after {} work unit(s)",
self.phase, self.units_done
)
}
}
impl std::error::Error for Cancelled {}
/// Cooperative cancellation handle. Cheap to clone (one `Arc`); a clone shares
/// the flag, so cancelling any clone cancels the convert.
///
/// `Default` is a **live** token that has not been cancelled — not an inert one —
/// so a token obtained by any route can always actually fire.
#[derive(Clone, Debug)]
pub struct CancelToken {
inner: Arc<Inner>,
}
#[derive(Debug)]
struct Inner {
/// The hot field: read once per work unit, never contended for writes.
flag: AtomicBool,
/// The FIRST checkpoint that observed `flag` set. Written only on the cancel
/// path (once), so the hot path never touches this lock. It exists because
/// the error the caller finally sees is decided by a teardown race inside
/// gatling; the site that actually stopped the run is not.
site: Mutex<Option<Cancelled>>,
}
impl Default for CancelToken {
fn default() -> Self {
Self::new()
}
}
impl CancelToken {
pub fn new() -> Self {
Self {
inner: Arc::new(Inner {
flag: AtomicBool::new(false),
site: Mutex::new(None),
}),
}
}
/// Request cancellation. Idempotent, callable from any thread.
pub fn cancel(&self) {
self.inner.flag.store(true, Ordering::Relaxed);
}
#[inline]
pub fn is_cancelled(&self) -> bool {
self.inner.flag.load(Ordering::Relaxed)
}
/// The checkpoint the first observer recorded — `None` until one fires.
pub fn stop_site(&self) -> Option<Cancelled> {
self.inner.site.lock().ok().and_then(|g| *g)
}
/// The checkpoint itself: `Err(Cancelled)` once the flag is set.
#[inline]
pub fn check(&self, phase: &'static str, units_done: u64) -> anyhow::Result<()> {
if self.is_cancelled() {
let c = Cancelled { phase, units_done };
if let Ok(mut g) = self.inner.site.lock() {
g.get_or_insert(c);
}
return Err(c.into());
}
Ok(())
}
/// Re-label a failed run as a cancellation when the flag is set.
///
/// Needed because gatling's collector and its dispatch loop race on the way
/// down: the sink's [`Cancelled`] kills the collector, and the dispatch loop
/// then usually fails first with its own generic "collector closed", masking
/// the real reason. Whichever error wins that race, the phase and unit count
/// reported are the ones the checkpoint that actually stopped the run
/// recorded ([`stop_site`](Self::stop_site)) — so the message is the same
/// either way, and `fallback_phase` is only used if nothing was recorded.
pub fn normalize<T>(
&self,
r: anyhow::Result<T>,
fallback_phase: &'static str,
) -> anyhow::Result<T> {
match r {
Ok(v) => Ok(v),
Err(e) if !self.is_cancelled() => Err(e),
Err(e) => {
let site = self.stop_site().unwrap_or(Cancelled {
phase: fallback_phase,
units_done: 0,
});
// `Cancelled` becomes the outermost cause, so it is what the
// phase log prints and what `downcast_ref` finds; the error the
// teardown race happened to produce is retained underneath, so a
// genuine failure that merely coincided with a shutdown is not
// erased.
Err(e.context(site))
}
}
}
}
/// True when `e` is (or was normalized to) a [`Cancelled`].
pub fn is_cancelled_error(e: &anyhow::Error) -> bool {
e.downcast_ref::<Cancelled>().is_some()
}
/// The one checkpoint implementation: a [`znippy_zoomies::gatling::TypedSink`] decorator that
/// tests the token once per work unit before delegating.
///
/// Borrows the inner sink, so the caller keeps it and can still call its own
/// inherent `finish()` after `run_typed` returns.
pub struct Checkpointed<'s, S: ?Sized> {
inner: &'s mut S,
cancel: CancelToken,
phase: &'static str,
units: u64,
}
impl<'s, S: ?Sized> Checkpointed<'s, S> {
pub fn new(inner: &'s mut S, cancel: &CancelToken, phase: &'static str) -> Self {
Self {
inner,
cancel: cancel.clone(),
phase,
units: 0,
}
}
/// Work units this checkpoint has passed through.
pub fn units(&self) -> u64 {
self.units
}
}
impl<T, S: znippy_zoomies::gatling::TypedSink<T> + ?Sized> znippy_zoomies::gatling::TypedSink<T>
for Checkpointed<'_, S>
{
#[inline]
fn process(&mut self, output: T, is_last: bool) -> anyhow::Result<()> {
self.cancel.check(self.phase, self.units)?;
self.units += 1;
self.inner.process(output, is_last)
}
fn finish(&mut self) -> anyhow::Result<()> {
self.cancel.check(self.phase, self.units)?;
self.inner.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
struct Counting(u64);
impl znippy_zoomies::gatling::TypedSink<u32> for Counting {
fn process(&mut self, _o: u32, _last: bool) -> anyhow::Result<()> {
self.0 += 1;
Ok(())
}
}
#[test]
fn checkpoint_stops_at_the_unit_the_flag_was_set_on() {
use znippy_zoomies::gatling::TypedSink as _;
let mut inner = Counting(0);
let tok = CancelToken::new();
let mut cp = Checkpointed::new(&mut inner, &tok, "pass1");
for _ in 0..3 {
cp.process(1, false).expect("not cancelled yet");
}
tok.cancel();
let e = cp.process(1, false).expect_err("must refuse after cancel");
let c = e.downcast_ref::<Cancelled>().expect("Cancelled");
assert_eq!(c.phase, "pass1");
assert_eq!(c.units_done, 3, "must report the units already forwarded");
assert_eq!(
inner.0, 3,
"the cancelled unit must NOT reach the inner sink"
);
}
#[test]
fn normalize_relabels_a_masked_error_only_when_cancelled() {
let tok = CancelToken::new();
let masked = || Err::<(), _>(anyhow::anyhow!("collector closed"));
// Not cancelled → the original error survives untouched.
let e = tok.normalize(masked(), "pass2").unwrap_err();
assert!(!is_cancelled_error(&e), "must not invent a cancellation");
assert_eq!(e.to_string(), "collector closed");
// Cancelled → the mask is replaced by the real reason, and the mask is
// kept underneath rather than thrown away.
tok.cancel();
let e = tok.normalize(masked(), "pass2").unwrap_err();
assert!(is_cancelled_error(&e));
assert_eq!(e.downcast_ref::<Cancelled>().unwrap().phase, "pass2");
assert_eq!(
e.to_string(),
"convert cancelled during pass2 after 0 work unit(s)"
);
assert!(
format!("{e:#}").contains("collector closed"),
"the raced-away error must survive as the cause: {e:#}"
);
}
/// The whole point of `stop_site`: the phase reported must be where the run
/// really stopped, not whichever generic error won gatling's teardown race.
#[test]
fn normalize_reports_the_checkpoint_that_actually_fired_not_the_racing_error() {
use znippy_zoomies::gatling::TypedSink as _;
let mut inner = Counting(0);
let tok = CancelToken::new();
tok.cancel();
{
let mut cp = Checkpointed::new(&mut inner, &tok, "pass1");
cp.process(1, false).unwrap_err();
}
// The caller instead sees gatling's generic dispatch-side error, and
// normalizes it at the outer `convert` level.
let raced = Err::<(), _>(anyhow::anyhow!("collector closed"));
let e = tok.normalize(raced, "convert").unwrap_err();
let c = e.downcast_ref::<Cancelled>().unwrap();
assert_eq!(
c.phase, "pass1",
"must report where it stopped, not the fallback"
);
assert_eq!(c.units_done, 0);
}
#[test]
fn normalize_never_turns_a_success_into_a_failure() {
let tok = CancelToken::new();
tok.cancel();
assert!(tok.normalize(Ok::<_, anyhow::Error>(7), "pass1").is_ok());
}
}