mz_compute/logging/

compute.rs

// 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.

//! Logging dataflows for events generated by clusterd.

use std::cell::RefCell;
use std::collections::{BTreeMap, BTreeSet};
use std::fmt::{Display, Write};
use std::rc::Rc;
use std::time::{Duration, Instant};

use columnar::{Columnar, Index, Ref};
use differential_dataflow::VecCollection;
use differential_dataflow::collection::AsCollection;
use differential_dataflow::trace::{BatchReader, Cursor};
use mz_compute_types::plan::LirId;
use mz_ore::cast::CastFrom;
use mz_repr::{Datum, Diff, GlobalId, Row, RowRef, Timestamp};
use mz_timely_util::columnar::batcher;
use mz_timely_util::columnar::builder::ColumnBuilder;
use mz_timely_util::columnar::{Col2ValBatcher, Column, columnar_exchange};
use mz_timely_util::replay::MzReplay;
use timely::dataflow::channels::pact::{ExchangeCore, Pipeline};
use timely::dataflow::operators::Operator;
use timely::dataflow::operators::generic::OutputBuilder;
use timely::dataflow::operators::generic::builder_rc::OperatorBuilder;
use timely::dataflow::operators::generic::operator::empty;
use timely::dataflow::{Scope, StreamVec};
use timely::scheduling::activate::{Activations, Activator};
use tracing::error;
use uuid::Uuid;

use crate::extensions::arrange::MzArrangeCore;
use crate::logging::{
    ComputeLog, EventQueue, LogCollection, LogVariant, OutputSessionColumnar, PermutedRowPacker,
    SharedLoggingState, Update,
};
use crate::typedefs::RowRowSpine;
use mz_row_spine::RowRowBuilder;

/// Type alias for a logger of compute events.
pub type Logger = timely::logging_core::Logger<ComputeEventBuilder>;
pub type ComputeEventBuilder = ColumnBuilder<(Duration, ComputeEvent)>;

/// A dataflow exports a global ID.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct Export {
    /// Identifier of the export.
    pub export_id: GlobalId,
    /// Timely worker index of the exporting dataflow.
    pub dataflow_index: usize,
}

/// The export for a global id was dropped.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct ExportDropped {
    /// Identifier of the export.
    pub export_id: GlobalId,
}

/// A peek event with a [`PeekType`], and an installation status.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct PeekEvent {
    /// The identifier of the view the peek targets.
    pub id: GlobalId,
    /// The logical timestamp requested.
    pub time: Timestamp,
    /// The ID of the peek.
    pub uuid: uuid::Bytes,
    /// The relevant _type_ of peek: index or persist.
    // Note that this is not stored on the Peek event for data-packing reasons only.
    pub peek_type: PeekType,
    /// True if the peek is being installed; false if it's being removed.
    pub installed: bool,
}

/// Frontier change event.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct Frontier {
    pub export_id: GlobalId,
    pub time: Timestamp,
    pub diff: i8,
}

/// An import frontier change.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct ImportFrontier {
    pub import_id: GlobalId,
    pub export_id: GlobalId,
    pub time: Timestamp,
    pub diff: i8,
}

/// A change in an arrangement's heap size.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct ArrangementHeapSize {
    /// Operator index
    pub operator_id: usize,
    /// Delta of the heap size in bytes of the arrangement.
    pub delta_size: isize,
}

/// A change in an arrangement's heap capacity.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct ArrangementHeapCapacity {
    /// Operator index
    pub operator_id: usize,
    /// Delta of the heap capacity in bytes of the arrangement.
    pub delta_capacity: isize,
}

/// A change in an arrangement's heap allocation count.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct ArrangementHeapAllocations {
    /// Operator index
    pub operator_id: usize,
    /// Delta of distinct heap allocations backing the arrangement.
    pub delta_allocations: isize,
}

/// Announcing an operator that manages an arrangement.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct ArrangementHeapSizeOperator {
    /// Operator index
    pub operator_id: usize,
    /// The address of the operator.
    pub address: Vec<usize>,
}

/// Drop event for an operator managing an arrangement.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct ArrangementHeapSizeOperatorDrop {
    /// Operator index
    pub operator_id: usize,
}

/// Dataflow shutdown event.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct DataflowShutdown {
    /// Timely worker index of the dataflow.
    pub dataflow_index: usize,
}

/// Error count update event.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct ErrorCount {
    /// Identifier of the export.
    pub export_id: GlobalId,
    /// The change in error count.
    pub diff: Diff,
}

/// An export is hydrated.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct Hydration {
    /// Identifier of the export.
    pub export_id: GlobalId,
}

/// An operator's hydration status changed.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct OperatorHydration {
    /// Identifier of the export.
    pub export_id: GlobalId,
    /// Identifier of the operator's LIR node.
    pub lir_id: LirId,
    /// Whether the operator is hydrated.
    pub hydrated: bool,
}

/// Announce a mapping of an LIR operator to a dataflow operator for a global ID.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct LirMapping {
    /// The `GlobalId` in which the LIR operator is rendered.
    ///
    /// NB a single a dataflow may have many `GlobalId`s inside it.
    /// A separate mapping (using `ComputeEvent::DataflowGlobal`)
    /// tracks the many-to-one relationship between `GlobalId`s and
    /// dataflows.
    pub global_id: GlobalId,
    /// The actual mapping.
    /// Represented this way to reduce the size of `ComputeEvent`.
    pub mapping: Vec<(LirId, LirMetadata)>,
}

/// Announce that a dataflow supports a specific global ID.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct DataflowGlobal {
    /// The identifier of the dataflow.
    pub dataflow_index: usize,
    /// A `GlobalId` that is rendered as part of this dataflow.
    pub global_id: GlobalId,
}

/// A logged compute event.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub enum ComputeEvent {
    /// A dataflow export was created.
    Export(Export),
    /// A dataflow export was dropped.
    ExportDropped(ExportDropped),
    /// Peek command.
    Peek(PeekEvent),
    /// Available frontier information for dataflow exports.
    Frontier(Frontier),
    /// Available frontier information for dataflow imports.
    ImportFrontier(ImportFrontier),
    /// Arrangement heap size update
    ArrangementHeapSize(ArrangementHeapSize),
    /// Arrangement heap size update
    ArrangementHeapCapacity(ArrangementHeapCapacity),
    /// Arrangement heap size update
    ArrangementHeapAllocations(ArrangementHeapAllocations),
    /// Arrangement size operator address
    ArrangementHeapSizeOperator(ArrangementHeapSizeOperator),
    /// Arrangement size operator dropped
    ArrangementHeapSizeOperatorDrop(ArrangementHeapSizeOperatorDrop),
    /// All operators of a dataflow have shut down.
    DataflowShutdown(DataflowShutdown),
    /// The number of errors in a dataflow export has changed.
    ErrorCount(ErrorCount),
    /// A dataflow export was hydrated.
    Hydration(Hydration),
    /// A dataflow operator's hydration status changed.
    OperatorHydration(OperatorHydration),
    /// An LIR operator was mapped to some particular dataflow operator.
    ///
    /// Cf. `ComputeLog::LirMaping`
    LirMapping(LirMapping),
    DataflowGlobal(DataflowGlobal),
}

