dial9-viewer 0.5.0-rc2

//! Sampling based, demand-driven aggregation
//!
//! Instead of batch-aggregating an entire window up front, a query folds
//! [source files] one at a time, in a deterministic pseudo-random [order key]
//! order that is uniform across host and time, so the first few files are a
//! representative spread of the scope. Results are served over whatever subset
//! has been folded so far, with a [coverage] report, and refined as more files
//! fold.
//!
//! Key properties (see `docs/adr/0003-folded-set-is-the-output-listing.md`):
//!
//! - **The output part-file's existence is the record that a file is folded.**
//!   There is no manifest and no skip-set. A zero-sample file still writes an
//!   empty part-file, so it is never re-fetched.
//! - **Folding is idempotent.** A source file folds to a deterministically
//!   named part-file (`samples/service=…/date=…/host=…/{blake3(source_key)}`),
//!   so re-folding writes the same key.
//! - **Aggregation reads part-files through the `StorageBackend`**, so it works
//!   identically over S3, the local FS, and the simulated S3 used in tests.
//!
//! [source files]: crate::ingest::aggregate
//! [order key]: order_key
//! [coverage]: Coverage

use crate::storage::{ObjectInfo, StorageBackend};
use arrow::array::Array;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::Semaphore;

use super::decode;
use super::parquet_writer;

/// Scheduling version baked into the [`order_key`] hash input. Bump to change
/// the fetch-order permutation. Lives ONLY here, never in an output path: the
/// folded samples are order-independent and must survive a bump untouched.
pub(crate) const ORDER_VERSION: u32 = 1;

/// Storage-format version baked into the output key path
/// (`{output_prefix}/v{N}/…`). Bump when changing *what* we persist and we want
/// a deliberate recompute; reads/writes then target a fresh empty tree that
/// repopulates lazily. The value is a monotonic cache namespace, not a schema
/// revision, so skipped values are expected. The old tree is abandoned and
/// GC'd out-of-band.
pub const SAMPLES_FORMAT_VERSION: u32 = 7;

/// Default raw-trace segment duration, in seconds. A source file covers
/// `[epoch, epoch + segment_duration)`; the [`Scope`] time filter pads by this
/// so a file that *started* just before the window but runs into it is not
/// dropped. Configurable per deployment (the producer's rotation period).
pub(crate) const DEFAULT_SEGMENT_DURATION_SECS: i64 = 60;

/// The deterministic pseudo-random total order over source files:
/// `BLAKE3(ORDER_VERSION_le ++ source_key)`. Uniform across host and time, so
/// the first K files in this order are a representative spread of a scope
/// rather than one host's earliest minutes.
fn order_key(source_key: &str) -> [u8; 16] {
    let mut hasher = blake3::Hasher::new();
    hasher.update(&ORDER_VERSION.to_le_bytes());
    hasher.update(source_key.as_bytes());
    let mut id = [0u8; 16];
    id.copy_from_slice(&hasher.finalize().as_bytes()[..16]);
    id
}

/// Content-addressed leaf name for a source file's output part-files: the first
/// 128 bits of `BLAKE3(source_key)`, hex-encoded (32 chars). Stable + unique per
/// source file, so folding is idempotent; the part-files are named after this
/// leaf, so it is also the membership token of the folded set. 128 bits is
/// ample collision resistance for any realistic file count (same width as
/// `stack_id`), and the shorter leaf keeps the partitioned key well within
/// filesystem path-component limits.
pub(crate) fn part_leaf_of(source_key: &str) -> String {
    let hash = blake3::hash(source_key.as_bytes());
    hash.to_hex()[..32].to_string()
}

/// `{output_prefix}/v{SAMPLES_FORMAT_VERSION}/bucket={source_bucket}` — the root
/// of the current storage-format generation for one source bucket. The output
/// store is namespaced by source bucket so that bring-your-own-credentials
/// sources fold into isolated, independently-prunable/GC-able trees and the
/// folded-set LIST never mixes buckets.
fn versioned_root(output_prefix: &str, source_bucket: &str) -> String {
    let p = output_prefix.trim_end_matches('/');
    debug_assert!(!p.is_empty(), "output_prefix must not be empty");
    let bucket = bucket_segment(source_bucket);
    format!("{p}/v{SAMPLES_FORMAT_VERSION}/bucket={bucket}")
}

/// The source bucket a `source_key` came from: the `{bucket}` in
/// `s3://{bucket}/{key}`, or `"local"` for a bare (local-FS) key. This is what
/// namespaces the output path, so it is derived from the key itself — the
/// per-file path functions never need it threaded in separately.
fn parse_source_bucket(source_key: &str) -> String {
    if let Some(rest) = source_key.strip_prefix("s3://") {
        rest.split_once('/')
            .map(|(b, _)| b)
            .unwrap_or(rest)
            .to_string()
    } else {
        "local".to_string()
    }
}

/// Sanitize a bucket name for use as a path segment. S3 bucket names are already
/// path-safe (lowercase, digits, `-`, `.`), but guard against `/` defensively.
fn bucket_segment(source_bucket: &str) -> String {
    source_bucket.replace('/', "_")
}

/// Output key for a source file's samples part-file: a Hive-partitioned path
/// (`bucket=…/samples/service=…/date=…/host=…/{hash}.parquet`) so the folded-set
/// LIST is scope-prunable and DataFusion-style partition pruning is possible.
/// The hash is only the leaf; the source bucket is derived from `source_key`.
fn samples_part_key(output_prefix: &str, source_key: &str) -> String {
    let (date, service, host) = parse_scope_fields(source_key);
    format!(
        "{root}/samples/service={service}/date={date}/host={host}/{leaf}.parquet",
        root = versioned_root(output_prefix, &parse_source_bucket(source_key)),
        leaf = part_leaf_of(source_key),
    )
}

/// Output key for a source file's stacks-dictionary part-file. Content-addressed
/// stack_ids dedup naturally across files when the dicts are merged.
fn dict_part_key(output_prefix: &str, source_key: &str) -> String {
    format!(
        "{root}/dict/stacks/{leaf}.parquet",
        root = versioned_root(output_prefix, &parse_source_bucket(source_key)),
        leaf = part_leaf_of(source_key),
    )
}

fn polls_part_key(output_prefix: &str, source_key: &str) -> String {
    format!(
        "{root}/polls/{leaf}.parquet",
        root = versioned_root(output_prefix, &parse_source_bucket(source_key)),
        leaf = part_leaf_of(source_key),
    )
}

fn spans_part_key(output_prefix: &str, source_key: &str) -> String {
    let (date, service, host) = parse_scope_fields(source_key);
    format!(
        "{root}/spans/service={service}/date={date}/host={host}/{leaf}.parquet",
        root = versioned_root(output_prefix, &parse_source_bucket(source_key)),
        leaf = part_leaf_of(source_key),
    )
}

/// Public accessor for `spans_part_key` used by span_stats endpoint.
pub(crate) fn spans_part_key_pub(output_prefix: &str, source_key: &str) -> String {
    spans_part_key(output_prefix, source_key)
}

/// The `samples/` prefix for one source bucket under the versioned root — the
/// folded-set LIST target. Pruned to a single source bucket.
fn samples_prefix(output_prefix: &str, source_bucket: &str) -> String {
    format!("{}/samples/", versioned_root(output_prefix, source_bucket))
}

/// A query's selection: an optional wall-clock time range (epoch nanoseconds)
/// and optional service / host filters. Translated to the matched set.
#[derive(Debug, Clone, Default)]
pub(crate) struct Scope {
    pub start_ns: Option<i64>,
    pub end_ns: Option<i64>,
    /// Exact service match (the `service=` path component).
    pub service: Option<String>,
    /// Host filter. Empty = all hosts. Non-empty = the host path component must
    /// equal one of these (a *set*, because a heatmap box selection spans many
    /// hosts). A single entry behaves like an exact single-host filter.
    pub hosts: Vec<String>,
}

/// Parse `(date, service, host)` from a source key, anchored on the
/// `YYYY-MM-DD` date component so a leading prefix (e.g. `traces/`) does not
/// shift the positions. Layout: `…/{date}/{HHMM}/{service}/{host}/{boot}/{file}`.
fn parse_scope_fields(key: &str) -> (String, String, String) {
    let path = strip_s3(key);
    let parts: Vec<&str> = path.split('/').collect();
    if let Some(anchor) = parts.iter().position(|p| is_date(p)) {
        let date = parts.get(anchor).copied().unwrap_or("").to_string();
        let service = parts.get(anchor + 2).copied().unwrap_or("").to_string();
        let host = parts.get(anchor + 3).copied().unwrap_or("").to_string();
        (date, service, host)
    } else {
        // No date anchor — fall back to the legacy fixed-index parse.
        let date = parts.first().copied().unwrap_or("").to_string();
        let service = parts.get(2).copied().unwrap_or("").to_string();
        let host = parts.get(3).copied().unwrap_or("").to_string();
        (date, service, host)
    }
}

/// The host path component of a source key (empty when the key has no host
/// segment). Used to count distinct hosts for the coverage's fleet-spread badge.
pub(crate) fn host_of(key: &str) -> String {
    parse_scope_fields(key).2
}

/// Parse the file start time (epoch SECONDS) from the filename `{ts}-{i}.bin.gz`.
fn parse_epoch_secs(key: &str) -> Option<i64> {
    let file = key.rsplit('/').next()?;
    let stem = file.split('.').next()?; // strip .bin.gz
    let ts = stem.split('-').next()?; // {ts}-{i}
    ts.parse::<i64>().ok()
}

fn strip_s3(key: &str) -> &str {
    if let Some(rest) = key.strip_prefix("s3://") {
        rest.split_once('/').map_or(rest, |(_, p)| p)
    } else {
        key
    }
}

fn is_date(s: &str) -> bool {
    let b = s.as_bytes();
    b.len() == 10
        && b[4] == b'-'
        && b[7] == b'-'
        && b[..4].iter().all(u8::is_ascii_digit)
        && b[5..7].iter().all(u8::is_ascii_digit)
        && b[8..].iter().all(u8::is_ascii_digit)
}

/// True if a raw source key is a trace segment (not our own Parquet output).
pub(crate) fn is_trace_segment(key: &str) -> bool {
    (key.ends_with(".bin.gz") || key.ends_with(".bin"))
        && !key.contains("/samples/")
        && !key.contains("/dict/")
        && !key.contains("/flamegraph-data/")
}

/// Filter a raw source listing down to the files a [`Scope`] selects, then sort
/// by [`order_key`]. The result is the ordered matched set: the coverage
/// denominator and the fold order.
fn matched_and_ordered(
    objects: Vec<ObjectInfo>,
    scope: &Scope,
    segment_duration_secs: i64,
) -> Vec<ObjectInfo> {
    let mut matched: Vec<ObjectInfo> = objects
        .into_iter()
        .filter(|o| is_trace_segment(&o.key))
        .filter(|o| scope_matches(&o.key, scope, segment_duration_secs))
        .collect();
    matched.sort_by_key(|o| order_key(&o.key));
    matched
}

fn scope_matches(key: &str, scope: &Scope, segment_duration_secs: i64) -> bool {
    let (_date, service, host) = parse_scope_fields(key);
    if let Some(want) = &scope.service
        && &service != want
    {
        return false;
    }
    if !scope.hosts.is_empty() && !scope.hosts.iter().any(|h| h == &host) {
        return false;
    }
    // Interval-overlap on wall-clock time, padding by the segment duration so a
    // file that started before the window but runs into it is kept.
    // Uses half-open interval semantics: file range [start, end) overlaps query
    // [start_ns, end_ns) iff file_start < query_end && file_end > query_start.
    if scope.start_ns.is_some() || scope.end_ns.is_some() {
        let Some(epoch_secs) = parse_epoch_secs(key) else {
            // Can't place it in time — keep it rather than silently drop data.
            return true;
        };
        let file_start_ns = epoch_secs.saturating_mul(1_000_000_000);
        let file_end_ns = (epoch_secs + segment_duration_secs).saturating_mul(1_000_000_000);
        if let Some(start) = scope.start_ns
            && file_end_ns <= start
        {
            return false;
        }
        if let Some(end) = scope.end_ns
            && file_start_ns >= end
        {
            return false;
        }
    }
    true
}

