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;
use differential_dataflow::Collection;
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, Timestamp};
use mz_timely_util::containers::{Column, ColumnBuilder, ProvidedBuilder};
use mz_timely_util::replay::MzReplay;
use timely::dataflow::channels::pact::Pipeline;
use timely::dataflow::operators::Operator;
use timely::dataflow::operators::core::Map;
use timely::dataflow::operators::generic::builder_rc::OperatorBuilder;
use timely::dataflow::{Scope, Stream};
use timely::scheduling::Scheduler;
use timely::{Container, Data};
use tracing::error;
use uuid::Uuid;

use crate::extensions::arrange::MzArrange;
use crate::logging::{
    ComputeLog, EventQueue, LogCollection, LogVariant, OutputSessionColumnar, OutputSessionVec,
    PermutedRowPacker, SharedLoggingState, Update,
};
use crate::row_spine::{RowRowBatcher, RowRowBuilder};
use crate::typedefs::RowRowSpine;

/// 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 [`Peek`], a [`PeekType`], and an installation status.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct PeekEvent {
    /// The data for the peek itself.
    pub peek: Peek,
    /// 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 {
    pub export_id: GlobalId,
}

/// 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),
    /// 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",
        }
    }
}

/// A logged peek event.
#[derive(Debug, Clone, PartialOrd, PartialEq, Columnar)]
pub struct Peek {
    /// The identifier of the view the peek targets.
    id: GlobalId,
    /// The logical timestamp requested.
    time: Timestamp,
    /// The ID of the peek.
    uuid: uuid::Bytes,
}

impl Peek {
    /// Create a new peek from its arguments.
    pub fn new(id: GlobalId, time: Timestamp, uuid: Uuid) -> Self {
        let uuid = uuid.into_bytes();
        Self { id, time, uuid }
    }
}