/// A peek type distinguishing between index and persist peeks.
#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Debug, Columnar)]
pub enum PeekType {
    /// A peek against an index.
    Index,
    /// A peek against persist.
    Persist,
}

impl PeekType {
    /// A human-readable name for a peek type.
    fn name(self) -> &'static str {
        match self {
            PeekType::Index => "index",
            PeekType::Persist => "persist",
        }
    }
}

/// Metadata for LIR operators.
#[derive(Clone, Debug, PartialEq, PartialOrd, Columnar)]
pub struct LirMetadata {
    /// The LIR operator, as a string (see `FlatPlanNode::humanize`).
    operator: String,
    /// The LIR identifier of the parent (if any).
    parent_lir_id: Option<LirId>,
    /// How nested the operator is (for nice indentation).
    nesting: u8,
    /// The dataflow operator ids, given as start (inclusive) and end (exclusive).
    /// If `start == end`, then no operators were used.
    operator_span: (usize, usize),
}

impl LirMetadata {
    /// Construct a new LIR metadata object.
    pub fn new(
        operator: String,
        parent_lir_id: Option<LirId>,
        nesting: u8,
        operator_span: (usize, usize),
    ) -> Self {
        Self {
            operator,
            parent_lir_id,
            nesting,
            operator_span,
        }
    }
}

/// The return type of the [`construct`] function.
pub(super) struct Return {
    /// Collections returned by [`construct`].
    pub collections: BTreeMap<LogVariant, LogCollection>,
}

/// Constructs the logging dataflow fragment for compute logs.
///
/// Params
/// * `scope`: The Timely scope hosting the log analysis dataflow.
/// * `scheduler`: The timely scheduler to obtainer activators.
/// * `config`: Logging configuration.
/// * `event_queue`: The source to read compute log events from.
/// * `compute_event_streams`: Additional compute event streams to absorb.
/// * `shared_state`: Shared state between logging dataflow fragments.
pub(super) fn construct<'scope>(
    scope: Scope<'scope, Timestamp>,
    activations: Rc<RefCell<Activations>>,
    config: &mz_compute_client::logging::LoggingConfig,
    event_queue: EventQueue<Column<(Duration, ComputeEvent)>>,
    shared_state: Rc<RefCell<SharedLoggingState>>,
) -> Return {
    let logging_interval_ms = std::cmp::max(1, config.interval.as_millis());

    scope.scoped("compute logging", move |scope| {
        let enable_logging = config.enable_logging;
        let (logs, token) = if enable_logging {
            event_queue.links.mz_replay(
                scope,
                "compute logs",
                config.interval,
                event_queue.activator,
            )
        } else {
            let token: Rc<dyn std::any::Any> = Rc::new(Box::new(()));
            (empty(scope), token)
        };

        // Build a demux operator that splits the replayed event stream up into the separate
        // logging streams.
        let mut demux = OperatorBuilder::new("Compute Logging Demux".to_string(), scope.clone());
        let mut input = demux.new_input(logs, Pipeline);
        let (export_out, export) = demux.new_output();
        let mut export_out = OutputBuilder::from(export_out);
        let (frontier_out, frontier) = demux.new_output();
        let mut frontier_out = OutputBuilder::from(frontier_out);
        let (import_frontier_out, import_frontier) = demux.new_output();
        let mut import_frontier_out = OutputBuilder::from(import_frontier_out);
        let (peek_out, peek) = demux.new_output();
        let mut peek_out = OutputBuilder::from(peek_out);
        let (peek_duration_out, peek_duration) = demux.new_output();
        let mut peek_duration_out = OutputBuilder::from(peek_duration_out);
        let (arrangement_heap_size_out, arrangement_heap_size) = demux.new_output();
        let mut arrangement_heap_size_out = OutputBuilder::from(arrangement_heap_size_out);
        let (arrangement_heap_capacity_out, arrangement_heap_capacity) = demux.new_output();
        let mut arrangement_heap_capacity_out = OutputBuilder::from(arrangement_heap_capacity_out);
        let (arrangement_heap_allocations_out, arrangement_heap_allocations) = demux.new_output();
        let mut arrangement_heap_allocations_out =
            OutputBuilder::from(arrangement_heap_allocations_out);
        let (error_count_out, error_count) = demux.new_output();
        let mut error_count_out = OutputBuilder::from(error_count_out);
        let (hydration_time_out, hydration_time) = demux.new_output();
        let mut hydration_time_out = OutputBuilder::from(hydration_time_out);
        let (operator_hydration_status_out, operator_hydration_status) = demux.new_output();
        let mut operator_hydration_status_out = OutputBuilder::from(operator_hydration_status_out);
        let (lir_mapping_out, lir_mapping) = demux.new_output();
        let mut lir_mapping_out = OutputBuilder::from(lir_mapping_out);
        let (dataflow_global_ids_out, dataflow_global_ids) = demux.new_output();
        let mut dataflow_global_ids_out = OutputBuilder::from(dataflow_global_ids_out);

        let mut demux_state = DemuxState::new(activations, scope.index());
        demux.build(move |_capability| {
            move |_frontiers| {
                let mut export = export_out.activate();
                let mut frontier = frontier_out.activate();
                let mut import_frontier = import_frontier_out.activate();
                let mut peek = peek_out.activate();
                let mut peek_duration = peek_duration_out.activate();
                let mut arrangement_heap_size = arrangement_heap_size_out.activate();
                let mut arrangement_heap_capacity = arrangement_heap_capacity_out.activate();
                let mut arrangement_heap_allocations = arrangement_heap_allocations_out.activate();
                let mut error_count = error_count_out.activate();
                let mut hydration_time = hydration_time_out.activate();
                let mut operator_hydration_status = operator_hydration_status_out.activate();
                let mut lir_mapping = lir_mapping_out.activate();
                let mut dataflow_global_ids = dataflow_global_ids_out.activate();

                input.for_each(|cap, data| {
                    let mut output_sessions = DemuxOutput {
                        export: export.session_with_builder(&cap),
                        frontier: frontier.session_with_builder(&cap),
                        import_frontier: import_frontier.session_with_builder(&cap),
                        peek: peek.session_with_builder(&cap),
                        peek_duration: peek_duration.session_with_builder(&cap),
                        arrangement_heap_allocations: arrangement_heap_allocations
                            .session_with_builder(&cap),
                        arrangement_heap_capacity: arrangement_heap_capacity
                            .session_with_builder(&cap),
                        arrangement_heap_size: arrangement_heap_size.session_with_builder(&cap),
                        error_count: error_count.session_with_builder(&cap),
                        hydration_time: hydration_time.session_with_builder(&cap),
                        operator_hydration_status: operator_hydration_status
                            .session_with_builder(&cap),
                        lir_mapping: lir_mapping.session_with_builder(&cap),
                        dataflow_global_ids: dataflow_global_ids.session_with_builder(&cap),
                    };

                    let shared_state = &mut shared_state.borrow_mut();
                    for (time, event) in data.borrow().into_index_iter() {
                        DemuxHandler {
                            state: &mut demux_state,
                            shared_state,
                            output: &mut output_sessions,
                            logging_interval_ms,
                            time,
                        }
                        .handle(event);
                    }
                });
            }
        });

        use ComputeLog::*;
        let logs = [
            (ArrangementHeapAllocations, arrangement_heap_allocations),
            (ArrangementHeapCapacity, arrangement_heap_capacity),
            (ArrangementHeapSize, arrangement_heap_size),
            (DataflowCurrent, export),
            (DataflowGlobal, dataflow_global_ids),
            (ErrorCount, error_count),
            (FrontierCurrent, frontier),
            (HydrationTime, hydration_time),
            (ImportFrontierCurrent, import_frontier),
            (LirMapping, lir_mapping),
            (OperatorHydrationStatus, operator_hydration_status),
            (PeekCurrent, peek),
            (PeekDuration, peek_duration),
        ];

        // Build the output arrangements.
        let mut collections = BTreeMap::new();
        for (variant, stream) in logs {
            let variant = LogVariant::Compute(variant);
            if config.index_logs.contains_key(&variant) {
                let exchange = ExchangeCore::<ColumnBuilder<_>, _>::new_core(
                    columnar_exchange::<Row, Row, Timestamp, Diff>,
                );
                let trace = stream
                    .mz_arrange_core::<
                        _,
                        batcher::Chunker<_>,
                        Col2ValBatcher<_, _, _, _>,
                        RowRowBuilder<_, _>,
                        RowRowSpine<_, _>,
                    >(exchange, &format!("Arrange {variant:?}"))
                    .trace;
                let collection = LogCollection {
                    trace,
                    token: Rc::clone(&token),
                };
                collections.insert(variant, collection);
            }
        }

        Return { collections }
    })
}

