mz_compute/logging/
compute.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Logging dataflows for events generated by clusterd.
11
12use std::cell::RefCell;
13use std::collections::{BTreeMap, BTreeSet};
14use std::fmt::{Display, Write};
15use std::rc::Rc;
16use std::time::{Duration, Instant};
17
18use chrono::DateTime;
19use columnar::{Columnar, Index, Ref};
20use differential_dataflow::VecCollection;
21use differential_dataflow::collection::AsCollection;
22use differential_dataflow::trace::{Cursor, Navigable, Span};
23use mz_compute_types::plan::LirId;
24use mz_ore::cast::CastFrom;
25use mz_repr::adt::timestamp::CheckedTimestamp;
26use mz_repr::{Datum, Diff, GlobalId, Row, RowRef, Timestamp};
27use mz_row_spine::RowRowBuilder;
28use mz_timely_util::columnar::batcher;
29use mz_timely_util::columnar::builder::ColumnBuilder;
30use mz_timely_util::columnar::{Col2ValBatcher, Column, columnar_exchange};
31use mz_timely_util::replay::MzReplay;
32use timely::dataflow::channels::pact::{ExchangeCore, Pipeline};
33use timely::dataflow::operators::Operator;
34use timely::dataflow::operators::generic::OutputBuilder;
35use timely::dataflow::operators::generic::builder_rc::OperatorBuilder;
36use timely::dataflow::operators::generic::operator::empty;
37use timely::dataflow::{Scope, StreamVec};
38use timely::scheduling::activate::{Activations, Activator};
39use tracing::error;
40use uuid::Uuid;
41
42use crate::extensions::arrange::MzArrangeCore;
43use crate::logging::{
44    ComputeLog, EventQueue, LogCollection, LogVariant, OutputSessionColumnar, PermutedRowPacker,
45    SharedLoggingState, Update,
46};
47use crate::typedefs::RowRowSpine;
48
49/// Type alias for a logger of compute events.
50pub type Logger = timely::logging_core::Logger<ComputeEventBuilder>;
51pub type ComputeEventBuilder = ColumnBuilder<(Duration, ComputeEvent)>;
52
53/// A dataflow exports a global ID.
54#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
55pub struct Export {
56    /// Identifier of the export.
57    pub export_id: GlobalId,
58    /// Timely worker index of the exporting dataflow.
59    pub dataflow_index: usize,
60}
61
62/// The export for a global id was dropped.
63#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
64pub struct ExportDropped {
65    /// Identifier of the export.
66    pub export_id: GlobalId,
67}
68
69/// A peek event with a [`PeekType`], and an installation status.
70#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
71pub struct PeekEvent {
72    /// The identifier of the view the peek targets.
73    pub id: GlobalId,
74    /// The logical timestamp requested.
75    pub time: Timestamp,
76    /// The ID of the peek.
77    pub uuid: uuid::Bytes,
78    /// The relevant _type_ of peek: index or persist.
79    // Note that this is not stored on the Peek event for data-packing reasons only.
80    pub peek_type: PeekType,
81    /// True if the peek is being installed; false if it's being removed.
82    pub installed: bool,
83}
84
85/// Frontier change event.
86#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
87pub struct Frontier {
88    pub export_id: GlobalId,
89    pub time: Timestamp,
90    pub diff: i8,
91}
92
93/// An import frontier change.
94#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
95pub struct ImportFrontier {
96    pub import_id: GlobalId,
97    pub export_id: GlobalId,
98    pub time: Timestamp,
99    pub diff: i8,
100}
101
102/// A change in an arrangement's heap size.
103#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
104pub struct ArrangementHeapSize {
105    /// Operator index
106    pub operator_id: usize,
107    /// Delta of the heap size in bytes of the arrangement.
108    pub delta_size: isize,
109}
110
111/// A change in an arrangement's heap capacity.
112#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
113pub struct ArrangementHeapCapacity {
114    /// Operator index
115    pub operator_id: usize,
116    /// Delta of the heap capacity in bytes of the arrangement.
117    pub delta_capacity: isize,
118}
119
120/// A change in an arrangement's heap allocation count.
121#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
122pub struct ArrangementHeapAllocations {
123    /// Operator index
124    pub operator_id: usize,
125    /// Delta of distinct heap allocations backing the arrangement.
126    pub delta_allocations: isize,
127}
128
129/// Announcing an operator that manages an arrangement.
130#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
131pub struct ArrangementHeapSizeOperator {
132    /// Operator index
133    pub operator_id: usize,
134    /// The address of the operator.
135    pub address: Vec<usize>,
136}
137
138/// Drop event for an operator managing an arrangement.
139#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
140pub struct ArrangementHeapSizeOperatorDrop {
141    /// Operator index
142    pub operator_id: usize,
143}
144
145/// Dataflow shutdown event.
146#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
147pub struct DataflowShutdown {
148    /// Timely worker index of the dataflow.
149    pub dataflow_index: usize,
150}
151
152/// Error count update event.
153#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
154pub struct ErrorCount {
155    /// Identifier of the export.
156    pub export_id: GlobalId,
157    /// The change in error count.
158    pub diff: Diff,
159}
160
161/// An export started hydrating.
162#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
163pub struct HydrationStart {
164    /// Identifier of the export.
165    pub export_id: GlobalId,
166}
167
168/// An export is hydrated.
169#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
170pub struct Hydration {
171    /// Identifier of the export.
172    pub export_id: GlobalId,
173}
174
175/// An operator's hydration status changed.
176#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
177pub struct OperatorHydration {
178    /// Identifier of the export.
179    pub export_id: GlobalId,
180    /// Identifier of the operator's LIR node.
181    pub lir_id: LirId,
182    /// Whether the operator is hydrated.
183    pub hydrated: bool,
184}
185
186/// Announce a mapping of an LIR operator to a dataflow operator for a global ID.
187#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
188pub struct LirMapping {
189    /// The `GlobalId` in which the LIR operator is rendered.
190    ///
191    /// NB a single a dataflow may have many `GlobalId`s inside it.
192    /// A separate mapping (using `ComputeEvent::DataflowGlobal`)
193    /// tracks the many-to-one relationship between `GlobalId`s and
194    /// dataflows.
195    pub global_id: GlobalId,
196    /// The actual mapping.
197    /// Represented this way to reduce the size of `ComputeEvent`.
198    pub mapping: Vec<(LirId, LirMetadata)>,
199}
200
201/// Announce that a dataflow supports a specific global ID.
202#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
203pub struct DataflowGlobal {
204    /// The identifier of the dataflow.
205    pub dataflow_index: usize,
206    /// A `GlobalId` that is rendered as part of this dataflow.
207    pub global_id: GlobalId,
208}
209
210/// A logged compute event.
211#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
212pub enum ComputeEvent {
213    /// A dataflow export was created.
214    Export(Export),
215    /// A dataflow export was dropped.
216    ExportDropped(ExportDropped),
217    /// Peek command.
218    Peek(PeekEvent),
219    /// Available frontier information for dataflow exports.
220    Frontier(Frontier),
221    /// Available frontier information for dataflow imports.
222    ImportFrontier(ImportFrontier),
223    /// Arrangement heap size update
224    ArrangementHeapSize(ArrangementHeapSize),
225    /// Arrangement heap size update
226    ArrangementHeapCapacity(ArrangementHeapCapacity),
227    /// Arrangement heap size update
228    ArrangementHeapAllocations(ArrangementHeapAllocations),
229    /// Arrangement size operator address
230    ArrangementHeapSizeOperator(ArrangementHeapSizeOperator),
231    /// Arrangement size operator dropped
232    ArrangementHeapSizeOperatorDrop(ArrangementHeapSizeOperatorDrop),
233    /// All operators of a dataflow have shut down.
234    DataflowShutdown(DataflowShutdown),
235    /// The number of errors in a dataflow export has changed.
236    ErrorCount(ErrorCount),
237    /// A dataflow export started hydrating, i.e. its dataflow was unsuspended.
238    HydrationStart(HydrationStart),
239    /// A dataflow export was hydrated.
240    Hydration(Hydration),
241    /// A dataflow operator's hydration status changed.
242    OperatorHydration(OperatorHydration),
243    /// An LIR operator was mapped to some particular dataflow operator.
244    ///
245    /// Cf. `ComputeLog::LirMaping`
246    LirMapping(LirMapping),
247    DataflowGlobal(DataflowGlobal),
248}
249
250/// A peek type distinguishing between index and persist peeks.
251#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug, Columnar)]
252pub enum PeekType {
253    /// A peek against an index.
254    Index,
255    /// A peek against persist.
256    Persist,
257}
258
259impl PeekType {
260    /// A human-readable name for a peek type.
261    fn name(self) -> &'static str {
262        match self {
263            PeekType::Index => "index",
264            PeekType::Persist => "persist",
265        }
266    }
267}
268
269/// Metadata for LIR operators.
270#[derive(Clone, Debug, PartialEq, PartialOrd, Columnar)]
271pub struct LirMetadata {
272    /// The LIR operator, as a string (see `FlatPlanNode::humanize`).
273    operator: String,
274    /// The LIR identifier of the parent (if any).
275    parent_lir_id: Option<LirId>,
276    /// How nested the operator is (for nice indentation).
277    nesting: u8,
278    /// The dataflow operator ids, given as start (inclusive) and end (exclusive).
279    /// If `start == end`, then no operators were used.
280    operator_span: (usize, usize),
281}
282
283impl LirMetadata {
284    /// Construct a new LIR metadata object.
285    pub fn new(
286        operator: String,
287        parent_lir_id: Option<LirId>,
288        nesting: u8,
289        operator_span: (usize, usize),
290    ) -> Self {
291        Self {
292            operator,
293            parent_lir_id,
294            nesting,
295            operator_span,
296        }
297    }
298}
299
300/// The return type of the [`construct`] function.
301pub(super) struct Return {
302    /// Collections returned by [`construct`].
303    pub collections: BTreeMap<LogVariant, LogCollection>,
304}
305
306/// Constructs the logging dataflow fragment for compute logs.
307///
308/// Params
309/// * `scope`: The Timely scope hosting the log analysis dataflow.
310/// * `scheduler`: The timely scheduler to obtainer activators.
311/// * `config`: Logging configuration.
312/// * `event_queue`: The source to read compute log events from.
313/// * `compute_event_streams`: Additional compute event streams to absorb.
314/// * `shared_state`: Shared state between logging dataflow fragments.
315pub(super) fn construct<'scope>(
316    scope: Scope<'scope, Timestamp>,
317    activations: Rc<RefCell<Activations>>,
318    config: &mz_compute_client::logging::LoggingConfig,
319    event_queue: EventQueue<Column<(Duration, ComputeEvent)>>,
320    shared_state: Rc<RefCell<SharedLoggingState>>,
321) -> Return {
322    let logging_interval_ms = std::cmp::max(1, config.interval.as_millis());
323
324    scope.scoped("compute logging", move |scope| {
325        let enable_logging = config.enable_logging;
326        let (logs, token) = if enable_logging {
327            event_queue.links.mz_replay(
328                scope,
329                "compute logs",
330                config.interval,
331                event_queue.activator,
332            )
333        } else {
334            let token: Rc<dyn std::any::Any> = Rc::new(Box::new(()));
335            (empty(scope), token)
336        };
337
338        // Build a demux operator that splits the replayed event stream up into the separate
339        // logging streams.
340        let mut demux = OperatorBuilder::new("Compute Logging Demux".to_string(), scope.clone());
341        let mut input = demux.new_input(logs, Pipeline);
342        let (export_out, export) = demux.new_output();
343        let mut export_out = OutputBuilder::from(export_out);
344        let (frontier_out, frontier) = demux.new_output();
345        let mut frontier_out = OutputBuilder::from(frontier_out);
346        let (import_frontier_out, import_frontier) = demux.new_output();
347        let mut import_frontier_out = OutputBuilder::from(import_frontier_out);
348        let (peek_out, peek) = demux.new_output();
349        let mut peek_out = OutputBuilder::from(peek_out);
350        let (peek_duration_out, peek_duration) = demux.new_output();
351        let mut peek_duration_out = OutputBuilder::from(peek_duration_out);
352        let (arrangement_heap_size_out, arrangement_heap_size) = demux.new_output();
353        let mut arrangement_heap_size_out = OutputBuilder::from(arrangement_heap_size_out);
354        let (arrangement_heap_capacity_out, arrangement_heap_capacity) = demux.new_output();
355        let mut arrangement_heap_capacity_out = OutputBuilder::from(arrangement_heap_capacity_out);
356        let (arrangement_heap_allocations_out, arrangement_heap_allocations) = demux.new_output();
357        let mut arrangement_heap_allocations_out =
358            OutputBuilder::from(arrangement_heap_allocations_out);
359        let (error_count_out, error_count) = demux.new_output();
360        let mut error_count_out = OutputBuilder::from(error_count_out);
361        let (hydration_time_out, hydration_time) = demux.new_output();
362        let mut hydration_time_out = OutputBuilder::from(hydration_time_out);
363        let (operator_hydration_status_out, operator_hydration_status) = demux.new_output();
364        let mut operator_hydration_status_out = OutputBuilder::from(operator_hydration_status_out);
365        let (lir_mapping_out, lir_mapping) = demux.new_output();
366        let mut lir_mapping_out = OutputBuilder::from(lir_mapping_out);
367        let (dataflow_global_ids_out, dataflow_global_ids) = demux.new_output();
368        let mut dataflow_global_ids_out = OutputBuilder::from(dataflow_global_ids_out);
369
370        let mut demux_state = DemuxState::new(activations, scope.index());
371        demux.build(move |_capability| {
372            move |_frontiers| {
373                let mut export = export_out.activate();
374                let mut frontier = frontier_out.activate();
375                let mut import_frontier = import_frontier_out.activate();
376                let mut peek = peek_out.activate();
377                let mut peek_duration = peek_duration_out.activate();
378                let mut arrangement_heap_size = arrangement_heap_size_out.activate();
379                let mut arrangement_heap_capacity = arrangement_heap_capacity_out.activate();
380                let mut arrangement_heap_allocations = arrangement_heap_allocations_out.activate();
381                let mut error_count = error_count_out.activate();
382                let mut hydration_time = hydration_time_out.activate();
383                let mut operator_hydration_status = operator_hydration_status_out.activate();
384                let mut lir_mapping = lir_mapping_out.activate();
385                let mut dataflow_global_ids = dataflow_global_ids_out.activate();
386
387                input.for_each(|cap, data| {
388                    let mut output_sessions = DemuxOutput {
389                        export: export.session_with_builder(&cap),
390                        frontier: frontier.session_with_builder(&cap),
391                        import_frontier: import_frontier.session_with_builder(&cap),
392                        peek: peek.session_with_builder(&cap),
393                        peek_duration: peek_duration.session_with_builder(&cap),
394                        arrangement_heap_allocations: arrangement_heap_allocations
395                            .session_with_builder(&cap),
396                        arrangement_heap_capacity: arrangement_heap_capacity
397                            .session_with_builder(&cap),
398                        arrangement_heap_size: arrangement_heap_size.session_with_builder(&cap),
399                        error_count: error_count.session_with_builder(&cap),
400                        hydration_time: hydration_time.session_with_builder(&cap),
401                        operator_hydration_status: operator_hydration_status
402                            .session_with_builder(&cap),
403                        lir_mapping: lir_mapping.session_with_builder(&cap),
404                        dataflow_global_ids: dataflow_global_ids.session_with_builder(&cap),
405                    };
406
407                    let shared_state = &mut shared_state.borrow_mut();
408                    for (time, event) in data.borrow().into_index_iter() {
409                        DemuxHandler {
410                            state: &mut demux_state,
411                            shared_state,
412                            output: &mut output_sessions,
413                            logging_interval_ms,
414                            time,
415                        }
416                        .handle(event);
417                    }
418                });
419            }
420        });
421
422        use ComputeLog::*;
423        let logs = [
424            (ArrangementHeapAllocations, arrangement_heap_allocations),
425            (ArrangementHeapCapacity, arrangement_heap_capacity),
426            (ArrangementHeapSize, arrangement_heap_size),
427            (DataflowCurrent, export),
428            (DataflowGlobal, dataflow_global_ids),
429            (ErrorCount, error_count),
430            (FrontierCurrent, frontier),
431            (HydrationTime, hydration_time),
432            (ImportFrontierCurrent, import_frontier),
433            (LirMapping, lir_mapping),
434            (OperatorHydrationStatus, operator_hydration_status),
435            (PeekCurrent, peek),
436            (PeekDuration, peek_duration),
437        ];
438
439        // Build the output arrangements.
440        let mut collections = BTreeMap::new();
441        for (variant, stream) in logs {
442            let variant = LogVariant::Compute(variant);
443            if config.index_logs.contains_key(&variant) {
444                let exchange = ExchangeCore::<ColumnBuilder<_>, _>::new_core(
445                    columnar_exchange::<Row, Row, Timestamp, Diff>,
446                );
447                let trace = stream
448                    .mz_arrange_core::<
449                        _,
450                        Col2ValBatcher<_, _, _, _, batcher::Chunker<_>, RowRowBuilder<_, _>>,
451                        RowRowSpine<_, _>,
452                    >(exchange, &format!("Arrange {variant:?}"))
453                    .trace;
454                let collection = LogCollection {
455                    trace,
456                    token: Rc::clone(&token),
457                };
458                collections.insert(variant, collection);
459            }
460        }
461
462        Return { collections }
463    })
464}
465
466/// Format the given value and pack it into a `Datum::String`.
467///
468/// The `scratch` buffer is used to perform the string conversion without an allocation.
469/// Callers should not assume anything about the contents of this buffer after this function
470/// returns.
471fn make_string_datum<V>(value: V, scratch: &mut String) -> Datum<'_>
472where
473    V: Display,
474{
475    scratch.clear();
476    write!(scratch, "{}", value).expect("writing to a `String` can't fail");
477    Datum::String(scratch)
478}
479
480/// Pack an offset from the Unix epoch into a `Datum::TimestampTz`.
481///
482/// The offset is rounded to microseconds, the maximum resolution of `timestamptz`. Rounding is
483/// monotone, so it cannot reorder two offsets that were ordered before.
484fn epoch_offset_datum(offset: Duration) -> Datum<'static> {
485    let secs = i64::try_from(offset.as_secs()).expect("must fit");
486    let datetime = DateTime::from_timestamp(secs, offset.subsec_nanos())
487        .expect("epoch offset is in range for `DateTime`");
488    let timestamp = CheckedTimestamp::try_from(datetime)
489        .expect("epoch offset is a valid timestamp")
490        .round_to_precision(None)
491        .expect("epoch offset is far from the maximum timestamp");
492    Datum::TimestampTz(timestamp)
493}
494
495/// State maintained by the demux operator.
496struct DemuxState {
497    /// The timely activations handle.
498    activations: Rc<RefCell<Activations>>,
499    /// The index of this worker.
500    worker_id: usize,
501    /// A reusable scratch string for formatting IDs.
502    scratch_string_a: String,
503    /// A reusable scratch string for formatting IDs.
504    scratch_string_b: String,
505    /// State tracked per dataflow export.
506    exports: BTreeMap<GlobalId, ExportState>,
507    /// Maps pending peeks to their installation time.
508    peek_stash: BTreeMap<Uuid, Duration>,
509    /// Arrangement size stash.
510    arrangement_size: BTreeMap<usize, ArrangementSizeState>,
511    /// LIR -> operator span mapping.
512    lir_mapping: BTreeMap<GlobalId, BTreeMap<LirId, LirMetadata>>,
513    /// Dataflow -> `GlobalId` mapping (many-to-one).
514    dataflow_global_ids: BTreeMap<usize, BTreeSet<GlobalId>>,
515    /// A row packer for the arrangement heap allocations output.
516    arrangement_heap_allocations_packer: PermutedRowPacker,
517    /// A row packer for the arrangement heap capacity output.
518    arrangement_heap_capacity_packer: PermutedRowPacker,
519    /// A row packer for the arrangement heap size output.
520    arrangement_heap_size_packer: PermutedRowPacker,
521    /// A row packer for the dataflow global output.
522    dataflow_global_packer: PermutedRowPacker,
523    /// A row packer for the error count output.
524    error_count_packer: PermutedRowPacker,
525    /// A row packer for the exports output.
526    export_packer: PermutedRowPacker,
527    /// A row packer for the frontier output.
528    frontier_packer: PermutedRowPacker,
529    /// A row packer for the exports output.
530    import_frontier_packer: PermutedRowPacker,
531    /// A row packer for the LIR mapping output.
532    lir_mapping_packer: PermutedRowPacker,
533    /// A row packer for the operator hydration status output.
534    operator_hydration_status_packer: PermutedRowPacker,
535    /// A row packer for the peek durations output.
536    peek_duration_packer: PermutedRowPacker,
537    /// A row packer for the peek output.
538    peek_packer: PermutedRowPacker,
539    /// A row packer for the hydration time output.
540    hydration_time_packer: PermutedRowPacker,
541}
542
543impl DemuxState {
544    fn new(activations: Rc<RefCell<Activations>>, worker_id: usize) -> Self {
545        Self {
546            activations,
547            worker_id,
548            scratch_string_a: String::new(),
549            scratch_string_b: String::new(),
550            exports: Default::default(),
551            peek_stash: Default::default(),
552            arrangement_size: Default::default(),
553            lir_mapping: Default::default(),
554            dataflow_global_ids: Default::default(),
555            arrangement_heap_allocations_packer: PermutedRowPacker::new(
556                ComputeLog::ArrangementHeapAllocations,
557            ),
558            arrangement_heap_capacity_packer: PermutedRowPacker::new(
559                ComputeLog::ArrangementHeapCapacity,
560            ),
561            arrangement_heap_size_packer: PermutedRowPacker::new(ComputeLog::ArrangementHeapSize),
562            dataflow_global_packer: PermutedRowPacker::new(ComputeLog::DataflowGlobal),
563            error_count_packer: PermutedRowPacker::new(ComputeLog::ErrorCount),
564            export_packer: PermutedRowPacker::new(ComputeLog::DataflowCurrent),
565            frontier_packer: PermutedRowPacker::new(ComputeLog::FrontierCurrent),
566            hydration_time_packer: PermutedRowPacker::new(ComputeLog::HydrationTime),
567            import_frontier_packer: PermutedRowPacker::new(ComputeLog::ImportFrontierCurrent),
568            lir_mapping_packer: PermutedRowPacker::new(ComputeLog::LirMapping),
569            operator_hydration_status_packer: PermutedRowPacker::new(
570                ComputeLog::OperatorHydrationStatus,
571            ),
572            peek_duration_packer: PermutedRowPacker::new(ComputeLog::PeekDuration),
573            peek_packer: PermutedRowPacker::new(ComputeLog::PeekCurrent),
574        }
575    }
576
577    /// Pack an arrangement heap allocations update key-value for the given operator.
578    fn pack_arrangement_heap_allocations_update(
579        &mut self,
580        operator_id: usize,
581    ) -> (&RowRef, &RowRef) {
582        self.arrangement_heap_allocations_packer.pack_slice(&[
583            Datum::UInt64(operator_id.try_into().expect("operator_id too big")),
584            Datum::UInt64(u64::cast_from(self.worker_id)),
585        ])
586    }
587
588    /// Pack an arrangement heap capacity update key-value for the given operator.
589    fn pack_arrangement_heap_capacity_update(&mut self, operator_id: usize) -> (&RowRef, &RowRef) {
590        self.arrangement_heap_capacity_packer.pack_slice(&[
591            Datum::UInt64(operator_id.try_into().expect("operator_id too big")),
592            Datum::UInt64(u64::cast_from(self.worker_id)),
593        ])
594    }
595
596    /// Pack an arrangement heap size update key-value for the given operator.
597    fn pack_arrangement_heap_size_update(&mut self, operator_id: usize) -> (&RowRef, &RowRef) {
598        self.arrangement_heap_size_packer.pack_slice(&[
599            Datum::UInt64(operator_id.try_into().expect("operator_id too big")),
600            Datum::UInt64(u64::cast_from(self.worker_id)),
601        ])
602    }
603
604    /// Pack a dataflow global update key-value for the given dataflow index and global ID.
605    fn pack_dataflow_global_update(
606        &mut self,
607        dataflow_index: usize,
608        global_id: GlobalId,
609    ) -> (&RowRef, &RowRef) {
610        self.dataflow_global_packer.pack_slice(&[
611            Datum::UInt64(u64::cast_from(dataflow_index)),
612            Datum::UInt64(u64::cast_from(self.worker_id)),
613            make_string_datum(global_id, &mut self.scratch_string_a),
614        ])
615    }
616
617    /// Pack an error count update key-value for the given export ID and count.
618    fn pack_error_count_update(&mut self, export_id: GlobalId, count: Diff) -> (&RowRef, &RowRef) {
619        // Normally we would use DD's diff field to encode counts, but in this case we can't: The total
620        // per-worker error count might be negative and at the SQL level having negative multiplicities
621        // is treated as an error.
622        self.error_count_packer.pack_slice(&[
623            make_string_datum(export_id, &mut self.scratch_string_a),
624            Datum::UInt64(u64::cast_from(self.worker_id)),
625            Datum::Int64(count.into_inner()),
626        ])
627    }
628
629    /// Pack an export update key-value for the given export ID and dataflow index.
630    fn pack_export_update(
631        &mut self,
632        export_id: GlobalId,
633        dataflow_index: usize,
634    ) -> (&RowRef, &RowRef) {
635        self.export_packer.pack_slice(&[
636            make_string_datum(export_id, &mut self.scratch_string_a),
637            Datum::UInt64(u64::cast_from(self.worker_id)),
638            Datum::UInt64(u64::cast_from(dataflow_index)),
639        ])
640    }
641
642    /// Pack a hydration time update key-value for the given export ID, hydration time, and
643    /// hydration lifecycle timestamps.
644    fn pack_hydration_time_update(
645        &mut self,
646        export_id: GlobalId,
647        time_ns: Option<u64>,
648        timestamps: &HydrationTimestamps,
649    ) -> (&RowRef, &RowRef) {
650        self.hydration_time_packer.pack_slice(&[
651            make_string_datum(export_id, &mut self.scratch_string_a),
652            Datum::UInt64(u64::cast_from(self.worker_id)),
653            Datum::from(time_ns),
654            epoch_offset_datum(timestamps.installed_at),
655            timestamps
656                .started_at
657                .map_or(Datum::Null, epoch_offset_datum),
658            timestamps
659                .hydrated_at
660                .map_or(Datum::Null, epoch_offset_datum),
661        ])
662    }
663
664    /// Pack an import frontier update key-value for the given export ID and dataflow index.
665    fn pack_import_frontier_update(
666        &mut self,
667        export_id: GlobalId,
668        import_id: GlobalId,
669        time: Timestamp,
670    ) -> (&RowRef, &RowRef) {
671        self.import_frontier_packer.pack_slice(&[
672            make_string_datum(export_id, &mut self.scratch_string_a),
673            make_string_datum(import_id, &mut self.scratch_string_b),
674            Datum::UInt64(u64::cast_from(self.worker_id)),
675            Datum::MzTimestamp(time),
676        ])
677    }
678
679    /// Pack an LIR mapping update key-value for the given LIR operator metadata.
680    fn pack_lir_mapping_update(
681        &mut self,
682        global_id: GlobalId,
683        lir_id: LirId,
684        operator: String,
685        parent_lir_id: Option<LirId>,
686        nesting: u8,
687        operator_span: (usize, usize),
688    ) -> (&RowRef, &RowRef) {
689        self.lir_mapping_packer.pack_slice(&[
690            make_string_datum(global_id, &mut self.scratch_string_a),
691            Datum::UInt64(lir_id.into()),
692            Datum::UInt64(u64::cast_from(self.worker_id)),
693            make_string_datum(operator, &mut self.scratch_string_b),
694            parent_lir_id.map_or(Datum::Null, |lir_id| Datum::UInt64(lir_id.into())),
695            Datum::UInt16(u16::cast_from(nesting)),
696            Datum::UInt64(u64::cast_from(operator_span.0)),
697            Datum::UInt64(u64::cast_from(operator_span.1)),
698        ])
699    }
700
701    /// Pack an operator hydration status update key-value for the given export ID, LIR ID, and
702    /// hydration status.
703    fn pack_operator_hydration_status_update(
704        &mut self,
705        export_id: GlobalId,
706        lir_id: LirId,
707        hydrated: bool,
708    ) -> (&RowRef, &RowRef) {
709        self.operator_hydration_status_packer.pack_slice(&[
710            make_string_datum(export_id, &mut self.scratch_string_a),
711            Datum::UInt64(lir_id.into()),
712            Datum::UInt64(u64::cast_from(self.worker_id)),
713            Datum::from(hydrated),
714        ])
715    }
716
717    /// Pack a peek duration update key-value for the given peek and peek type.
718    fn pack_peek_duration_update(
719        &mut self,
720        peek_type: PeekType,
721        bucket: u128,
722    ) -> (&RowRef, &RowRef) {
723        self.peek_duration_packer.pack_slice(&[
724            Datum::UInt64(u64::cast_from(self.worker_id)),
725            Datum::String(peek_type.name()),
726            Datum::UInt64(bucket.try_into().expect("bucket too big")),
727        ])
728    }
729
730    /// Pack a peek update key-value for the given peek and peek type.
731    fn pack_peek_update(
732        &mut self,
733        id: GlobalId,
734        time: Timestamp,
735        uuid: Uuid,
736        peek_type: PeekType,
737    ) -> (&RowRef, &RowRef) {
738        self.peek_packer.pack_slice(&[
739            Datum::Uuid(uuid),
740            Datum::UInt64(u64::cast_from(self.worker_id)),
741            make_string_datum(id, &mut self.scratch_string_a),
742            Datum::String(peek_type.name()),
743            Datum::MzTimestamp(time),
744        ])
745    }
746
747    /// Pack a frontier update key-value for the given export ID and time.
748    fn pack_frontier_update(&mut self, export_id: GlobalId, time: Timestamp) -> (&RowRef, &RowRef) {
749        self.frontier_packer.pack_slice(&[
750            make_string_datum(export_id, &mut self.scratch_string_a),
751            Datum::UInt64(u64::cast_from(self.worker_id)),
752            Datum::MzTimestamp(time),
753        ])
754    }
755}
756
757/// Wallclock instants of an export's hydration lifecycle, as durations since the Unix epoch.
758///
759/// These are sampled from compute logging event times, which advance off an `Instant` anchored to
760/// the epoch once per worker when logging is initialized. They are therefore monotone within a
761/// worker and unaffected by system clock steps, but they carry that worker's anchor, so comparing
762/// them across workers absorbs anchor skew.
763///
764/// The invariant `installed_at <= started_at <= hydrated_at` holds over the non-NULL values, and no
765/// later stage is ever set without its predecessors.
766#[derive(Debug, Clone, Copy)]
767struct HydrationTimestamps {
768    /// When this export's dataflow was installed, still suspended. Also the start of queueing.
769    installed_at: Duration,
770    /// When hydration work began, that is, when the dataflow was unsuspended.
771    ///
772    /// Equals `installed_at` for a dataflow that was never suspended, since such a dataflow starts
773    /// running the moment it is built.
774    started_at: Option<Duration>,
775    /// When this export's output frontier passed the dataflow as-of.
776    hydrated_at: Option<Duration>,
777}
778
779/// State tracked for each dataflow export.
780struct ExportState {
781    /// The ID of the dataflow maintaining this export.
782    dataflow_index: usize,
783    /// Number of errors in this export.
784    ///
785    /// This must be a signed integer, since per-worker error counts can be negative, only the
786    /// cross-worker total has to sum up to a non-negative value.
787    error_count: Diff,
788    /// When this export was created.
789    created_at: Instant,
790    /// Whether the exported collection is hydrated.
791    hydration_time_ns: Option<u64>,
792    /// Wallclock instants of this export's hydration lifecycle.
793    ///
794    /// Distinct from `created_at`/`hydration_time_ns`, which remain the authoritative per-worker
795    /// duration at nanosecond precision. These instants exist to identify and bound hydration
796    /// episodes, which a duration cannot do.
797    hydration_timestamps: HydrationTimestamps,
798    /// Hydration status of operators feeding this export.
799    operator_hydration: BTreeMap<LirId, bool>,
800}
801
802impl ExportState {
803    fn new(dataflow_index: usize, installed_at: Duration) -> Self {
804        Self {
805            dataflow_index,
806            error_count: Diff::ZERO,
807            created_at: Instant::now(),
808            hydration_time_ns: None,
809            hydration_timestamps: HydrationTimestamps {
810                installed_at,
811                started_at: None,
812                hydrated_at: None,
813            },
814            operator_hydration: BTreeMap::new(),
815        }
816    }
817}
818
819/// State for tracking arrangement sizes.
820#[derive(Default, Debug)]
821struct ArrangementSizeState {
822    size: isize,
823    capacity: isize,
824    count: isize,
825}
826
827/// Bundled output sessions used by the demux operator.
828struct DemuxOutput<'a, 'b> {
829    export: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
830    frontier: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
831    import_frontier: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
832    peek: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
833    peek_duration: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
834    arrangement_heap_allocations: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
835    arrangement_heap_capacity: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
836    arrangement_heap_size: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
837    hydration_time: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
838    operator_hydration_status: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
839    error_count: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
840    lir_mapping: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
841    dataflow_global_ids: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
842}
843
844/// Event handler of the demux operator.
845struct DemuxHandler<'a, 'b, 'c> {
846    /// State kept by the demux operator.
847    state: &'a mut DemuxState,
848    /// State shared across log receivers.
849    shared_state: &'a mut SharedLoggingState,
850    /// Demux output sessions.
851    output: &'a mut DemuxOutput<'b, 'c>,
852    /// The logging interval specifying the time granularity for the updates.
853    logging_interval_ms: u128,
854    /// The current event time.
855    time: Duration,
856}
857
858impl DemuxHandler<'_, '_, '_> {
859    /// Return the timestamp associated with the current event, based on the event time and the
860    /// logging interval.
861    fn ts(&self) -> Timestamp {
862        let time_ms = self.time.as_millis();
863        let interval = self.logging_interval_ms;
864        let rounded = (time_ms / interval + 1) * interval;
865        rounded.try_into().expect("must fit")
866    }
867
868    /// Handle the given compute event.
869    fn handle(&mut self, event: Ref<'_, ComputeEvent>) {
870        use ComputeEventReference::*;
871        match event {
872            Export(export) => self.handle_export(export),
873            ExportDropped(export_dropped) => self.handle_export_dropped(export_dropped),
874            Peek(peek) if peek.installed => self.handle_peek_install(peek),
875            Peek(peek) => self.handle_peek_retire(peek),
876            Frontier(frontier) => self.handle_frontier(frontier),
877            ImportFrontier(import_frontier) => self.handle_import_frontier(import_frontier),
878            ArrangementHeapSize(inner) => self.handle_arrangement_heap_size(inner),
879            ArrangementHeapCapacity(inner) => self.handle_arrangement_heap_capacity(inner),
880            ArrangementHeapAllocations(inner) => self.handle_arrangement_heap_allocations(inner),
881            ArrangementHeapSizeOperator(inner) => self.handle_arrangement_heap_size_operator(inner),
882            ArrangementHeapSizeOperatorDrop(inner) => {
883                self.handle_arrangement_heap_size_operator_dropped(inner)
884            }
885            DataflowShutdown(shutdown) => self.handle_dataflow_shutdown(shutdown),
886            ErrorCount(error_count) => self.handle_error_count(error_count),
887            HydrationStart(hydration) => self.handle_hydration_start(hydration),
888            Hydration(hydration) => self.handle_hydration(hydration),
889            OperatorHydration(hydration) => self.handle_operator_hydration(hydration),
890            LirMapping(mapping) => self.handle_lir_mapping(mapping),
891            DataflowGlobal(global) => self.handle_dataflow_global(global),
892        }
893    }
894
895    fn handle_export(
896        &mut self,
897        ExportReference {
898            export_id,
899            dataflow_index,
900        }: Ref<'_, Export>,
901    ) {
902        let export_id = Columnar::into_owned(export_id);
903        let ts = self.ts();
904        let datum = self.state.pack_export_update(export_id, dataflow_index);
905        self.output.export.give((datum, ts, Diff::ONE));
906
907        // Stamp the event time, not `ts`, which is rounded up to the logging interval. The rounding
908        // then only delays when an update becomes visible, rather than skewing recorded instants.
909        let installed_at = self.time;
910
911        let existing = self
912            .state
913            .exports
914            .insert(export_id, ExportState::new(dataflow_index, installed_at));
915        if existing.is_some() {
916            error!(%export_id, "export already registered");
917        }
918
919        // Insert hydration time logging for this export.
920        let timestamps = HydrationTimestamps {
921            installed_at,
922            started_at: None,
923            hydrated_at: None,
924        };
925        let datum = self
926            .state
927            .pack_hydration_time_update(export_id, None, &timestamps);
928        self.output.hydration_time.give((datum, ts, Diff::ONE));
929    }
930
931    fn handle_export_dropped(
932        &mut self,
933        ExportDroppedReference { export_id }: Ref<'_, ExportDropped>,
934    ) {
935        let export_id = Columnar::into_owned(export_id);
936        let Some(export) = self.state.exports.remove(&export_id) else {
937            error!(%export_id, "missing exports entry at time of export drop");
938            return;
939        };
940
941        let ts = self.ts();
942        let dataflow_index = export.dataflow_index;
943
944        let datum = self.state.pack_export_update(export_id, dataflow_index);
945        self.output.export.give((datum, ts, Diff::MINUS_ONE));
946
947        // Remove error count logging for this export.
948        if export.error_count != Diff::ZERO {
949            let datum = self
950                .state
951                .pack_error_count_update(export_id, export.error_count);
952            self.output.error_count.give((datum, ts, Diff::MINUS_ONE));
953        }
954
955        // Remove hydration time logging for this export.
956        let datum = self.state.pack_hydration_time_update(
957            export_id,
958            export.hydration_time_ns,
959            &export.hydration_timestamps,
960        );
961        self.output
962            .hydration_time
963            .give((datum, ts, Diff::MINUS_ONE));
964
965        // Remove operator hydration logging for this export.
966        for (lir_id, hydrated) in export.operator_hydration {
967            let datum = self
968                .state
969                .pack_operator_hydration_status_update(export_id, lir_id, hydrated);
970            self.output
971                .operator_hydration_status
972                .give((datum, ts, Diff::MINUS_ONE));
973        }
974    }
975
976    fn handle_dataflow_shutdown(
977        &mut self,
978        DataflowShutdownReference { dataflow_index }: Ref<'_, DataflowShutdown>,
979    ) {
980        let ts = self.ts();
981
982        // We deal with any `GlobalId` based mappings in this event.
983        if let Some(global_ids) = self.state.dataflow_global_ids.remove(&dataflow_index) {
984            for global_id in global_ids {
985                // Remove dataflow/`GlobalID` mapping.
986                let datum = self
987                    .state
988                    .pack_dataflow_global_update(dataflow_index, global_id);
989                self.output
990                    .dataflow_global_ids
991                    .give((datum, ts, Diff::MINUS_ONE));
992
993                // Remove LIR mapping.
994                if let Some(mappings) = self.state.lir_mapping.remove(&global_id) {
995                    for (
996                        lir_id,
997                        LirMetadata {
998                            operator,
999                            parent_lir_id,
1000                            nesting,
1001                            operator_span,
1002                        },
1003                    ) in mappings
1004                    {
1005                        let datum = self.state.pack_lir_mapping_update(
1006                            global_id,
1007                            lir_id,
1008                            operator,
1009                            parent_lir_id,
1010                            nesting,
1011                            operator_span,
1012                        );
1013                        self.output.lir_mapping.give((datum, ts, Diff::MINUS_ONE));
1014                    }
1015                }
1016            }
1017        }
1018    }
1019
1020    fn handle_error_count(&mut self, ErrorCountReference { export_id, diff }: Ref<'_, ErrorCount>) {
1021        let ts = self.ts();
1022        let export_id = Columnar::into_owned(export_id);
1023
1024        let Some(export) = self.state.exports.get_mut(&export_id) else {
1025            // The export might have already been dropped, in which case we are no longer
1026            // interested in its errors.
1027            return;
1028        };
1029
1030        let old_count = export.error_count;
1031        let new_count = old_count + diff;
1032        export.error_count = new_count;
1033
1034        if old_count != Diff::ZERO {
1035            let datum = self.state.pack_error_count_update(export_id, old_count);
1036            self.output.error_count.give((datum, ts, Diff::MINUS_ONE));
1037        }
1038        if new_count != Diff::ZERO {
1039            let datum = self.state.pack_error_count_update(export_id, new_count);
1040            self.output.error_count.give((datum, ts, Diff::ONE));
1041        }
1042    }
1043
1044    fn handle_hydration_start(
1045        &mut self,
1046        HydrationStartReference { export_id }: Ref<'_, HydrationStart>,
1047    ) {
1048        let ts = self.ts();
1049        // Stamp the event time rather than `ts`, as in `handle_export`.
1050        let started_at = self.time;
1051        let export_id = Columnar::into_owned(export_id);
1052
1053        let Some(export) = self.state.exports.get_mut(&export_id) else {
1054            error!(%export_id, "hydration start event for unknown export");
1055            return;
1056        };
1057        if export.hydration_timestamps.started_at.is_some() {
1058            // A `Schedule` command is re-sent for dataflows retained across reconciliation, and is
1059            // sent even for dataflows that were never suspended and have hydrated already. Ignore
1060            // the repeats, mirroring the guard in `handle_hydration`.
1061            return;
1062        }
1063
1064        let old_timestamps = export.hydration_timestamps;
1065        export.hydration_timestamps.started_at = Some(started_at);
1066        let new_timestamps = export.hydration_timestamps;
1067        let time_ns = export.hydration_time_ns;
1068
1069        let retraction = self
1070            .state
1071            .pack_hydration_time_update(export_id, time_ns, &old_timestamps);
1072        self.output
1073            .hydration_time
1074            .give((retraction, ts, Diff::MINUS_ONE));
1075        let insertion = self
1076            .state
1077            .pack_hydration_time_update(export_id, time_ns, &new_timestamps);
1078        self.output.hydration_time.give((insertion, ts, Diff::ONE));
1079    }
1080
1081    fn handle_hydration(&mut self, HydrationReference { export_id }: Ref<'_, Hydration>) {
1082        let ts = self.ts();
1083        // Stamp the event time rather than `ts`, as in `handle_export`.
1084        let hydrated_at = self.time;
1085        let export_id = Columnar::into_owned(export_id);
1086
1087        let Some(export) = self.state.exports.get_mut(&export_id) else {
1088            error!(%export_id, "hydration event for unknown export");
1089            return;
1090        };
1091        if export.hydration_time_ns.is_some() {
1092            // Hydration events for already hydrated dataflows can occur when a dataflow is reused
1093            // after reconciliation. We can simply ignore these.
1094            return;
1095        }
1096
1097        let duration = export.created_at.elapsed();
1098        let nanos = u64::try_from(duration.as_nanos()).expect("must fit");
1099        export.hydration_time_ns = Some(nanos);
1100
1101        let old_timestamps = export.hydration_timestamps;
1102        export.hydration_timestamps.hydrated_at = Some(hydrated_at);
1103        // A dataflow can reach hydration before its `Schedule` arrives, and not only when it has
1104        // no imports to suspend: an index over an already-hydrated arrangement reports hydration
1105        // while still suspended, which happens for a handful of `mz_catalog_server` indexes on
1106        // every bootstrap. So this is a normal path, not a repair for an exotic one.
1107        //
1108        // Stamp `started_at` from `installed_at`, which keeps `installed_at <= started_at <=
1109        // hydrated_at` total and reports the queueing interval as zero. Stamping `hydrated_at`
1110        // instead would invert it, charging the whole life to queueing and reporting zero
1111        // hydration time for a dataflow that only ever hydrated.
1112        if export.hydration_timestamps.started_at.is_none() {
1113            export.hydration_timestamps.started_at = Some(export.hydration_timestamps.installed_at);
1114        }
1115        let new_timestamps = export.hydration_timestamps;
1116
1117        let retraction = self
1118            .state
1119            .pack_hydration_time_update(export_id, None, &old_timestamps);
1120        self.output
1121            .hydration_time
1122            .give((retraction, ts, Diff::MINUS_ONE));
1123        let insertion =
1124            self.state
1125                .pack_hydration_time_update(export_id, Some(nanos), &new_timestamps);
1126        self.output.hydration_time.give((insertion, ts, Diff::ONE));
1127    }
1128
1129    fn handle_operator_hydration(
1130        &mut self,
1131        OperatorHydrationReference {
1132            export_id,
1133            lir_id,
1134            hydrated,
1135        }: Ref<'_, OperatorHydration>,
1136    ) {
1137        let ts = self.ts();
1138        let export_id = Columnar::into_owned(export_id);
1139        let lir_id = Columnar::into_owned(lir_id);
1140        let hydrated = Columnar::into_owned(hydrated);
1141
1142        let Some(export) = self.state.exports.get_mut(&export_id) else {
1143            // The export might have already been dropped, in which case we are no longer
1144            // interested in its operator hydration events.
1145            return;
1146        };
1147
1148        let old_status = export.operator_hydration.get(&lir_id).copied();
1149        export.operator_hydration.insert(lir_id, hydrated);
1150
1151        if let Some(hydrated) = old_status {
1152            let retraction = self
1153                .state
1154                .pack_operator_hydration_status_update(export_id, lir_id, hydrated);
1155            self.output
1156                .operator_hydration_status
1157                .give((retraction, ts, Diff::MINUS_ONE));
1158        }
1159
1160        let insertion = self
1161            .state
1162            .pack_operator_hydration_status_update(export_id, lir_id, hydrated);
1163        self.output
1164            .operator_hydration_status
1165            .give((insertion, ts, Diff::ONE));
1166    }
1167
1168    fn handle_peek_install(
1169        &mut self,
1170        PeekEventReference {
1171            id,
1172            time,
1173            uuid,
1174            peek_type,
1175            installed: _,
1176        }: Ref<'_, PeekEvent>,
1177    ) {
1178        let id = Columnar::into_owned(id);
1179        let uuid = Uuid::from_bytes(uuid::Bytes::into_owned(uuid));
1180        let ts = self.ts();
1181        let datum = self.state.pack_peek_update(id, time, uuid, peek_type);
1182        self.output.peek.give((datum, ts, Diff::ONE));
1183
1184        let existing = self.state.peek_stash.insert(uuid, self.time);
1185        if existing.is_some() {
1186            error!(%uuid, "peek already registered");
1187        }
1188    }
1189
1190    fn handle_peek_retire(
1191        &mut self,
1192        PeekEventReference {
1193            id,
1194            time,
1195            uuid,
1196            peek_type,
1197            installed: _,
1198        }: Ref<'_, PeekEvent>,
1199    ) {
1200        let id = Columnar::into_owned(id);
1201        let uuid = Uuid::from_bytes(uuid::Bytes::into_owned(uuid));
1202        let ts = self.ts();
1203        let datum = self.state.pack_peek_update(id, time, uuid, peek_type);
1204        self.output.peek.give((datum, ts, Diff::MINUS_ONE));
1205
1206        if let Some(start) = self.state.peek_stash.remove(&uuid) {
1207            let elapsed_ns = self.time.saturating_sub(start).as_nanos();
1208            let bucket = elapsed_ns.next_power_of_two();
1209            let datum = self.state.pack_peek_duration_update(peek_type, bucket);
1210            self.output.peek_duration.give((datum, ts, Diff::ONE));
1211        } else {
1212            error!(%uuid, "peek not yet registered");
1213        }
1214    }
1215
1216    fn handle_frontier(
1217        &mut self,
1218        FrontierReference {
1219            export_id,
1220            time,
1221            diff,
1222        }: Ref<'_, Frontier>,
1223    ) {
1224        let export_id = Columnar::into_owned(export_id);
1225        let diff = Diff::from(*diff);
1226        let ts = self.ts();
1227        let time = Columnar::into_owned(time);
1228        let datum = self.state.pack_frontier_update(export_id, time);
1229        self.output.frontier.give((datum, ts, diff));
1230    }
1231
1232    fn handle_import_frontier(
1233        &mut self,
1234        ImportFrontierReference {
1235            import_id,
1236            export_id,
1237            time,
1238            diff,
1239        }: Ref<'_, ImportFrontier>,
1240    ) {
1241        let import_id = Columnar::into_owned(import_id);
1242        let export_id = Columnar::into_owned(export_id);
1243        let diff = Diff::from(*diff);
1244        let ts = self.ts();
1245        let time = Columnar::into_owned(time);
1246        let datum = self
1247            .state
1248            .pack_import_frontier_update(export_id, import_id, time);
1249        self.output.import_frontier.give((datum, ts, diff));
1250    }
1251
1252    /// Update the allocation size for an arrangement.
1253    fn handle_arrangement_heap_size(
1254        &mut self,
1255        ArrangementHeapSizeReference {
1256            operator_id,
1257            delta_size,
1258        }: Ref<'_, ArrangementHeapSize>,
1259    ) {
1260        let ts = self.ts();
1261        let Some(state) = self.state.arrangement_size.get_mut(&operator_id) else {
1262            return;
1263        };
1264
1265        state.size += delta_size;
1266
1267        let datum = self.state.pack_arrangement_heap_size_update(operator_id);
1268        let diff = Diff::cast_from(delta_size);
1269        self.output.arrangement_heap_size.give((datum, ts, diff));
1270    }
1271
1272    /// Update the allocation capacity for an arrangement.
1273    fn handle_arrangement_heap_capacity(
1274        &mut self,
1275        ArrangementHeapCapacityReference {
1276            operator_id,
1277            delta_capacity,
1278        }: Ref<'_, ArrangementHeapCapacity>,
1279    ) {
1280        let ts = self.ts();
1281        let Some(state) = self.state.arrangement_size.get_mut(&operator_id) else {
1282            return;
1283        };
1284
1285        state.capacity += delta_capacity;
1286
1287        let datum = self
1288            .state
1289            .pack_arrangement_heap_capacity_update(operator_id);
1290        let diff = Diff::cast_from(delta_capacity);
1291        self.output
1292            .arrangement_heap_capacity
1293            .give((datum, ts, diff));
1294    }
1295
1296    /// Update the allocation count for an arrangement.
1297    fn handle_arrangement_heap_allocations(
1298        &mut self,
1299        ArrangementHeapAllocationsReference {
1300            operator_id,
1301            delta_allocations,
1302        }: Ref<'_, ArrangementHeapAllocations>,
1303    ) {
1304        let ts = self.ts();
1305        let Some(state) = self.state.arrangement_size.get_mut(&operator_id) else {
1306            return;
1307        };
1308
1309        state.count += delta_allocations;
1310
1311        let datum = self
1312            .state
1313            .pack_arrangement_heap_allocations_update(operator_id);
1314        let diff = Diff::cast_from(delta_allocations);
1315        self.output
1316            .arrangement_heap_allocations
1317            .give((datum, ts, diff));
1318    }
1319
1320    /// Indicate that a new arrangement exists, start maintaining the heap size state.
1321    fn handle_arrangement_heap_size_operator(
1322        &mut self,
1323        ArrangementHeapSizeOperatorReference {
1324            operator_id,
1325            address,
1326        }: Ref<'_, ArrangementHeapSizeOperator>,
1327    ) {
1328        let activator = Activator::new(
1329            address.into_iter().collect(),
1330            Rc::clone(&self.state.activations),
1331        );
1332        let existing = self
1333            .state
1334            .arrangement_size
1335            .insert(operator_id, Default::default());
1336        if existing.is_some() {
1337            error!(%operator_id, "arrangement size operator already registered");
1338        }
1339        let existing = self
1340            .shared_state
1341            .arrangement_size_activators
1342            .insert(operator_id, activator);
1343        if existing.is_some() {
1344            error!(%operator_id, "arrangement size activator already registered");
1345        }
1346    }
1347
1348    /// Indicate that an arrangement has been dropped and we can cleanup the heap size state.
1349    fn handle_arrangement_heap_size_operator_dropped(
1350        &mut self,
1351        event: Ref<'_, ArrangementHeapSizeOperatorDrop>,
1352    ) {
1353        let operator_id = event.operator_id;
1354        if let Some(state) = self.state.arrangement_size.remove(&operator_id) {
1355            let ts = self.ts();
1356            let allocations = self
1357                .state
1358                .pack_arrangement_heap_allocations_update(operator_id);
1359            let diff = -Diff::cast_from(state.count);
1360            self.output
1361                .arrangement_heap_allocations
1362                .give((allocations, ts, diff));
1363
1364            let capacity = self
1365                .state
1366                .pack_arrangement_heap_capacity_update(operator_id);
1367            let diff = -Diff::cast_from(state.capacity);
1368            self.output
1369                .arrangement_heap_capacity
1370                .give((capacity, ts, diff));
1371
1372            let size = self.state.pack_arrangement_heap_size_update(operator_id);
1373            let diff = -Diff::cast_from(state.size);
1374            self.output.arrangement_heap_size.give((size, ts, diff));
1375        }
1376        self.shared_state
1377            .arrangement_size_activators
1378            .remove(&operator_id);
1379    }
1380
1381    /// Indicate that a new LIR operator exists; record the dataflow address it maps to.
1382    fn handle_lir_mapping(
1383        &mut self,
1384        LirMappingReference { global_id, mapping }: Ref<'_, LirMapping>,
1385    ) {
1386        let global_id = Columnar::into_owned(global_id);
1387        // record the state (for the later drop)
1388        let mappings = || mapping.into_iter().map(Columnar::into_owned);
1389        self.state
1390            .lir_mapping
1391            .entry(global_id)
1392            .and_modify(|existing_mapping| existing_mapping.extend(mappings()))
1393            .or_insert_with(|| mappings().collect());
1394
1395        // send the datum out
1396        let ts = self.ts();
1397        for (lir_id, meta) in mapping.into_iter() {
1398            let datum = self.state.pack_lir_mapping_update(
1399                global_id,
1400                Columnar::into_owned(lir_id),
1401                Columnar::into_owned(meta.operator),
1402                Columnar::into_owned(meta.parent_lir_id),
1403                Columnar::into_owned(meta.nesting),
1404                Columnar::into_owned(meta.operator_span),
1405            );
1406            self.output.lir_mapping.give((datum, ts, Diff::ONE));
1407        }
1408    }
1409
1410    fn handle_dataflow_global(
1411        &mut self,
1412        DataflowGlobalReference {
1413            dataflow_index,
1414            global_id,
1415        }: Ref<'_, DataflowGlobal>,
1416    ) {
1417        let global_id = Columnar::into_owned(global_id);
1418        self.state
1419            .dataflow_global_ids
1420            .entry(dataflow_index)
1421            .and_modify(|globals| {
1422                // NB BTreeSet::insert() returns `false` when the element was already in the set
1423                if !globals.insert(global_id) {
1424                    error!(%dataflow_index, %global_id, "dataflow mapping already knew about this GlobalId");
1425                }
1426            })
1427            .or_insert_with(|| BTreeSet::from([global_id]));
1428
1429        let ts = self.ts();
1430        let datum = self
1431            .state
1432            .pack_dataflow_global_update(dataflow_index, global_id);
1433        self.output.dataflow_global_ids.give((datum, ts, Diff::ONE));
1434    }
1435}
1436
1437/// Logging state maintained for a compute collection.
1438///
1439/// This type is used to produce appropriate log events in response to changes of logged collection
1440/// state, e.g. frontiers, and to produce cleanup events when a collection is dropped.
1441pub struct CollectionLogging {
1442    export_id: GlobalId,
1443    logger: Logger,
1444
1445    logged_frontier: Option<Timestamp>,
1446    logged_import_frontiers: BTreeMap<GlobalId, Timestamp>,
1447}
1448
1449impl CollectionLogging {
1450    /// Create new logging state for the identified collection and emit initial logging events.
1451    pub fn new(
1452        export_id: GlobalId,
1453        logger: Logger,
1454        dataflow_index: usize,
1455        import_ids: impl Iterator<Item = GlobalId>,
1456    ) -> Self {
1457        logger.log(&ComputeEvent::Export(Export {
1458            export_id,
1459            dataflow_index,
1460        }));
1461
1462        let mut self_ = Self {
1463            export_id,
1464            logger,
1465            logged_frontier: None,
1466            logged_import_frontiers: Default::default(),
1467        };
1468
1469        // Initialize frontier logging.
1470        let initial_frontier = Some(Timestamp::MIN);
1471        self_.set_frontier(initial_frontier);
1472        import_ids.for_each(|id| self_.set_import_frontier(id, initial_frontier));
1473
1474        self_
1475    }
1476
1477    /// Set the collection frontier to the given new time and emit corresponding logging events.
1478    pub fn set_frontier(&mut self, new_time: Option<Timestamp>) {
1479        let old_time = self.logged_frontier;
1480        self.logged_frontier = new_time;
1481
1482        if old_time != new_time {
1483            let export_id = self.export_id;
1484            let retraction = old_time.map(|time| {
1485                ComputeEvent::Frontier(Frontier {
1486                    export_id,
1487                    time,
1488                    diff: -1,
1489                })
1490            });
1491            let insertion = new_time.map(|time| {
1492                ComputeEvent::Frontier(Frontier {
1493                    export_id,
1494                    time,
1495                    diff: 1,
1496                })
1497            });
1498            let events = retraction.as_ref().into_iter().chain(insertion.as_ref());
1499            self.logger.log_many(events);
1500        }
1501    }
1502
1503    /// Set the frontier of the given import to the given new time and emit corresponding logging
1504    /// events.
1505    pub fn set_import_frontier(&mut self, import_id: GlobalId, new_time: Option<Timestamp>) {
1506        let old_time = self.logged_import_frontiers.remove(&import_id);
1507        if let Some(time) = new_time {
1508            self.logged_import_frontiers.insert(import_id, time);
1509        }
1510
1511        if old_time != new_time {
1512            let export_id = self.export_id;
1513            let retraction = old_time.map(|time| {
1514                ComputeEvent::ImportFrontier(ImportFrontier {
1515                    import_id,
1516                    export_id,
1517                    time,
1518                    diff: -1,
1519                })
1520            });
1521            let insertion = new_time.map(|time| {
1522                ComputeEvent::ImportFrontier(ImportFrontier {
1523                    import_id,
1524                    export_id,
1525                    time,
1526                    diff: 1,
1527                })
1528            });
1529            let events = retraction.as_ref().into_iter().chain(insertion.as_ref());
1530            self.logger.log_many(events);
1531        }
1532    }
1533
1534    /// Record that the collection's dataflow was unsuspended, so hydration work has begun.
1535    ///
1536    /// Repeated calls are ignored by the demux, so callers need not track whether the dataflow was
1537    /// already unsuspended.
1538    pub fn set_hydration_start(&self) {
1539        self.logger
1540            .log(&ComputeEvent::HydrationStart(HydrationStart {
1541                export_id: self.export_id,
1542            }));
1543    }
1544
1545    /// Set the collection as hydrated.
1546    pub fn set_hydrated(&self) {
1547        self.logger.log(&ComputeEvent::Hydration(Hydration {
1548            export_id: self.export_id,
1549        }));
1550    }
1551
1552    /// The export global ID of this collection.
1553    pub fn export_id(&self) -> GlobalId {
1554        self.export_id
1555    }
1556}
1557
1558impl Drop for CollectionLogging {
1559    fn drop(&mut self) {
1560        // Emit retraction events to clean up events previously logged.
1561        self.set_frontier(None);
1562
1563        let import_ids: Vec<_> = self.logged_import_frontiers.keys().copied().collect();
1564        for import_id in import_ids {
1565            self.set_import_frontier(import_id, None);
1566        }
1567
1568        self.logger.log(&ComputeEvent::ExportDropped(ExportDropped {
1569            export_id: self.export_id,
1570        }));
1571    }
1572}
1573
1574/// Extension trait to attach `ComputeEvent::DataflowError` logging operators to collections and
1575/// batch streams.
1576pub(crate) trait LogDataflowErrors {
1577    fn log_dataflow_errors(self, logger: Logger, export_id: GlobalId) -> Self;
1578}
1579
1580impl<'scope, T, D> LogDataflowErrors for VecCollection<'scope, T, D, Diff>
1581where
1582    T: timely::progress::Timestamp,
1583    D: Clone + 'static,
1584{
1585    fn log_dataflow_errors(self, logger: Logger, export_id: GlobalId) -> Self {
1586        self.inner
1587            .unary(Pipeline, "LogDataflowErrorsCollection", |_cap, _info| {
1588                move |input, output| {
1589                    input.for_each(|cap, data| {
1590                        let diff = data.iter().map(|(_d, _t, r)| *r).sum::<Diff>();
1591                        logger.log(&ComputeEvent::ErrorCount(ErrorCount { export_id, diff }));
1592
1593                        output.session(&cap).give_container(data);
1594                    });
1595                }
1596            })
1597            .as_collection()
1598    }
1599}
1600
1601impl<'scope, T, B> LogDataflowErrors for StreamVec<'scope, T, Span<T, B>>
1602where
1603    T: timely::progress::Timestamp,
1604    B: Navigable + Clone + 'static,
1605    for<'a> B::Cursor: Cursor<DiffGat<'a> = &'a Diff>,
1606{
1607    fn log_dataflow_errors(self, logger: Logger, export_id: GlobalId) -> Self {
1608        self.unary(Pipeline, "LogDataflowErrorsStream", |_cap, _info| {
1609            move |input, output| {
1610                input.for_each(|cap, data| {
1611                    let diff = data
1612                        .iter()
1613                        .filter_map(|span| span.inner.as_ref())
1614                        .map(sum_batch_diffs)
1615                        .sum::<Diff>();
1616                    logger.log(&ComputeEvent::ErrorCount(ErrorCount { export_id, diff }));
1617
1618                    output.session(&cap).give_container(data);
1619                });
1620            }
1621        })
1622    }
1623}
1624
1625/// Return the sum of all diffs within the given batch.
1626///
1627/// Note that this operation can be expensive: Its runtime is O(N) with N being the number of
1628/// unique (key, value, time) tuples. We only use it on error streams, which are expected to
1629/// contain only a small number of records, so this doesn't matter much. But avoid using it when
1630/// batches might become large.
1631fn sum_batch_diffs<B>(batch: &B) -> Diff
1632where
1633    B: Navigable,
1634    for<'a> B::Cursor: Cursor<DiffGat<'a> = &'a Diff>,
1635{
1636    let mut sum = Diff::ZERO;
1637    let mut cursor = batch.cursor();
1638
1639    while cursor.key_valid(batch) {
1640        while cursor.val_valid(batch) {
1641            cursor.map_times(batch, |_t, r| sum += r);
1642            cursor.step_val(batch);
1643        }
1644        cursor.step_key(batch);
1645    }
1646
1647    sum
1648}
1649
1650#[cfg(test)]
1651mod tests {
1652    use super::*;
1653
1654    #[mz_ore::test]
1655    fn test_compute_event_size() {
1656        // This could be a static assertion, but we don't use those yet in this crate.
1657        assert_eq!(56, std::mem::size_of::<ComputeEvent>())
1658    }
1659}
mz_compute/logging/compute.rs

mz_compute/logging/
compute.rs