/// The canonical source key recorded in part-files and passed to the decoder
/// (which derives host/service/date): a bare key for a local source,
/// `s3://{bucket}/{key}` for S3.
fn full_source_key(source_is_local: bool, source_bucket: &str, key: &str) -> String {
    if source_is_local {
        key.to_string()
    } else {
        format!("s3://{source_bucket}/{key}")
    }
}

/// Process-global concurrency limits for the demand-driven fold pipeline.
///
/// These are shared across all in-flight `/api/flamegraph` requests (held in
/// `AppState` and cloned per request — the inner `Arc<Semaphore>`s are shared),
/// so total fold work is bounded *application-wide* rather than per request.
/// Without this, N concurrent polls each running their own bounded batch could
/// still oversubscribe the box by a factor of N.
///
/// The stages are bounded independently because they bottleneck on different
/// resources:
///
/// - [`fetch`](Self::fetch): network-bound source GETs (~37–50 MB compressed
///   each). Mostly waiting on the network; bounded by `inflight` in practice
///   (a fetch can't start without an inflight permit).
/// - [`cpu`](Self::cpu): gunzip + decode + parquet-encode, run on blocking
///   threads. This is the dominant memory consumer — decoding one segment
///   expands to >1 GB of transient structures — so it is capped by a small
///   ABSOLUTE number ([`MAX_DECODE_CONCURRENCY`]), NOT by core count. Scaling it
///   with cores is what OOM-killed a 64-core box.
/// - [`inflight`](Self::inflight): the total number of folds that may hold a
///   fetched-but-not-yet-written segment buffer at once. This is the memory
///   backstop: [`fold_one`] releases the fetch permit before acquiring the CPU
///   permit (so network and CPU overlap), which means a decoded-slower-than-
///   fetched work-list would otherwise let downloaded buffers pile up without
///   bound between the two stages — thousands of them for a broad scope,
///   OOM-killing the process. Holding one inflight permit across the *whole*
///   fold caps the resident segment buffers regardless of work-list size.
///
/// Part-file writes (small, network-bound) run ungated: they are cheap relative
/// to the fetch and we don't want to hold a CPU permit across their I/O.
#[derive(Clone)]
pub(crate) struct FoldLimits {
    /// Bounds concurrent source fetches (network-bound).
    pub fetch: Arc<Semaphore>,
    /// Bounds concurrent decode/encode work (CPU-bound).
    pub cpu: Arc<Semaphore>,
    /// Bounds concurrently in-flight folds (memory backstop; held fetch→write).
    pub inflight: Arc<Semaphore>,
}

impl FoldLimits {
    /// Construct with explicit permit counts. Each is clamped to at least 1 so a
    /// zero never deadlocks the pipeline. `inflight` bounds resident segment
    /// buffers and must be ≥ `fetch` (a fetch can't proceed without an inflight
    /// permit, so a smaller inflight would cap effective fetch concurrency).
    pub(crate) fn new(fetch_permits: usize, cpu_permits: usize, inflight_permits: usize) -> Self {
        Self {
            fetch: Arc::new(Semaphore::new(fetch_permits.max(1))),
            cpu: Arc::new(Semaphore::new(cpu_permits.max(1))),
            inflight: Arc::new(Semaphore::new(inflight_permits.max(1))),
        }
    }

    /// Default sizing.
    ///
    /// Memory — not cores — is the binding constraint on a fold, so the
    /// memory-bound stages use small ABSOLUTE caps rather than scaling with
    /// parallelism (an earlier core-scaled sizing OOM-killed a 64-core box:
    /// 64 concurrent decodes × ~1.5 GB each ≈ 96 GB):
    ///
    /// - `cpu` (decode/encode) is the real memory lever: decoding one segment
    ///   expands ~1.3M events to well over 1 GB of transient structures. Capped
    ///   at [`MAX_DECODE_CONCURRENCY`] so peak decode memory is bounded to roughly
    ///   that many segments regardless of core count, then further limited to the
    ///   available parallelism on small boxes.
    /// - `inflight` (whole-fold, memory backstop) bounds how many fetched
    ///   segment buffers are resident; a small multiple of the decode cap gives
    ///   fetch a little runway ahead of decode without unbounded pile-up.
    /// - `fetch` (network) can safely exceed the memory caps since a GET only
    ///   holds a compressed (~40 MB) buffer, but it never needs to exceed
    ///   `inflight` (a fetch can't start without an inflight permit).
    pub(crate) fn from_available_parallelism() -> Self {
        let par = std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(4);
        let cpu = par.min(MAX_DECODE_CONCURRENCY);
        let inflight = cpu * 2;
        let fetch = inflight;
        Self::new(fetch, cpu, inflight)
    }
}

/// Hard ceiling on concurrent segment decodes. Each decode transiently holds
/// over a gigabyte (a ~1.3M-event segment expanded in memory), so this count
/// times ~1.5 GB is roughly the fold pipeline's peak decode memory. Deliberately
/// small and core-independent: adding cores speeds each decode but must not
/// multiply the memory high-water mark.
const MAX_DECODE_CONCURRENCY: usize = 6;

impl Default for FoldLimits {
    fn default() -> Self {
        Self::from_available_parallelism()
    }
}

/// Encoded part-file buffers produced by the CPU stage of a fold, ready to write.
struct EncodedParts {
    samples_buf: Vec<u8>,
    dict_buf: Vec<u8>,
    polls_buf: Vec<u8>,
    spans_buf: Vec<u8>,
    /// CPU-stage timing/count breakdown, threaded out to the per-file metric.
    cpu_stats: CpuStageStats,
}

/// Timing and size breakdown of the CPU stage ([`decode_and_encode`]), surfaced
/// to the per-file [`FoldFileMetrics`] so the fold's cost is attributable to
/// gunzip vs. the individual decode phases vs. parquet encode.
#[derive(Default)]
struct CpuStageStats {
    decompressed_bytes: u64,
    gunzip: std::time::Duration,
    parquet_encode: std::time::Duration,
    decode: crate::ingest::decode::DecodeStats,
}

/// CPU-bound stage of a fold: gunzip + decode + parquet-encode over
/// already-fetched bytes. Pure and synchronous so the caller can run it on a
/// blocking thread under a CPU concurrency permit, decoupled from the
/// network-bound fetch and write stages. (The parquet encode previously ran on
/// the async executor thread; moving it here keeps CPU work off the runtime.)
fn decode_and_encode(bytes: &[u8], full_key: &str) -> anyhow::Result<EncodedParts> {
    use std::time::Instant;
    let mut cpu_stats = CpuStageStats::default();

    let t_gunzip = Instant::now();
    let raw = maybe_gunzip(bytes);
    cpu_stats.gunzip = t_gunzip.elapsed();
    cpu_stats.decompressed_bytes = raw.len() as u64;

    let ((samples, stacks, polls, spans), decode_stats) =
        decode::decode_samples_with_stats(&raw, full_key)
            .map_err(|e| anyhow::anyhow!("decode {full_key}: {e}"))?;
    cpu_stats.decode = decode_stats;

    // Always encode the samples part-file, even with zero rows: its existence is
    // the record that this file is folded, so it is never re-fetched.
    let t_encode = Instant::now();
    let metadata = HashMap::new();
    let mut samples_buf = Vec::new();
    parquet_writer::write_samples(&mut samples_buf, &samples, &metadata)?;

    let stacks_map: HashMap<[u8; 16], Vec<String>> = stacks.into_iter().collect();
    let mut dict_buf = Vec::new();
    parquet_writer::write_stacks_dict(&mut dict_buf, &stacks_map)?;

    let mut polls_buf = Vec::new();
    parquet_writer::write_polls(&mut polls_buf, &polls)?;

    // Write spans part-file (may be empty for files with no tracing spans).
    let mut spans_buf = Vec::new();
    parquet_writer::write_spans(&mut spans_buf, &spans)?;
    cpu_stats.parquet_encode = t_encode.elapsed();

    Ok(EncodedParts {
        samples_buf,
        dict_buf,
        polls_buf,
        spans_buf,
        cpu_stats,
    })
}

/// Write stage of a fold: the (small) part-file PUTs.
///
/// The `samples/` part is the durable record of "this file is folded"
/// ([`list_folded_leaves`] lists it; see ADR-0003), so it MUST be written LAST,
/// only after the dict, polls, and spans parts have landed. Writing it
/// concurrently would let a mid-write failure — or a cancelled fold task (the
/// streaming endpoints abort in-flight folds whenever the client disconnects) —
/// commit a file as folded while its dict/polls/spans parts are missing,
/// permanently: a folded file is never re-folded, so the gap would never heal.
/// Orphaned dict/polls/spans parts from the reverse interleaving are harmless —
/// the file stays unfolded and a later re-fold idempotently overwrites the same
/// keys.
async fn write_parts(
    output: &dyn StorageBackend,
    output_bucket: &str,
    output_prefix: &str,
    full_key: &str,
    encoded: EncodedParts,
) -> anyhow::Result<()> {
    let part_key = samples_part_key(output_prefix, full_key);
    let dict_key = dict_part_key(output_prefix, full_key);
    let polls_key = polls_part_key(output_prefix, full_key);
    let spans_key = spans_part_key(output_prefix, full_key);
    // Write dict, polls, and spans concurrently — all before the samples commit marker.
    let (dict_res, polls_res, spans_res) = tokio::join!(
        output.put_object(output_bucket, &dict_key, encoded.dict_buf),
        output.put_object(output_bucket, &polls_key, encoded.polls_buf),
        output.put_object(output_bucket, &spans_key, encoded.spans_buf),
    );
    dict_res.map_err(|e| anyhow::anyhow!("write dict {dict_key}: {e}"))?;
    polls_res.map_err(|e| anyhow::anyhow!("write polls {polls_key}: {e}"))?;
    spans_res.map_err(|e| anyhow::anyhow!("write spans {spans_key}: {e}"))?;
    // Samples part LAST — its presence is the folded-set record (ADR-0003).
    output
        .put_object(output_bucket, &part_key, encoded.samples_buf)
        .await
        .map_err(|e| anyhow::anyhow!("write samples {part_key}: {e}"))?;
    Ok(())
}