/// Format the given value and pack it into a `Datum::String`.
///
/// The `scratch` buffer is used to perform the string conversion without an allocation.
/// Callers should not assume anything about the contents of this buffer after this function
/// returns.
fn make_string_datum<V>(value: V, scratch: &mut String) -> Datum<'_>
where
    V: Display,
{
    scratch.clear();
    write!(scratch, "{}", value).expect("writing to a `String` can't fail");
    Datum::String(scratch)
}

/// State maintained by the demux operator.
struct DemuxState {
    /// The timely activations handle.
    activations: Rc<RefCell<Activations>>,
    /// The index of this worker.
    worker_id: usize,
    /// A reusable scratch string for formatting IDs.
    scratch_string_a: String,
    /// A reusable scratch string for formatting IDs.
    scratch_string_b: String,
    /// State tracked per dataflow export.
    exports: BTreeMap<GlobalId, ExportState>,
    /// Maps pending peeks to their installation time.
    peek_stash: BTreeMap<Uuid, Duration>,
    /// Arrangement size stash.
    arrangement_size: BTreeMap<usize, ArrangementSizeState>,
    /// LIR -> operator span mapping.
    lir_mapping: BTreeMap<GlobalId, BTreeMap<LirId, LirMetadata>>,
    /// Dataflow -> `GlobalId` mapping (many-to-one).
    dataflow_global_ids: BTreeMap<usize, BTreeSet<GlobalId>>,
    /// A row packer for the arrangement heap allocations output.
    arrangement_heap_allocations_packer: PermutedRowPacker,
    /// A row packer for the arrangement heap capacity output.
    arrangement_heap_capacity_packer: PermutedRowPacker,
    /// A row packer for the arrangement heap size output.
    arrangement_heap_size_packer: PermutedRowPacker,
    /// A row packer for the dataflow global output.
    dataflow_global_packer: PermutedRowPacker,
    /// A row packer for the error count output.
    error_count_packer: PermutedRowPacker,
    /// A row packer for the exports output.
    export_packer: PermutedRowPacker,
    /// A row packer for the frontier output.
    frontier_packer: PermutedRowPacker,
    /// A row packer for the exports output.
    import_frontier_packer: PermutedRowPacker,
    /// A row packer for the LIR mapping output.
    lir_mapping_packer: PermutedRowPacker,
    /// A row packer for the operator hydration status output.
    operator_hydration_status_packer: PermutedRowPacker,
    /// A row packer for the peek durations output.
    peek_duration_packer: PermutedRowPacker,
    /// A row packer for the peek output.
    peek_packer: PermutedRowPacker,
    /// A row packer for the hydration time output.
    hydration_time_packer: PermutedRowPacker,
}

impl DemuxState {
    fn new(activations: Rc<RefCell<Activations>>, worker_id: usize) -> Self {
        Self {
            activations,
            worker_id,
            scratch_string_a: String::new(),
            scratch_string_b: String::new(),
            exports: Default::default(),
            peek_stash: Default::default(),
            arrangement_size: Default::default(),
            lir_mapping: Default::default(),
            dataflow_global_ids: Default::default(),
            arrangement_heap_allocations_packer: PermutedRowPacker::new(
                ComputeLog::ArrangementHeapAllocations,
            ),
            arrangement_heap_capacity_packer: PermutedRowPacker::new(
                ComputeLog::ArrangementHeapCapacity,
            ),
            arrangement_heap_size_packer: PermutedRowPacker::new(ComputeLog::ArrangementHeapSize),
            dataflow_global_packer: PermutedRowPacker::new(ComputeLog::DataflowGlobal),
            error_count_packer: PermutedRowPacker::new(ComputeLog::ErrorCount),
            export_packer: PermutedRowPacker::new(ComputeLog::DataflowCurrent),
            frontier_packer: PermutedRowPacker::new(ComputeLog::FrontierCurrent),
            hydration_time_packer: PermutedRowPacker::new(ComputeLog::HydrationTime),
            import_frontier_packer: PermutedRowPacker::new(ComputeLog::ImportFrontierCurrent),
            lir_mapping_packer: PermutedRowPacker::new(ComputeLog::LirMapping),
            operator_hydration_status_packer: PermutedRowPacker::new(
                ComputeLog::OperatorHydrationStatus,
            ),
            peek_duration_packer: PermutedRowPacker::new(ComputeLog::PeekDuration),
            peek_packer: PermutedRowPacker::new(ComputeLog::PeekCurrent),
        }
    }

