compute: serve fast-path peeks on the interactive runtime

src/compute/src/compute_state/error_scan.rs

@@ -8,27 +8,60 @@
 
 use std::time::{Duration, Instant};
 
-use differential_dataflow::trace::{Cursor, TraceReader};
+use differential_dataflow::trace::{Cursor, Navigable, TraceReader};
 use mz_compute_client::protocol::response::PeekError;
 use mz_repr::{Diff, GlobalId, Timestamp};
 use timely::order::PartialOrder;
 use tracing::error;
 
 use crate::arrangement::manager::PaddedTrace;
 use crate::compute_state::{PeekRowIterationTracker, peek_result_iterator};
+use crate::render::errors::DataflowErrorSer;
 use crate::typedefs::ErrAgent;
 
 /// The error trace of an index, as
 /// [`TraceBundle::errs_mut`](crate::arrangement::manager::TraceBundle::errs_mut) hands it out.
 pub(super) type ErrsHandle = PaddedTrace<ErrAgent<Timestamp, Diff>>;
 
+/// A trace an index peek's error walk can read.
+///
+/// The bound is spelled once here, so the walk and everything that carries it name the shape
+/// rather than restate it.
+pub(super) trait PeekErrsTrace:
+    TraceReader<
+        Time = Timestamp,
+        Batch: Navigable<
+            Cursor: for<'a> Cursor<
+                Key<'a> = &'a DataflowErrorSer,
+                TimeGat<'a>: PartialOrder<Timestamp>,
+                DiffGat<'a> = &'a Diff,
+            >,
+        >,
+    >
+{
+}
+
+impl<Tr> PeekErrsTrace for Tr where
+    Tr: TraceReader<
+            Time = Timestamp,
+            Batch: Navigable<
+                Cursor: for<'a> Cursor<
+                    Key<'a> = &'a DataflowErrorSer,
+                    TimeGat<'a>: PartialOrder<Timestamp>,
+                    DiffGat<'a> = &'a Diff,
+                >,
+            >,
+        >
+{
+}
+
 /// A walk over an index peek's error trace, suspendable between cursor positions.
 ///
 /// Holds nothing of the ok trace or of the rows a peek returns. A peek reaches those only once
 /// this walk reports [`ErrorScanStep::Finished`] with an `Ok`.