/// Fold one source file: fetch, decode, and write its samples + stacks-dict
/// part-files to deterministic keys, gating the network fetch and the CPU
/// decode/encode on two separate, caller-supplied (process-global) semaphores.
///
/// The fetch permit is released before CPU work starts, and the CPU permit is
/// released before the part-file writes, so each global limit bounds only the
/// stage it is meant to bound. A zero-sample file still writes an empty samples
/// part-file (the "folded" record). Idempotent: re-folding writes the same keys.
pub(crate) async fn fold_one(
    agg: &AggContext,
    raw_key: &str,
    limits: &FoldLimits,
) -> anyhow::Result<()> {
    use std::time::Instant;
    // One per-file metric per fold: assemble phase timings as we go, emit once on
    // the way out (success or failure). See `FoldFileMetrics`.
    let mut metric = crate::server::metrics::FoldFileMetricsBuilder::new();
    let t_total = Instant::now();
    // Emit-on-exit helper so every early return still publishes what we measured.
    let emit = |mut metric: crate::server::metrics::FoldFileMetricsBuilder,
                total: std::time::Duration,
                failed: bool| {
        metric.total(total).failed(failed);
        metric.emit();
    };

    // Memory backstop — held across the WHOLE fold (fetch → decode → write), so
    // the number of resident ~40 MB segment buffers is bounded regardless of how
    // many files the work-list spawned. Without this, fetch (fast, network) races
    // ahead of decode (slow, CPU) and downloaded buffers pile up between the two
    // stages, OOM-killing the process on a broad scope. Acquired before the fetch
    // permit so a fold doesn't occupy fetch concurrency while waiting for memory.
    let _inflight = limits
        .inflight
        .acquire()
        .await
        .expect("inflight semaphore is never closed");

    // Stage 1 — fetch (network-bound). Permit held only for the duration of the GET.
    let t_fetch = Instant::now();
    let bytes = {
        let _permit = limits
            .fetch
            .acquire()
            .await
            .expect("fetch semaphore is never closed");
        match agg.source.get_object(&agg.source_bucket, raw_key).await {
            Ok(b) => b,
            Err(e) => {
                emit(metric, t_total.elapsed(), true);
                return Err(anyhow::anyhow!("fetch {raw_key}: {e}"));
            }
        }
    };
    metric.fetch(t_fetch.elapsed());
    metric.source_bytes(bytes.len() as u64);

    // Stage 2 — decode + encode (CPU-bound) on a blocking thread, gated so
    // concurrent folds don't oversubscribe the cores.
    let full_key = full_source_key(agg.source_is_local, &agg.source_bucket, raw_key);
    let decode_key = full_key.clone();
    let encoded = {
        let _permit = limits
            .cpu
            .acquire()
            .await
            .expect("cpu semaphore is never closed");
        match tokio::task::spawn_blocking(move || decode_and_encode(&bytes, &decode_key)).await {
            Ok(Ok(encoded)) => encoded,
            Ok(Err(e)) => {
                emit(metric, t_total.elapsed(), true);
                return Err(e);
            }
            Err(e) => {
                emit(metric, t_total.elapsed(), true);
                return Err(anyhow::anyhow!("decode task panicked: {e}"));
            }
        }
    };
    // Fold the CPU-stage breakdown into the metric before `encoded` is consumed
    // by the writer.
    let cpu = &encoded.cpu_stats;
    metric
        .gunzip(cpu.gunzip)
        .decompressed_bytes(cpu.decompressed_bytes)
        .decode_phases(&cpu.decode)
        .parquet_encode(cpu.parquet_encode);

    // Stage 3 — write part-files (small, network-bound). Ungated.
    let t_write = Instant::now();
    let write_result = write_parts(
        &*agg.output,
        &agg.output_bucket,
        &agg.output_prefix,
        &full_key,
        encoded,
    )
    .await;
    metric.write_parts(t_write.elapsed());

    let failed = write_result.is_err();
    emit(metric, t_total.elapsed(), failed);
    write_result
}

/// LIST the folded set for one source bucket and optional service: the
/// source-file leaf hashes that already have a samples part-file under that
/// partitioned root. A BYOC source's folded set never mixes with another's, and
/// a service scope does not enumerate sibling services.
pub(crate) async fn list_folded_leaves(
    output: &dyn StorageBackend,
    output_bucket: &str,
    output_prefix: &str,
    source_bucket: &str,
    service: Option<&str>,
) -> HashSet<String> {
    let prefix = folded_set_prefix(output_prefix, source_bucket, service);
    let objects = output
        .list_objects_all(output_bucket, &prefix)
        .await
        .unwrap_or_else(|e| {
            // Treating an error as "nothing folded" makes the refinement loop
            // re-fold files it already processed, wasting the whole budget on
            // redundant work. We can't cheaply propagate (callers expect a set),
            // but we must not swallow it silently.
            tracing::warn!(
                bucket = %output_bucket,
                prefix = %prefix,
                error = %e,
                "list_folded_leaves: failed to list folded set; treating as empty \
                 (already-folded files may be re-folded this round)"
            );
            Vec::new()
        });
    objects
        .iter()
        .filter_map(|o| {
            let name = o.key.rsplit('/').next()?;
            name.strip_suffix(".parquet").map(|s| s.to_string())
        })
        .collect()
}

fn folded_set_prefix(output_prefix: &str, source_bucket: &str, service: Option<&str>) -> String {
    let prefix = samples_prefix(output_prefix, source_bucket);
    match service {
        Some(service) => format!("{prefix}service={service}/"),
        None => prefix,
    }
}

/// How much of a scope has been folded so far. Reported on every query.
#[derive(Debug, Clone, serde::Serialize)]
pub(crate) struct Coverage {
    pub files_matched: usize,
    pub files_folded: usize,
    /// Deterministic digest of the successfully represented folded leaf set.
    /// Clients use this to ensure partial refreshes patch statistics produced
    /// from exactly the same files, not merely the same file count. Omitted by
    /// endpoints that do not track their represented set precisely.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub folded_set_id: Option<String>,
    /// Deterministic digest of the full cached seed set this stream is working
    /// toward. Seed-only exemplar refreshes expose it on every cumulative event,
    /// allowing clients to preview subsets only when the final target matches
    /// the catalog they are patching.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub target_folded_set_id: Option<String>,
    /// Number of matched files in the deterministic prefix eligible for new
    /// folding in this request. Cached coverage may exceed this value.
    pub fold_work_cap: usize,
    pub samples_folded: usize,
    /// Total bytes of all matched source files in the scope.
    pub total_bytes: u64,
    /// Distinct hosts across the matched set (the scope's fleet breadth).
    pub hosts_matched: usize,
    /// Distinct hosts among the folded files (how much of that breadth the
    /// current sample actually spans), so the UI can show fleet-representativeness
    /// e.g. "8 / 40 hosts".
    pub hosts_folded: usize,
    /// Number of files whose fold FAILED this stream (fetch/decode/write error —
    /// e.g. an unwritable output bucket → `PutObject` AccessDenied). Non-zero
    /// means the tree may be incomplete for a reason other than the sampling cap,
    /// so the UI surfaces it instead of showing a silent empty/partial result.
    pub fold_errors: usize,
    /// A representative fold error message (the most recent), for the UI to
    /// display. `None` when `fold_errors == 0`. Truncated to keep the event small.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub fold_error_sample: Option<String>,
}

/// Wire value of the `CpuProfile` CPU-sample source (periodic on-CPU sample).
const SOURCE_CPU_PROFILE: u8 = 0;
/// Wire value of the `SchedEvent` CPU-sample source (context switch, off-CPU).
/// Mirrors `CpuSampleSource::SchedEvent` and the JS `source === 1` check.
const SOURCE_SCHED_EVENT: u8 = 1;

// ─── Generic facet system ────────────────────────────────────────────────────
//
// Each facet is defined once in [`FACETS`]. The read loop extracts facet values
// generically, records them for the toolbar, and applies an optional exact-match
// filter. Adding a new facet requires only one new entry here (+ the Parquet
// column from ingest).

/// A facet definition: how to read a column and produce a string value.
#[derive(Clone)]
pub(crate) struct FacetDef {
    /// Query parameter / filter key name (e.g. `"source"`, `"thread_class"`).
    pub name: &'static str,
    /// Human label for the toolbar selector.
    pub label: &'static str,
    /// How to extract this facet's value from a row. Virtual facets derive from
    /// non-string columns; direct facets just read a nullable Utf8 column.
    pub kind: FacetKind,
    /// Default filter value when the param is absent. `"cpu"` for source (the
    /// on-CPU default), empty string for all others (= no constraint).
    pub default_filter: &'static str,
}

#[derive(Clone)]
pub(crate) enum FacetKind {
    /// Read from a UInt8 column and map wire values to labels.
    MappedU8 {
        column: &'static str,
        map: &'static [(u8, &'static str)],
        /// Fallback value when the column is absent (backwards compat with older
        /// part-files that predate this column).
        absent_value: &'static str,
    },
    /// Derived from a nullable column: `"worker"` if non-null, `"off-worker"` if null.
    NullDerived {
        column: &'static str,
        present_label: &'static str,
        absent_label: &'static str,
        /// Label when the entire column is missing from an old part-file.
        missing_column_label: &'static str,
    },
    /// Read directly from a nullable Utf8 column. Null rows produce no value
    /// (excluded from facet set and never match a filter).
    DirectString { column: &'static str },
}

/// The facet registry. Order here is the toolbar display order.
pub(crate) const FACETS: &[FacetDef] = &[
    FacetDef {
        name: "source",
        label: "Source",
        kind: FacetKind::MappedU8 {
            column: "source",
            map: &[(SOURCE_CPU_PROFILE, "cpu"), (SOURCE_SCHED_EVENT, "sched")],
            absent_value: "cpu",
        },
        default_filter: "cpu",
    },
    FacetDef {
        name: "thread_class",
        label: "Thread",
        kind: FacetKind::NullDerived {
            column: "worker_id",
            present_label: "worker",
            absent_label: "off-worker",
            missing_column_label: "worker",
        },
        default_filter: "",
    },
    FacetDef {
        name: "host",
        label: "Host",
        kind: FacetKind::DirectString { column: "host" },
        default_filter: "",
    },
    FacetDef {
        name: "spawn_location",
        label: "Task",
        kind: FacetKind::DirectString {
            column: "spawn_location",
        },
        default_filter: "",
    },
];

/// One facet's response: name + label + sorted distinct values.
#[derive(Debug, Clone, serde::Serialize)]
pub(crate) struct FacetResult {
    pub name: &'static str,
    pub label: &'static str,
    pub values: Vec<String>,
}

/// A generic set of active facet filters: name → required value. An entry with
/// an empty string means "no constraint" (match all). Entries are matched by
/// exact equality against the facet value extracted for each row.
pub(crate) type FacetFilters = HashMap<&'static str, String>;

/// Accumulates distinct facet values across part-files. One `HashSet<String>`
/// per facet definition.
#[derive(Clone)]
struct FacetAccum {
    /// Distinct values per facet (indexed same as [`FACETS`]).
    sets: Vec<HashSet<String>>,
    /// Distinct hosts that passed ALL filters (for the "N hosts" badge).
    matched_hosts: HashSet<String>,
}

impl FacetAccum {
    fn new() -> Self {
        Self {
            sets: FACETS.iter().map(|_| HashSet::new()).collect(),
            matched_hosts: HashSet::new(),
        }
    }

    /// Snapshot the accumulated facet values without consuming, so a streaming
    /// query can produce a fresh [`FacetResult`] set after every merged file.
    fn results(&self) -> Vec<FacetResult> {
        FACETS
            .iter()
            .zip(&self.sets)
            .map(|(def, set)| {
                let mut values: Vec<String> = set.iter().cloned().collect();
                values.sort();
                FacetResult {
                    name: def.name,
                    label: def.label,
                    values,
                }
            })
            .collect()
    }
}

/// The combined per-query filter: time range + poll-duration band + per-facet
/// exact-match filters + span-type filter.
#[derive(Debug, Clone, Default)]
pub(crate) struct SampleFilter {
    /// Optional time range filter (epoch nanoseconds, half-open: [start, end)).
    pub start_ns: Option<i64>,
    pub end_ns: Option<i64>,
    /// Optional poll-duration band (nanoseconds, inclusive: [min, max]). Keeps
    /// only samples attributed to a poll whose duration falls in the band — the
    /// "why are the slow polls slow" slice. A sample with no `poll_duration_ns`
    /// (off-worker / between polls) is excluded whenever either bound is set,
    /// since a poll-duration question only concerns in-poll samples.
    ///
    /// NOTE: this is *poll* duration (PollStart→PollEnd), not request/span
    /// latency — the decoder does not yet capture request spans. A future
    /// request-latency band would be a separate pair of fields, not a reuse of
    /// these.
    pub min_poll_ns: Option<i64>,
    pub max_poll_ns: Option<i64>,
    /// Per-facet filters. Key = facet name, value = required value. Empty string
    /// or absent = no constraint. For "source", the default is "cpu" (set by the
    /// endpoint when the param is absent).
    pub facets: FacetFilters,
    /// Span type UID filter (raw 16 bytes). When Some, only samples whose
    /// `enclosing_spans` column contains a matching type UID pass. Reads
    /// the new v4 `enclosing_spans` column; old part-files without it pass all
    /// samples when no span filter is active (backwards compatible).
    pub span_type_uid: Option<[u8; 16]>,
    /// Minimum span elapsed_ns for span filtering (inclusive).
    pub min_span_ns: Option<i64>,
    /// Maximum span elapsed_ns for span filtering (inclusive).
    pub max_span_ns: Option<i64>,
}