    /// Pack an arrangement heap allocations update key-value for the given operator.
    fn pack_arrangement_heap_allocations_update(
        &mut self,
        operator_id: usize,
    ) -> (&RowRef, &RowRef) {
        self.arrangement_heap_allocations_packer.pack_slice(&[
            Datum::UInt64(operator_id.try_into().expect("operator_id too big")),
            Datum::UInt64(u64::cast_from(self.worker_id)),
        ])
    }

    /// Pack an arrangement heap capacity update key-value for the given operator.
    fn pack_arrangement_heap_capacity_update(&mut self, operator_id: usize) -> (&RowRef, &RowRef) {
        self.arrangement_heap_capacity_packer.pack_slice(&[
            Datum::UInt64(operator_id.try_into().expect("operator_id too big")),
            Datum::UInt64(u64::cast_from(self.worker_id)),
        ])
    }

    /// Pack an arrangement heap size update key-value for the given operator.
    fn pack_arrangement_heap_size_update(&mut self, operator_id: usize) -> (&RowRef, &RowRef) {
        self.arrangement_heap_size_packer.pack_slice(&[
            Datum::UInt64(operator_id.try_into().expect("operator_id too big")),
            Datum::UInt64(u64::cast_from(self.worker_id)),
        ])
    }

    /// Pack a dataflow global update key-value for the given dataflow index and global ID.
    fn pack_dataflow_global_update(
        &mut self,
        dataflow_index: usize,
        global_id: GlobalId,
    ) -> (&RowRef, &RowRef) {
        self.dataflow_global_packer.pack_slice(&[
            Datum::UInt64(u64::cast_from(dataflow_index)),
            Datum::UInt64(u64::cast_from(self.worker_id)),
            make_string_datum(global_id, &mut self.scratch_string_a),
        ])
    }

    /// Pack an error count update key-value for the given export ID and count.
    fn pack_error_count_update(&mut self, export_id: GlobalId, count: Diff) -> (&RowRef, &RowRef) {
        // Normally we would use DD's diff field to encode counts, but in this case we can't: The total
        // per-worker error count might be negative and at the SQL level having negative multiplicities
        // is treated as an error.
        self.error_count_packer.pack_slice(&[
            make_string_datum(export_id, &mut self.scratch_string_a),
            Datum::UInt64(u64::cast_from(self.worker_id)),
            Datum::Int64(count.into_inner()),
        ])
    }

    /// Pack an export update key-value for the given export ID and dataflow index.
    fn pack_export_update(
        &mut self,
        export_id: GlobalId,
        dataflow_index: usize,
    ) -> (&RowRef, &RowRef) {
        self.export_packer.pack_slice(&[
            make_string_datum(export_id, &mut self.scratch_string_a),
            Datum::UInt64(u64::cast_from(self.worker_id)),
            Datum::UInt64(u64::cast_from(dataflow_index)),
        ])
    }

    /// Pack a hydration time update key-value for the given export ID and hydration time.
    fn pack_hydration_time_update(
        &mut self,
        export_id: GlobalId,
        time_ns: Option<u64>,
    ) -> (&RowRef, &RowRef) {
        self.hydration_time_packer.pack_slice(&[
            make_string_datum(export_id, &mut self.scratch_string_a),
            Datum::UInt64(u64::cast_from(self.worker_id)),
            Datum::from(time_ns),
        ])
    }

    /// Pack an import frontier update key-value for the given export ID and dataflow index.
    fn pack_import_frontier_update(
        &mut self,
        export_id: GlobalId,
        import_id: GlobalId,
        time: Timestamp,
    ) -> (&RowRef, &RowRef) {
        self.import_frontier_packer.pack_slice(&[
            make_string_datum(export_id, &mut self.scratch_string_a),
            make_string_datum(import_id, &mut self.scratch_string_b),
            Datum::UInt64(u64::cast_from(self.worker_id)),
            Datum::MzTimestamp(time),
        ])
    }

    /// Pack an LIR mapping update key-value for the given LIR operator metadata.
    fn pack_lir_mapping_update(
        &mut self,
        global_id: GlobalId,
        lir_id: LirId,
        operator: String,
        parent_lir_id: Option<LirId>,
        nesting: u8,
        operator_span: (usize, usize),
    ) -> (&RowRef, &RowRef) {
        self.lir_mapping_packer.pack_slice(&[
            make_string_datum(global_id, &mut self.scratch_string_a),
            Datum::UInt64(lir_id.into()),
            Datum::UInt64(u64::cast_from(self.worker_id)),
            make_string_datum(operator, &mut self.scratch_string_b),
            parent_lir_id.map_or(Datum::Null, |lir_id| Datum::UInt64(lir_id.into())),
            Datum::UInt16(u16::cast_from(nesting)),
            Datum::UInt64(u64::cast_from(operator_span.0)),
            Datum::UInt64(u64::cast_from(operator_span.1)),
        ])
    }

    /// Pack an operator hydration status update key-value for the given export ID, LIR ID, and
    /// hydration status.
    fn pack_operator_hydration_status_update(
        &mut self,
        export_id: GlobalId,
        lir_id: LirId,
        hydrated: bool,
    ) -> (&RowRef, &RowRef) {
        self.operator_hydration_status_packer.pack_slice(&[
            make_string_datum(export_id, &mut self.scratch_string_a),
            Datum::UInt64(lir_id.into()),
            Datum::UInt64(u64::cast_from(self.worker_id)),
            Datum::from(hydrated),
        ])
    }

    /// Pack a peek duration update key-value for the given peek and peek type.
    fn pack_peek_duration_update(
        &mut self,
        peek_type: PeekType,
        bucket: u128,
    ) -> (&RowRef, &RowRef) {
        self.peek_duration_packer.pack_slice(&[
            Datum::UInt64(u64::cast_from(self.worker_id)),
            Datum::String(peek_type.name()),
            Datum::UInt64(bucket.try_into().expect("bucket too big")),
        ])
    }