-pub(super) struct ErrorScan {
-    cursor: peek_result_iterator::TraceCursor<ErrsHandle>,
-    storage: peek_result_iterator::TraceStorage<ErrsHandle>,
+pub(super) struct ErrorScan<Tr: PeekErrsTrace> {
+    cursor: peek_result_iterator::TraceCursor<Tr>,
+    storage: peek_result_iterator::TraceStorage<Tr>,
     /// The limit spans this walk and the ok scan after it, so the count accrued here is handed
     /// on with [`ErrorScanStep::Finished`].
     row_iteration_tracker: PeekRowIterationTracker,
@@ -48,12 +81,12 @@ pub(super) enum ErrorScanStep {
     OutOfFuel,
 }
 
-impl ErrorScan {
+impl<Tr: PeekErrsTrace> ErrorScan<Tr> {
     /// Opens a walk over `errs`.
     ///
     /// The walk starts without a row-iteration limit. The limit in effect is the caller's to
     /// supply through [`ErrorScan::set_row_iteration_limit`] before each step.
-    pub(super) fn new(errs: &mut ErrsHandle) -> Self {
+    pub(super) fn new(errs: &mut Tr) -> Self {
         let scan_start = Instant::now();
         let (cursor, storage) = errs.cursor();
         let mut scan = Self::from_cursor(cursor, storage);
@@ -63,8 +96,8 @@ impl ErrorScan {
 
     /// Opens a walk over an already-opened cursor.
     pub(super) fn from_cursor(
-        cursor: peek_result_iterator::TraceCursor<ErrsHandle>,
-        storage: peek_result_iterator::TraceStorage<ErrsHandle>,
+        cursor: peek_result_iterator::TraceCursor<Tr>,
+        storage: peek_result_iterator::TraceStorage<Tr>,
     ) -> Self {
         Self {
             cursor,

src/compute/src/compute_state/error_scan/tests.rs

@@ -51,7 +51,10 @@ pub(crate) fn error_batch(
 
 /// Builds a walk over a single-batch error trace holding `updates`, bounded by
 /// `row_iteration_limit`.
-pub(crate) fn error_scan(updates: ErrorUpdates, row_iteration_limit: Option<usize>) -> ErrorScan {
+pub(crate) fn error_scan(
+    updates: ErrorUpdates,
+    row_iteration_limit: Option<usize>,
+) -> ErrorScan<ErrsHandle> {
     let storage = vec![error_batch(updates)];
     let cursor = CursorList::new(vec![storage[0].cursor()], &storage);
     let mut scan = ErrorScan::from_cursor(cursor, storage);
@@ -79,7 +82,10 @@ pub(crate) fn holding(error: &DataflowErrorSer) -> ErrorUpdates {
 
 /// Runs `scan` to an answer in slices of `fuel_per_step` units, and returns that answer, the
 /// fuel the walk spent, and the number of calls it took.
-fn run_sliced(scan: &mut ErrorScan, fuel_per_step: usize) -> (ErrorScanStep, usize, usize) {
+fn run_sliced(
+    scan: &mut ErrorScan<ErrsHandle>,
+    fuel_per_step: usize,
+) -> (ErrorScanStep, usize, usize) {
     let mut consumed = 0;
     // Bounded so that a walk which restarts from the first key on each resumption fails the
     // test instead of hanging it.

src/compute/src/compute_state/index_peek_tests.rs

@@ -27,6 +27,7 @@ use super::error_scan::tests::{
     ErrorUpdates, PEEK_TIMESTAMP, cancelling, error, error_batch, holding,
 };
 use super::*;
+use crate::arrangement::manager::TraceBundle;
 
 /// The collection the peeks in these tests read.
 pub(crate) const TARGET_ID: GlobalId = GlobalId::User(1);
@@ -198,10 +199,10 @@ impl TestMetrics {
         }
     }
 
-    /// How often each metric that `collect_finished_data` can observe into was observed.
+    /// How often each metric that a walk can observe into was observed.
     ///
     /// The two histograms the enclosing `seek_fulfillment` owns are left out, because the
-    /// tests that read this call `collect_finished_data` directly.
+    /// tests that read this call [`collect`] directly.
     fn observations(&self) -> BTreeMap<&'static str, u64> {
         let metrics = &self.metrics;
         BTreeMap::from([
@@ -282,8 +283,8 @@ enum Answer {
     Ready(PeekResponse),
 }
 
-impl From<PeekStatus> for Answer {
-    fn from(status: PeekStatus) -> Self {
+impl From<PeekStatus<IndexPeekScan>> for Answer {
+    fn from(status: PeekStatus<IndexPeekScan>) -> Self {
         match status {
             PeekStatus::NotReady => Answer::NotReady,
             // The scan an offload carries has no comparison of its own. What is comparable
@@ -295,10 +296,35 @@ impl From<PeekStatus> for Answer {
 }
 
 /// An index peek of `peek` over an index holding `keys` and `errors`.
+/// Walks `subject` without the frontier gate, so the observations are the walk's alone.
+fn collect(
+    subject: &mut IndexPeek,
+    max_result_size: u64,
+    stash: StashBounds,
+    row_iteration_limit: Option<usize>,
+    fuel: &mut usize,
+    metrics: &IndexPeekMetrics<'_>,
+) -> PeekStatus<IndexPeekScan> {
+    let (oks, errs) = subject
+        .traces
+        .resolve(subject.peek.target.id())
+        .expect("local traces resolve");
+    IndexPeek::walk_traces(
+        &subject.peek,
+        oks,
+        errs,
+        max_result_size,
+        stash,
+        row_iteration_limit,
+        fuel,
+        metrics,
+    )
+}
+
 fn index_peek_over(peek: Peek, keys: &[Row], errors: ErrorUpdates) -> IndexPeek {
     IndexPeek {
         peek,
-        trace_bundle: trace_bundle(keys, errors),
+        traces: IndexTraces::Local(trace_bundle(keys, errors)),
         span: tracing::Span::none(),
     }
 }
@@ -330,7 +356,8 @@ fn a_completed_scan_answers_with_rows_and_reports_every_phase() {
     );
     let metrics = TestMetrics::new();
 
-    let answer = subject.collect_finished_data(
+    let answer = collect(
+        &mut subject,
         u64::MAX,
         NO_STASH,
         None,
@@ -356,7 +383,8 @@ fn an_error_answered_peek_reports_no_phase_timers() {
     let mut subject = index_peek_over(index_peek(trivial_finishing(), None), &keys, errors);
     let metrics = TestMetrics::new();
 
-    let answer = subject.collect_finished_data(
+    let answer = collect(
+        &mut subject,
         u64::MAX,
         NO_STASH,
         None,
@@ -390,7 +418,8 @@ fn a_scan_that_fills_a_batch_leaves_the_worker_with_fuel_to_spare() {
     // A threshold of zero bytes is crossed by the first row, so the scan fills a batch well
     // before the trace runs out and well before unbounded fuel could run out.
     let mut fuel = unbounded_fuel();
-    let answer = subject.collect_finished_data(
+    let answer = collect(
+        &mut subject,
         u64::MAX,
         STASH_EVERYTHING,
         None,
@@ -420,7 +449,8 @@ fn a_batch_ready_suspension_out_of_fuel_is_offloaded_too() {
     // suspends holding a full batch and out of fuel, with both causes of a suspension in force
     // at once.
     let mut fuel = 1;
-    let answer = subject.collect_finished_data(
+    let answer = collect(
+        &mut subject,
         u64::MAX,
         STASH_EVERYTHING,
         None,
@@ -451,8 +481,14 @@ fn a_scan_that_outruns_its_fuel_leaves_the_worker_reporting_nothing() {
     // An empty error trace is walked out within a position or two, so this fuel is spent
     // inside the ok walk with most of the six keys still ahead of it.
     let mut fuel = 2;
-    let answer =
-        subject.collect_finished_data(u64::MAX, NO_STASH, None, &mut fuel, &metrics.as_metrics());
+    let answer = collect(
+        &mut subject,
+        u64::MAX,
+        NO_STASH,
+        None,
+        &mut fuel,
+        &metrics.as_metrics(),
+    );
 
     assert_eq!(Answer::from(answer), Answer::Offload);
     assert_eq!(
@@ -482,7 +518,8 @@ fn an_ok_phase_failure_reports_the_phases_the_walk_reached() {
     // A ceiling of one byte is crossed by the first row the ok walk produces, so the peek
     // fails inside that walk rather than in the error walk before it.
     let max_result_size = 1;
-    let answer = subject.collect_finished_data(
+    let answer = collect(
+        &mut subject,
         max_result_size,
         NO_STASH,
         None,

src/compute/src/compute_state/index_traces.rs

@@ -0,0 +1,125 @@
+// Copyright Materialize, Inc. and contributors. All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+//! The traces an index peek reads, from an arrangement this runtime maintains or one the sharing
+//! registry publishes.
+
+use differential_dataflow::trace::TraceReader;
+use mz_repr::{Diff, GlobalId, Timestamp};
+use timely::progress::frontier::AntichainRef;
+
+use crate::arrangement::manager::{PaddedTrace, TraceBundle};
+use crate::compute_state::error_scan::ErrsHandle;
+use crate::shared_trace::{SharedErrsHandle, SharedOksHandle};
+use crate::sharing::ArrangementSharingRegistry;
+use crate::typedefs::{ErrSpine, RowRowAgent, RowRowSpine};
+
+/// Where an index peek finds the traces that answer it.
+pub(super) enum IndexTraces {
+    /// Traces this runtime maintains, pinned for the peek's life.
+    Local(TraceBundle),
+    /// An arrangement the sharing registry publishes, resolved on every attempt. A parked peek
+    /// holds nothing of the arrangement, so an unpublished slot registers no hold at the minimum.
+    Shared {
+        registry: ArrangementSharingRegistry,
+        worker_index: usize,
+    },
+}
+
+impl IndexTraces {
+    /// Handles on the traces of `id` for one attempt, or `None` while a shared index is
+    /// unpublished.
+    ///
+    /// Both variants hand out owned handles so the scan can carry them off the worker. A local
+    /// handle is a clone of the pinned one, which registers a hold the pinned one already keeps.
+    pub(super) fn resolve(&mut self, id: GlobalId) -> Option<(PeekOks, PeekErrs)> {
+        match self {
+            IndexTraces::Local(bundle) => {
+                let (oks, errs) = bundle.oks_errs_mut();
+                Some((PeekOks::Local(oks.clone()), PeekErrs::Local(errs.clone())))
+            }
+            IndexTraces::Shared {
+                registry,
+                worker_index,
+            } => registry
+                .handles(&id, *worker_index)
+                .map(|(oks, errs)| (PeekOks::Shared(oks), PeekErrs::Shared(errs))),
+        }
+    }
+}
+
+/// The ok trace an index peek reads.
+pub(super) enum PeekOks {
+    Local(PaddedTrace<RowRowAgent<Timestamp, Diff>>),
+    Shared(SharedOksHandle),
+}
+
+/// The error trace an index peek reads.
+pub(super) enum PeekErrs {
+    Local(ErrsHandle),
+    Shared(SharedErrsHandle),
+}
+
+/// Both variants read the same batch type, so the enum is a `TraceReader` by delegation.
+macro_rules! delegate_trace_reader {
+    ($ty:ident, $spine:ty) => {
+        impl TraceReader for $ty {
+            type Time = Timestamp;
+            type Batch = <$spine as TraceReader>::Batch;
+
+            fn set_logical_compaction(&mut self, frontier: AntichainRef<Timestamp>) {
+                match self {
+                    $ty::Local(trace) => trace.set_logical_compaction(frontier),
+                    $ty::Shared(trace) => trace.set_logical_compaction(frontier),
+                }
+            }
+
+            fn get_logical_compaction(&mut self) -> AntichainRef<'_, Timestamp> {
+                match self {
+                    $ty::Local(trace) => trace.get_logical_compaction(),
+                    $ty::Shared(trace) => trace.get_logical_compaction(),
+                }
+            }
+
+            fn set_physical_compaction(&mut self, frontier: AntichainRef<Timestamp>) {
+                match self {
+                    $ty::Local(trace) => trace.set_physical_compaction(frontier),
+                    $ty::Shared(trace) => trace.set_physical_compaction(frontier),
+                }
+            }
+
+            fn get_physical_compaction(&mut self) -> AntichainRef<'_, Timestamp> {
+                match self {
+                    $ty::Local(trace) => trace.get_physical_compaction(),
+                    $ty::Shared(trace) => trace.get_physical_compaction(),
+                }
+            }
+
+            fn map_batches<F: FnMut(&Self::Batch)>(&self, f: F) {
+                match self {
+                    $ty::Local(trace) => trace.map_batches(f),
+                    $ty::Shared(trace) => trace.map_batches(f),
+                }
+            }
+
+            fn batches_through(
+                &mut self,
+                upper: AntichainRef<Timestamp>,
+            ) -> Option<Vec<Self::Batch>> {
+                match self {
+                    $ty::Local(trace) => trace.batches_through(upper),
+                    $ty::Shared(trace) => trace.batches_through(upper),
+                }
+            }
+        }
+    };
+}
+
+delegate_trace_reader!(PeekOks, RowRowSpine<Timestamp, Diff>);
+delegate_trace_reader!(PeekErrs, ErrSpine<Timestamp, Diff>);