/// One bar of the poll-duration histogram: a log-scale duration bucket
/// `[lo_ns, hi_ns)` and the number of samples whose enclosing poll fell in it.
/// Sample-weighted, so bar height == the samples you'd get by selecting this
/// band. Emitted as explicit ns ranges so the UI needs no bucketing math. The
/// *filter* is never bucketed (`min/max_poll_ns` are exact ns); only the display
/// histogram is, and only because we can't ship one bar per distinct duration.
#[derive(Debug, Clone, serde::Serialize)]
pub(crate) struct PollDurationBucket {
    pub lo_ns: i64,
    pub hi_ns: i64,
    pub samples: u64,
}

/// Sub-octave subdivisions per power of two. `1` = plain log₂ (each bar 2× the
/// last); `4` = quarter-octave (each bar ≈1.19×), 4× finer while still bounded
/// (a fixed number of bars per decade). Bump for a finer histogram.
const POLL_HIST_SUBDIV: u32 = 4;

/// Log-scale bucket index for a positive poll duration, at [`POLL_HIST_SUBDIV`]
/// bins per octave: `floor(SUBDIV · log₂(ns))`. `None` for non-positive input.
/// Bucket `k` covers `[2^(k/SUBDIV), 2^((k+1)/SUBDIV)) ns` (see [`bucket_edge_ns`]).
fn poll_bucket(ns: i64) -> Option<u32> {
    if ns <= 0 {
        return None;
    }
    // floor(log₂) is the integer part; add the fractional octave via log2 of the
    // mantissa so sub-octave bins are uniform in log space.
    let log2 = (ns as f64).log2();
    Some((log2 * POLL_HIST_SUBDIV as f64).floor() as u32)
}

/// The lower ns edge of sub-octave bucket `k`: `2^(k/SUBDIV)`, rounded to a whole
/// ns. Monotonic in `k`, so `bucket_edge_ns(k+1)` is the bar's upper edge.
fn bucket_edge_ns(k: u32) -> i64 {
    2f64.powf(k as f64 / POLL_HIST_SUBDIV as f64).round() as i64
}

/// Sample-weighted poll-duration histogram: sub-octave bucket index → sample
/// count. Accumulated over rows that pass the time + facet filters but BEFORE the
/// poll band, so the bars always describe the full distribution the band selects
/// from.
type PollHist = HashMap<u32, u64>;

/// Convert the raw bucket map into sorted, explicit-range bars for the response.
/// Adjacent bars share an edge (`hi_ns` of bar `k` == `lo_ns` of bar `k+1`), so
/// a UI brush maps cleanly to a contiguous ns band.
fn poll_hist_bars(hist: &PollHist) -> Vec<PollDurationBucket> {
    let mut buckets: Vec<u32> = hist.keys().copied().collect();
    buckets.sort_unstable();
    buckets
        .into_iter()
        .map(|k| PollDurationBucket {
            lo_ns: bucket_edge_ns(k),
            hi_ns: bucket_edge_ns(k + 1),
            samples: hist[&k],
        })
        .collect()
}

/// A borrowed snapshot of a [`FlamegraphAccum`] at a point during streaming: the
/// dictionary is borrowed (never cloned per emit — it grows monotonically and is
/// the large part), while `stack_counts` and `facets` are materialized because
/// the caller iterates them to build the tree and toolbar.
pub(crate) struct AggSnapshot<'a> {
    pub stack_counts: Vec<(Vec<u8>, u64)>,
    pub stacks_dict: &'a HashMap<Vec<u8>, Vec<String>>,
    pub total_samples: usize,
    pub hosts: usize,
    pub min_ts: Option<i64>,
    pub max_ts: Option<i64>,
    /// Generic facet results: each facet's distinct values seen so far.
    pub facets: Vec<FacetResult>,
    /// Sample-weighted poll-duration histogram (the minimap over the band picker).
    pub poll_duration_histogram: Vec<PollDurationBucket>,
}

/// Incremental flamegraph accumulator: merge folded part-files one at a time
/// under a fixed [`SampleFilter`], so a streaming query can emit a fresh
/// [`snapshot`](Self::snapshot) after every file rather than re-reading the whole
/// folded set each poll. The shared accumulators live here and are merged into
/// serially, so no locking is needed even when part-file GETs run concurrently.
pub(crate) struct FlamegraphAccum {
    filter: SampleFilter,
    counts: HashMap<[u8; 16], u64>,
    dict: HashMap<Vec<u8>, Vec<String>>,
    facets: FacetAccum,
    total_samples: usize,
    min_ts: Option<i64>,
    max_ts: Option<i64>,
    /// Sample-weighted poll-duration histogram, accumulated pre-band (see
    /// [`PollHist`]).
    poll_hist: PollHist,
}

impl FlamegraphAccum {
    pub(crate) fn new(filter: SampleFilter) -> Self {
        Self {
            filter,
            counts: HashMap::new(),
            dict: HashMap::new(),
            facets: FacetAccum::new(),
            total_samples: 0,
            min_ts: None,
            max_ts: None,
            poll_hist: HashMap::new(),
        }
    }

    /// Merge one folded file's samples part-file (and its optional stacks dict)
    /// into the running totals **transactionally**: both the sample counts and the
    /// dict are staged into temporaries and committed only when the full merge
    /// (samples + dict) succeeds. A failure in either step leaves `self` unchanged.
    pub(crate) fn merge(&mut self, samples: Vec<u8>, dict: Option<Vec<u8>>) -> anyhow::Result<()> {
        // Stage: clone the mutable state that read_samples_part and read_dict_part
        // would mutate, so a failure in either step leaves self unchanged.
        let mut staged_counts = self.counts.clone();
        let mut staged_dict = self.dict.clone();
        let mut staged_facets = self.facets.clone();
        let mut staged_total = self.total_samples;
        let mut staged_min_ts = self.min_ts;
        let mut staged_max_ts = self.max_ts;
        let mut staged_poll_hist = self.poll_hist.clone();

        // Apply samples into the staged state.
        self.read_samples_part_into(
            samples,
            &mut staged_counts,
            &mut staged_dict,
            &mut staged_facets,
            &mut staged_total,
            &mut staged_min_ts,
            &mut staged_max_ts,
            &mut staged_poll_hist,
        )?;

        // Apply dict into the staged state.
        if let Some(dict) = dict {
            read_dict_part(dict, &mut staged_dict)?;
        }

        // Commit: both succeeded — swap staged state into self.
        self.counts = staged_counts;
        self.dict = staged_dict;
        self.facets = staged_facets;
        self.total_samples = staged_total;
        self.min_ts = staged_min_ts;
        self.max_ts = staged_max_ts;
        self.poll_hist = staged_poll_hist;
        Ok(())
    }

    /// Parse a single samples part-file and merge its rows into the provided
    /// staged accumulators, applying time/facet/band filters. Used by the
    /// transactional [`merge`](Self::merge) to stage mutations without touching
    /// `self` until both samples and dict succeed.
    #[allow(clippy::too_many_arguments)]
    fn read_samples_part_into(
        &self,
        data: Vec<u8>,
        counts: &mut HashMap<[u8; 16], u64>,
        _dict: &mut HashMap<Vec<u8>, Vec<String>>,
        facets: &mut FacetAccum,
        total_samples: &mut usize,
        min_ts: &mut Option<i64>,
        max_ts: &mut Option<i64>,
        poll_hist: &mut PollHist,
    ) -> anyhow::Result<()> {
        // `Bytes::from(Vec<u8>)` reuses the allocation (no copy); threading the
        // owned buffer in from the caller avoids the round-trip through `&[u8]`.
        let reader = ::parquet::arrow::arrow_reader::ParquetRecordBatchReader::try_new(
            bytes::Bytes::from(data),
            4096,
        )?;
        for batch in reader {
            let batch = batch?;
            let stack_col = batch.column_by_name("stack_id").and_then(|c| {
                c.as_any()
                    .downcast_ref::<arrow::array::FixedSizeBinaryArray>()
            });
            let Some(stack_arr) = stack_col else { continue };
            let ts_arr = batch
                .column_by_name("timestamp_ns")
                .and_then(|c| c.as_any().downcast_ref::<arrow::array::Int64Array>());
            // Poll-duration column, for the latency-band filter. Nullable and absent
            // from old part-files; either case means "no poll duration for this row".
            let poll_arr = batch
                .column_by_name("poll_duration_ns")
                .and_then(|c| c.as_any().downcast_ref::<arrow::array::Int64Array>());
            let poll_band = self.filter.min_poll_ns.is_some() || self.filter.max_poll_ns.is_some();

            // Pre-resolve column references for each facet in this batch.
            let facet_cols: Vec<ResolvedFacetCol> = FACETS
                .iter()
                .map(|def| resolve_facet_col(&batch, def))
                .collect();

            for i in 0..batch.num_rows() {
                // Time range filter.
                if let Some(ts) = ts_arr {
                    let v = ts.value(i);
                    if self.filter.start_ns.is_some_and(|start| v < start) {
                        continue;
                    }
                    if self.filter.end_ns.is_some_and(|end| v >= end) {
                        continue;
                    }
                }

                // Extract facet values for this row and record them (pre-filter).
                let mut row_values: Vec<Option<String>> = Vec::with_capacity(FACETS.len());
                for (fi, col) in facet_cols.iter().enumerate() {
                    let val = extract_facet_value(col, i);
                    if let Some(ref v) = val {
                        facets.sets[fi].insert(v.clone());
                    }
                    row_values.push(val);
                }

                // Apply facet filters: every active filter must match.
                let mut passes = true;
                for (fi, def) in FACETS.iter().enumerate() {
                    if let Some(wanted) = self.filter.facets.get(def.name) {
                        if wanted.is_empty() {
                            continue;
                        }
                        match &row_values[fi] {
                            Some(v) if v == wanted => {}
                            _ => {
                                passes = false;
                                break;
                            }
                        }
                    }
                }
                if !passes {
                    continue;
                }

                // This row passed the time + facet filters. Its enclosing-poll
                // duration (if any) feeds two things: the pre-band histogram (so the
                // minimap shows the full distribution the band selects from) and the
                // band filter itself.
                let poll_dur = poll_arr.and_then(|a| (!a.is_null(i)).then(|| a.value(i)));

                // Sample-weighted poll-duration histogram, BEFORE the band filter.
                // Off-poll rows (no duration) don't fall in any log₂ bucket, so they
                // simply don't contribute — matching the band's exclusion of them.
                if let Some(k) = poll_dur.and_then(poll_bucket) {
                    *poll_hist.entry(k).or_insert(0) += 1;
                }

                // Poll-duration band filter. A row with no poll duration (null column,
                // or the column absent in an old part-file) is excluded whenever a
                // band is set — the slice is inherently about in-poll samples.
                if poll_band {
                    match poll_dur {
                        Some(d) => {
                            if self.filter.min_poll_ns.is_some_and(|min| d < min) {
                                continue;
                            }
                            if self.filter.max_poll_ns.is_some_and(|max| d > max) {
                                continue;
                            }
                        }
                        None => continue,
                    }
                }

                // Span-type filter: when active, only samples whose
                // enclosing_spans list contains a matching span_type_uid pass.
                // Old part-files without the column pass all rows when no span
                // filter is active (backwards compatible per design §7).
                if let Some(ref wanted_uid) = self.filter.span_type_uid
                    && !span_filter_matches(
                        &batch,
                        i,
                        wanted_uid,
                        self.filter.min_span_ns,
                        self.filter.max_span_ns,
                    )
                {
                    continue;
                }

                // Count this sample.
                let mut id = [0u8; 16];
                id.copy_from_slice(stack_arr.value(i));
                *counts.entry(id).or_insert(0) += 1;
                *total_samples += 1;
                if let Some(ts) = ts_arr {
                    let v = ts.value(i);
                    *min_ts = Some(min_ts.map_or(v, |m| m.min(v)));
                    *max_ts = Some(max_ts.map_or(v, |m| m.max(v)));
                }
                // Track matched hosts for the "N hosts" badge.
                if let Some(ref h) = row_values[host_facet_index()] {
                    facets.matched_hosts.insert(h.clone());
                }
            }
        }
        Ok(())
    }