    /// Pack a peek update key-value for the given peek and peek type.
    fn pack_peek_update(
        &mut self,
        id: GlobalId,
        time: Timestamp,
        uuid: Uuid,
        peek_type: PeekType,
    ) -> (&RowRef, &RowRef) {
        self.peek_packer.pack_slice(&[
            Datum::Uuid(uuid),
            Datum::UInt64(u64::cast_from(self.worker_id)),
            make_string_datum(id, &mut self.scratch_string_a),
            Datum::String(peek_type.name()),
            Datum::MzTimestamp(time),
        ])
    }

    /// Pack a frontier update key-value for the given export ID and time.
    fn pack_frontier_update(&mut self, export_id: GlobalId, time: Timestamp) -> (&RowRef, &RowRef) {
        self.frontier_packer.pack_slice(&[
            make_string_datum(export_id, &mut self.scratch_string_a),
            Datum::UInt64(u64::cast_from(self.worker_id)),
            Datum::MzTimestamp(time),
        ])
    }
}

/// State tracked for each dataflow export.
struct ExportState {
    /// The ID of the dataflow maintaining this export.
    dataflow_index: usize,
    /// Number of errors in this export.
    ///
    /// This must be a signed integer, since per-worker error counts can be negative, only the
    /// cross-worker total has to sum up to a non-negative value.
    error_count: Diff,
    /// When this export was created.
    created_at: Instant,
    /// Whether the exported collection is hydrated.
    hydration_time_ns: Option<u64>,
    /// Hydration status of operators feeding this export.
    operator_hydration: BTreeMap<LirId, bool>,
}

impl ExportState {
    fn new(dataflow_index: usize) -> Self {
        Self {
            dataflow_index,
            error_count: Diff::ZERO,
            created_at: Instant::now(),
            hydration_time_ns: None,
            operator_hydration: BTreeMap::new(),
        }
    }
}

/// State for tracking arrangement sizes.
#[derive(Default, Debug)]
struct ArrangementSizeState {
    size: isize,
    capacity: isize,
    count: isize,
}

/// Bundled output sessions used by the demux operator.
struct DemuxOutput<'a, 'b> {
    export: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    frontier: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    import_frontier: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    peek: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    peek_duration: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    arrangement_heap_allocations: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    arrangement_heap_capacity: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    arrangement_heap_size: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    hydration_time: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    operator_hydration_status: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    error_count: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    lir_mapping: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
    dataflow_global_ids: OutputSessionColumnar<'a, 'b, Update<(Row, Row)>>,
}

/// Event handler of the demux operator.
struct DemuxHandler<'a, 'b, 'c> {
    /// State kept by the demux operator.
    state: &'a mut DemuxState,
    /// State shared across log receivers.
    shared_state: &'a mut SharedLoggingState,
    /// Demux output sessions.
    output: &'a mut DemuxOutput<'b, 'c>,
    /// The logging interval specifying the time granularity for the updates.
    logging_interval_ms: u128,
    /// The current event time.
    time: Duration,
}

