diff --git a/Cargo.lock b/Cargo.lock index b7df075307db4..37b2af7ea4eb0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6784,6 +6784,7 @@ dependencies = [ "scopeguard", "serde", "smallvec", + "static_assertions", "timely", "tokio", "tracing", diff --git a/doc/developer/design/20260825_peek_execution.md b/doc/developer/design/20260825_peek_execution.md new file mode 100644 index 0000000000000..b95e9cd556c41 --- /dev/null +++ b/doc/developer/design/20260825_peek_execution.md @@ -0,0 +1,283 @@ +# One execution path for index peeks + +- Associated: + [CPU-217](https://linear.app/materializeinc/issue/CPU-217), + [CPU-195](https://linear.app/materializeinc/issue/CPU-195) + +Scope. This covers how a fast-path index peek's arrangement walk is executed: +where it runs, how it is bounded, how it is cancelled, and how it hands results +to the peek response stash. It does not change what a peek returns, how peeks +are planned, or how the adapter routes them. Persist fast-path peeks keep their +own path, because they read a different substrate. + +## The problem + +A fast-path index peek walks an arrangement cursor to collect its result. Today +that walk runs inline on the single-threaded timely worker that received the +peek, to completion, with no preemption point. A large scan therefore holds the +worker for its full duration: dataflows are not scheduled, commands are not +handled, peeks queued behind it wait, and the peek cannot observe its own +cancellation. This is the defect that three separate efforts have attacked from +three directions, and the reason to settle on one execution path before adding a +fourth. + +The direct cost is head-of-line blocking, and it is measured. A point lookup +behind three concurrent scans reached 5783.7 ms; the same lookup with the walk +moved off the worker reached 180.4 ms. A skewed point lookup over a hot key was +slow on 58 of 261 samples inline, and 0 of 261 with the walk moved. Peeks +broadcast to every worker and retire only when all have answered, so one worker +holding its thread delays every peek in flight, not just the peeks behind it on +that worker. + +The indirect cost is that the workarounds have multiplied. There are now three +`PendingPeek` states for one target type, two independent accumulation loops, +and two places where a walk is thrown away and restarted. Each was a reasonable +local fix. Together they are more mechanism than the problem needs, and the +combination is what this document replaces. + +## What exists today + +Four code paths serve or propose to serve an index peek. + +**Inline walk.** `IndexPeek::seek_fulfillment` gates on frontiers, then +`collect_finished_data` scans the error trace and `collect_ok_finished_data` +walks the oks trace to completion. No budget, no preemption. + +**Stash diversion.** When accumulated bytes cross +`peek_response_stash_threshold_bytes`, the walk returns +`PeekStatus::UsePeekStash` and **discards everything it accumulated** +(`compute_state.rs`, the `total_size > peek_stash_threshold_bytes` branch). +`process_peek` then builds a fresh iterator over the same trace and starts +again, which the code acknowledges: "A fresh walk over the same trace: the +iterator that produced `UsePeekStash` was consumed deciding that the result is +too big to return inline." The restarted walk becomes a `PendingPeek::Stash`, +whose rows the worker pumps across activations into a tokio upload task through +a bounded channel. + +**Offloaded walk (#38429).** Takes owned cursors and walks them on a tokio +blocking task, so the serving worker is free. It pays the same restart: +`snapshot_for_offload` takes a *second* cursor up front, held for the life of +every stash-eligible peek, purely so the walk can start over on diversion. + +**Cooperative walk (#38040).** Makes the inline walk resumable. +`PeekResultIterator::step(&mut fuel) -> Step` returns `Step::OutOfFuel` from +inside the scan loop, charging one unit per cursor position. This is the only +one of the four that bounds a scan against an adversarial input, and it is the +piece this design keeps. + +Two observations make the consolidation possible. + +**The discarded prefix was always reusable.** Stash eligibility requires +`RowSetFinishing::is_streamable`, which is `order_by.is_empty() && project == +identity`. The accumulation loop only thins when `num_rows_needed()` is `Some`, +and for empty `order_by` that branch truncates and returns a *complete* answer +rather than continuing. So at the moment of diversion the accumulated rows have +never been sorted, thinned, or reordered. They are an in-order prefix of the +stream, and handing them onward is correct. The restart buys nothing. + +**The reason for the worker pump no longer holds.** `StashingPeek` documents its +channel as necessary because "the underlying trace reader is not Send/Sync". +That is false. `TraceReader::cursor` returns its batches by value, and since the +Arc-backed production spines those batches are `Arc`s, so the pair it returns +owns what it reads and crosses threads as it stands. `PeekResultIterator<..>: +Send` is asserted in `peek_result_iterator.rs`, and `spawn_offloaded_walk` +already ships production cursors to another thread. + +## Design + +One scan type, driven in two placements, with the stash as a state transition +rather than a restart. + +```mermaid +flowchart TD + A[Peek arrives] --> B[Pending, holds no cursor] + B -->|worker activation| C[Take cursor, run one inline slice] + C -->|Complete| D[Answer Rows] + C -->|Failed| E[Answer Error] + C -->|Suspended| F[Queue for a permit] + F -->|permit acquired| G[Tokio task drives the same scan] + F -->|cancelled| H[Drop scan, release cursor] + G -->|Suspended, batch| I[Write batch to stash, continue] + I --> G + G -->|Complete| J[Answer Rows or Stashed] + G -->|cancelled at slice boundary| H +``` + +### The scan + +`PeekScan` owns the oks cursor, the errs cursor, the accumulated rows, the size +accounting, and the literal state. It performs no IO and never awaits. Its +`step` returns: + +```rust +enum ScanOutcome { + /// Stopped with work left. `batch` is present when accumulation crossed the + /// stash threshold: an in-order prefix the driver must take, because the + /// scan cannot both hold it and keep going. + Suspended { batch: Option }, + Complete(PeekResponse), + Failed(PeekError), +} +``` + +`Suspended` is one state with an independent payload rather than two variants, +because a scan can run out of budget and cross the stash threshold on the same +step. The two facts a driver needs are "it stopped" and "is there a batch to +dispose of", and those are orthogonal. + +### The two drivers + +**Inline.** Runs from the worker's peek processing, one slice per peek per +activation. It never inspects `batch` and never performs IO: `Complete` and +`Failed` answer directly, `Suspended` moves the whole scan to the tokio queue. +The invariant that the inline driver performs no IO is what keeps the design +free of async colouring, and it is load-bearing rather than incidental. + +**Tokio.** Holds a permit, drives the same `step` in a loop, writes a batch to +the peek stash when one is produced, checks for cancellation at each slice +boundary, and yields. Because the scan survives the diversion, the stash upload +is fed by the walk that is already running rather than by a second walk. + +### Placement policy + +A peek runs its first slice inline with a small budget, sized so that point +lookups finish there and nothing else does. If it completes, the peek never +leaves the worker and its latency is what it is today. If it suspends, it is +measured to be expensive and moves off the critical path. Cost is measured +rather than predicted, which is what makes a skewed point lookup over a hot key +behave correctly without being special-cased: it enters as a point lookup, +overruns the inline budget, and offloads. + +### Bounding concurrency + +The tokio side is bounded by a replica-wide semaphore. A scan acquires a permit +before running and holds it until it completes or is dropped. Excess scans queue +in the semaphore rather than running, so a peek storm costs queue entries rather +than threads and retained batches. Permits release on drop, including on panic, +so no expiry or renewal is needed. + +Two queues exist and they have different costs, which the implementation must +keep distinct. A peek waiting for its first inline slice holds only the `Peek` +itself and is genuinely free. A scan waiting for a permit holds its cursor, and +therefore `Arc` handles that pin batches against physical compaction, plus its +accumulated prefix. Taking the cursor at dispatch rather than at arrival is what +keeps the first queue free. + +### Budgets + +All budgets are counted in **consumed values**, meaning cursor positions +visited, not rows returned. This is not a stylistic choice. `PeekResultIterator` +loops internally when the MFP rejects a row, stepping the cursor and continuing +without returning, so a selective filter over a large arrangement produces no +rows while walking arbitrarily far. Counting returned rows yields a budget such a +peek never spends. `rows_processed` already increments per position visited, +before extraction, and #38040's fuel uses the same unit. + +Time-based budgets are deliberately excluded. They are nondeterministic, make +behaviour irreproducible between runs and under load, and require a clock-read +granularity hack to be affordable. A count of consumed values is deterministic +and is the quantity that actually bounds the work. + +Three parameters: + +* **Inline budget**, default 1024 consumed values. How far a peek may walk on + the worker before it is offloaded. Sized for point lookups, not for scans. +* **Yield granularity**, default 10000 consumed values. How often a promoted + scan checks for cancellation and yields. At a plausible 100ns to 1us per + position this bounds cancellation latency to single-digit milliseconds, while + keeping the yield overhead immaterial. +* **Permit count.** How many scans may hold a tokio slot at once. This bounds + retained batches and runtime threads, and is the value the queue backs up + behind. + +All three are read through handles rather than values, so a configuration change +reaches scans already in flight without discarding work they have done. This +follows #38158, which needed the same property for the same reason. + +### Cancellation + +Cancellation must work in five states, and the fifth is the one that needs more +than a check. + +1. **Pending, pre-slice.** Removed from the pending map. Holds nothing. +2. **Mid inline slice.** The slice is bounded by the inline budget, so it + finishes; cancellation is observed before promotion. +3. **Queued for a permit.** Leaves the queue, dropping the scan and releasing + its cursor and prefix. +4. **Running on tokio.** Observed at a slice boundary. + `handle_cancel_peek` removes the `PendingPeek`, which drops the result + channel's receiver, so the scan sees a closed sender without any additional + mechanism. +5. **Mid stash upload.** Partial batches already written to the stash shard must + be cleaned up, or they leak blobs. This is what `StashingPeek`'s abort handle + does today and it is the part of that type which survives. + +### Literal constraints + +`Literals::seek_next_literal_key` performs one `seek_key` per literal that has +no matching key, and is called both from `Literals::new` and from `step_key`. +The call from `new` is outside any budget, and the call from `step_key` is +charged one unit for the whole loop, so an `IN` list of mostly-absent values can +walk far outside its budget in both places. The fix is to pass fuel into that +loop and let it suspend mid-way. `new` then does not seek at all, because the +first `step` does it under budget. + +## What this deletes + +* Both restarts: `PeekStatus::UsePeekStash` as a control-flow return, and + `snapshot_for_offload`'s spare `oks_stash` cursor. +* `PendingPeek::Stash`, `StashingPeek::start_upload`, `StashingPeek::pump_rows`, + and its `peek_iterator` and `rows_tx` fields. +* `PEEK_STASH_NUM_BATCHES`, which is worker-pump granularity and meaningless + once nothing pumps. +* `OffloadSnapshot` and `IndexOffloadPeek` from #38429, replaced by the scan plus + a permit. +* The `time:` half of the yielding configuration. + +`PendingPeek` ends with fewer index-peek states than it has today, not more. + +## Relationship to the open PRs + +This work happens in the scope of #38429. #38040 and #38158 are not merged +first and then subsumed. Their commits are absorbed into this branch with +attribution, because merging work we would immediately rewrite costs review +effort twice and leaves the tree carrying mechanism that never had a user. + +**From #38040 (Aljoscha Krettek).** The budget-aware iterator: +`PeekResultIterator::step(&mut fuel) -> Step`, fuel charged per cursor position, +`Step::OutOfFuel` returned from inside the scan loop, and `next` reimplemented +over `step`. This is the load-bearing piece and it is kept close to as written. +The placement policy, the per-activation budget accounting, and the time-based +half of `YieldSpec` are not carried over. + +**From #38158.** The consumed-values counter and the structured `PeekError` with +its SQLSTATE reporting and bincode mirror type. That PR deliberately stopped at +the peek stash because "a stashed peek restarts its scan ... and the restart then +charges the same rows twice". Deleting the restart is exactly what this design +does, so the work it deferred becomes trivial: the count continues because the +scan does. + +Sequencing note: #38158 changes `PeekResponse::Error` into a structured type and +rewrites `merge_peek_responses` precedence. This design touches the same enum and +the same function. The protocol-level change should be settled first, and this +design's `Failed(PeekError)` assumes it. + +## What is unknown + +**Whether the inline budget is right.** 1024 consumed values is a starting +estimate, not a measurement. Too low and ordinary peeks pay a promotion they did +not need; too high and the worker stalls it was meant to prevent. The cheapest +experiment is a latency histogram of inline-completed peeks against consumed +values, taken on a real workload, to see where the point-lookup population ends. + +**Whether peek CPU needs its own bound.** Today one blocked worker implicitly +caps peek CPU at one core per worker. Promoted scans remove that: N concurrent +scans use N cores, and the permit count is then doing load-shedding duty rather +than only memory duty. Whether that starves dataflow maintenance under a peek +storm is not known and should be measured before the permit default is chosen. + +**The regression band.** A peek that overruns the inline budget by a little now +pays promotion plus a retirement step where today it finishes inline. It is +bounded by one slice plus one activation, but it is a real regression for peeks +sitting just past the budget, and its width should be measured rather than +argued. diff --git a/doc/user/data/metrics.yml b/doc/user/data/metrics.yml index fee34ab44368d..06003a6174610 100644 --- a/doc/user/data/metrics.yml +++ b/doc/user/data/metrics.yml @@ -1095,6 +1095,13 @@ metrics: help: Total time processing index peeks, from process_peek entry to response. Excluding peeks that use the peek response stash. source: src/compute/src/metrics.rs visibility: internal +- name: mz_index_peek_walks_total + help: 'The total number of fast-path index peek walks, by the substrate that ran them: inline, offload, or capped (wanted the offload, ran inline because the worker was at its in-flight cap).' + labels: + - substrate + - worker_id + source: src/compute/src/metrics.rs + visibility: internal - name: mz_kafka_partition_offset_max help: High watermark offset on broker for partition labels: diff --git a/misc/python/materialize/mzcompose/__init__.py b/misc/python/materialize/mzcompose/__init__.py index 0731ed7ea5c6e..e1ec1b0ce71a2 100644 --- a/misc/python/materialize/mzcompose/__init__.py +++ b/misc/python/materialize/mzcompose/__init__.py @@ -298,6 +298,11 @@ def get_variable_system_parameters( "true", ["true", "false"], ), + VariableSystemParameter( + "enable_index_peek_offload", + "true", + ["true", "false"], + ), VariableSystemParameter( "enable_union_cancellation_after_relation_cse", "true", @@ -590,6 +595,7 @@ def get_default_system_parameters( # all. Only add it in UNINTERESTING_SYSTEM_PARAMETERS if none of the above # apply. UNINTERESTING_SYSTEM_PARAMETERS = [ + "index_peek_offload_max_inflight", "enable_compute_half_join2", "enable_mz_join_core", "linear_join_yielding", diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index 868907a41cc8c..88f661833ae1a 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -3023,6 +3023,8 @@ def __init__( BOOLEAN_FLAG_VALUES ) self.flags_with_values["enable_upsert_v2"] = BOOLEAN_FLAG_VALUES + self.flags_with_values["enable_index_peek_offload"] = BOOLEAN_FLAG_VALUES + self.flags_with_values["index_peek_offload_max_inflight"] = ["1", "16", "64"] self.flags_with_values["enable_coalesce_case_transform"] = BOOLEAN_FLAG_VALUES self.flags_with_values["enable_compute_sync_mv_sink"] = BOOLEAN_FLAG_VALUES self.flags_with_values["enable_column_paged_batcher"] = BOOLEAN_FLAG_VALUES diff --git a/src/compute-types/src/dyncfgs.rs b/src/compute-types/src/dyncfgs.rs index ff03b29b64aed..c2620b6a98f7b 100644 --- a/src/compute-types/src/dyncfgs.rs +++ b/src/compute-types/src/dyncfgs.rs @@ -578,6 +578,55 @@ pub const PEEK_RESPONSE_STASH_READ_MEMORY_BUDGET_BYTES: Config = Config:: ParameterScope::Environment, ); +/// Whether to walk a fast-path index peek's cursor on a blocking task instead of inline on the +/// timely worker that received it. +/// +/// The serving worker still takes the snapshot, which costs a mutex and a handful of `Arc` clones, +/// and then dispatches. +/// +/// What it buys: a long scan no longer delays the peeks queued behind it on the serving worker. +/// Without it, a worker that serves peeks inline blocks every peek behind the longest walk it is +/// running. +/// +/// What it costs: the walk now runs concurrently with the serving worker rather than instead of +/// it, so on a CPU-saturated replica it competes for cores with the work that worker went on to +/// do. Each in-flight walk also pins the batches its cursor covers, bounded by +/// `index_peek_offload_max_inflight`. +/// +/// Applies to peeks the peek response stash could take as well. The offloaded walk makes the same +/// size-based diversion partway through, and drives the upload from its own thread rather than +/// handing rows back to the worker to pump. +pub const ENABLE_INDEX_PEEK_OFFLOAD: Config = Config::new( + "enable_index_peek_offload", + false, + "Walk fast-path index peeks on a blocking task rather than on the serving timely worker.", + ParameterScope::Replica, +); + +/// How many offloaded index-peek walks one worker may have in flight before it falls back to +/// walking inline. +/// +/// Each in-flight walk retains the batches its snapshot covers, so unbounded concurrency trades +/// memory for latency. It does not hold the trace back from compacting: the walk owns `Arc` +/// batches and the dispatching path drops its trace handle before the walk starts. Serving inline +/// bounds the retained set implicitly at one walk per worker, and this is the explicit form of +/// that bound. +/// +/// Counted per worker, so a replica retains up to `workers * this` snapshots at once and occupies +/// that many blocking-pool threads. Size it against the replica's worker count, not against the +/// replica. +/// +/// NOTE: a walk that diverts to the peek response stash holds its slot and its blocking thread +/// across the persist upload, not just the cursor walk, so under a stash-heavy workload slots turn +/// over on network latency rather than on walk cost. `mz_index_peek_walks_total{substrate="capped"}` +/// is what shows the cap being reached. +pub const INDEX_PEEK_OFFLOAD_MAX_INFLIGHT: Config = Config::new( + "index_peek_offload_max_inflight", + 16, + "Maximum offloaded index-peek walks in flight per worker before falling back to an inline walk.", + ParameterScope::Replica, +); + /// The number of batches to pump from the peek result iterator when stashing peek responses. pub const PEEK_STASH_NUM_BATCHES: Config = Config::new( "compute_peek_stash_num_batches", @@ -678,6 +727,8 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet { .add(&PEEK_RESPONSE_STASH_BATCH_MAX_RUNS) .add(&PEEK_RESPONSE_STASH_READ_BATCH_SIZE_BYTES) .add(&PEEK_RESPONSE_STASH_READ_MEMORY_BUDGET_BYTES) + .add(&ENABLE_INDEX_PEEK_OFFLOAD) + .add(&INDEX_PEEK_OFFLOAD_MAX_INFLIGHT) .add(&PEEK_STASH_NUM_BATCHES) .add(&PEEK_STASH_BATCH_SIZE) .add(&COMPUTE_PROMETHEUS_INTROSPECTION_SCRAPE_INTERVAL) diff --git a/src/compute/Cargo.toml b/src/compute/Cargo.toml index 9300ebd57da87..a0f98aba0936e 100644 --- a/src/compute/Cargo.toml +++ b/src/compute/Cargo.toml @@ -60,6 +60,7 @@ criterion.workspace = true mz-storage-types = { path = "../storage-types", features = ["proptest"] } proptest.workspace = true rand.workspace = true +static_assertions.workspace = true [features] default = [] diff --git a/src/compute/src/compute_state.rs b/src/compute/src/compute_state.rs index 3a2f7b42e2697..645b8511e13ff 100644 --- a/src/compute/src/compute_state.rs +++ b/src/compute/src/compute_state.rs @@ -12,6 +12,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::num::NonZeroUsize; use std::rc::Rc; use std::sync::Arc; +use std::sync::atomic::{self, AtomicUsize}; use std::time::{Duration, Instant}; use bytesize::ByteSize; @@ -30,8 +31,9 @@ use mz_compute_client::protocol::response::{ }; use mz_compute_types::dataflows::DataflowDescription; use mz_compute_types::dyncfgs::{ - ENABLE_PEEK_RESPONSE_STASH, PEEK_RESPONSE_STASH_BATCH_MAX_RUNS, - PEEK_RESPONSE_STASH_THRESHOLD_BYTES, PEEK_STASH_BATCH_SIZE, PEEK_STASH_NUM_BATCHES, + ENABLE_INDEX_PEEK_OFFLOAD, ENABLE_PEEK_RESPONSE_STASH, INDEX_PEEK_OFFLOAD_MAX_INFLIGHT, + PEEK_RESPONSE_STASH_BATCH_MAX_RUNS, PEEK_RESPONSE_STASH_THRESHOLD_BYTES, PEEK_STASH_BATCH_SIZE, + PEEK_STASH_NUM_BATCHES, }; use mz_compute_types::plan::render_plan::RenderPlan; use mz_dyncfg::ConfigSet; @@ -68,13 +70,15 @@ use tokio::sync::{oneshot, watch}; use tracing::{Level, debug, error, info, span, trace, warn}; use uuid::Uuid; -use crate::arrangement::manager::{TraceBundle, TraceManager}; +use crate::arrangement::manager::{PaddedTrace, TraceBundle, TraceManager}; +use crate::compute_state::peek_result_iterator::{TraceCursor, TraceStorage}; use crate::logging; use crate::logging::compute::{CollectionLogging, ComputeEvent, PeekEvent}; use crate::logging::initialize::LoggingTraces; use crate::metrics::{CollectionMetrics, WorkerMetrics}; use crate::render::{LinearJoinSpec, StartSignal}; use crate::server::{ComputeInstanceContext, ResponseSender}; +use crate::typedefs::{ErrAgent, RowRowAgent}; mod peek_result_iterator; mod peek_stash; @@ -107,6 +111,19 @@ pub struct ComputeState { pub copy_to_response_buffer: Rc>>, /// Peek commands that are awaiting fulfillment. pub pending_peeks: BTreeMap, + /// How many offloaded index-peek walks this worker has in flight. + /// + /// Each one retains the `Arc` batches its cursor covers for the length of the walk, and for the + /// upload too if it diverts to the stash, and occupies a blocking-pool thread throughout. It + /// does *not* hold the trace back from compacting: both dispatch paths drop their trace handles + /// before the walk starts, so the cost is retained memory and a thread, not a compaction hold. + /// `INDEX_PEEK_OFFLOAD_MAX_INFLIGHT` caps it and peeks over the cap walk inline instead. + /// + /// Owned by an `InFlightOffload` guard rather than adjusted by hand, so that every way a peek can + /// leave `pending_peeks` (retired, cancelled, dropped by reconciliation) decrements it. Hand + /// accounting leaked on the cancel and reconciliation paths and silently disabled the offload + /// once the cap was reached. + pub in_flight_offloaded_peeks: Arc, /// The persist location where we can stash large peek results. pub peek_stash_persist_location: Option, /// The logger, from Timely's logging framework, if logs are enabled. @@ -195,6 +212,7 @@ impl ComputeState { subscribe_response_buffer: Default::default(), copy_to_response_buffer: Default::default(), pending_peeks: Default::default(), + in_flight_offloaded_peeks: Arc::new(AtomicUsize::new(0)), peek_stash_persist_location: None, compute_logger: None, persist_clients, @@ -1024,6 +1042,53 @@ impl<'a> ActiveComputeState<'a> { } } + /// Decides where this worker walks the next fast-path index peek, claiming a slot of the + /// in-flight budget if the walk is to be offloaded. + /// + /// Declines at the cap because each in-flight walk retains the batches its cursor covers, and + /// a walk that diverts to the stash retains them for the upload as well. The inline walk is + /// always correct, so declining costs latency and nothing else. + /// + /// Checking and claiming are one step so the count cannot drift from the number of live + /// guards. Only the worker that owns this counter ever touches it, which is what makes the + /// unsynchronized check-then-claim sound. + fn offload_placement(&self) -> WalkPlacement { + if !ENABLE_INDEX_PEEK_OFFLOAD.get(&self.compute_state.worker_config) { + return WalkPlacement::Inline; + } + let max_inflight = INDEX_PEEK_OFFLOAD_MAX_INFLIGHT.get(&self.compute_state.worker_config); + let counter = Arc::clone(&self.compute_state.in_flight_offloaded_peeks); + if counter.load(atomic::Ordering::SeqCst) >= max_inflight { + return WalkPlacement::Capped; + } + counter.fetch_add(1, atomic::Ordering::SeqCst); + WalkPlacement::Offloaded(InFlightOffload(counter)) + } + + /// The stash configuration an offloaded walk needs, or `None` when the stash cannot take this + /// peek and the walk must resolve inline. + fn offload_stash( + &self, + peek_stash_usable: bool, + threshold_bytes: usize, + ) -> Option { + if !peek_stash_usable { + return None; + } + Some(OffloadStash { + persist_clients: Arc::clone(&self.compute_state.persist_clients), + persist_location: self + .compute_state + .peek_stash_persist_location + .clone() + .expect("peek stash usability implies a configured location"), + threshold_bytes, + batch_max_runs: PEEK_RESPONSE_STASH_BATCH_MAX_RUNS + .get(&self.compute_state.worker_config), + batch_size: PEEK_STASH_BATCH_SIZE.get(&self.compute_state.worker_config), + }) + } + /// Either complete the peek (and send the response) or put it in the pending set. fn process_peek(&mut self, upper: &mut Antichain, mut peek: PendingPeek) { let response = match &mut peek { @@ -1048,71 +1113,134 @@ impl<'a> ActiveComputeState<'a> { let peek_stash_threshold_bytes = PEEK_RESPONSE_STASH_THRESHOLD_BYTES.get(&self.compute_state.worker_config); - let metrics = IndexPeekMetrics { - seek_fulfillment_seconds: &self - .compute_state - .metrics - .index_peek_seek_fulfillment_seconds, - frontier_check_seconds: &self - .compute_state - .metrics - .index_peek_frontier_check_seconds, - error_scan_seconds: &self.compute_state.metrics.index_peek_error_scan_seconds, - cursor_setup_seconds: &self - .compute_state - .metrics - .index_peek_cursor_setup_seconds, - row_iteration_seconds: &self - .compute_state - .metrics - .index_peek_row_iteration_seconds, - row_iteration_rows: &self.compute_state.metrics.index_peek_row_iteration_rows, - result_sort_seconds: &self.compute_state.metrics.index_peek_result_sort_seconds, - result_sort_rows: &self.compute_state.metrics.index_peek_result_sort_rows, - row_collection_seconds: &self - .compute_state - .metrics - .index_peek_row_collection_seconds, - }; - - let status = peek.seek_fulfillment( - upper, - self.compute_state.max_result_size, - peek_stash_enabled && peek_stash_eligible, - peek_stash_threshold_bytes, - &metrics, - ); - - self.compute_state - .metrics - .index_peek_total_seconds - .observe(start.elapsed().as_secs_f64()); - - match status { - PeekStatus::Ready(result) => Some(result), - PeekStatus::NotReady => None, - PeekStatus::UsePeekStash => { - let _span = - span!(parent: &peek.span, Level::DEBUG, "process_stash_peek").entered(); + let peek_stash_usable = peek_stash_enabled && peek_stash_eligible; + + let placement = self.offload_placement(); + if let WalkPlacement::Offloaded(in_flight) = placement { + let stash = self.offload_stash(peek_stash_usable, peek_stash_threshold_bytes); + match peek.snapshot_for_offload(upper, stash.is_some()) { + OffloadSnapshot::NotReady => None, + OffloadSnapshot::Response(response) => Some(response), + OffloadSnapshot::Ready { + oks, + errs, + oks_stash, + } => { + let offloaded = spawn_offloaded_walk::< + PaddedTrace>, + ErrAgent, + >( + peek.peek.clone(), + self.compute_state.max_result_size, + oks, + errs, + oks_stash, + stash, + in_flight, + peek.span.clone(), + self.timely_worker.sync_activator_for([].into()), + ); + self.compute_state + .metrics + .index_peek_walks_offload_total + .inc(); + self.compute_state + .pending_peeks + .insert(offloaded.peek.uuid, PendingPeek::IndexOffload(offloaded)); + return; + } + } + } else { + // The walk runs here either way, but which counter it lands in is what + // separates "the offload is off" from "the offload gave up". + match placement { + WalkPlacement::Capped => self + .compute_state + .metrics + .index_peek_walks_capped_total + .inc(), + _ => self + .compute_state + .metrics + .index_peek_walks_inline_total + .inc(), + } + let metrics = IndexPeekMetrics { + seek_fulfillment_seconds: &self + .compute_state + .metrics + .index_peek_seek_fulfillment_seconds, + frontier_check_seconds: &self + .compute_state + .metrics + .index_peek_frontier_check_seconds, + error_scan_seconds: &self + .compute_state + .metrics + .index_peek_error_scan_seconds, + cursor_setup_seconds: &self + .compute_state + .metrics + .index_peek_cursor_setup_seconds, + row_iteration_seconds: &self + .compute_state + .metrics + .index_peek_row_iteration_seconds, + row_iteration_rows: &self + .compute_state + .metrics + .index_peek_row_iteration_rows, + result_sort_seconds: &self + .compute_state + .metrics + .index_peek_result_sort_seconds, + result_sort_rows: &self.compute_state.metrics.index_peek_result_sort_rows, + row_collection_seconds: &self + .compute_state + .metrics + .index_peek_row_collection_seconds, + }; + + let status = peek.seek_fulfillment( + upper, + self.compute_state.max_result_size, + peek_stash_usable, + peek_stash_threshold_bytes, + &metrics, + ); - let peek_stash_batch_max_runs = PEEK_RESPONSE_STASH_BATCH_MAX_RUNS - .get(&self.compute_state.worker_config); + self.compute_state + .metrics + .index_peek_total_seconds + .observe(start.elapsed().as_secs_f64()); + + match status { + PeekStatus::Ready(result) => Some(result), + PeekStatus::NotReady => None, + PeekStatus::UsePeekStash => { + let _span = + span!(parent: &peek.span, Level::DEBUG, "process_stash_peek") + .entered(); + + let peek_stash_batch_max_runs = PEEK_RESPONSE_STASH_BATCH_MAX_RUNS + .get(&self.compute_state.worker_config); + + let stash_task = peek_stash::StashingPeek::start_upload( + Arc::clone(&self.compute_state.persist_clients), + self.compute_state + .peek_stash_persist_location + .as_ref() + .expect("verified above"), + peek.peek.clone(), + peek.trace_bundle.clone(), + peek_stash_batch_max_runs, + ); - let stash_task = peek_stash::StashingPeek::start_upload( - Arc::clone(&self.compute_state.persist_clients), self.compute_state - .peek_stash_persist_location - .as_ref() - .expect("verified above"), - peek.peek.clone(), - peek.trace_bundle.clone(), - peek_stash_batch_max_runs, - ); - - self.compute_state - .pending_peeks - .insert(peek.peek.uuid, PendingPeek::Stash(stash_task)); - return; + .pending_peeks + .insert(peek.peek.uuid, PendingPeek::Stash(stash_task)); + return; + } } } } @@ -1123,6 +1251,35 @@ impl<'a> ActiveComputeState<'a> { .observe(duration.as_secs_f64()); result }), + PendingPeek::IndexOffload(peek) => { + match peek.result.try_recv() { + Ok((result, duration)) => { + // The walk's own duration, which is the whole peek's cost for an offloaded + // peek. The dispatching path deliberately does not observe this histogram, + // so there is one observation per peek either way. The in-flight slot is + // returned by the guard when this `PendingPeek` drops. + self.compute_state + .metrics + .index_peek_total_seconds + .observe(duration.as_secs_f64()); + Some(result) + } + Err(oneshot::error::TryRecvError::Empty) => None, + // The sender is gone without a result, which means the walk's thread panicked: + // tokio catches a panic in a blocking closure and drops the closure's state. The + // peek must be answered, otherwise it sits in `pending_peeks` forever and the + // client waits indefinitely. + Err(oneshot::error::TryRecvError::Closed) => { + soft_panic_or_log!( + "offloaded index peek {} lost its walk without a result", + peek.peek.uuid, + ); + Some(PeekResponse::Error( + "offloaded index peek walk failed".to_string(), + )) + } + } + } PendingPeek::Stash(stashing_peek) => { let num_batches = PEEK_STASH_NUM_BATCHES.get(&self.compute_state.worker_config); let batch_size = PEEK_STASH_BATCH_SIZE.get(&self.compute_state.worker_config); @@ -1304,6 +1461,10 @@ pub enum PendingPeek { /// A peek against an index that is being stashed in the peek stash by an /// async background task. Stash(peek_stash::StashingPeek), + /// A peek against an index whose cursor walk runs on a blocking task, off the timely worker + /// that received it. Produced by either peek path, since both can hand an owned cursor to + /// another thread. + IndexOffload(IndexOffloadPeek), } impl PendingPeek { @@ -1425,6 +1586,7 @@ impl PendingPeek { PendingPeek::Index(p) => &p.span, PendingPeek::Persist(p) => &p.span, PendingPeek::Stash(p) => &p.span, + PendingPeek::IndexOffload(p) => &p.span, } } @@ -1433,10 +1595,57 @@ impl PendingPeek { PendingPeek::Index(p) => &p.peek, PendingPeek::Persist(p) => &p.peek, PendingPeek::Stash(p) => &p.peek, + PendingPeek::IndexOffload(p) => &p.peek, } } } +/// A fast-path index peek whose cursor walk is running on a blocking task. +/// +/// Note that this intentionally does not implement or derive `Clone`, as each pending peek is +/// meant to be dropped after it's responded to. +/// Holds one slot of the offload's in-flight budget for as long as the peek exists. +/// +/// A started `spawn_blocking` closure cannot be aborted, so dropping this does not stop the walk. It +/// only returns the slot, which is what keeps the budget honest across cancellation and +/// reconciliation. +/// Where a fast-path index peek's cursor walk runs, and why. +/// +/// The two declining variants both walk inline. They are distinct because they need different +/// responses: `Inline` is the configured state, `Capped` is the offload failing to engage under +/// load, and a counter that merged them could not tell one from the other. +enum WalkPlacement { + /// `enable_index_peek_offload` is off, so the walk runs on the serving worker. + Inline, + /// The offload is on, but this worker is at `index_peek_offload_max_inflight`, so the walk + /// runs on the serving worker anyway. + Capped, + /// A slot in the in-flight budget is claimed and held by the guard; the walk runs on a + /// blocking task. + Offloaded(InFlightOffload), +} + +pub struct InFlightOffload(Arc); + +impl Drop for InFlightOffload { + fn drop(&mut self) { + self.0.fetch_sub(1, atomic::Ordering::SeqCst); + } +} + +pub struct IndexOffloadPeek { + pub(crate) peek: Peek, + /// Returns this walk's slot in the in-flight budget when the peek goes away, by any route. + _in_flight: InFlightOffload, + /// The walk. Dropping this cannot abort a `spawn_blocking` closure that already started, so the + /// walk runs to completion and discards its result. + _abort_handle: AbortOnDropHandle<()>, + /// The result of the walk, eventually. + result: oneshot::Receiver<(PeekResponse, Duration)>, + /// The `tracing::Span` tracking this peek's operation. + span: tracing::Span, +} + /// An in-progress Persist peek. /// /// Note that `PendingPeek` intentionally does not implement or derive `Clone`, @@ -1600,6 +1809,44 @@ pub(crate) struct IndexPeekMetrics<'a> { } impl IndexPeek { + /// Takes owned, `Send` cursors over the oks and errs traces if the peek is ready to be + /// answered, so the walk can move to another thread. + /// + /// Applies the same two gates as the inline walk, in the same order: the traces must be sealed + /// beyond the peek time, and their compaction frontier must not have passed it. + fn snapshot_for_offload( + &mut self, + upper: &mut Antichain, + want_stash: bool, + ) -> OffloadSnapshot>, ErrAgent> { + self.trace_bundle.oks_mut().read_upper(upper); + if upper.less_equal(&self.peek.timestamp) { + return OffloadSnapshot::NotReady; + } + self.trace_bundle.errs_mut().read_upper(upper); + if upper.less_equal(&self.peek.timestamp) { + return OffloadSnapshot::NotReady; + } + + let read_frontier = self.trace_bundle.compaction_frontier(); + if !read_frontier.less_equal(&self.peek.timestamp) { + return OffloadSnapshot::Response(PeekResponse::Error(format!( + "Arrangement compaction frontier ({:?}) is beyond the time of the attempted read ({})", + read_frontier.elements(), + self.peek.timestamp, + ))); + } + + // `TraceReader::cursor` hands back the batches by value, so what it returns already owns + // the `Arc`s it reads and crosses to another thread as it is. The trace handles are only + // borrowed for the length of these calls, which is why the walk holds nothing back. + OffloadSnapshot::Ready { + oks: self.trace_bundle.oks_mut().cursor(), + errs: self.trace_bundle.errs_mut().cursor(), + oks_stash: want_stash.then(|| self.trace_bundle.oks_mut().cursor()), + } + } + /// Attempts to fulfill the peek and reports success. /// /// To produce output at `peek.timestamp`, we must be certain that @@ -1671,44 +1918,75 @@ impl IndexPeek { // Check if there exist any errors and, if so, return whatever one we // find first. - let (mut cursor, storage) = self.trace_bundle.errs_mut().cursor(); + let (cursor, storage) = self.trace_bundle.errs_mut().cursor(); + if let Some(error) = Self::scan_errs_for_error::>>( + self.peek.target.id(), + self.peek.timestamp, + cursor, + storage, + ) { + return PeekStatus::Ready(error); + } + + metrics + .error_scan_seconds + .observe(error_scan_start.elapsed().as_secs_f64()); + + Self::collect_ok_finished_data( + &self.peek, + self.trace_bundle.oks_mut(), + max_result_size, + peek_stash_eligible, + peek_stash_threshold_bytes, + metrics, + ) + } + + /// Scans an errs cursor for any error at or before `peek_timestamp`, returning the first one + /// found (or `None`). + /// + /// Shared between the inline errs scan in `collect_finished_data` (cursor borrowed live off a + /// local `TraceBundle`) and the offloaded walk in `offloaded_response` (cursor owned by a + /// snapshot): a `(cursor, storage)` pair looks the same to this scan either way. + fn scan_errs_for_error( + target_id: GlobalId, + peek_timestamp: Timestamp, + mut cursor: peek_result_iterator::TraceCursor, + storage: peek_result_iterator::TraceStorage, + ) -> Option + where + Tr: TraceReader, + for<'a> BatchCursor: Cursor< + Key<'a>: std::fmt::Display, + TimeGat<'a>: PartialOrder, + DiffGat<'a> = &'a Diff, + >, + { while cursor.key_valid(&storage) { let mut copies = Diff::ZERO; cursor.map_times(&storage, |time, diff| { - if time.less_equal(&self.peek.timestamp) { + if time.less_equal(&peek_timestamp) { copies += diff; } }); if copies.is_negative() { let error = cursor.key(&storage); error!( - target = %self.peek.target.id(), diff = %copies, %error, + target = %target_id, diff = %copies, %error, "index peek encountered negative multiplicities in error trace", ); - return PeekStatus::Ready(PeekResponse::Error(format!( + return Some(PeekResponse::Error(format!( "Invalid data in source errors, \ saw retractions ({}) for row that does not exist: {}", -copies, error, ))); } if copies.is_positive() { - return PeekStatus::Ready(PeekResponse::Error(cursor.key(&storage).to_string())); + return Some(PeekResponse::Error(cursor.key(&storage).to_string())); } cursor.step_key(&storage); } - - metrics - .error_scan_seconds - .observe(error_scan_start.elapsed().as_secs_f64()); - - Self::collect_ok_finished_data( - &self.peek, - self.trace_bundle.oks_mut(), - max_result_size, - peek_stash_eligible, - peek_stash_threshold_bytes, - metrics, - ) + None } /// Collects data for a known-complete peek from the ok stream. @@ -1730,15 +2008,12 @@ impl IndexPeek { DiffGat<'a> = &'a Diff, >, { - let max_result_size = usize::cast_from(max_result_size); - let count_byte_size = size_of::(); - // Cursor setup timing let cursor_setup_start = Instant::now(); // We clone `literal_constraints` here because we don't want to move the constraints // out of the peek struct, and don't want to modify in-place. - let mut peek_iterator = peek_result_iterator::PeekResultIterator::new( + let peek_iterator = peek_result_iterator::PeekResultIterator::new( peek.target.id().clone(), peek.map_filter_project.clone(), peek.timestamp, @@ -1750,6 +2025,44 @@ impl IndexPeek { .cursor_setup_seconds .observe(cursor_setup_start.elapsed().as_secs_f64()); + Self::drain_ok_iterator( + peek_iterator, + peek, + max_result_size, + peek_stash_eligible, + peek_stash_threshold_bytes, + Some(metrics), + ) + } + + /// Drains a [`peek_result_iterator::PeekResultIterator`] into a [`PeekStatus`], sorting and + /// truncating per `peek.finishing`. + /// + /// Shared between the inline walk (`collect_ok_finished_data`, iterator borrowed live off a + /// local `TraceBundle`) and the offloaded walk (`offloaded_response`, iterator over an owned + /// snapshot): the accumulation logic is identical either way. Only the inline walk records + /// per-phase metrics; the offloaded walk passes `metrics` as `None`. + fn drain_ok_iterator( + mut peek_iterator: peek_result_iterator::PeekResultIterator, + peek: &Peek, + max_result_size: u64, + peek_stash_eligible: bool, + peek_stash_threshold_bytes: usize, + metrics: Option<&IndexPeekMetrics<'_>>, + ) -> PeekStatus + where + Tr: TraceReader, + for<'a> BatchCursor: Cursor< + Key<'a>: ExtendDatums + Eq, + KeyContainer: BatchContainer, + Val<'a>: ExtendDatums, + TimeGat<'a>: PartialOrder, + DiffGat<'a> = &'a Diff, + >, + { + let max_result_size = usize::cast_from(max_result_size); + let count_byte_size = size_of::(); + // Accumulated `Vec<(row, count)>` results that we are likely to return. let mut results = Vec::new(); let mut total_size: usize = 0; @@ -1812,23 +2125,27 @@ impl IndexPeek { { if peek.finishing.order_by.is_empty() { results.truncate(max_results); - metrics - .row_iteration_seconds - .observe(row_iteration_start.elapsed().as_secs_f64()); - metrics - .row_iteration_rows - .observe(f64::cast_lossy(peek_iterator.rows_processed())); - metrics - .result_sort_seconds - .observe(sort_time_accum.as_secs_f64()); - metrics - .result_sort_rows - .observe(f64::cast_lossy(rows_sorted)); + if let Some(metrics) = metrics { + metrics + .row_iteration_seconds + .observe(row_iteration_start.elapsed().as_secs_f64()); + metrics + .row_iteration_rows + .observe(f64::cast_lossy(peek_iterator.rows_processed())); + metrics + .result_sort_seconds + .observe(sort_time_accum.as_secs_f64()); + metrics + .result_sort_rows + .observe(f64::cast_lossy(rows_sorted)); + } let row_collection_start = Instant::now(); let collection = RowCollection::new(results, &peek.finishing.order_by); - metrics - .row_collection_seconds - .observe(row_collection_start.elapsed().as_secs_f64()); + if let Some(metrics) = metrics { + metrics + .row_collection_seconds + .observe(row_collection_start.elapsed().as_secs_f64()); + } return PeekStatus::Ready(PeekResponse::Rows(vec![collection])); } else { // We can sort `results` and then truncate to `max_results`. @@ -1860,28 +2177,417 @@ impl IndexPeek { } } - metrics - .row_iteration_seconds - .observe(row_iteration_start.elapsed().as_secs_f64()); - metrics - .row_iteration_rows - .observe(f64::cast_lossy(peek_iterator.rows_processed())); - metrics - .result_sort_seconds - .observe(sort_time_accum.as_secs_f64()); - metrics - .result_sort_rows - .observe(f64::cast_lossy(rows_sorted)); + if let Some(metrics) = metrics { + metrics + .row_iteration_seconds + .observe(row_iteration_start.elapsed().as_secs_f64()); + metrics + .row_iteration_rows + .observe(f64::cast_lossy(peek_iterator.rows_processed())); + metrics + .result_sort_seconds + .observe(sort_time_accum.as_secs_f64()); + metrics + .result_sort_rows + .observe(f64::cast_lossy(rows_sorted)); + } let row_collection_start = Instant::now(); let collection = RowCollection::new(results, &peek.finishing.order_by); - metrics - .row_collection_seconds - .observe(row_collection_start.elapsed().as_secs_f64()); + if let Some(metrics) = metrics { + metrics + .row_collection_seconds + .observe(row_collection_start.elapsed().as_secs_f64()); + } PeekStatus::Ready(PeekResponse::Rows(vec![collection])) } } +/// Spawns the cursor walk for `peek` on a blocking task and returns the pending peek that +/// collects its result. +/// +/// The caller has already taken the owned cursors, so nothing here touches a trace. That is the +/// whole point: the batches the cursors cover are `Arc`-backed, so the serving worker can keep +/// stepping (and the publishing worker can keep merging) while this walk runs elsewhere. +/// +/// `spawn_blocking`, not `spawn`: a walk over a large arrangement is CPU-bound and would occupy +/// an async runtime thread for its whole duration. +fn spawn_offloaded_walk( + peek: Peek, + max_result_size: u64, + oks: (TraceCursor, TraceStorage), + errs: (TraceCursor, TraceStorage), + oks_stash: Option<(TraceCursor, TraceStorage)>, + stash: Option, + in_flight: InFlightOffload, + span: tracing::Span, + activator: timely::scheduling::activate::SyncActivator, +) -> IndexOffloadPeek +where + OksTr: TraceReader + 'static, + TraceCursor: Send, + TraceStorage: Send, + for<'a> BatchCursor: Cursor< + Key<'a>: ExtendDatums + Eq, + KeyContainer: BatchContainer, + Val<'a>: ExtendDatums, + TimeGat<'a>: PartialOrder, + DiffGat<'a> = &'a Diff, + >, + ErrsTr: TraceReader + 'static, + TraceCursor: Send, + TraceStorage: Send, + for<'a> BatchCursor: Cursor< + Key<'a>: std::fmt::Display, + TimeGat<'a>: PartialOrder, + DiffGat<'a> = &'a Diff, + >, +{ + let (result_tx, result_rx) = oneshot::channel(); + let task_peek = peek.clone(); + let peek_uuid = peek.uuid; + + let task_handle = mz_ore::task::spawn_blocking( + || "index_peek::offload", + move || { + let start = Instant::now(); + let response = offloaded_response::( + &task_peek, + max_result_size, + oks, + errs, + oks_stash, + stash, + ); + match result_tx.send((response, start.elapsed())) { + Ok(()) => {} + Err((_response, elapsed)) => { + debug!(duration = ?elapsed, "dropping result for cancelled peek {peek_uuid}") + } + } + // Wake the serving worker so `process_peeks` retires this peek. Without it the + // result sits in the oneshot until the worker happens to step for another reason. + if activator.activate().is_err() { + debug!("unable to wake timely after completed offloaded index peek {peek_uuid}"); + } + }, + ); + + IndexOffloadPeek { + peek, + _in_flight: in_flight, + _abort_handle: task_handle.abort_on_drop(), + result: result_rx, + span, + } +} + +/// What an offloaded walk needs to divert to the peek response stash partway through. +struct OffloadStash { + persist_clients: Arc, + persist_location: PersistLocation, + threshold_bytes: usize, + batch_max_runs: usize, + batch_size: usize, +} + +/// Builds a peek response from owned cursors, off the serving worker's thread. +/// +/// Mirrors the inline walk: scan the errors first and report one if present, then drain the ok +/// rows. A result that grows past the stash threshold diverts to persist from this thread, so the +/// serving worker is not involved in either outcome. +fn offloaded_response( + peek: &Peek, + max_result_size: u64, + oks: (TraceCursor, TraceStorage), + errs: (TraceCursor, TraceStorage), + oks_stash: Option<(TraceCursor, TraceStorage)>, + stash: Option, +) -> PeekResponse +where + OksTr: TraceReader, + for<'a> BatchCursor: Cursor< + Key<'a>: ExtendDatums + Eq, + KeyContainer: BatchContainer, + Val<'a>: ExtendDatums, + TimeGat<'a>: PartialOrder, + DiffGat<'a> = &'a Diff, + >, + ErrsTr: TraceReader, + for<'a> BatchCursor: Cursor< + Key<'a>: std::fmt::Display, + TimeGat<'a>: PartialOrder, + DiffGat<'a> = &'a Diff, + >, +{ + let target_id = peek.target.id(); + let (errs_cursor, errs_storage) = errs; + if let Some(error) = IndexPeek::scan_errs_for_error::( + target_id, + peek.timestamp, + errs_cursor, + errs_storage, + ) { + return error; + } + + let (oks_cursor, oks_storage) = oks; + let peek_iterator = peek_result_iterator::PeekResultIterator::::new_over_cursor( + target_id, + peek.map_filter_project.clone(), + peek.timestamp, + peek.literal_constraints.clone().as_deref_mut(), + oks_cursor, + oks_storage, + ); + + let (stash_eligible, threshold_bytes) = match &stash { + Some(stash) => (true, stash.threshold_bytes), + None => (false, 0), + }; + + // `NotReady` is unreachable: it comes from the frontier checks the caller already passed, not + // from draining an owned cursor. + match IndexPeek::drain_ok_iterator( + peek_iterator, + peek, + max_result_size, + stash_eligible, + threshold_bytes, + None, + ) { + PeekStatus::Ready(response) => response, + PeekStatus::UsePeekStash => { + let stash = stash.expect("diversion only reported when the stash is configured"); + let (oks_cursor, oks_storage) = + oks_stash.expect("spare cursor taken whenever the stash is configured"); + let peek_iterator = peek_result_iterator::PeekResultIterator::::new_over_cursor( + target_id, + peek.map_filter_project.clone(), + peek.timestamp, + peek.literal_constraints.clone().as_deref_mut(), + oks_cursor, + oks_storage, + ); + peek_stash::StashingPeek::upload_blocking( + stash.persist_clients, + &stash.persist_location, + peek, + peek_iterator, + stash.batch_max_runs, + stash.batch_size, + ) + } + PeekStatus::NotReady => { + unreachable!("an offloaded index peek always resolves to a response") + } + } +} +#[cfg(test)] +mod index_peek_tests { + use std::rc::Rc; + + use differential_dataflow::operators::arrange::TraceAgent; + use differential_dataflow::trace::{Builder, Description, Trace}; + use mz_expr::{MapFilterProject, RowSetFinishing}; + use mz_repr::{Datum, RelationDesc, SqlScalarType}; + use mz_row_spine::{RowRowBuilder, RowRowSpine}; + use mz_timely_util::columnation::ColumnationStack; + use timely::container::PushInto; + use timely::dataflow::operators::generic::OperatorInfo; + use timely::progress::Timestamp as _; + use uuid::Uuid; + + use super::*; + use crate::server::ComputeRuntimeRole; + use crate::typedefs::{ErrAgent, ErrBuilder, ErrSpine, RowRowAgent}; + + fn row(x: i64) -> Row { + Row::pack_slice(&[Datum::Int64(x)]) + } + + /// Builds a one-batch `[0, upper)` oks trace with `rows`, wrapped exactly like a real + /// index's `TraceBundle.oks` (a `PaddedTrace>`), but constructed directly + /// (bypassing rendering a dataflow) for test purposes. + /// + /// The batch is inserted through the `TraceWriter` (not `Trace::insert` on the bare spine + /// directly), because the writer tracks its own idea of the trace's current upper and + /// asserts new batches are contiguous with it; inserting straight into the spine before + /// wrapping desyncs that bookkeeping, and the writer's `Drop` (which seals the trace to the + /// empty frontier) then panics. Closing the trace this way is fine for a test snapshot: an + /// empty (fully closed) upper is readable at any finite peek timestamp. + fn oks_trace_with_rows( + upper: Timestamp, + rows: Vec<((Row, Row), Timestamp, Diff)>, + ) -> PaddedTrace> { + let spine: RowRowSpine = + Trace::new(OperatorInfo::new(0, 0, Rc::from(vec![0])), None, None); + let (agent, mut writer) = + TraceAgent::new(spine, OperatorInfo::new(1, 0, Rc::from(vec![0])), None); + + let description = Description::new( + Antichain::from_elem(Timestamp::minimum()), + Antichain::from_elem(upper), + Antichain::from_elem(Timestamp::minimum()), + ); + let mut chunk = ColumnationStack::default(); + for row in rows { + chunk.push_into(row); + } + let batch = RowRowBuilder::::seal(&mut vec![chunk], description); + writer.insert(batch, Some(Timestamp::minimum())); + + agent.into() + } + + /// Builds a one-batch `[0, upper)` errs trace with no errors, wrapped like a real index's + /// `TraceBundle.errs`. + fn errs_trace_empty(upper: Timestamp) -> PaddedTrace> { + let spine: ErrSpine = + Trace::new(OperatorInfo::new(2, 0, Rc::from(vec![0])), None, None); + let (agent, mut writer) = + TraceAgent::new(spine, OperatorInfo::new(3, 0, Rc::from(vec![0])), None); + + let description = Description::new( + Antichain::from_elem(Timestamp::minimum()), + Antichain::from_elem(upper), + Antichain::from_elem(Timestamp::minimum()), + ); + let chunk = ColumnationStack::default(); + let batch = ErrBuilder::::seal(&mut vec![chunk], description); + writer.insert(batch, Some(Timestamp::minimum())); + + agent.into() + } + + fn test_metrics(registry: &mz_ore::metrics::MetricsRegistry) -> crate::metrics::ComputeMetrics { + crate::metrics::ComputeMetrics::register_with(registry, ComputeRuntimeRole::Maintenance) + } + + fn make_peek(timestamp: Timestamp) -> Peek { + let result_desc = RelationDesc::builder() + .with_column("k", SqlScalarType::Int64.nullable(false)) + .with_column("v", SqlScalarType::Int64.nullable(false)) + .finish(); + Peek { + target: PeekTarget::Index { + id: GlobalId::User(1), + }, + result_desc, + literal_constraints: None, + uuid: Uuid::new_v4(), + timestamp, + finishing: RowSetFinishing::trivial(2), + map_filter_project: MapFilterProject::new(2) + .into_plan() + .expect("identity MFP plans") + .into_nontemporal() + .expect("identity MFP has no temporal filters"), + otel_ctx: OpenTelemetryContext::empty(), + } + } + + fn index_metrics(metrics: &crate::metrics::WorkerMetrics) -> IndexPeekMetrics<'_> { + IndexPeekMetrics { + seek_fulfillment_seconds: &metrics.index_peek_seek_fulfillment_seconds, + frontier_check_seconds: &metrics.index_peek_frontier_check_seconds, + error_scan_seconds: &metrics.index_peek_error_scan_seconds, + cursor_setup_seconds: &metrics.index_peek_cursor_setup_seconds, + row_iteration_seconds: &metrics.index_peek_row_iteration_seconds, + row_iteration_rows: &metrics.index_peek_row_iteration_rows, + result_sort_seconds: &metrics.index_peek_result_sort_seconds, + result_sort_rows: &metrics.index_peek_result_sort_rows, + row_collection_seconds: &metrics.index_peek_row_collection_seconds, + } + } + + /// Publishes `rows` (at time 0, sealed to 1) as a real index arrangement into a fresh registry + /// under `id` on worker 0 of 1, mirroring how a maintained index publishes on the maintenance + /// runtime. + #[mz_ore::test] + #[cfg_attr(miri, ignore)] // differential-dataflow's Columnation isn't miri-clean + fn offloaded_walk_matches_the_inline_walk() { + let registry = mz_ore::metrics::MetricsRegistry::new(); + let metrics = test_metrics(®istry).for_worker(0); + let index_metrics = index_metrics(&metrics); + + let kv = [(row(1), row(10)), (row(2), row(20)), (row(3), row(30))]; + let peek_ts = Timestamp::new(0); + let trace_upper = Timestamp::new(1); + + let bundle = || { + TraceBundle::new( + oks_trace_with_rows( + trace_upper, + kv.iter() + .cloned() + .map(|(k, v)| ((k, v), peek_ts, Diff::ONE)) + .collect(), + ), + errs_trace_empty(trace_upper), + ) + }; + + // Local source, inline. + let mut inline_peek = IndexPeek { + peek: make_peek(peek_ts), + trace_bundle: bundle(), + span: tracing::Span::none(), + }; + let mut upper = Antichain::new(); + let inline_local = + match inline_peek.seek_fulfillment(&mut upper, u64::MAX, false, 0, &index_metrics) { + PeekStatus::Ready(response) => response, + _ => panic!("inline local walk must resolve directly"), + }; + + // Local source, offloaded. + let mut offload_peek = IndexPeek { + peek: make_peek(peek_ts), + trace_bundle: bundle(), + span: tracing::Span::none(), + }; + let mut upper = Antichain::new(); + let offloaded_local = match offload_peek.snapshot_for_offload(&mut upper, false) { + OffloadSnapshot::Ready { + oks, + errs, + oks_stash, + } => offloaded_response::< + PaddedTrace>, + ErrAgent, + >(&make_peek(peek_ts), u64::MAX, oks, errs, oks_stash, None), + _ => panic!("local snapshot must be ready"), + }; + assert_eq!( + inline_local, offloaded_local, + "offloaded walk over a local trace diverged from the inline walk" + ); + } +} + +/// The result of trying to take owned cursors for an offloaded index-peek walk. +/// +/// Generic over the trace source so that any cursor an owned snapshot can produce feeds the same +/// walk. +enum OffloadSnapshot, ErrsTr: TraceReader> { + /// The traces are not sealed through the peek time yet. Keep the peek pending. + NotReady, + /// The peek resolves without a walk, for example because compaction passed its read time. + Response(PeekResponse), + /// Owned cursors covering the read, ready to be walked on another thread. + Ready { + oks: (TraceCursor, TraceStorage), + errs: (TraceCursor, TraceStorage), + /// A second cursor over the same read, present only when the stash could take this peek. + /// + /// Diversion to the stash is decided partway through a walk, and the walk consumes its + /// cursor. The inline path re-walks from a fresh iterator for the same reason. Taking the + /// spare up front keeps that possible without holding the trace handle across threads. + oks_stash: Option<(TraceCursor, TraceStorage)>, + }, +} + /// For keeping track of the state of pending or ready peeks, and managing /// control flow. enum PeekStatus { diff --git a/src/compute/src/compute_state/peek_result_iterator.rs b/src/compute/src/compute_state/peek_result_iterator.rs index 8e3a2967693fb..5b0aaa713eee9 100644 --- a/src/compute/src/compute_state/peek_result_iterator.rs +++ b/src/compute/src/compute_state/peek_result_iterator.rs @@ -14,10 +14,10 @@ use differential_dataflow::trace::implementations::BatchContainer; use differential_dataflow::trace::{Cursor, Navigable, TraceReader}; /// The merged cursor a [`TraceReader::cursor`] hands out over all of a trace's batches: a -/// [`CursorList`] over the per-batch cursors. -type TraceCursor = CursorList>; +/// `CursorList` over the per-batch cursors. +pub(crate) type TraceCursor = CursorList>; /// Backing storage for a [`TraceCursor`]: the batches the cursor borrows from. -type TraceStorage = Vec<::Batch>; +pub(crate) type TraceStorage = Vec<::Batch>; use mz_ore::result::ResultExt; use mz_repr::fixed_length::ExtendDatums; use mz_repr::{DatumVec, Diff, GlobalId, Row, RowArena}; @@ -150,6 +150,35 @@ where self.rows_processed } + /// Builds a [`PeekResultIterator`] over an already-owned cursor and its backing storage. + /// + /// Unlike [`Self::new`], this takes the `(cursor, storage)` pair directly rather than a live + /// `&mut Tr`, so a caller that already holds an owned cursor can feed it without borrowing a + /// trace for the walk. + pub fn new_over_cursor( + target_id: GlobalId, + map_filter_project: mz_expr::SafeMfpPlan, + peek_timestamp: mz_repr::Timestamp, + literal_constraints: Option<&mut [Row]>, + mut cursor: TraceCursor, + storage: TraceStorage, + ) -> Self { + let literals = literal_constraints + .map(|constraints| Literals::new(constraints, &mut cursor, &storage)); + + Self { + target_id, + cursor, + storage, + map_filter_project, + peek_timestamp, + row_builder: Row::default(), + datum_vec: DatumVec::new(), + literals, + rows_processed: 0, + } + } + /// Returns `true` if the iterator has no more literals to process, or if there are no literals at all. fn literals_exhausted(&self) -> bool { self.literals.as_ref().map_or(false, Literals::is_exhausted) @@ -325,3 +354,17 @@ where false } } + +#[cfg(test)] +mod tests { + use mz_repr::{Diff, Timestamp}; + use static_assertions::assert_impl_all; + + use crate::typedefs::RowRowSpine; + + // `TraceReader::cursor` returns its batches by value, and since #38396 those batches are + // `Arc`-backed, so an iterator built over them owns everything it reads. That is what lets a + // peek's walk run on a thread other than the worker that owns the trace, and it is a property + // of the cursor rather than of any wrapper, so assert it directly. + assert_impl_all!(super::PeekResultIterator>: Send); +} diff --git a/src/compute/src/compute_state/peek_stash.rs b/src/compute/src/compute_state/peek_stash.rs index 214d9f371f64f..96dea33476d79 100644 --- a/src/compute/src/compute_state/peek_stash.rs +++ b/src/compute/src/compute_state/peek_stash.rs @@ -123,6 +123,74 @@ impl StashingPeek { } } + /// Stashes `peek_iterator`'s rows and blocks until the upload finishes, returning the response. + /// + /// For a walk that already runs off the serving worker. The worker-pumped path exists because a + /// trace cursor is not `Send` and so cannot be given to the upload task. A caller that owns a + /// `Send` snapshot has no such problem: it drives the same upload directly and never involves + /// the worker. + /// + /// Must not be called from an async context. It blocks the calling thread for the length of the + /// upload, which is why it belongs on a blocking task and why the offload's in-flight cap also + /// bounds how many blocking threads this can occupy. + pub fn upload_blocking( + persist_clients: Arc, + persist_location: &PersistLocation, + peek: &Peek, + peek_iterator: impl Iterator>, + batch_max_runs: usize, + batch_size: usize, + ) -> PeekResponse { + let (rows_tx, rows_rx) = tokio::sync::mpsc::channel(10); + let persist_location = persist_location.clone(); + let peek_uuid = peek.uuid; + let relation_desc = peek.result_desc.clone(); + let rows_needed_by_finishing = peek.finishing.num_rows_needed(); + + let upload = mz_ore::task::spawn( + || format!("peek_stash::stash_peek_response({peek_uuid})"), + async move { + Self::do_upload( + &persist_clients, + persist_location, + batch_max_runs, + peek_uuid, + relation_desc, + rows_needed_by_finishing, + rows_rx, + ) + .await + }, + ); + + let mut peek_iterator = peek_iterator.peekable(); + loop { + let rows: Result, _> = peek_iterator.by_ref().take(batch_size).collect(); + let (rows, done) = match rows { + Ok(rows) if rows.is_empty() => break, + Ok(rows) => { + let done = peek_iterator.peek().is_none(); + (Ok(rows), done) + } + Err(e) => (Err(e), true), + }; + // A send error means the upload stopped reading, which it does once the finishing's + // row bound is met. That is a normal early exit, not a failure. + if rows_tx.blocking_send(rows).is_err() { + break; + } + if done { + break; + } + } + drop(rows_tx); + + match tokio::runtime::Handle::current().block_on(upload) { + Ok(response) => response, + Err(e) => PeekResponse::Error(e), + } + } + async fn do_upload( persist_clients: &PersistClientCache, persist_location: PersistLocation, diff --git a/src/compute/src/metrics.rs b/src/compute/src/metrics.rs index 16861f1b378e3..85a5e2b755b1c 100644 --- a/src/compute/src/metrics.rs +++ b/src/compute/src/metrics.rs @@ -61,6 +61,9 @@ pub struct ComputeMetrics { stashed_peek_seconds: HistogramVec, handle_command_duration_seconds: HistogramVec, + // peek walk substrate + index_peek_walks_total: raw::IntCounterVec, + // Index peek timing phases (per-cluster, no worker label) index_peek_total_seconds: Histogram, index_peek_seek_fulfillment_seconds: Histogram, @@ -172,6 +175,13 @@ impl ComputeMetrics { help: "The total number of dataflows that were replaced during compute reconciliation.", var_labels: ["worker_id", "reason"], ), role)), + index_peek_walks_total: registry.register(with_role(metric!( + name: "mz_index_peek_walks_total", + help: "The total number of fast-path index peek walks, by the substrate that ran \ + them: inline, offload, or capped (wanted the offload, ran inline because \ + the worker was at its in-flight cap).", + var_labels: ["worker_id", "substrate"], + ), role)), arrangement_maintenance_seconds_total: registry.register(with_role(metric!( name: "mz_arrangement_maintenance_seconds_total", help: "The total time spent maintaining arrangements.", @@ -310,6 +320,15 @@ impl ComputeMetrics { self.handle_command_duration_seconds .with_label_values(&[worker.as_ref(), typ]) }); + let index_peek_walks_inline_total = self + .index_peek_walks_total + .with_label_values(&[&worker, "inline"]); + let index_peek_walks_offload_total = self + .index_peek_walks_total + .with_label_values(&[&worker, "offload"]); + let index_peek_walks_capped_total = self + .index_peek_walks_total + .with_label_values(&[&worker, "capped"]); let index_peek_total_seconds = self.index_peek_total_seconds.clone(); let index_peek_seek_fulfillment_seconds = self.index_peek_seek_fulfillment_seconds.clone(); let index_peek_error_scan_seconds = self.index_peek_error_scan_seconds.clone(); @@ -339,6 +358,9 @@ impl ComputeMetrics { persist_peek_seconds, stashed_peek_seconds, handle_command_duration_seconds, + index_peek_walks_inline_total, + index_peek_walks_offload_total, + index_peek_walks_capped_total, index_peek_total_seconds, index_peek_seek_fulfillment_seconds, index_peek_error_scan_seconds, @@ -378,6 +400,25 @@ pub struct WorkerMetrics { pub(crate) stashed_peek_seconds: Histogram, /// Histogram of command handling durations. pub(crate) handle_command_duration_seconds: CommandMetrics, + /// Fast-path index peek walks run inline on this worker because the offload is off. + /// + /// The three walk counters partition every fast-path index peek walk, so their sum is the + /// total and each one alone is a rate of the whole. + pub(crate) index_peek_walks_inline_total: IntCounter, + /// Fast-path index peek walks this worker dispatched to a blocking task. + /// + /// Zero while `enable_index_peek_offload` is off. With it on, a peek the stash could take is + /// offloaded too and diverts from the walking thread, so a flat counter does mean the flag did + /// not reach this worker. + pub(crate) index_peek_walks_offload_total: IntCounter, + /// Fast-path index peek walks that wanted to offload but ran inline at the in-flight cap. + /// + /// NOTE: cap saturation is self-reinforcing. Slots return when walks finish, so a worker whose + /// walks run long accumulates in-flight walks, reaches the cap, and falls back to the inline + /// walk that blocks its step loop, which makes it accumulate faster still. A replica where + /// this climbs while `offload` flattens is losing the offload exactly where it was needed, and + /// `index_peek_offload_max_inflight` is the knob, not the flag. + pub(crate) index_peek_walks_capped_total: IntCounter, /// Histogram of total index peek durations. pub(crate) index_peek_total_seconds: Histogram, /// Histogram of index peek seek_fulfillment durations.