/// 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<S: Scheduler + 'static, G: Scope<Timestamp = Timestamp>>(
    mut scope: G,
    scheduler: S,
    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) = event_queue.links.mz_replay::<_, ProvidedBuilder<_>, _>(
            scope,
            "compute logs",
            config.interval,
            event_queue.activator,
            move |mut session, mut data| {
                // If logging is disabled, we still need to install the indexes, but we can leave them
                // empty. We do so by immediately filtering all logs events.
                if enable_logging {
                    session.give_container(data.to_mut())
                }
            },
        );

        // 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 (mut export_out, export) = demux.new_output();
        let (mut frontier_out, frontier) = demux.new_output();
        let (mut import_frontier_out, import_frontier) = demux.new_output();
        let (mut peek_out, peek) = demux.new_output();
        let (mut peek_duration_out, peek_duration) = demux.new_output();
        let (mut shutdown_duration_out, shutdown_duration) = demux.new_output();
        let (mut arrangement_heap_size_out, arrangement_heap_size) = demux.new_output();
        let (mut arrangement_heap_capacity_out, arrangement_heap_capacity) = demux.new_output();
        let (mut arrangement_heap_allocations_out, arrangement_heap_allocations) =
            demux.new_output();
        let (mut error_count_out, error_count) = demux.new_output();
        let (mut hydration_time_out, hydration_time) = demux.new_output();
        let (mut lir_mapping_out, lir_mapping) = demux.new_output();
        let (mut dataflow_global_ids_out, dataflow_global_ids) = demux.new_output();

        let mut demux_state = DemuxState::new(scheduler);
        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 shutdown_duration = shutdown_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 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(&cap),
                        frontier: frontier.session(&cap),
                        import_frontier: import_frontier.session(&cap),
                        peek: peek.session(&cap),
                        peek_duration: peek_duration.session(&cap),
                        shutdown_duration: shutdown_duration.session(&cap),
                        arrangement_heap_size: arrangement_heap_size.session(&cap),
                        arrangement_heap_capacity: arrangement_heap_capacity.session(&cap),
                        arrangement_heap_allocations: arrangement_heap_allocations.session(&cap),
                        error_count: error_count.session(&cap),
                        hydration_time: hydration_time.session(&cap),
                        lir_mapping: lir_mapping.session_with_builder(&cap),
                        dataflow_global_ids: dataflow_global_ids.session(&cap),
                    };

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

        let worker_id = scope.index();

        // Encode the contents of each logging stream into its expected `Row` format.
        let mut packer = PermutedRowPacker::new(ComputeLog::DataflowCurrent);
        let dataflow_current = export.as_collection().map({
            let mut scratch = String::new();
            move |datum| {
                packer.pack_slice_owned(&[
                    make_string_datum(datum.export_id, &mut scratch),
                    Datum::UInt64(u64::cast_from(worker_id)),
                    Datum::UInt64(u64::cast_from(datum.dataflow_index)),
                ])
            }
        });
        let mut packer = PermutedRowPacker::new(ComputeLog::FrontierCurrent);
        let frontier_current = frontier.as_collection().map({
            let mut scratch = String::new();
            move |datum| {
                packer.pack_slice_owned(&[
                    make_string_datum(datum.export_id, &mut scratch),
                    Datum::UInt64(u64::cast_from(worker_id)),
                    Datum::MzTimestamp(datum.time),
                ])
            }
        });
        let mut packer = PermutedRowPacker::new(ComputeLog::ImportFrontierCurrent);
        let import_frontier_current = import_frontier.as_collection().map({
            let mut scratch1 = String::new();
            let mut scratch2 = String::new();
            move |datum| {
                packer.pack_slice_owned(&[
                    make_string_datum(datum.export_id, &mut scratch1),
                    make_string_datum(datum.import_id, &mut scratch2),
                    Datum::UInt64(u64::cast_from(worker_id)),
                    Datum::MzTimestamp(datum.time),
                ])
            }
        });
        let mut packer = PermutedRowPacker::new(ComputeLog::PeekCurrent);
        let peek_current = peek.as_collection().map({
            let mut scratch = String::new();
            move |PeekDatum { peek, peek_type }| {
                packer.pack_slice_owned(&[
                    Datum::Uuid(Uuid::from_bytes(peek.uuid)),
                    Datum::UInt64(u64::cast_from(worker_id)),
                    make_string_datum(peek.id, &mut scratch),
                    Datum::String(peek_type.name()),
                    Datum::MzTimestamp(peek.time),
                ])
            }
        });
        let mut packer = PermutedRowPacker::new(ComputeLog::PeekDuration);
        let peek_duration =
            peek_duration
                .as_collection()
                .map(move |PeekDurationDatum { peek_type, bucket }| {
                    packer.pack_slice_owned(&[
                        Datum::UInt64(u64::cast_from(worker_id)),
                        Datum::String(peek_type.name()),
                        Datum::UInt64(bucket.try_into().expect("bucket too big")),
                    ])
                });
        let mut packer = PermutedRowPacker::new(ComputeLog::ShutdownDuration);
        let shutdown_duration = shutdown_duration.as_collection().map(move |bucket| {
            packer.pack_slice_owned(&[
                Datum::UInt64(u64::cast_from(worker_id)),
                Datum::UInt64(bucket.try_into().expect("bucket too big")),
            ])
        });

        let arrangement_heap_datum_to_row =
            move |packer: &mut PermutedRowPacker, ArrangementHeapDatum { operator_id }| {
                packer.pack_slice_owned(&[
                    Datum::UInt64(operator_id.try_into().expect("operator_id too big")),
                    Datum::UInt64(u64::cast_from(worker_id)),
                ])
            };

        let mut packer = PermutedRowPacker::new(ComputeLog::ArrangementHeapSize);
        let arrangement_heap_size = arrangement_heap_size
            .as_collection()
            .map(move |d| arrangement_heap_datum_to_row(&mut packer, d));

        let mut packer = PermutedRowPacker::new(ComputeLog::ArrangementHeapCapacity);
        let arrangement_heap_capacity = arrangement_heap_capacity
            .as_collection()
            .map(move |d| arrangement_heap_datum_to_row(&mut packer, d));

        let mut packer = PermutedRowPacker::new(ComputeLog::ArrangementHeapSize);
        let arrangement_heap_allocations = arrangement_heap_allocations
            .as_collection()
            .map(move |d| arrangement_heap_datum_to_row(&mut packer, d));

        let mut packer = PermutedRowPacker::new(ComputeLog::ErrorCount);
        let error_count = error_count.as_collection().map({
            let mut scratch = String::new();
            move |datum| {
                packer.pack_slice_owned(&[
                    make_string_datum(datum.export_id, &mut scratch),
                    Datum::UInt64(u64::cast_from(worker_id)),
                    Datum::Int64(datum.count.into_inner()),
                ])
            }
        });

        let mut packer = PermutedRowPacker::new(ComputeLog::HydrationTime);
        let hydration_time = hydration_time.as_collection().map({
            let mut scratch = String::new();
            move |datum| {
                packer.pack_slice_owned(&[
                    make_string_datum(datum.export_id, &mut scratch),
                    Datum::UInt64(u64::cast_from(worker_id)),
                    Datum::from(datum.time_ns),
                ])
            }
        });

        let mut scratch1 = String::new();
        let mut scratch2 = String::new();
        let mut packer = PermutedRowPacker::new(ComputeLog::LirMapping);
        let lir_mapping = lir_mapping
            .map(move |(datum, time, diff)| {
                let row = packer.pack_slice_owned(&[
                    make_string_datum(GlobalId::into_owned(datum.global_id), &mut scratch1),
                    Datum::UInt64(<LirId as Columnar>::into_owned(datum.lir_id).into()),
                    Datum::UInt64(u64::cast_from(worker_id)),
                    make_string_datum(datum.operator, &mut scratch2),
                    datum
                        .parent_lir_id
                        .map(|lir_id| Datum::UInt64(LirId::into_owned(lir_id).into()))
                        .unwrap_or_else(|| Datum::Null),
                    Datum::UInt16(u16::cast_from(*datum.nesting)),
                    Datum::UInt64(u64::cast_from(datum.operator_span.0)),
                    Datum::UInt64(u64::cast_from(datum.operator_span.1)),
                ]);
                (row, Timestamp::into_owned(time), diff)
            })
            .as_collection();

        let mut packer = PermutedRowPacker::new(ComputeLog::DataflowGlobal);
        let dataflow_global_ids = dataflow_global_ids.as_collection().map({
            let mut scratch = String::new();
            move |datum| {
                packer.pack_slice_owned(&[
                    Datum::UInt64(u64::cast_from(datum.dataflow_index)),
                    Datum::UInt64(u64::cast_from(worker_id)),
                    make_string_datum(datum.global_id, &mut scratch),
                ])
            }
        });

        use ComputeLog::*;
        let logs = [
            (DataflowCurrent, dataflow_current),
            (FrontierCurrent, frontier_current),
            (ImportFrontierCurrent, import_frontier_current),
            (PeekCurrent, peek_current),
            (PeekDuration, peek_duration),
            (ShutdownDuration, shutdown_duration),
            (ArrangementHeapSize, arrangement_heap_size),
            (ArrangementHeapCapacity, arrangement_heap_capacity),
            (ArrangementHeapAllocations, arrangement_heap_allocations),
            (ErrorCount, error_count),
            (HydrationTime, hydration_time),
            (LirMapping, lir_mapping),
            (DataflowGlobal, dataflow_global_ids),
        ];

        // Build the output arrangements.
        let mut collections = BTreeMap::new();
        for (variant, collection) in logs {
            let variant = LogVariant::Compute(variant);
            if config.index_logs.contains_key(&variant) {
                let trace = collection
                    .mz_arrange::<RowRowBatcher<_, _>, RowRowBuilder<_, _>, RowRowSpine<_, _>>(
                        &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<A> {
    /// The timely scheduler.
    scheduler: A,
    /// State tracked per dataflow export.
    exports: BTreeMap<GlobalId, ExportState>,
    /// Maps live dataflows to counts of their exports.
    dataflow_export_counts: BTreeMap<usize, u32>,
    /// Maps dropped dataflows to their drop time.
    dataflow_drop_times: BTreeMap<usize, Duration>,
    /// Contains dataflows that have shut down but not yet been dropped.
    shutdown_dataflows: BTreeSet<usize>,
    /// 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>>,
}

impl<A: Scheduler> DemuxState<A> {
    fn new(scheduler: A) -> Self {
        Self {
            scheduler,
            exports: Default::default(),
            dataflow_export_counts: Default::default(),
            dataflow_drop_times: Default::default(),
            shutdown_dataflows: Default::default(),
            peek_stash: Default::default(),
            arrangement_size: Default::default(),
            lir_mapping: Default::default(),
            dataflow_global_ids: Default::default(),
        }
    }
}

/// 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>,
}

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

/// 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> {
    export: OutputSessionVec<'a, Update<ExportDatum>>,
    frontier: OutputSessionVec<'a, Update<FrontierDatum>>,
    import_frontier: OutputSessionVec<'a, Update<ImportFrontierDatum>>,
    peek: OutputSessionVec<'a, Update<PeekDatum>>,
    peek_duration: OutputSessionVec<'a, Update<PeekDurationDatum>>,
    shutdown_duration: OutputSessionVec<'a, Update<u128>>,
    arrangement_heap_size: OutputSessionVec<'a, Update<ArrangementHeapDatum>>,
    arrangement_heap_capacity: OutputSessionVec<'a, Update<ArrangementHeapDatum>>,
    arrangement_heap_allocations: OutputSessionVec<'a, Update<ArrangementHeapDatum>>,
    hydration_time: OutputSessionVec<'a, Update<HydrationTimeDatum>>,
    error_count: OutputSessionVec<'a, Update<ErrorCountDatum>>,
    lir_mapping: OutputSessionColumnar<'a, Update<LirMappingDatum>>,
    dataflow_global_ids: OutputSessionVec<'a, Update<DataflowGlobalDatum>>,
}

#[derive(Clone)]
struct ExportDatum {
    export_id: GlobalId,
    dataflow_index: usize,
}

#[derive(Clone)]
struct FrontierDatum {
    export_id: GlobalId,
    time: Timestamp,
}

#[derive(Clone)]
struct ImportFrontierDatum {
    export_id: GlobalId,
    import_id: GlobalId,
    time: Timestamp,
}

#[derive(Clone)]
struct PeekDatum {
    peek: Peek,
    peek_type: PeekType,
}

#[derive(Clone)]
struct PeekDurationDatum {
    peek_type: PeekType,
    bucket: u128,
}

#[derive(Clone, Copy)]
struct ArrangementHeapDatum {
    operator_id: usize,
}

#[derive(Clone)]
struct HydrationTimeDatum {
    export_id: GlobalId,
    time_ns: Option<u64>,
}

#[derive(Clone)]
struct ErrorCountDatum {
    export_id: GlobalId,
    // 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.
    count: Diff,
}

#[derive(Clone, Columnar)]
struct LirMappingDatum {
    global_id: GlobalId,
    lir_id: LirId,
    operator: String,
    parent_lir_id: Option<LirId>,
    nesting: u8,
    operator_span: (usize, usize),
}

#[derive(Clone)]
struct DataflowGlobalDatum {
    dataflow_index: usize,
    global_id: GlobalId,
}

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

impl<A: Scheduler> DemuxHandler<'_, '_, A> {
    /// 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: <ComputeEvent as Columnar>::Ref<'_>) {
        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),
            LirMapping(mapping) => self.handle_lir_mapping(mapping),
            DataflowGlobal(global) => self.handle_dataflow_global(global),
        }
    }

    fn handle_export(
        &mut self,
        ExportReference {
            export_id,
            dataflow_index,
        }: <Export as Columnar>::Ref<'_>,
    ) {
        let export_id = Columnar::into_owned(export_id);
        let ts = self.ts();
        let datum = ExportDatum {
            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");
        }

        *self
            .state
            .dataflow_export_counts
            .entry(dataflow_index)
            .or_default() += 1;

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

    fn handle_export_dropped(
        &mut self,
        ExportDroppedReference { export_id }: <ExportDropped as Columnar>::Ref<'_>,
    ) {
        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 = ExportDatum {
            export_id,
            dataflow_index,
        };
        self.output.export.give((datum, ts, Diff::MINUS_ONE));

        match self.state.dataflow_export_counts.get_mut(&dataflow_index) {
            entry @ Some(0) | entry @ None => {
                error!(
                    %export_id,
                    %dataflow_index,
                    "invalid dataflow_export_counts entry at time of export drop: {entry:?}",
                );
            }
            Some(1) => self.handle_dataflow_dropped(dataflow_index),
            Some(count) => *count -= 1,
        }

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

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

    fn handle_dataflow_dropped(&mut self, dataflow_index: usize) {
        self.state.dataflow_export_counts.remove(&dataflow_index);

        if self.state.shutdown_dataflows.remove(&dataflow_index) {
            // Dataflow has already shut down before it was dropped.
            self.output
                .shutdown_duration
                .give((0, self.ts(), Diff::ONE));
        } else {
            // Dataflow has not yet shut down.
            let existing = self
                .state
                .dataflow_drop_times
                .insert(dataflow_index, self.time);
            if existing.is_some() {
                error!(%dataflow_index, "dataflow already dropped");
            }
        }
    }

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

        if let Some(start) = self.state.dataflow_drop_times.remove(&dataflow_index) {
            // Dataflow has already been dropped.
            let elapsed_ns = self.time.saturating_sub(start).as_nanos();
            let elapsed_pow = elapsed_ns.next_power_of_two();
            self.output
                .shutdown_duration
                .give((elapsed_pow, ts, Diff::ONE));
        } else {
            // Dataflow has not yet been dropped.
            let was_new = self.state.shutdown_dataflows.insert(dataflow_index);
            if !was_new {
                error!(%dataflow_index, "dataflow already shutdown");
            }
        }

        // 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 = DataflowGlobalDatum {
                    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 = LirMappingDatum {
                            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 }: <ErrorCount as Columnar>::Ref<'_>,
    ) {
        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;

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

        export.error_count = new_count;
    }

    fn handle_hydration(
        &mut self,
        HydrationReference { export_id }: <Hydration as Columnar>::Ref<'_>,
    ) {
        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");

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

        export.hydration_time_ns = Some(nanos);
    }

    fn handle_peek_install(
        &mut self,
        PeekEventReference {
            peek,
            peek_type,
            installed: _,
        }: <PeekEvent as Columnar>::Ref<'_>,
    ) {
        let peek = Peek::into_owned(peek);
        let uuid = Uuid::from_bytes(peek.uuid);
        let ts = self.ts();
        self.output
            .peek
            .give((PeekDatum { peek, peek_type }, 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 {
            peek,
            peek_type,
            installed: _,
        }: <PeekEvent as Columnar>::Ref<'_>,
    ) {
        let peek = Peek::into_owned(peek);
        let uuid = Uuid::from_bytes(peek.uuid);
        let ts = self.ts();
        self.output
            .peek
            .give((PeekDatum { peek, peek_type }, 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();
            self.output.peek_duration.give((
                PeekDurationDatum { peek_type, bucket },
                ts,
                Diff::ONE,
            ));
        } else {
            error!(%uuid, "peek not yet registered");
        }
    }

    fn handle_frontier(
        &mut self,
        FrontierReference {
            export_id,
            time,
            diff,
        }: <Frontier as Columnar>::Ref<'_>,
    ) {
        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 = FrontierDatum { export_id, time };
        self.output.frontier.give((datum, ts, diff));
    }

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

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

        let datum = ArrangementHeapDatum { operator_id };
        self.output
            .arrangement_heap_size
            .give((datum, ts, Diff::cast_from(delta_size)));

        state.size += delta_size;
    }

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

        let datum = ArrangementHeapDatum { operator_id };
        self.output
            .arrangement_heap_capacity
            .give((datum, ts, Diff::cast_from(delta_capacity)));

        state.capacity += delta_capacity;
    }

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

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

        state.count += delta_allocations;
    }

    /// Indicate that a new arrangement exists, start maintaining the heap size state.
    fn handle_arrangement_heap_size_operator(
        &mut self,
        ArrangementHeapSizeOperatorReference {
            operator_id,
            address,
        }: <ArrangementHeapSizeOperator as Columnar>::Ref<'_>,
    ) {
        let activator = self
            .state
            .scheduler
            .activator_for(address.into_iter().collect());
        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,
        ArrangementHeapSizeOperatorDropReference { operator_id }: <ArrangementHeapSizeOperatorDrop as Columnar>::Ref<'_>,
    ) {
        if let Some(state) = self.state.arrangement_size.remove(&operator_id) {
            let ts = self.ts();
            let datum = ArrangementHeapDatum { operator_id };

            let diff = -Diff::cast_from(state.size);
            self.output.arrangement_heap_size.give((datum, ts, diff));

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

            let diff = -Diff::cast_from(state.count);
            self.output
                .arrangement_heap_allocations
                .give((datum, 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 }: <LirMapping as Columnar>::Ref<'_>,
    ) {
        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 = LirMappingDatumReference {
                global_id,
                lir_id,
                operator: meta.operator,
                parent_lir_id: meta.parent_lir_id,
                nesting: meta.nesting,
                operator_span: meta.operator_span,
            };
            self.output.lir_mapping.give((datum, ts, Diff::ONE));
        }
    }

    fn handle_dataflow_global(
        &mut self,
        DataflowGlobalReference {
            dataflow_index,
            global_id,
        }: <DataflowGlobal as Columnar>::Ref<'_>,
    ) {
        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 = DataflowGlobalDatum {
            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,
        }));
    }
}

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<G, D> LogDataflowErrors for Collection<G, D, Diff>
where
    G: Scope,
    D: Data,
{
    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<G, B> LogDataflowErrors for Stream<G, B>
where
    G: Scope,
    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>())
    }
}