impl DemuxHandler<'_, '_, '_> {
    /// Return the timestamp associated with the current event, based on the event time and the
    /// logging interval.
    fn ts(&self) -> Timestamp {
        let time_ms = self.time.as_millis();
        let interval = self.logging_interval_ms;
        let rounded = (time_ms / interval + 1) * interval;
        rounded.try_into().expect("must fit")
    }

    /// Handle the given compute event.
    fn handle(&mut self, event: Ref<'_, ComputeEvent>) {
        use ComputeEventReference::*;
        match event {
            Export(export) => self.handle_export(export),
            ExportDropped(export_dropped) => self.handle_export_dropped(export_dropped),
            Peek(peek) if peek.installed => self.handle_peek_install(peek),
            Peek(peek) => self.handle_peek_retire(peek),
            Frontier(frontier) => self.handle_frontier(frontier),
            ImportFrontier(import_frontier) => self.handle_import_frontier(import_frontier),
            ArrangementHeapSize(inner) => self.handle_arrangement_heap_size(inner),
            ArrangementHeapCapacity(inner) => self.handle_arrangement_heap_capacity(inner),
            ArrangementHeapAllocations(inner) => self.handle_arrangement_heap_allocations(inner),
            ArrangementHeapSizeOperator(inner) => self.handle_arrangement_heap_size_operator(inner),
            ArrangementHeapSizeOperatorDrop(inner) => {
                self.handle_arrangement_heap_size_operator_dropped(inner)
            }
            DataflowShutdown(shutdown) => self.handle_dataflow_shutdown(shutdown),
            ErrorCount(error_count) => self.handle_error_count(error_count),
            Hydration(hydration) => self.handle_hydration(hydration),
            OperatorHydration(hydration) => self.handle_operator_hydration(hydration),
            LirMapping(mapping) => self.handle_lir_mapping(mapping),
            DataflowGlobal(global) => self.handle_dataflow_global(global),
        }
    }

    fn handle_export(
        &mut self,
        ExportReference {
            export_id,
            dataflow_index,
        }: Ref<'_, Export>,
    ) {
        let export_id = Columnar::into_owned(export_id);
        let ts = self.ts();
        let datum = self.state.pack_export_update(export_id, dataflow_index);
        self.output.export.give((datum, ts, Diff::ONE));

        let existing = self
            .state
            .exports
            .insert(export_id, ExportState::new(dataflow_index));
        if existing.is_some() {
            error!(%export_id, "export already registered");
        }

        // Insert hydration time logging for this export.
        let datum = self.state.pack_hydration_time_update(export_id, None);
        self.output.hydration_time.give((datum, ts, Diff::ONE));
    }

    fn handle_export_dropped(
        &mut self,
        ExportDroppedReference { export_id }: Ref<'_, ExportDropped>,
    ) {
        let export_id = Columnar::into_owned(export_id);
        let Some(export) = self.state.exports.remove(&export_id) else {
            error!(%export_id, "missing exports entry at time of export drop");
            return;
        };

        let ts = self.ts();
        let dataflow_index = export.dataflow_index;

        let datum = self.state.pack_export_update(export_id, dataflow_index);
        self.output.export.give((datum, ts, Diff::MINUS_ONE));

        // Remove error count logging for this export.
        if export.error_count != Diff::ZERO {
            let datum = self
                .state
                .pack_error_count_update(export_id, export.error_count);
            self.output.error_count.give((datum, ts, Diff::MINUS_ONE));
        }

        // Remove hydration time logging for this export.
        let datum = self
            .state
            .pack_hydration_time_update(export_id, export.hydration_time_ns);
        self.output
            .hydration_time
            .give((datum, ts, Diff::MINUS_ONE));

        // Remove operator hydration logging for this export.
        for (lir_id, hydrated) in export.operator_hydration {
            let datum = self
                .state
                .pack_operator_hydration_status_update(export_id, lir_id, hydrated);
            self.output
                .operator_hydration_status
                .give((datum, ts, Diff::MINUS_ONE));
        }
    }

    fn handle_dataflow_shutdown(
        &mut self,
        DataflowShutdownReference { dataflow_index }: Ref<'_, DataflowShutdown>,
    ) {
        let ts = self.ts();

        // We deal with any `GlobalId` based mappings in this event.
        if let Some(global_ids) = self.state.dataflow_global_ids.remove(&dataflow_index) {
            for global_id in global_ids {
                // Remove dataflow/`GlobalID` mapping.
                let datum = self
                    .state
                    .pack_dataflow_global_update(dataflow_index, global_id);
                self.output
                    .dataflow_global_ids
                    .give((datum, ts, Diff::MINUS_ONE));

                // Remove LIR mapping.
                if let Some(mappings) = self.state.lir_mapping.remove(&global_id) {
                    for (
                        lir_id,
                        LirMetadata {
                            operator,
                            parent_lir_id,
                            nesting,
                            operator_span,
                        },
                    ) in mappings
                    {
                        let datum = self.state.pack_lir_mapping_update(
                            global_id,
                            lir_id,
                            operator,
                            parent_lir_id,
                            nesting,
                            operator_span,
                        );
                        self.output.lir_mapping.give((datum, ts, Diff::MINUS_ONE));
                    }
                }
            }
        }
    }

    fn handle_error_count(&mut self, ErrorCountReference { export_id, diff }: Ref<'_, ErrorCount>) {
        let ts = self.ts();
        let export_id = Columnar::into_owned(export_id);

        let Some(export) = self.state.exports.get_mut(&export_id) else {
            // The export might have already been dropped, in which case we are no longer
            // interested in its errors.
            return;
        };

        let old_count = export.error_count;
        let new_count = old_count + diff;
        export.error_count = new_count;

        if old_count != Diff::ZERO {
            let datum = self.state.pack_error_count_update(export_id, old_count);
            self.output.error_count.give((datum, ts, Diff::MINUS_ONE));
        }
        if new_count != Diff::ZERO {
            let datum = self.state.pack_error_count_update(export_id, new_count);
            self.output.error_count.give((datum, ts, Diff::ONE));
        }
    }

    fn handle_hydration(&mut self, HydrationReference { export_id }: Ref<'_, Hydration>) {
        let ts = self.ts();
        let export_id = Columnar::into_owned(export_id);

        let Some(export) = self.state.exports.get_mut(&export_id) else {
            error!(%export_id, "hydration event for unknown export");
            return;
        };
        if export.hydration_time_ns.is_some() {
            // Hydration events for already hydrated dataflows can occur when a dataflow is reused
            // after reconciliation. We can simply ignore these.
            return;
        }

        let duration = export.created_at.elapsed();
        let nanos = u64::try_from(duration.as_nanos()).expect("must fit");
        export.hydration_time_ns = Some(nanos);

        let retraction = self.state.pack_hydration_time_update(export_id, None);
        self.output
            .hydration_time
            .give((retraction, ts, Diff::MINUS_ONE));
        let insertion = self
            .state
            .pack_hydration_time_update(export_id, Some(nanos));
        self.output.hydration_time.give((insertion, ts, Diff::ONE));
    }

    fn handle_operator_hydration(
        &mut self,
        OperatorHydrationReference {
            export_id,
            lir_id,
            hydrated,
        }: Ref<'_, OperatorHydration>,
    ) {
        let ts = self.ts();
        let export_id = Columnar::into_owned(export_id);
        let lir_id = Columnar::into_owned(lir_id);
        let hydrated = Columnar::into_owned(hydrated);

        let Some(export) = self.state.exports.get_mut(&export_id) else {
            // The export might have already been dropped, in which case we are no longer
            // interested in its operator hydration events.
            return;
        };

        let old_status = export.operator_hydration.get(&lir_id).copied();
        export.operator_hydration.insert(lir_id, hydrated);

        if let Some(hydrated) = old_status {
            let retraction = self
                .state
                .pack_operator_hydration_status_update(export_id, lir_id, hydrated);
            self.output
                .operator_hydration_status
                .give((retraction, ts, Diff::MINUS_ONE));
        }

        let insertion = self
            .state
            .pack_operator_hydration_status_update(export_id, lir_id, hydrated);
        self.output
            .operator_hydration_status
            .give((insertion, ts, Diff::ONE));
    }

    fn handle_peek_install(
        &mut self,
        PeekEventReference {
            id,
            time,
            uuid,
            peek_type,
            installed: _,
        }: Ref<'_, PeekEvent>,
    ) {
        let id = Columnar::into_owned(id);
        let uuid = Uuid::from_bytes(uuid::Bytes::into_owned(uuid));
        let ts = self.ts();
        let datum = self.state.pack_peek_update(id, time, uuid, peek_type);
        self.output.peek.give((datum, ts, Diff::ONE));

        let existing = self.state.peek_stash.insert(uuid, self.time);
        if existing.is_some() {
            error!(%uuid, "peek already registered");
        }
    }

    fn handle_peek_retire(
        &mut self,
        PeekEventReference {
            id,
            time,
            uuid,
            peek_type,
            installed: _,
        }: Ref<'_, PeekEvent>,
    ) {
        let id = Columnar::into_owned(id);
        let uuid = Uuid::from_bytes(uuid::Bytes::into_owned(uuid));
        let ts = self.ts();
        let datum = self.state.pack_peek_update(id, time, uuid, peek_type);
        self.output.peek.give((datum, ts, Diff::MINUS_ONE));

        if let Some(start) = self.state.peek_stash.remove(&uuid) {
            let elapsed_ns = self.time.saturating_sub(start).as_nanos();
            let bucket = elapsed_ns.next_power_of_two();
            let datum = self.state.pack_peek_duration_update(peek_type, bucket);
            self.output.peek_duration.give((datum, ts, Diff::ONE));
        } else {
            error!(%uuid, "peek not yet registered");
        }
    }

    fn handle_frontier(
        &mut self,
        FrontierReference {
            export_id,
            time,
            diff,
        }: Ref<'_, Frontier>,
    ) {
        let export_id = Columnar::into_owned(export_id);
        let diff = Diff::from(*diff);
        let ts = self.ts();
        let time = Columnar::into_owned(time);
        let datum = self.state.pack_frontier_update(export_id, time);
        self.output.frontier.give((datum, ts, diff));
    }

    fn handle_import_frontier(
        &mut self,
        ImportFrontierReference {
            import_id,
            export_id,
            time,
            diff,
        }: Ref<'_, ImportFrontier>,
    ) {
        let import_id = Columnar::into_owned(import_id);
        let export_id = Columnar::into_owned(export_id);
        let diff = Diff::from(*diff);
        let ts = self.ts();
        let time = Columnar::into_owned(time);
        let datum = self
            .state
            .pack_import_frontier_update(export_id, import_id, time);
        self.output.import_frontier.give((datum, ts, diff));
    }

    /// Update the allocation size for an arrangement.
    fn handle_arrangement_heap_size(
        &mut self,
        ArrangementHeapSizeReference {
            operator_id,
            delta_size,
        }: Ref<'_, ArrangementHeapSize>,
    ) {
        let ts = self.ts();
        let Some(state) = self.state.arrangement_size.get_mut(&operator_id) else {
            return;
        };

        state.size += delta_size;

        let datum = self.state.pack_arrangement_heap_size_update(operator_id);
        let diff = Diff::cast_from(delta_size);
        self.output.arrangement_heap_size.give((datum, ts, diff));
    }

    /// Update the allocation capacity for an arrangement.
    fn handle_arrangement_heap_capacity(
        &mut self,
        ArrangementHeapCapacityReference {
            operator_id,
            delta_capacity,
        }: Ref<'_, ArrangementHeapCapacity>,
    ) {
        let ts = self.ts();
        let Some(state) = self.state.arrangement_size.get_mut(&operator_id) else {
            return;
        };

        state.capacity += delta_capacity;

        let datum = self
            .state
            .pack_arrangement_heap_capacity_update(operator_id);
        let diff = Diff::cast_from(delta_capacity);
        self.output
            .arrangement_heap_capacity
            .give((datum, ts, diff));
    }

    /// Update the allocation count for an arrangement.
    fn handle_arrangement_heap_allocations(
        &mut self,
        ArrangementHeapAllocationsReference {
            operator_id,
            delta_allocations,
        }: Ref<'_, ArrangementHeapAllocations>,
    ) {
        let ts = self.ts();
        let Some(state) = self.state.arrangement_size.get_mut(&operator_id) else {
            return;
        };

        state.count += delta_allocations;

        let datum = self
            .state
            .pack_arrangement_heap_allocations_update(operator_id);
        let diff = Diff::cast_from(delta_allocations);
        self.output
            .arrangement_heap_allocations
            .give((datum, ts, diff));
    }

    /// Indicate that a new arrangement exists, start maintaining the heap size state.
    fn handle_arrangement_heap_size_operator(
        &mut self,
        ArrangementHeapSizeOperatorReference {
            operator_id,
            address,
        }: Ref<'_, ArrangementHeapSizeOperator>,
    ) {
        let activator = Activator::new(
            address.into_iter().collect(),
            Rc::clone(&self.state.activations),
        );
        let existing = self
            .state
            .arrangement_size
            .insert(operator_id, Default::default());
        if existing.is_some() {
            error!(%operator_id, "arrangement size operator already registered");
        }
        let existing = self
            .shared_state
            .arrangement_size_activators
            .insert(operator_id, activator);
        if existing.is_some() {
            error!(%operator_id, "arrangement size activator already registered");
        }
    }

    /// Indicate that an arrangement has been dropped and we can cleanup the heap size state.
    fn handle_arrangement_heap_size_operator_dropped(
        &mut self,
        event: Ref<'_, ArrangementHeapSizeOperatorDrop>,
    ) {
        let operator_id = event.operator_id;
        if let Some(state) = self.state.arrangement_size.remove(&operator_id) {
            let ts = self.ts();
            let allocations = self
                .state
                .pack_arrangement_heap_allocations_update(operator_id);
            let diff = -Diff::cast_from(state.count);
            self.output
                .arrangement_heap_allocations
                .give((allocations, ts, diff));

            let capacity = self
                .state
                .pack_arrangement_heap_capacity_update(operator_id);
            let diff = -Diff::cast_from(state.capacity);
            self.output
                .arrangement_heap_capacity
                .give((capacity, ts, diff));

            let size = self.state.pack_arrangement_heap_size_update(operator_id);
            let diff = -Diff::cast_from(state.size);
            self.output.arrangement_heap_size.give((size, ts, diff));
        }
        self.shared_state
            .arrangement_size_activators
            .remove(&operator_id);
    }

    /// Indicate that a new LIR operator exists; record the dataflow address it maps to.
    fn handle_lir_mapping(
        &mut self,
        LirMappingReference { global_id, mapping }: Ref<'_, LirMapping>,
    ) {
        let global_id = Columnar::into_owned(global_id);
        // record the state (for the later drop)
        let mappings = || mapping.into_iter().map(Columnar::into_owned);
        self.state
            .lir_mapping
            .entry(global_id)
            .and_modify(|existing_mapping| existing_mapping.extend(mappings()))
            .or_insert_with(|| mappings().collect());

        // send the datum out
        let ts = self.ts();
        for (lir_id, meta) in mapping.into_iter() {
            let datum = self.state.pack_lir_mapping_update(
                global_id,
                Columnar::into_owned(lir_id),
                Columnar::into_owned(meta.operator),
                Columnar::into_owned(meta.parent_lir_id),
                Columnar::into_owned(meta.nesting),
                Columnar::into_owned(meta.operator_span),
            );
            self.output.lir_mapping.give((datum, ts, Diff::ONE));
        }
    }

    fn handle_dataflow_global(
        &mut self,
        DataflowGlobalReference {
            dataflow_index,
            global_id,
        }: Ref<'_, DataflowGlobal>,
    ) {
        let global_id = Columnar::into_owned(global_id);
        self.state
            .dataflow_global_ids
            .entry(dataflow_index)
            .and_modify(|globals| {
                // NB BTreeSet::insert() returns `false` when the element was already in the set
                if !globals.insert(global_id) {
                    error!(%dataflow_index, %global_id, "dataflow mapping already knew about this GlobalId");
                }
            })
            .or_insert_with(|| BTreeSet::from([global_id]));

        let ts = self.ts();
        let datum = self
            .state
            .pack_dataflow_global_update(dataflow_index, global_id);
        self.output.dataflow_global_ids.give((datum, ts, Diff::ONE));
    }
}

