osm-katana/src/xml_scan.rs
//! **Bounded-memory streaming HARVEST over compressed OSM XML.**
//!
//! [`crate::relations::stream_admin_region_names`] and its geometry sibling
//! [`crate::relations::stream_admin_region_polygon`] proved the shape: walk the
//! planet once per pass, keep only the handful of bytes the caller asked for, and
//! never materialise a node store. Both, however, could only walk a **PBF** —
//! while the planet on this box is `planet-260427.osm.bz2`, 173 GB of bzip2'd OSM
//! XML. So Sweden's own border sat inside a file the crate could already read at
//! ~170 MB/s and the admin scan could not reach, purely because the two halves had
//! never been wired together.
//!
//! This module is that wire. It is **not a second bz2 reader**: the decode, the
//! slot split, the block seek and the Gatling config are all
//! [`crate::xml_to_pbf`]'s — the exact path `osm-katana convert` takes over the
//! same file. What is new here is only the *payload*: instead of encoding each
//! decoded XML run into PBF blobs, each Gatling worker runs the caller's
//! `harvest` closure over the run's elements and keeps whatever the closure
//! returns.
//!
//! FIVE-ZEROS: fan-out is `gatling` (never rayon) — one worker per bz2 block
//! range, each decoding into its own slot and DROPPING it after harvesting, so
//! peak memory is bounded by what the closure keeps, not by the file. The
//! collector stitches the elements that straddle a block boundary
//! ([`crate::xml_to_pbf::xml_middle_span`]) exactly as the PBF sink does, so an
//! element on a segment edge is harvested once — not twice, not never.
use std::path::Path;
use anyhow::Result;
use znippy_zoomies::gatling;
use crate::xml_to_pbf::{
Bz2Prepared, Bz2Skip, ElemRef, bz2_prepare, bz2_split, mk_bz2_cfg, vtd_parse_serial,
xml_middle_span,
};
/// A harvest closure: given one run of decoded XML and the elements found in it
/// (zero-copy — every [`ElemRef`] indexes back into `xml`), return the items to
/// keep. Runs on a Gatling worker, so it must be `Sync`.
pub(crate) trait Harvest: Sync {
/// What one element contributes. `Send + 'static` so a worker can hand it to
/// the collector.
type Item: Send + 'static;
/// Harvest one decoded XML run. Called once per segment per worker, plus once
/// per stitched segment boundary on the collector.
fn harvest(&self, xml: &[u8], elems: &[ElemRef]) -> Vec<Self::Item>;
/// Whether this harvest needs the run parsed into elements.
///
/// `false` turns the sweep into a **decompression-only probe**: the worker
/// decodes its block range, hands the WHOLE decoded run to `harvest` with an
/// empty element slice, and skips both the VTD scan and the boundary stitch.
/// That is the arm a bench needs to say how much of a pass is bzip2 and how
/// much is everything else — see [`bz2_decode_probe`].
fn needs_elements(&self) -> bool {
true
}
}
/// Blanket impl so a plain closure is a [`Harvest`] — callers write
/// `harvest_xml_bz2(path, skip, workers, |xml, elems| { … })`.
impl<T, F> Harvest for F
where
T: Send + 'static,
F: Fn(&[u8], &[ElemRef]) -> Vec<T> + Sync,
{
type Item = T;
#[inline]
fn harvest(&self, xml: &[u8], elems: &[ElemRef]) -> Vec<T> {
self(xml, elems)
}
}
/// Which consumer of a shared pass produced an item.
///
/// See the `(A, B)` [`Harvest`] impl: two extractions ride ONE sweep and the
/// collector tells their outputs apart by this tag.
pub(crate) enum Of2<A, B> {
A(A),
B(B),
}
/// **Two consumers, one pass over the planet.**
///
/// This is the plan-shaped half of the design, for a *stream*. The PBF side
/// already has one: [`crate::reader::Pass2Plan`] is a list of blob INDICES, so
/// one pass can serve several extractions by naming which blobs to revisit. That
/// concept cannot cross to bz2 — a bzip2 stream has no addressable units to name,
/// and inventing addresses for it would mean seeking a file we decompress
/// end-to-end anyway. The stream analogue of "what does this pass do" is not a
/// list of places, it is a list of CONSUMERS: run both harvests over the same
/// decoded elements, once.
///
/// Composes, so `((a, b), c)` rides one pass with three consumers. The region
/// catalog + one region's boundary use it together
/// ([`crate::relations::stream_admin_catalog_and_polygon_bz2`]); road or POI
/// extraction would be the next consumer, at no extra decompression.
impl<X: Harvest, Y: Harvest> Harvest for (X, Y) {
type Item = Of2<X::Item, Y::Item>;
fn harvest(&self, xml: &[u8], elems: &[ElemRef]) -> Vec<Self::Item> {
let mut out: Vec<Self::Item> = self.0.harvest(xml, elems).into_iter().map(Of2::A).collect();
out.extend(self.1.harvest(xml, elems).into_iter().map(Of2::B));
out
}
/// Elements are parsed if EITHER consumer wants them.
fn needs_elements(&self) -> bool {
self.0.needs_elements() || self.1.needs_elements()
}
}
/// Split a shared pass's tagged output back into its two consumers' lists.
pub(crate) fn unzip_of2<A, B>(items: Vec<Of2<A, B>>) -> (Vec<A>, Vec<B>) {
let mut a = Vec::new();
let mut b = Vec::new();
for it in items {
match it {
Of2::A(x) => a.push(x),
Of2::B(y) => b.push(y),
}
}
(a, b)
}
/// One worker's result for one decoded segment: the harvested items plus the
/// partial elements at each edge, which only the in-order collector can stitch.
struct HarvestSegment<T> {
items: Vec<T>,
left_stub: Vec<u8>,
right_stub: Vec<u8>,
}
impl<T> Default for HarvestSegment<T> {
fn default() -> Self {
Self {
items: Vec::new(),
left_stub: Vec::new(),
right_stub: Vec::new(),
}
}
}
/// bz2 [`gatling::TypedCodec`] whose `transform` decodes one block range and
/// harvests it. Decode + parse + harvest all happen on the worker; the collector
/// only stitches boundaries.
struct Bz2HarvestCodec<'a, H: Harvest> {
max_blocksize: u32,
harvest: &'a H,
}
impl<H: Harvest> gatling::TypedCodec for Bz2HarvestCodec<'_, H> {
type Seg = (u64, u64);
type Output = HarvestSegment<H::Item>;
fn split(
&self,
data: &[u8],
n_workers: usize,
is_last: bool,
) -> Option<gatling::Split<(u64, u64)>> {
bz2_split(data, n_workers, self.max_blocksize, is_last)
}
fn transform(&self, data: &[u8], &(start_bit, end_bit): &(u64, u64)) -> Self::Output {
let decoded = lbzip2::chunk::decode_segment(data, start_bit, end_bit, self.max_blocksize);
if decoded.is_empty() {
return HarvestSegment::default();
}
if !self.harvest.needs_elements() {
// Decompression-only probe: the decode is the measurement, so no VTD
// scan and no stubs (there is no element to lose — nothing is parsed).
return HarvestSegment {
items: self.harvest.harvest(&decoded, &[]),
..HarvestSegment::default()
};
}
let (first, last) = xml_middle_span(&decoded);
let middle = &decoded[first..last];
let elems = vtd_parse_serial(middle);
HarvestSegment {
items: self.harvest.harvest(middle, &elems),
left_stub: decoded[..first].to_vec(),
right_stub: decoded[last..].to_vec(),
}
}
}
/// In-order collector: appends each segment's items and harvests the elements
/// that straddled a segment boundary (`prev.right_stub + next.left_stub`), the
/// same stitch [`crate::xml_to_pbf`]'s PBF sink performs.
struct HarvestSink<'a, H: Harvest> {
items: Vec<H::Item>,
prev_stub: Vec<u8>,
harvest: &'a H,
}
impl<H: Harvest> HarvestSink<'_, H> {
/// Harvest a stitched boundary run. Parsed WITHOUT re-trimming: the stub is
/// already exactly the element(s) the two segments split, so trimming it
/// again would throw away the very element this stitch exists to recover.
fn flush_stub(&mut self, stub: &[u8]) {
if stub.is_empty() {
return;
}
let elems = vtd_parse_serial(stub);
if elems.is_empty() {
return;
}
self.items.extend(self.harvest.harvest(stub, &elems));
}
}
impl<H: Harvest> gatling::TypedSink<HarvestSegment<H::Item>> for HarvestSink<'_, H> {
fn process(&mut self, seg: HarvestSegment<H::Item>, _is_last: bool) -> Result<()> {
let boundary = [self.prev_stub.as_slice(), seg.left_stub.as_slice()].concat();
self.flush_stub(&boundary);
self.items.extend(seg.items);
self.prev_stub = seg.right_stub;
Ok(())
}
}
/// **Harvest a `.osm.bz2` in ONE streaming, gatling-parallel pass.**
///
/// `skip` picks where the pass starts, using the block seeks already proven by
/// `convert` ([`Bz2Skip::Changesets`] to the first data block, [`Bz2Skip::Nodes`]
/// past the whole node section). OSM XML is ordered nodes → ways → relations, so
/// a relation- or way-only harvest costs a fraction of the file.
///
/// Returns every item the closure kept, in stream order. Peak memory is the
/// harvest's own output plus the in-flight slots — a whole-planet run keeps
/// nothing it was not asked for.
pub(crate) fn harvest_xml_bz2<H: Harvest>(
path: &Path,
skip: Bz2Skip<'_>,
n_workers: usize,
harvest: &H,
) -> Result<Vec<H::Item>> {
let Bz2Prepared {
reader,
max_blocksize,
initial_carry,
..
} = bz2_prepare(path, skip)?;
let codec = Bz2HarvestCodec {
max_blocksize,
harvest,
};
let mut sink = HarvestSink {
items: Vec::new(),
prev_stub: Vec::new(),
harvest,
};
gatling::run_typed(
reader,
codec,
&mut sink,
n_workers,
mk_bz2_cfg(initial_carry),
)?;
// The trailing stub is the file's last partial run; harvest it too. Without
// this the LAST element of the file is silently dropped — the classic
// block-boundary tail bug, and exactly what
// `bz2_harvest_sees_every_element_a_serial_decode_does` catches.
let tail = std::mem::take(&mut sink.prev_stub);
sink.flush_stub(&tail);
Ok(sink.items)
}
/// **Decompression-only probe over a `.osm.bz2`** — decode every block on a
/// gatling worker, discard the bytes, return how many there were.
///
/// This is the bz2 primitive with the payload removed, so a bench can attribute
/// a pass's wall time: run it at `n_workers = 1` for the serial decode this box
/// used to be stuck with (osmium's single-stream bzip2 reader, ~8 h projected on
/// this planet — `bench_history.json`, 2026-06-18) and at `n_workers = cores` for
/// the block-parallel one, over the same bytes.
///
/// `skip_to_data` seeks past the changeset prologue exactly as `convert` does.
pub fn bz2_decode_probe(path: &Path, n_workers: usize, skip_to_data: bool) -> Result<u64> {
/// Counts decoded bytes and keeps nothing.
struct DecodeOnly;
impl Harvest for DecodeOnly {
type Item = u64;
fn harvest(&self, decoded: &[u8], _elems: &[ElemRef]) -> Vec<u64> {
vec![decoded.len() as u64]
}
fn needs_elements(&self) -> bool {
false
}
}
let skip = if skip_to_data {
Bz2Skip::Changesets(crate::xml_to_pbf::DATA_NEEDLES)
} else {
Bz2Skip::None
};
Ok(harvest_xml_bz2(path, skip, n_workers, &DecodeOnly)?
.into_iter()
.sum())
}
/// **Sample real node coordinates out of a `.osm.bz2`** — every `stride`-th node,
/// as `(lon, lat)`.
///
/// A clip's cost is decided by where the planet's nodes actually ARE: how many
/// land inside the region's bounding box and pay for a ray-cast, and how many of
/// those sit near the boundary. Neither can be answered by generated points, and
/// re-cutting a planet to find out is the expensive thing being measured. This
/// gives a bench (and the mask-equivalence guard) a corpus of real coordinates at
/// bounded memory: `stride` thins the stream, so a whole planet can be sampled
/// into a few hundred MB.
///
/// The pass starts at the first data block and is the usual block-parallel sweep;
/// `stride` is applied per decoded segment, so the sample is spread across the
/// file rather than taken from its head.
pub fn sample_node_coords(path: &Path, n_workers: usize, stride: usize) -> Result<Vec<(f64, f64)>> {
use crate::xml_vtd::ElemKind;
let stride = stride.max(1);
let harvest = move |_xml: &[u8], elems: &[ElemRef]| -> Vec<(f64, f64)> {
elems
.iter()
.filter(|e| e.kind == ElemKind::Node)
.step_by(stride)
.map(|e| (e.lon_e7 as f64 / 1e7, e.lat_e7 as f64 / 1e7))
.collect()
};
harvest_xml_bz2(
path,
Bz2Skip::Changesets(crate::xml_to_pbf::DATA_NEEDLES),
n_workers,
&harvest,
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::xml_vtd::ElemKind;
/// A REAL bzip2'd OSM XML extract that is many bz2 blocks long — the only
/// input on which a block-boundary bug can show at all. Lives in the
/// `lbzip2` sibling crate's test data (5.2 MB, ~60 MB decoded).
fn liechtenstein_bz2() -> std::path::PathBuf {
std::path::PathBuf::from(concat!(
env!("CARGO_MANIFEST_DIR"),
"/../lbzip2/test_data/liechtenstein.osm.bz2"
))
}
/// Decode a whole `.bz2` through the BLOCK-PARALLEL engine and concatenate
/// the segments in stream order — the bytes the harvest workers actually see.
/// Uses the production split ([`bz2_split`]) and the production per-segment
/// decoder, so what it proves is what production does.
fn block_parallel_decode(path: &std::path::Path, n_workers: usize) -> Result<Vec<u8>> {
use crate::xml_to_pbf::{Bz2Prepared, Bz2Skip, bz2_prepare, bz2_split, mk_bz2_cfg};
struct RawCodec {
max_blocksize: u32,
}
impl gatling::TypedCodec for RawCodec {
type Seg = (u64, u64);
type Output = Vec<u8>;
fn split(
&self,
data: &[u8],
n: usize,
is_last: bool,
) -> Option<gatling::Split<(u64, u64)>> {
bz2_split(data, n, self.max_blocksize, is_last)
}
fn transform(&self, data: &[u8], &(s, e): &(u64, u64)) -> Vec<u8> {
lbzip2::chunk::decode_segment(data, s, e, self.max_blocksize)
}
}
struct Cat(Vec<u8>);
impl gatling::TypedSink<Vec<u8>> for Cat {
fn process(&mut self, out: Vec<u8>, _is_last: bool) -> Result<()> {
self.0.extend_from_slice(&out);
Ok(())
}
}
let Bz2Prepared {
reader,
max_blocksize,
initial_carry,
..
} = bz2_prepare(path, Bz2Skip::None)?;
let mut sink = Cat(Vec::new());
gatling::run_typed(
reader,
RawCodec { max_blocksize },
&mut sink,
n_workers,
mk_bz2_cfg(initial_carry),
)?;
Ok(sink.0)
}
/// **LAW 2 — the block-boundary guard.** A block-parallel decode must be
/// BYTE-IDENTICAL to a serial one.
///
/// bzip2 blocks are bit-aligned (`1AY&SY`, not byte-aligned) and the parallel
/// path finds their boundaries itself, so the failure mode is silent: a
/// dropped last block, an off-by-one bit, a segment decoded twice. Every one
/// of those still produces plausible XML — the only check that can see them
/// is the whole byte stream against an independent SERIAL decoder
/// (`lbzip2::stream::decompress`, a different implementation, not ours
/// re-run with one worker).
///
/// RED-PROVEN 2026-08-02 on oden: truncating the parallel output by its last
/// segment fails on the length assert (`60_121_XXX vs 60_125_XXX`), and
/// dropping the FIRST segment fails on the byte-mismatch assert. The equality
/// is therefore load-bearing, not decorative.
#[test]
fn block_parallel_bz2_decode_is_byte_identical_to_serial() {
let path = liechtenstein_bz2();
let raw = std::fs::read(&path).expect("liechtenstein.osm.bz2 fixture (lbzip2/test_data)");
let serial = lbzip2::stream::decompress(&raw).expect("serial bz2 decode");
assert!(
serial.len() > 40_000_000,
"fixture decodes to a real extract, got {}",
serial.len()
);
for workers in [1usize, 8, 31] {
let par = block_parallel_decode(&path, workers).expect("block-parallel decode");
assert_eq!(
par.len(),
serial.len(),
"block-parallel decode with {workers} workers lost/duplicated bytes: \
{} vs serial {}",
par.len(),
serial.len(),
);
let first_diff = par.iter().zip(serial.iter()).position(|(a, b)| a != b);
assert!(
first_diff.is_none(),
"block-parallel decode with {workers} workers differs from serial at byte {:?}",
first_diff,
);
}
}
/// **LAW 2 — the harvest must see every element the serial decode contains.**
///
/// Byte-identity above proves the DECODE. This proves the layer that splits
/// those bytes into elements: an element that straddles a bz2 block boundary
/// is carried as a stub and stitched by the collector, and the LAST element of
/// the file is only harvested if the trailing stub is flushed. Both are easy
/// to get silently wrong, and neither changes the decoded bytes at all.
///
/// The reference is the serial decode parsed in one piece — no stubs, no
/// segments, nothing of the parallel path involved.
///
/// RED-PROVEN 2026-08-02 on oden, and one honest gap:
/// * removing the segment-boundary stitch in `process` → **RED**, 395 778
/// elements found vs 395 807 in the serial decode (29 lost at 31 workers);
/// * dropping the last segment of every slot in `bz2_split` → **RED** here
/// and in the byte-identity guard above;
/// * removing the TRAILING `flush_stub` in [`harvest_xml_bz2`] → still
/// green. That line is **not covered**: a well-formed `.osm` ends with
/// `</osm>` after its last complete element, so the final stub holds no
/// element to lose. It is kept because a truncated stream does not have
/// that courtesy — but no test here proves it, and pretending otherwise
/// would be exactly the hollow guard this comment exists to avoid.
#[test]
fn bz2_harvest_sees_every_element_a_serial_decode_does() {
let path = liechtenstein_bz2();
let raw = std::fs::read(&path).expect("liechtenstein.osm.bz2 fixture (lbzip2/test_data)");
let serial = lbzip2::stream::decompress(&raw).expect("serial bz2 decode");
// Reference: parse the whole decoded file in one piece.
let want: Vec<(u8, i64)> = vtd_parse_serial(&serial)
.iter()
.map(|e| (e.kind as u8, e.id))
.collect();
assert!(
want.len() > 100_000,
"fixture holds a real element count, got {}",
want.len()
);
let harvest = |xml: &[u8], elems: &[ElemRef]| -> Vec<(u8, i64)> {
let _ = xml;
elems.iter().map(|e| (e.kind as u8, e.id)).collect()
};
for workers in [1usize, 31] {
let got =
harvest_xml_bz2(&path, Bz2Skip::None, workers, &harvest).expect("bz2 harvest");
assert_eq!(
got.len(),
want.len(),
"block-parallel harvest with {workers} workers found {} elements, \
serial decode holds {} — a boundary or tail element was dropped",
got.len(),
want.len(),
);
assert_eq!(
got, want,
"block-parallel harvest with {workers} workers disagrees with the serial \
decode on element identity/order",
);
}
}
/// The changeset-prologue seek must not cost elements: skipping to the first
/// data block yields exactly the node/way/relation elements, and the same set
/// at 1 worker and at 31.
#[test]
fn bz2_harvest_skip_to_data_keeps_every_osm_element() {
let path = liechtenstein_bz2();
let harvest = |_xml: &[u8], elems: &[ElemRef]| -> Vec<(u8, i64)> {
elems.iter().map(|e| (e.kind as u8, e.id)).collect()
};
let whole =
harvest_xml_bz2(&path, Bz2Skip::None, 31, &harvest).expect("whole-file harvest");
let skipped = harvest_xml_bz2(
&path,
Bz2Skip::Changesets(crate::xml_to_pbf::DATA_NEEDLES),
31,
&harvest,
)
.expect("skip-to-data harvest");
assert_eq!(
whole,
skipped,
"the changeset seek dropped OSM elements ({} whole vs {} skipped)",
whole.len(),
skipped.len(),
);
assert!(
whole.iter().any(|&(k, _)| k == ElemKind::Relation as u8),
"the fixture must carry relations, or the admin scan proves nothing",
);
}
/// The decode probe measures the SAME bytes the serial decoder produces —
/// so a throughput number taken with it is a number about real work.
#[test]
fn bz2_decode_probe_counts_the_serial_byte_total() {
let path = liechtenstein_bz2();
let raw = std::fs::read(&path).expect("liechtenstein.osm.bz2 fixture");
let serial = lbzip2::stream::decompress(&raw).expect("serial bz2 decode");
for workers in [1usize, 31] {
let n = bz2_decode_probe(&path, workers, false).expect("decode probe");
assert_eq!(
n,
serial.len() as u64,
"decode probe at {workers} workers counted {n} bytes, serial decode is {}",
serial.len(),
);
}
}
}