    /// A borrowed snapshot of the current totals for emitting one SSE event.
    pub(crate) fn snapshot(&self) -> AggSnapshot<'_> {
        let stack_counts: Vec<(Vec<u8>, u64)> =
            self.counts.iter().map(|(k, v)| (k.to_vec(), *v)).collect();
        AggSnapshot {
            stack_counts,
            stacks_dict: &self.dict,
            total_samples: self.total_samples,
            hosts: self.facets.matched_hosts.len().max(1),
            min_ts: self.min_ts,
            max_ts: self.max_ts,
            facets: self.facets.results(),
            poll_duration_histogram: poll_hist_bars(&self.poll_hist),
        }
    }
}

/// Concurrency for samples/dict part-file GETs, matching [`POLLS_READ_CONCURRENCY`].
const SAMPLES_READ_CONCURRENCY: usize = 24;

/// Fetch one folded file's samples + stacks-dict part-file bytes, issuing the
/// two GETs concurrently. Returns `None` when the samples part is missing (the
/// file is not yet readable); the dict is optional (`None` if its GET fails).
pub(crate) async fn fetch_sample_parts(
    output: &dyn StorageBackend,
    bucket: &str,
    output_prefix: &str,
    source_key: &str,
) -> Option<(Vec<u8>, Option<Vec<u8>>)> {
    let part_key = samples_part_key(output_prefix, source_key);
    let dict_key = dict_part_key(output_prefix, source_key);
    let (samples, dict_data) = tokio::join!(
        output.get_object(bucket, &part_key),
        output.get_object(bucket, &dict_key),
    );
    Some((samples.ok()?, dict_data.ok()))
}

/// Fetch the samples + dict part-files for each explicit source key,
/// concurrently (`buffer_unordered`). Returns `(leaf, Result)` pairs keyed by
/// [`part_leaf_of`] so callers can identify exactly which leaves succeeded or
/// failed regardless of completion order. Used by the fold-stream driver to
/// seed a [`FlamegraphAccum`] one bounded batch at a time.
pub(crate) async fn fetch_folded_sample_parts(
    output: &dyn StorageBackend,
    bucket: &str,
    output_prefix: &str,
    source_keys: &[String],
) -> Vec<(String, Result<(Vec<u8>, Option<Vec<u8>>), String>)> {
    use futures::stream::StreamExt;
    futures::stream::iter(source_keys.iter().cloned())
        .map(|sk| async move {
            let leaf = part_leaf_of(&sk);
            match fetch_sample_parts(output, bucket, output_prefix, &sk).await {
                Some(parts) => (leaf, Ok(parts)),
                None => (
                    leaf,
                    Err(format!(
                        "{}: sample parts GET failed",
                        sk.rsplit('/').next().unwrap_or(&sk)
                    )),
                ),
            }
        })
        .buffer_unordered(SAMPLES_READ_CONCURRENCY)
        .collect()
        .await
}

/// Concurrency for polls part-file GETs. A GET is a single round-trip with a
/// small body (one file's polls), so this can run wide.
const POLLS_READ_CONCURRENCY: usize = 24;

/// Fetch the `polls/` part-file bytes for each folded source key in
/// `source_keys`, concurrently. Returns `(raw_source_key, polls_bytes)` for the
/// keys whose part-file is both folded and present; not-yet-folded or missing
/// ones are skipped. The caller decodes the Parquet itself (the polls schema is
/// the tokio-stats endpoint's concern, not the aggregator's), so the part-key
/// path scheme stays private to this module.
pub(crate) async fn read_polls_parts(
    output: &dyn StorageBackend,
    bucket: &str,
    output_prefix: &str,
    source_keys: &[(String, String)],
    folded: &HashSet<String>,
) -> Vec<(String, Vec<u8>)> {
    use futures::stream::StreamExt;
    let fetches: Vec<(String, String)> = source_keys
        .iter()
        .filter(|(_, full)| folded.contains(&part_leaf_of(full)))
        .map(|(raw, full)| (raw.clone(), polls_part_key(output_prefix, full)))
        .collect();

    futures::stream::iter(fetches)
        .map(|(raw_key, polls_key)| async move {
            output
                .get_object(bucket, &polls_key)
                .await
                .ok()
                .map(|data| (raw_key, data))
        })
        .buffer_unordered(POLLS_READ_CONCURRENCY)
        .filter_map(|x| async { x })
        .collect()
        .await
}

/// Fetch one folded file's `polls/` part-file bytes. `None` when the part is
/// missing (not yet readable). The streaming tokio-stats path uses this to read
/// each newly-folded file as it lands, rather than re-reading the whole folded
/// set every poll.
pub(crate) async fn fetch_polls_part(
    output: &dyn StorageBackend,
    bucket: &str,
    output_prefix: &str,
    full_key: &str,
) -> Option<Vec<u8>> {
    let polls_key = polls_part_key(output_prefix, full_key);
    output.get_object(bucket, &polls_key).await.ok()
}

/// Index of the "host" facet in [`FACETS`]. A missing "host" facet is a
/// developer error (the facet table is a compile-time constant), so this panics
/// rather than silently picking the wrong column.
fn host_facet_index() -> usize {
    FACETS
        .iter()
        .position(|f| f.name == "host")
        .expect("FACETS must define a \"host\" facet")
}

enum ResolvedFacetCol<'a> {
    MappedU8 {
        arr: Option<&'a arrow::array::UInt8Array>,
        map: &'static [(u8, &'static str)],
        absent_value: &'static str,
    },
    NullDerived {
        arr: Option<&'a arrow::array::UInt32Array>,
        present_label: &'static str,
        absent_label: &'static str,
        missing_column_label: &'static str,
    },
    DirectString {
        arr: Option<&'a arrow::array::StringArray>,
    },
}

fn resolve_facet_col<'a>(
    batch: &'a arrow::record_batch::RecordBatch,
    def: &FacetDef,
) -> ResolvedFacetCol<'a> {
    match &def.kind {
        FacetKind::MappedU8 {
            column,
            map,
            absent_value,
        } => {
            let arr = batch
                .column_by_name(column)
                .and_then(|c| c.as_any().downcast_ref::<arrow::array::UInt8Array>());
            ResolvedFacetCol::MappedU8 {
                arr,
                map,
                absent_value,
            }
        }
        FacetKind::NullDerived {
            column,
            present_label,
            absent_label,
            missing_column_label,
        } => {
            let arr = batch
                .column_by_name(column)
                .and_then(|c| c.as_any().downcast_ref::<arrow::array::UInt32Array>());
            ResolvedFacetCol::NullDerived {
                arr,
                present_label,
                absent_label,
                missing_column_label,
            }
        }
        FacetKind::DirectString { column } => {
            let arr = batch
                .column_by_name(column)
                .and_then(|c| c.as_any().downcast_ref::<arrow::array::StringArray>());
            ResolvedFacetCol::DirectString { arr }
        }
    }
}

fn extract_facet_value(col: &ResolvedFacetCol, i: usize) -> Option<String> {
    match col {
        ResolvedFacetCol::MappedU8 {
            arr,
            map,
            absent_value,
        } => {
            let label = match arr {
                Some(a) => {
                    let v = a.value(i);
                    map.iter().find(|(k, _)| *k == v).map_or("", |(_, l)| l)
                }
                None => absent_value,
            };
            if label.is_empty() {
                None
            } else {
                Some(label.to_string())
            }
        }
        ResolvedFacetCol::NullDerived {
            arr,
            present_label,
            absent_label,
            missing_column_label,
        } => {
            let label = match arr {
                Some(a) => {
                    if a.is_null(i) {
                        absent_label
                    } else {
                        present_label
                    }
                }
                None => missing_column_label,
            };
            Some(label.to_string())
        }
        ResolvedFacetCol::DirectString { arr } => match arr {
            Some(a) if !a.is_null(i) => Some(a.value(i).to_string()),
            _ => None,
        },
    }
}

fn read_dict_part(data: Vec<u8>, dict: &mut HashMap<Vec<u8>, Vec<String>>) -> anyhow::Result<()> {
    // `Bytes::from(Vec<u8>)` reuses the allocation (no copy).
    let reader = ::parquet::arrow::arrow_reader::ParquetRecordBatchReader::try_new(
        bytes::Bytes::from(data),
        4096,
    )?;
    for batch in reader {
        let batch = batch?;
        let stack_arr = batch.column_by_name("stack_id").and_then(|c| {
            c.as_any()
                .downcast_ref::<arrow::array::FixedSizeBinaryArray>()
        });
        let frames_arr = batch
            .column_by_name("frames")
            .and_then(|c| c.as_any().downcast_ref::<arrow::array::ListArray>());
        let (Some(stack_arr), Some(frames_arr)) = (stack_arr, frames_arr) else {
            continue;
        };
        for i in 0..batch.num_rows() {
            let id = stack_arr.value(i).to_vec();
            if dict.contains_key(&id) {
                continue;
            }
            let frame_list = frames_arr.value(i);
            if let Some(str_arr) = frame_list
                .as_any()
                .downcast_ref::<arrow::array::StringArray>()
            {
                let frames: Vec<String> = (0..str_arr.len())
                    .map(|j| str_arr.value(j).to_string())
                    .collect();
                dict.insert(id, frames);
            }
        }
    }
    Ok(())
}

fn maybe_gunzip(data: &[u8]) -> Vec<u8> {
    if data.len() >= 2 && data[0] == 0x1f && data[1] == 0x8b {
        use std::io::Read;
        let mut decoder = flate2::read::GzDecoder::new(data);
        let mut out = Vec::new();
        match decoder.read_to_end(&mut out) {
            Ok(_) => out,
            Err(_) => data.to_vec(),
        }
    } else {
        data.to_vec()
    }
}

/// The ordered matched-set keys (as `(raw_key, full_key)` pairs) for a scope,
/// given a raw source listing, plus the total bytes of all matched files (for
/// the coverage block). Used by the refinement loop.
pub(crate) fn ordered_full_keys_with_size(
    objects: Vec<ObjectInfo>,
    scope: &Scope,
    segment_duration_secs: i64,
    source_is_local: bool,
    source_bucket: &str,
) -> (Vec<(String, String)>, u64) {
    let matched = matched_and_ordered(objects, scope, segment_duration_secs);
    let total_bytes: u64 = matched.iter().map(|o| o.size.max(0) as u64).sum();
    let keys = matched
        .into_iter()
        .map(|o| {
            let full = full_source_key(source_is_local, source_bucket, &o.key);
            (o.key, full)
        })
        .collect();
    (keys, total_bytes)
}