/// Logging state maintained for a compute collection.
///
/// This type is used to produce appropriate log events in response to changes of logged collection
/// state, e.g. frontiers, and to produce cleanup events when a collection is dropped.
pub struct CollectionLogging {
    export_id: GlobalId,
    logger: Logger,

    logged_frontier: Option<Timestamp>,
    logged_import_frontiers: BTreeMap<GlobalId, Timestamp>,
}

impl CollectionLogging {
    /// Create new logging state for the identified collection and emit initial logging events.
    pub fn new(
        export_id: GlobalId,
        logger: Logger,
        dataflow_index: usize,
        import_ids: impl Iterator<Item = GlobalId>,
    ) -> Self {
        logger.log(&ComputeEvent::Export(Export {
            export_id,
            dataflow_index,
        }));

        let mut self_ = Self {
            export_id,
            logger,
            logged_frontier: None,
            logged_import_frontiers: Default::default(),
        };

        // Initialize frontier logging.
        let initial_frontier = Some(Timestamp::MIN);
        self_.set_frontier(initial_frontier);
        import_ids.for_each(|id| self_.set_import_frontier(id, initial_frontier));

        self_
    }

    /// Set the collection frontier to the given new time and emit corresponding logging events.
    pub fn set_frontier(&mut self, new_time: Option<Timestamp>) {
        let old_time = self.logged_frontier;
        self.logged_frontier = new_time;

        if old_time != new_time {
            let export_id = self.export_id;
            let retraction = old_time.map(|time| {
                ComputeEvent::Frontier(Frontier {
                    export_id,
                    time,
                    diff: -1,
                })
            });
            let insertion = new_time.map(|time| {
                ComputeEvent::Frontier(Frontier {
                    export_id,
                    time,
                    diff: 1,
                })
            });
            let events = retraction.as_ref().into_iter().chain(insertion.as_ref());
            self.logger.log_many(events);
        }
    }