/// Check if a sample row at index `i` has an enclosing span matching the
/// span-type filter (type UID + optional duration band). Reads the nested
/// `enclosing_spans` LIST<STRUCT> column. Returns `true` if ANY membership
/// matches.
///
/// FAIL CLOSED: If the column is absent (old v3 part-file without span data)
/// or malformed, returns `false` — no samples pass a span filter when the
/// membership data is unavailable. This prevents silently including unvetted
/// samples when a user explicitly requests span-scoped analysis.
pub(crate) fn span_filter_matches(
    batch: &arrow::record_batch::RecordBatch,
    row: usize,
    wanted_uid: &[u8; 16],
    min_span_ns: Option<i64>,
    max_span_ns: Option<i64>,
) -> bool {
    use arrow::array::{Array, AsArray};

    let Some(col) = batch.column_by_name("enclosing_spans") else {
        // Old part-file without enclosing_spans column — fail closed: the sample
        // cannot prove membership, so it does not pass the span filter.
        return false;
    };
    let list_arr = match col.as_list_opt::<i32>() {
        Some(a) => a,
        None => return false, // Malformed column — fail closed
    };

    if list_arr.is_null(row) {
        return false;
    }

    let offsets = list_arr.offsets();
    let start = offsets[row] as usize;
    let end = offsets[row + 1] as usize;
    if start == end {
        return false; // Empty list = no enclosing spans
    }

    let values = list_arr.values();
    let struct_arr = match values.as_struct_opt() {
        Some(a) => a,
        None => return false, // Malformed — fail closed
    };

    // Find span_type_uid and elapsed_ns columns in the struct.
    let type_uid_col = struct_arr.column_by_name("span_type_uid").and_then(|c| {
        c.as_any()
            .downcast_ref::<arrow::array::FixedSizeBinaryArray>()
    });
    let elapsed_col = struct_arr
        .column_by_name("elapsed_ns")
        .and_then(|c| c.as_any().downcast_ref::<arrow::array::Int64Array>());

    let Some(type_uid_arr) = type_uid_col else {
        return false; // Schema doesn't have the expected field — fail closed
    };

    // When duration bounds exist, the elapsed_ns column MUST be present and
    // valid. If it's absent, fail closed — we cannot verify the bound.
    let has_bounds = min_span_ns.is_some() || max_span_ns.is_some();
    if has_bounds && elapsed_col.is_none() {
        return false; // Cannot verify bounds without elapsed_ns — fail closed
    }

    for idx in start..end {
        // Fail closed on null list children (malformed struct entries).
        if struct_arr.is_null(idx) {
            continue;
        }
        if type_uid_arr.is_null(idx) {
            continue; // Null type UID — skip (fail closed for this child)
        }
        let uid = type_uid_arr.value(idx);
        if uid == wanted_uid.as_slice() {
            // Type matches. Check optional duration band.
            if has_bounds {
                let elapsed_arr = elapsed_col.unwrap(); // safe: checked above
                if elapsed_arr.is_null(idx) {
                    // Null elapsed when bounds are required — fail closed for this entry.
                    continue;
                }
                let elapsed = elapsed_arr.value(idx);
                if min_span_ns.is_some_and(|min| elapsed < min) {
                    continue;
                }
                if max_span_ns.is_some_and(|max| elapsed > max) {
                    continue;
                }
            }
            return true;
        }
    }
    false
}

/// Shared `Arc`-friendly handle bundle the server uses to run the refinement
/// loop without re-reading config each call.
#[derive(Clone)]
pub struct AggContext {
    pub source: Arc<dyn StorageBackend>,
    pub output: Arc<dyn StorageBackend>,
    pub source_bucket: String,
    pub source_is_local: bool,
    pub output_bucket: String,
    pub output_prefix: String,
    pub source_prefixes: Vec<String>,
    pub segment_duration_secs: i64,
}

#[cfg(test)]
mod tests {
    use std::future::Future;
    use std::pin::Pin;

    use super::*;
    use crate::storage::{BucketInfo, StorageError};

    /// A source backend that records the peak number of `get_object` calls
    /// running at once, so a test can assert the in-flight fold cap holds. Each
    /// GET holds the "in-flight" state briefly (a yield-heavy spin) so overlap is
    /// observable, then returns junk bytes (decode fails downstream — the test
    /// only cares about fetch-stage concurrency).
    #[derive(Default)]
    struct ConcurrencyProbe {
        in_flight: std::sync::atomic::AtomicUsize,
        peak: std::sync::atomic::AtomicUsize,
    }

    impl StorageBackend for ConcurrencyProbe {
        fn list_buckets(
            &self,
        ) -> Pin<Box<dyn Future<Output = Result<Vec<BucketInfo>, StorageError>> + Send + '_>>
        {
            Box::pin(async { Ok(vec![]) })
        }
        fn list_objects(
            &self,
            _bucket: &str,
            _prefix: &str,
            _cap: usize,
        ) -> Pin<Box<dyn Future<Output = Result<crate::storage::ListPage, StorageError>> + Send + '_>>
        {
            Box::pin(async { Err(StorageError::NotFound("unused".into())) })
        }
        fn list_objects_all(
            &self,
            _bucket: &str,
            _prefix: &str,
        ) -> Pin<Box<dyn Future<Output = Result<Vec<ObjectInfo>, StorageError>> + Send + '_>>
        {
            Box::pin(async { Ok(vec![]) })
        }
        fn list_prefixes(
            &self,
            _bucket: &str,
            _prefix: &str,
        ) -> Pin<Box<dyn Future<Output = Result<Vec<String>, StorageError>> + Send + '_>> {
            Box::pin(async { Ok(vec![]) })
        }
        fn get_object(
            &self,
            _bucket: &str,
            _key: &str,
        ) -> Pin<Box<dyn Future<Output = Result<Vec<u8>, StorageError>> + Send + '_>> {
            use std::sync::atomic::Ordering;
            Box::pin(async move {
                let cur = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
                self.peak.fetch_max(cur, Ordering::SeqCst);
                // Hold the in-flight window open across several executor turns so
                // concurrent folds actually overlap here.
                for _ in 0..50 {
                    tokio::task::yield_now().await;
                }
                self.in_flight.fetch_sub(1, Ordering::SeqCst);
                // Junk bytes: fold_one's decode stage fails, but only AFTER the
                // fetch — which is all this probe measures.
                Ok(vec![0u8; 8])
            })
        }
        fn put_object(
            &self,
            _bucket: &str,
            _key: &str,
            _data: Vec<u8>,
        ) -> Pin<Box<dyn Future<Output = Result<(), StorageError>> + Send + '_>> {
            Box::pin(async { Ok(()) })
        }
    }

    /// The in-flight cap bounds how many folds hold a fetched segment buffer at
    /// once — the memory backstop against a broad scope OOM-killing the process.
    /// Even with a high fetch permit count, concurrent `fold_one`s must never run
    /// more source GETs simultaneously than `inflight` allows.
    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn inflight_cap_bounds_concurrent_fetches() {
        use std::sync::atomic::Ordering;

        let probe = Arc::new(ConcurrencyProbe::default());
        let agg = AggContext {
            source: probe.clone() as Arc<dyn StorageBackend>,
            output: probe.clone() as Arc<dyn StorageBackend>,
            source_bucket: "src".to_string(),
            source_is_local: false,
            output_bucket: "out".to_string(),
            output_prefix: "flamegraph-data".to_string(),
            source_prefixes: vec![],
            segment_duration_secs: 60,
        };

        // fetch permits high (32) but inflight capped low (3): the inflight cap
        // must be the binding constraint on concurrent fetches.
        const INFLIGHT: usize = 3;
        let limits = FoldLimits::new(32, 8, INFLIGHT);

        let mut tasks = tokio::task::JoinSet::new();
        for i in 0..64 {
            let agg = agg.clone();
            let limits = limits.clone();
            tasks.spawn(async move {
                // Decode fails on junk bytes; we don't care about the result,
                // only that the fetch obeyed the in-flight cap.
                let _ = fold_one(&agg, &format!("raw-{i}"), &limits).await;
            });
        }
        while tasks.join_next().await.is_some() {}

        let peak = probe.peak.load(Ordering::SeqCst);
        assert!(peak > 0, "probe should have observed fetches");
        assert!(
            peak <= INFLIGHT,
            "peak concurrent fetches {peak} exceeded the in-flight cap {INFLIGHT}"
        );
    }

    /// Regression guard for the OOM: the default fold sizing must NOT scale the
    /// memory-bound decode stage with core count. On a big box the decode cap has
    /// to stay small and absolute — a 64-core machine running 64 concurrent
    /// ~1.5 GB decodes is ~96 GB and gets OOM-killed. `cpu` is the decode gate and
    /// must never exceed MAX_DECODE_CONCURRENCY regardless of parallelism.
    #[test]
    fn default_fold_limits_do_not_scale_decode_with_cores() {
        let limits = FoldLimits::from_available_parallelism();
        assert!(
            limits.cpu.available_permits() <= MAX_DECODE_CONCURRENCY,
            "decode concurrency {} must stay within the absolute cap {MAX_DECODE_CONCURRENCY}",
            limits.cpu.available_permits()
        );
        // The whole-fold (memory backstop) cap is a small multiple of the decode
        // cap, and fetch never needs to exceed inflight.
        assert!(
            limits.inflight.available_permits() <= MAX_DECODE_CONCURRENCY * 2,
            "inflight {} should stay a small multiple of the decode cap",
            limits.inflight.available_permits()
        );
        assert!(
            limits.fetch.available_permits() <= limits.inflight.available_permits(),
            "fetch must not exceed inflight (a fetch needs an inflight permit)"
        );
    }

    #[test]
    fn sample_filter_source_and_thread() {
        use std::collections::HashMap;

        // Default filter: source=cpu, thread_class="" (all), others empty.
        let def = SampleFilter {
            facets: HashMap::from([
                ("source", "cpu".to_string()),
                ("thread_class", String::new()),
            ]),
            ..Default::default()
        };
        // The filter works by exact match on extracted values. Verify the
        // data-driven filtering contract: non-empty = must match, empty = pass.
        assert_eq!(def.facets.get("source"), Some(&"cpu".to_string()));
        assert_eq!(def.facets.get("thread_class"), Some(&String::new()));

        // A filter with source=sched should be expressible.
        let sched = SampleFilter {
            facets: HashMap::from([("source", "sched".to_string())]),
            ..Default::default()
        };
        assert_eq!(sched.facets.get("source"), Some(&"sched".to_string()));

        // FacetDef registry has our expected facets.
        let names: Vec<&str> = FACETS.iter().map(|f| f.name).collect();
        assert!(names.contains(&"source"));
        assert!(names.contains(&"thread_class"));
        assert!(names.contains(&"host"));
        assert!(names.contains(&"spawn_location"));
    }

    /// Build a one-batch samples Parquet buffer from `(stack_byte, poll_ns)`
    /// pairs. `poll_ns == None` means the sample was not inside a poll.
    fn samples_parquet(rows: &[(u8, Option<u64>)]) -> Vec<u8> {
        use crate::ingest::decode::ResolvedSample;
        use crate::ingest::parquet_writer::write_samples;
        let samples: Vec<ResolvedSample> = rows
            .iter()
            .enumerate()
            .map(|(i, (stack, poll))| ResolvedSample {
                timestamp_ns: 1000 + i as u64,
                stack_id: [*stack; 16],
                worker_id: Some(1),
                source: SOURCE_CPU_PROFILE,
                source_key: "2026-06-19/1450/shale/myhost/boot-1/123-0.bin.gz".to_string(),
                host: "myhost".to_string(),
                service: "shale".to_string(),
                date: "2026-06-19".to_string(),
                poll_duration_ns: *poll,
                spawn_location: Some("src/main.rs:42".to_string()),
                enclosing_spans: Vec::new(),
            })
            .collect();
        let mut buf = Vec::new();
        write_samples(&mut buf, &samples, &HashMap::new()).unwrap();
        buf
    }

    /// Total samples kept after merging `parquet` under a poll-duration band.
    fn samples_kept(parquet: Vec<u8>, min_poll_ns: Option<i64>, max_poll_ns: Option<i64>) -> usize {
        let filter = SampleFilter {
            min_poll_ns,
            max_poll_ns,
            // source defaults to cpu in the endpoint; here match all sources so
            // the test isolates the poll-band behavior.
            facets: HashMap::from([("source", "cpu".to_string())]),
            ..Default::default()
        };
        let mut accum = FlamegraphAccum::new(filter);
        accum.merge(parquet, None).unwrap();
        accum.snapshot().total_samples
    }

    #[test]
    fn poll_band_filters_samples_by_duration() {
        // Three in-poll samples at 0.5ms / 5ms / 50ms, plus one with no poll.
        let rows = [
            (1u8, Some(500_000)),
            (2u8, Some(5_000_000)),
            (3u8, Some(50_000_000)),
            (4u8, None),
        ];
        let mk = || samples_parquet(&rows);

        // No band → every row (including the null-poll one) is kept.
        assert_eq!(samples_kept(mk(), None, None), 4);

        // Lower bound only: ≥ 5ms keeps the 5ms and 50ms rows; excludes the
        // 0.5ms row AND the null-poll row.
        assert_eq!(samples_kept(mk(), Some(5_000_000), None), 2);

        // Upper bound only: ≤ 1ms keeps just the 0.5ms row; excludes null-poll.
        assert_eq!(samples_kept(mk(), None, Some(1_000_000)), 1);

        // Band [1ms, 10ms] keeps only the 5ms row (bounds inclusive).
        assert_eq!(samples_kept(mk(), Some(1_000_000), Some(10_000_000)), 1);

        // A band that matches nothing keeps nothing (rather than erroring).
        assert_eq!(samples_kept(mk(), Some(100_000_000), None), 0);
    }

    #[test]
    fn poll_bucket_is_monotonic_and_subdivides_octaves() {
        assert_eq!(poll_bucket(0), None, "0 has no bucket");
        assert_eq!(poll_bucket(-5), None, "negative has no bucket");
        // Sub-octave: with SUBDIV=4 an octave (2×) spans 4 buckets, so values a
        // little apart within an octave land in DIFFERENT buckets (finer than
        // plain log₂, which would collapse them).
        let b1 = poll_bucket(1_000_000).unwrap();
        let b2 = poll_bucket(1_300_000).unwrap();
        assert!(
            b2 > b1,
            "1.0ms and 1.3ms fall in different sub-octave buckets"
        );
        assert!(b2 - b1 <= POLL_HIST_SUBDIV, "…but within one octave");
        // Monotonic: larger duration → same-or-higher bucket.
        assert!(poll_bucket(50_000_000).unwrap() > poll_bucket(500_000).unwrap());
        // Bucket edges bracket the value that produced them.
        let k = poll_bucket(500_000).unwrap();
        assert!(bucket_edge_ns(k) <= 500_000 && 500_000 < bucket_edge_ns(k + 1));
    }

    #[test]
    fn poll_histogram_is_sample_weighted_and_pre_band() {
        // Two samples at ~0.5ms, one at 50ms, one off-poll. The two 0.5ms samples
        // are close enough to share a sub-octave bucket; 50ms is far away.
        let rows = [
            (1u8, Some(500_000)),
            (2u8, Some(500_001)),
            (3u8, Some(50_000_000)),
            (4u8, None),
        ];
        let hist = |min_poll_ns, max_poll_ns| {
            let filter = SampleFilter {
                min_poll_ns,
                max_poll_ns,
                facets: HashMap::from([("source", "cpu".to_string())]),
                ..Default::default()
            };
            let mut accum = FlamegraphAccum::new(filter);
            accum.merge(samples_parquet(&rows), None).unwrap();
            accum.snapshot().poll_duration_histogram
        };

        // No band: two occupied buckets. The 0.5ms bucket holds 2 samples
        // (weighted, not 1 poll), the 50ms bucket holds 1. The off-poll row
        // contributes to neither. Bars are sorted ascending, edges bracket input.
        let bars = hist(None, None);
        assert_eq!(bars.len(), 2, "two occupied buckets");
        assert_eq!(bars[0].samples, 2, "0.5ms bucket is sample-weighted (2)");
        assert_eq!(bars[1].samples, 1, "50ms bucket holds the one slow sample");
        assert!(
            bars[0].lo_ns <= 500_000 && 500_000 < bars[0].hi_ns,
            "fast bar brackets 0.5ms"
        );
        assert!(
            bars[1].lo_ns <= 50_000_000 && 50_000_000 < bars[1].hi_ns,
            "slow bar brackets 50ms"
        );
        assert!(
            bars[0].hi_ns <= bars[1].lo_ns,
            "bars are disjoint and ascending"
        );

        // The histogram is accumulated PRE-band: narrowing the band to the slow
        // bucket must NOT change the bars (the minimap always shows the full
        // distribution you're selecting from).
        assert_eq!(
            hist(Some(10_000_000), None).len(),
            2,
            "band does not shrink the histogram"
        );
    }

    #[test]
    fn order_key_is_deterministic_and_versioned() {
        let a = order_key("2026-06-19/1300/shale/host-a/boot/1-0.bin.gz");
        let b = order_key("2026-06-19/1300/shale/host-a/boot/1-0.bin.gz");
        assert_eq!(a, b, "same key → same order");
        let c = order_key("2026-06-19/1300/shale/host-b/boot/1-0.bin.gz");
        assert_ne!(a, c, "different key → different order (almost surely)");
    }

    #[test]
    fn parse_scope_fields_handles_prefix() {
        let (d, s, h) = parse_scope_fields(
            "traces/2026-04-09/1910/checkout-api/us-east-1/abcd/1744224000-3.bin.gz",
        );
        assert_eq!(d, "2026-04-09");
        assert_eq!(s, "checkout-api");
        assert_eq!(h, "us-east-1");
    }

    #[test]
    fn parse_scope_fields_handles_s3_uri() {
        let (d, s, h) =
            parse_scope_fields("s3://bkt/2026-06-19/1300/shale/host-a/boot-1/1-0.bin.gz");
        assert_eq!(d, "2026-06-19");
        assert_eq!(s, "shale");
        assert_eq!(h, "host-a");
    }

    #[test]
    fn parse_epoch_from_filename() {
        assert_eq!(
            parse_epoch_secs("2026-04-09/1910/svc/host/boot/1744224000-3.bin.gz"),
            Some(1744224000)
        );
    }

    #[test]
    fn part_keys_are_partitioned_with_hash_leaf() {
        let sk = "s3://bkt/2026-06-19/1300/shale/host-a/boot-1/1-0.bin.gz";
        let pk = samples_part_key("flamegraph-data", sk);
        // Output is namespaced by source bucket, then partitioned by scope.
        // Reference the version constant so bumping it doesn't break this test.
        let expected_prefix = format!(
            "flamegraph-data/v{SAMPLES_FORMAT_VERSION}/bucket=bkt/samples/service=shale/date=2026-06-19/host=host-a/"
        );
        assert!(pk.starts_with(&expected_prefix));
        assert!(pk.ends_with(".parquet"));
        // Leaf is the content hash, idempotent across calls.
        assert_eq!(pk, samples_part_key("flamegraph-data", sk));
    }

    #[test]
    fn folded_set_prefix_is_pruned_to_service() {
        let prefix = folded_set_prefix("flamegraph-data", "bkt", Some("shale"));
        assert_eq!(
            prefix,
            format!("flamegraph-data/v{SAMPLES_FORMAT_VERSION}/bucket=bkt/samples/service=shale/")
        );
    }

    #[test]
    fn parse_source_bucket_from_key() {
        assert_eq!(
            parse_source_bucket("s3://my-bucket/2026-06-19/1300/svc/host/boot/1-0.bin.gz"),
            "my-bucket"
        );
        // Bare (local) keys have no bucket → the "local" namespace.
        assert_eq!(
            parse_source_bucket("2026-06-19/1300/svc/host/boot/1-0.bin.gz"),
            "local"
        );
    }

    #[test]
    fn output_namespaced_by_source_bucket_isolates_buckets() {
        // Same scope path, two different source buckets → different output roots,
        // so their folded sets and LISTs never mix.
        let a = samples_part_key("out", "s3://bucket-a/2026-06-19/1300/svc/h/b/1-0.bin.gz");
        let b = samples_part_key("out", "s3://bucket-b/2026-06-19/1300/svc/h/b/1-0.bin.gz");
        assert!(a.contains("/bucket=bucket-a/"));
        assert!(b.contains("/bucket=bucket-b/"));
        assert_ne!(a, b);
        // And the per-bucket LIST prefixes are disjoint.
        assert_ne!(
            samples_prefix("out", "bucket-a"),
            samples_prefix("out", "bucket-b")
        );
    }

    #[test]
    fn scope_filters_by_service_host_and_time() {
        let scope = Scope {
            start_ns: Some(1_744_224_000_000_000_000),
            end_ns: Some(1_744_224_100_000_000_000),
            service: Some("shale".to_string()),
            hosts: vec!["host-a".to_string()],
        };
        // In service, host, and time window.
        assert!(scope_matches(
            "2026-04-09/1910/shale/host-a/boot/1744224050-0.bin.gz",
            &scope,
            60
        ));
        // Wrong service.
        assert!(!scope_matches(
            "2026-04-09/1910/other/host-a/boot/1744224050-0.bin.gz",
            &scope,
            60
        ));
        // Wrong host (exact match, not substring): host-a must not match host-ab.
        assert!(!scope_matches(
            "2026-04-09/1910/shale/host-z/boot/1744224050-0.bin.gz",
            &scope,
            60
        ));
        // Far in the future — outside window.
        assert!(!scope_matches(
            "2026-04-09/1910/shale/host-a/boot/1744999999-0.bin.gz",
            &scope,
            60
        ));
    }

    #[test]
    fn scope_host_set_matches_union_exactly() {
        let scope = Scope {
            hosts: vec!["host-a".to_string(), "host-c".to_string()],
            ..Default::default()
        };
        assert!(scope_matches("d/h/svc/host-a/boot/1-0.bin.gz", &scope, 60));
        assert!(scope_matches("d/h/svc/host-c/boot/1-0.bin.gz", &scope, 60));
        // Not in the set.
        assert!(!scope_matches("d/h/svc/host-b/boot/1-0.bin.gz", &scope, 60));
        // Exact match, not prefix/substring: "host-a" must not match "host-aa".
        assert!(!scope_matches(
            "d/h/svc/host-aa/boot/1-0.bin.gz",
            &scope,
            60
        ));
        // Empty host set = all hosts.
        let all = Scope::default();
        assert!(scope_matches("d/h/svc/any-host/boot/1-0.bin.gz", &all, 60));
    }

    #[test]
    fn scope_time_overlap_keeps_boundary_file() {
        // File starts 30s before the window opens but runs into it (60s segment).
        let scope = Scope {
            start_ns: Some(1_744_224_000_000_000_000),
            end_ns: Some(1_744_224_100_000_000_000),
            ..Default::default()
        };
        assert!(
            scope_matches("d/h/svc/host/boot/1744223970-0.bin.gz", &scope, 60),
            "file [t-30, t+30) overlaps window opening at t"
        );
    }

    #[test]
    fn matched_set_is_ordered_by_order_key() {
        let objs: Vec<ObjectInfo> = (0..20)
            .map(|i| ObjectInfo {
                key: format!("2026-06-19/1300/shale/host-{i}/boot/1-0.bin.gz"),
                size: 1,
                last_modified: None,
            })
            .collect();
        let ordered = matched_and_ordered(objs.clone(), &Scope::default(), 60);
        assert_eq!(ordered.len(), 20);
        // The order must match sorting by order_key, and be a permutation (not
        // the original lexicographic order in general).
        let mut by_key = ordered.clone();
        by_key.sort_by_key(|o| order_key(&o.key));
        assert_eq!(
            ordered.iter().map(|o| &o.key).collect::<Vec<_>>(),
            by_key.iter().map(|o| &o.key).collect::<Vec<_>>()
        );
    }

    /// Verify span_filter_matches fails closed when enclosing_spans column is absent
    /// (old v3 schema). A span filter should exclude rather than include samples
    /// with no provable membership.
    #[test]
    fn span_filter_old_schema_fail_closed() {
        use crate::ingest::decode::ResolvedSample;
        use crate::ingest::parquet_writer::write_samples;

        // Create a sample WITH NO enclosing_spans data (simulating old schema)
        let samples = vec![ResolvedSample {
            timestamp_ns: 1000,
            stack_id: [1u8; 16],
            worker_id: Some(1),
            source: SOURCE_CPU_PROFILE,
            source_key: "2026-06-19/1450/shale/myhost/boot-1/123-0.bin.gz".to_string(),
            host: "myhost".to_string(),
            service: "shale".to_string(),
            date: "2026-06-19".to_string(),
            poll_duration_ns: Some(5_000_000),
            spawn_location: Some("src/main.rs:42".to_string()),
            enclosing_spans: Vec::new(), // empty = no membership
        }];
        let mut buf = Vec::new();
        write_samples(&mut buf, &samples, &HashMap::new()).unwrap();

        // A span filter with a specific span_type_uid should NOT match this sample.
        let filter = SampleFilter {
            span_type_uid: Some([42u8; 16]),
            facets: HashMap::from([("source", "cpu".to_string())]),
            ..Default::default()
        };
        let mut accum = FlamegraphAccum::new(filter);
        accum.merge(buf, None).unwrap();
        assert_eq!(
            accum.snapshot().total_samples,
            0,
            "span filter must fail closed: sample with no membership data must NOT pass"
        );
    }

    /// Verify span_filter_matches works with valid membership data.
    #[test]
    fn span_filter_matches_valid_membership() {
        use crate::ingest::decode::{EnclosingSpanSummary, ResolvedSample};
        use crate::ingest::parquet_writer::write_samples;

        let target_uid = [42u8; 16];
        let samples = vec![ResolvedSample {
            timestamp_ns: 1000,
            stack_id: [1u8; 16],
            worker_id: Some(1),
            source: SOURCE_CPU_PROFILE,
            source_key: "2026-06-19/1450/shale/myhost/boot-1/123-0.bin.gz".to_string(),
            host: "myhost".to_string(),
            service: "shale".to_string(),
            date: "2026-06-19".to_string(),
            poll_duration_ns: Some(5_000_000),
            spawn_location: Some("src/main.rs:42".to_string()),
            enclosing_spans: vec![EnclosingSpanSummary {
                span_uid: [1u8; 16],
                span_type_uid: target_uid,
                elapsed_ns: 10_000_000,
                details_complete: true,
            }],
        }];
        let mut buf = Vec::new();
        write_samples(&mut buf, &samples, &HashMap::new()).unwrap();

        // Filter matches the target uid
        let filter = SampleFilter {
            span_type_uid: Some(target_uid),
            facets: HashMap::from([("source", "cpu".to_string())]),
            ..Default::default()
        };
        let mut accum = FlamegraphAccum::new(filter);
        accum.merge(buf, None).unwrap();
        assert_eq!(
            accum.snapshot().total_samples,
            1,
            "span filter must match when membership is present"
        );
    }

    /// Verify that min_span_ns/max_span_ns bounds work on exact boundaries.
    #[test]
    fn span_filter_exact_window_boundaries() {
        use crate::ingest::decode::{EnclosingSpanSummary, ResolvedSample};
        use crate::ingest::parquet_writer::write_samples;

        let target_uid = [42u8; 16];
        let make_sample = |elapsed_ns: u64| -> ResolvedSample {
            ResolvedSample {
                timestamp_ns: 1000,
                stack_id: [1u8; 16],
                worker_id: Some(1),
                source: SOURCE_CPU_PROFILE,
                source_key: "2026-06-19/1450/shale/myhost/boot-1/123-0.bin.gz".to_string(),
                host: "myhost".to_string(),
                service: "shale".to_string(),
                date: "2026-06-19".to_string(),
                poll_duration_ns: None,
                spawn_location: None,
                enclosing_spans: vec![EnclosingSpanSummary {
                    span_uid: [1u8; 16],
                    span_type_uid: target_uid,
                    elapsed_ns,
                    details_complete: true,
                }],
            }
        };

        let samples = vec![
            make_sample(1_000_000),  // exactly at min boundary
            make_sample(5_000_000),  // in the middle
            make_sample(10_000_000), // exactly at max boundary
            make_sample(10_000_001), // just above max
            make_sample(999_999),    // just below min
        ];
        let mut buf = Vec::new();
        write_samples(&mut buf, &samples, &HashMap::new()).unwrap();

        // Band [1ms, 10ms] inclusive — should keep exactly 3 samples
        let filter = SampleFilter {
            span_type_uid: Some(target_uid),
            min_span_ns: Some(1_000_000),
            max_span_ns: Some(10_000_000),
            facets: HashMap::from([("source", "cpu".to_string())]),
            ..Default::default()
        };
        let mut accum = FlamegraphAccum::new(filter);
        accum.merge(buf, None).unwrap();
        assert_eq!(
            accum.snapshot().total_samples,
            3,
            "span filter boundaries must be inclusive: [min, max]"
        );
    }

    /// DELIVERABLE (span-explorer bug repro): the `span_type_uid` filter must keep
    /// ONLY the samples enclosed by a span of the requested type, and the surviving
    /// flamegraph must therefore contain ONLY that type's frames. This is the test
    /// that makes a broken filter obvious: two span types (A, B) enclose samples
    /// with clearly distinguishable stack frames ("frame_A_only" vs "frame_B_only").
    /// Filtering to A must yield frame_A_only and NOT frame_B_only, and vice versa.
    ///
    /// The count-only span-filter tests above use identical stack frames, so they
    /// would still pass if the filter mixed the wrong samples in. This test asserts
    /// on the actual frames present in the folded stacks dictionary, so a filter
    /// that lets the wrong span type's samples through fails loudly.
    #[test]
    fn span_filter_keeps_only_matching_span_types_frames() {
        use crate::ingest::decode::{EnclosingSpanSummary, ResolvedSample};
        use crate::ingest::parquet_writer::{write_samples, write_stacks_dict};

        let type_a = [0xAAu8; 16];
        let type_b = [0xBBu8; 16];

        // Two stacks with distinct, easily-greppable leaf frames.
        let stack_a = [0x0Au8; 16];
        let stack_b = [0x0Bu8; 16];
        let frames_a = vec!["frame_A_only".to_string(), "shared_root".to_string()];
        let frames_b = vec!["frame_B_only".to_string(), "shared_root".to_string()];

        let sample = |stack_id: [u8; 16], type_uid: [u8; 16], ts: u64| ResolvedSample {
            timestamp_ns: ts,
            stack_id,
            worker_id: Some(1),
            source: SOURCE_CPU_PROFILE,
            source_key: "2026-06-19/1450/shale/myhost/boot-1/123-0.bin.gz".to_string(),
            host: "myhost".to_string(),
            service: "shale".to_string(),
            date: "2026-06-19".to_string(),
            poll_duration_ns: None,
            spawn_location: None,
            enclosing_spans: vec![EnclosingSpanSummary {
                span_uid: [1u8; 16],
                span_type_uid: type_uid,
                elapsed_ns: 5_000_000,
                details_complete: true,
            }],
        };

        // 3 samples enclosed by type A (frame_A_only), 2 by type B (frame_B_only).
        let samples = vec![
            sample(stack_a, type_a, 1000),
            sample(stack_a, type_a, 1001),
            sample(stack_a, type_a, 1002),
            sample(stack_b, type_b, 1003),
            sample(stack_b, type_b, 1004),
        ];
        let mut samples_buf = Vec::new();
        write_samples(&mut samples_buf, &samples, &HashMap::new()).unwrap();

        let mut dict = HashMap::new();
        dict.insert(stack_a, frames_a.clone());
        dict.insert(stack_b, frames_b.clone());
        let mut dict_buf = Vec::new();
        write_stacks_dict(&mut dict_buf, &dict).unwrap();

        // Collect the distinct frame names present in the surviving (kept) stacks.
        let surviving_frames =
            |span_type_uid: Option<[u8; 16]>| -> std::collections::HashSet<String> {
                let filter = SampleFilter {
                    span_type_uid,
                    facets: HashMap::from([("source", "cpu".to_string())]),
                    ..Default::default()
                };
                let mut accum = FlamegraphAccum::new(filter);
                accum
                    .merge(samples_buf.clone(), Some(dict_buf.clone()))
                    .unwrap();
                let snap = accum.snapshot();
                let mut frames = std::collections::HashSet::new();
                for (stack_id, count) in &snap.stack_counts {
                    assert!(*count > 0, "a stack with zero count should not be present");
                    if let Some(fs) = snap.stacks_dict.get(stack_id) {
                        for f in fs {
                            frames.insert(f.clone());
                        }
                    }
                }
                frames
            };

        // Filter to type A: only frame_A_only survives.
        let a = surviving_frames(Some(type_a));
        assert!(
            a.contains("frame_A_only"),
            "type-A filter must keep frame_A_only, got {a:?}"
        );
        assert!(
            !a.contains("frame_B_only"),
            "type-A filter must NOT keep frame_B_only (broken filter leaks B), got {a:?}"
        );

        // Filter to type B: only frame_B_only survives.
        let b = surviving_frames(Some(type_b));
        assert!(
            b.contains("frame_B_only"),
            "type-B filter must keep frame_B_only, got {b:?}"
        );
        assert!(
            !b.contains("frame_A_only"),
            "type-B filter must NOT keep frame_A_only (broken filter leaks A), got {b:?}"
        );

        // No filter: both frames survive.
        let both = surviving_frames(None);
        assert!(
            both.contains("frame_A_only") && both.contains("frame_B_only"),
            "no span filter must keep both span types' frames, got {both:?}"
        );
    }

    /// FlamegraphAccum::merge is transactional: a dict parse failure after
    /// samples have been read must leave the accumulator unchanged.
    #[test]
    fn flamegraph_merge_dict_failure_leaves_accum_unchanged() {
        use crate::ingest::decode::ResolvedSample;
        use crate::ingest::parquet_writer::write_samples;

        let sample = ResolvedSample {
            timestamp_ns: 5000,
            stack_id: [42u8; 16],
            worker_id: Some(0),
            source: SOURCE_CPU_PROFILE,
            source_key: "2026-06-19/1450/shale/myhost/boot-1/123-0.bin.gz".to_string(),
            host: "myhost".to_string(),
            service: "shale".to_string(),
            date: "2026-06-19".to_string(),
            poll_duration_ns: None,
            spawn_location: None,
            enclosing_spans: Vec::new(),
        };
        let mut samples_buf = Vec::new();
        write_samples(&mut samples_buf, &[sample], &HashMap::new()).unwrap();

        // A valid dict: the merge should succeed.
        let filter = SampleFilter {
            facets: HashMap::from([("source", "cpu".to_string())]),
            ..Default::default()
        };
        let mut accum = FlamegraphAccum::new(filter.clone());
        accum.merge(samples_buf.clone(), None).unwrap();
        assert_eq!(accum.snapshot().total_samples, 1);

        // Now attempt a merge with garbage dict: must fail and leave accum
        // unchanged (transactional — samples should not be committed).
        let mut accum2 = FlamegraphAccum::new(filter);
        let result = accum2.merge(samples_buf, Some(b"not valid parquet".to_vec()));
        assert!(result.is_err(), "garbage dict must fail");
        assert_eq!(
            accum2.snapshot().total_samples,
            0,
            "transactional merge must leave accum unchanged on dict failure"
        );
    }

    /// FlamegraphAccum::merge is transactional: a samples parse failure must
    /// leave the accumulator unchanged.
    #[test]
    fn flamegraph_merge_samples_failure_leaves_accum_unchanged() {
        let filter = SampleFilter {
            facets: HashMap::from([("source", "cpu".to_string())]),
            ..Default::default()
        };
        let mut accum = FlamegraphAccum::new(filter);
        let result = accum.merge(b"not valid parquet".to_vec(), None);
        assert!(result.is_err(), "garbage samples must fail");
        assert_eq!(
            accum.snapshot().total_samples,
            0,
            "transactional merge must leave accum unchanged on samples failure"
        );
    }
}