    /// Set the frontier of the given import to the given new time and emit corresponding logging
    /// events.
    pub fn set_import_frontier(&mut self, import_id: GlobalId, new_time: Option<Timestamp>) {
        let old_time = self.logged_import_frontiers.remove(&import_id);
        if let Some(time) = new_time {
            self.logged_import_frontiers.insert(import_id, time);
        }

        if old_time != new_time {
            let export_id = self.export_id;
            let retraction = old_time.map(|time| {
                ComputeEvent::ImportFrontier(ImportFrontier {
                    import_id,
                    export_id,
                    time,
                    diff: -1,
                })
            });
            let insertion = new_time.map(|time| {
                ComputeEvent::ImportFrontier(ImportFrontier {
                    import_id,
                    export_id,
                    time,
                    diff: 1,
                })
            });
            let events = retraction.as_ref().into_iter().chain(insertion.as_ref());
            self.logger.log_many(events);
        }
    }

    /// Set the collection as hydrated.
    pub fn set_hydrated(&self) {
        self.logger.log(&ComputeEvent::Hydration(Hydration {
            export_id: self.export_id,
        }));
    }

    /// The export global ID of this collection.
    pub fn export_id(&self) -> GlobalId {
        self.export_id
    }
}

impl Drop for CollectionLogging {
    fn drop(&mut self) {
        // Emit retraction events to clean up events previously logged.
        self.set_frontier(None);

        let import_ids: Vec<_> = self.logged_import_frontiers.keys().copied().collect();
        for import_id in import_ids {
            self.set_import_frontier(import_id, None);
        }

        self.logger.log(&ComputeEvent::ExportDropped(ExportDropped {
            export_id: self.export_id,
        }));
    }
}

/// Extension trait to attach `ComputeEvent::DataflowError` logging operators to collections and
/// batch streams.
pub(crate) trait LogDataflowErrors {
    fn log_dataflow_errors(self, logger: Logger, export_id: GlobalId) -> Self;
}

impl<'scope, T, D> LogDataflowErrors for VecCollection<'scope, T, D, Diff>
where
    T: timely::progress::Timestamp,
    D: Clone + 'static,
{
    fn log_dataflow_errors(self, logger: Logger, export_id: GlobalId) -> Self {
        self.inner
            .unary(Pipeline, "LogDataflowErrorsCollection", |_cap, _info| {
                move |input, output| {
                    input.for_each(|cap, data| {
                        let diff = data.iter().map(|(_d, _t, r)| *r).sum::<Diff>();
                        logger.log(&ComputeEvent::ErrorCount(ErrorCount { export_id, diff }));

                        output.session(&cap).give_container(data);
                    });
                }
            })
            .as_collection()
    }
}

impl<'scope, T, B> LogDataflowErrors for StreamVec<'scope, T, B>
where
    T: timely::progress::Timestamp,
    for<'a> B: BatchReader<DiffGat<'a> = &'a Diff> + Clone + 'static,
{
    fn log_dataflow_errors(self, logger: Logger, export_id: GlobalId) -> Self {
        self.unary(Pipeline, "LogDataflowErrorsStream", |_cap, _info| {
            move |input, output| {
                input.for_each(|cap, data| {
                    let diff = data.iter().map(sum_batch_diffs).sum::<Diff>();
                    logger.log(&ComputeEvent::ErrorCount(ErrorCount { export_id, diff }));

                    output.session(&cap).give_container(data);
                });
            }
        })
    }
}

/// Return the sum of all diffs within the given batch.
///
/// Note that this operation can be expensive: Its runtime is O(N) with N being the number of
/// unique (key, value, time) tuples. We only use it on error streams, which are expected to
/// contain only a small number of records, so this doesn't matter much. But avoid using it when
/// batches might become large.
fn sum_batch_diffs<B>(batch: &B) -> Diff
where
    for<'a> B: BatchReader<DiffGat<'a> = &'a Diff>,
{
    let mut sum = Diff::ZERO;
    let mut cursor = batch.cursor();

    while cursor.key_valid(batch) {
        while cursor.val_valid(batch) {
            cursor.map_times(batch, |_t, r| sum += r);
            cursor.step_val(batch);
        }
        cursor.step_key(batch);
    }

    sum
}

#[cfg(test)]
mod tests {
    use super::*;

    #[mz_ore::test]
    fn test_compute_event_size() {
        // This could be a static assertion, but we don't use those yet in this crate.
        assert_eq!(56, std::mem::size_of::<ComputeEvent>())
    }
}