mz_compute/logging/

timely.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 timely dataflow.

use std::cell::RefCell;
use std::collections::BTreeMap;
use std::rc::Rc;
use std::time::Duration;

use columnar::{Columnar, Index};
use columnation::{Columnation, CopyRegion};
use mz_compute_client::logging::LoggingConfig;
use mz_ore::cast::CastFrom;
use mz_repr::{Datum, Diff, Timestamp};
use mz_timely_util::columnar::batcher;
use mz_timely_util::columnar::builder::ColumnBuilder;
use mz_timely_util::columnar::{Col2ValBatcher, columnar_exchange};
use mz_timely_util::columnation::ColumnationChunker;
use mz_timely_util::replay::MzReplay;
use timely::dataflow::Scope;
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::logging::{
    ChannelsEvent, MessagesEvent, OperatesEvent, ParkEvent, ScheduleEvent, ShutdownEvent,
    TimelyEvent,
};
use tracing::error;

use crate::extensions::arrange::MzArrangeCore;
use crate::logging::compute::{ComputeEvent, DataflowShutdown};
use crate::logging::{
    EventQueue, LogVariant, OutputSessionColumnar, OutputSessionVec, PermutedRowPacker, TimelyLog,
    Update,
};
use crate::logging::{LogCollection, SharedLoggingState, consolidate_and_pack};
use crate::typedefs::{KeyBatcher, KeyValBatcher, RowRowSpine};
use mz_row_spine::RowRowBuilder;

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

/// Constructs the logging dataflow fragment for timely logs.
///
/// Params
/// * `scope`: The Timely scope hosting the log analysis dataflow.
/// * `config`: Logging configuration
/// * `event_queue`: The source to read log events from.
pub(super) fn construct(
    scope: Scope<'_, Timestamp>,
    config: &LoggingConfig,
    event_queue: EventQueue<Vec<(Duration, TimelyEvent)>>,
    shared_state: Rc<RefCell<SharedLoggingState>>,
    storage_log_reader: Option<crate::server::StorageTimelyLogReader>,
) -> Return {
    scope.scoped("timely logging", move |scope| {
        let enable_logging = config.enable_logging;
        let (logs, token) = if enable_logging {
            event_queue.links.mz_replay(
                scope,
                "timely logs",
                config.interval,
                event_queue.activator,
            )
        } else {
            let token: Rc<dyn std::any::Any> = Rc::new(Box::new(()));
            (empty(scope), token)
        };

        // If we have a storage log reader, replay it and concatenate with compute logs.
        let (logs, storage_token) = if let Some(reader) = storage_log_reader {
            use mz_timely_util::activator::RcActivator;
            use timely::dataflow::operators::Concatenate;

            let activator = RcActivator::new("storage_timely_activator".to_string(), 128);
            let (storage_logs, s_token) =
                [reader].mz_replay(scope, "storage timely logs", config.interval, activator);
            let merged = scope.concatenate([logs, storage_logs]);
            (merged, Some(s_token))
        } else {
            (logs, None)
        };
        // Convert storage token to Rc<dyn Any> so we can stash it alongside the compute token.
        let storage_token: Rc<dyn std::any::Any> = match storage_token {
            Some(t) => t,
            None => Rc::new(()),
        };

        // Build a demux operator that splits the replayed event stream up into the separate
        // logging streams.
        let mut demux = OperatorBuilder::new("Timely Logging Demux".to_string(), scope.clone());
        let mut input = demux.new_input(logs, Pipeline);
        let (operates_out, operates) = demux.new_output();
        let mut operates_out = OutputBuilder::from(operates_out);
        let (channels_out, channels) = demux.new_output();
        let mut channels_out = OutputBuilder::from(channels_out);
        let (addresses_out, addresses) = demux.new_output();
        let mut addresses_out = OutputBuilder::from(addresses_out);
        let (parks_out, parks) = demux.new_output();
        let mut parks_out = OutputBuilder::from(parks_out);
        let (messages_sent_out, messages_sent) = demux.new_output();
        let mut messages_sent_out = OutputBuilder::from(messages_sent_out);
        let (messages_received_out, messages_received) = demux.new_output();
        let mut messages_received_out = OutputBuilder::from(messages_received_out);
        let (schedules_duration_out, schedules_duration) = demux.new_output();
        let mut schedules_duration_out = OutputBuilder::from(schedules_duration_out);
        let (schedules_histogram_out, schedules_histogram) = demux.new_output();
        let mut schedules_histogram_out = OutputBuilder::from(schedules_histogram_out);
        let (batches_sent_out, batches_sent) = demux.new_output();
        let mut batches_sent_out = OutputBuilder::from(batches_sent_out);
        let (batches_received_out, batches_received) = demux.new_output();
        let mut batches_received_out = OutputBuilder::from(batches_received_out);

        let worker_id = scope.index();
        let mut demux_state = DemuxState::default();
        demux.build(|_capability| {
            let peers = scope.peers();
            let logging_interval_ms = std::cmp::max(1, config.interval.as_millis());
            move |_frontiers| {
                let mut operates = operates_out.activate();
                let mut channels = channels_out.activate();
                let mut addresses = addresses_out.activate();
                let mut parks = parks_out.activate();
                let mut messages_sent = messages_sent_out.activate();
                let mut messages_received = messages_received_out.activate();
                let mut batches_sent = batches_sent_out.activate();
                let mut batches_received = batches_received_out.activate();
                let mut schedules_duration = schedules_duration_out.activate();
                let mut schedules_histogram = schedules_histogram_out.activate();

                input.for_each(|cap, data| {
                    let mut output_buffers = DemuxOutput {
                        operates: operates.session_with_builder(&cap),
                        channels: channels.session_with_builder(&cap),
                        addresses: addresses.session_with_builder(&cap),
                        parks: parks.session_with_builder(&cap),
                        messages_sent: messages_sent.session_with_builder(&cap),
                        messages_received: messages_received.session_with_builder(&cap),
                        schedules_duration: schedules_duration.session_with_builder(&cap),
                        schedules_histogram: schedules_histogram.session_with_builder(&cap),
                        batches_sent: batches_sent.session_with_builder(&cap),
                        batches_received: batches_received.session_with_builder(&cap),
                    };

                    for (time, event) in data.drain(..) {
                        if let TimelyEvent::Messages(msg) = &event {
                            match msg.is_send {
                                true => assert_eq!(msg.source, worker_id),
                                false => assert_eq!(msg.target, worker_id),
                            }
                        }

                        DemuxHandler {
                            state: &mut demux_state,
                            shared_state: &mut shared_state.borrow_mut(),
                            output: &mut output_buffers,
                            logging_interval_ms,
                            peers,
                            time,
                        }
                        .handle(event);
                    }
                });
            }
        });

        // Encode the contents of each logging stream into its expected `Row` format.
        // We pre-arrange the logging streams to force a consolidation and reduce the amount of
        // updates that reach `Row` encoding.

        let operates = consolidate_and_pack::<
            ColumnationChunker<_>,
            KeyValBatcher<_, _, _, _>,
            ColumnBuilder<_>,
            _,
            _,
            _,
        >(
            operates,
            TimelyLog::Operates,
            move |data, packer, session| {
                for ((id, name), time, diff) in data.iter() {
                    let data = packer.pack_slice(&[
                        Datum::UInt64(u64::cast_from(*id)),
                        Datum::UInt64(u64::cast_from(worker_id)),
                        Datum::String(name),
                    ]);
                    session.give((data, time, diff));
                }
            },
        );

        // TODO: `consolidate_and_pack` requires columnation, which `ChannelDatum` does not
        // implement. Consider consolidating here once we support columnar.
        let channels = channels.unary::<ColumnBuilder<_>, _, _, _>(
            Pipeline,
            "ToRow Channels",
            |_cap, _info| {
                let mut packer = PermutedRowPacker::new(TimelyLog::Channels);
                move |input, output| {
                    input.for_each_time(|time, data| {
                        let mut session = output.session_with_builder(&time);
                        for d in data.flat_map(|c| c.borrow().into_index_iter()) {
                            let ((datum, ()), time, diff) = d;
                            let (source_node, source_port) = datum.source;
                            let (target_node, target_port) = datum.target;
                            let data = packer.pack_slice(&[
                                Datum::UInt64(u64::cast_from(datum.id)),
                                Datum::UInt64(u64::cast_from(worker_id)),
                                Datum::UInt64(u64::cast_from(source_node)),
                                Datum::UInt64(u64::cast_from(source_port)),
                                Datum::UInt64(u64::cast_from(target_node)),
                                Datum::UInt64(u64::cast_from(target_port)),
                                Datum::String(
                                    std::str::from_utf8(datum.typ).expect("valid string"),
                                ),
                            ]);
                            session.give((data, time, diff));
                        }
                    });
                }
            },
        );

        // Types to make rustfmt happy.
        type KVB<K, V, T, D> = KeyValBatcher<K, V, T, D>;
        type KB<K, T, D> = KeyBatcher<K, T, D>;

        let addresses = consolidate_and_pack::<
            ColumnationChunker<_>,
            KVB<_, _, _, _>,
            ColumnBuilder<_>,
            _,
            _,
            _,
        >(
            addresses,
            TimelyLog::Addresses,
            move |data, packer, session| {
                for ((id, address), time, diff) in data.iter() {
                    let data = packer.pack_by_index(|packer, index| match index {
                        0 => packer.push(Datum::UInt64(u64::cast_from(*id))),
                        1 => packer.push(Datum::UInt64(u64::cast_from(worker_id))),
                        2 => {
                            let list = address.iter().map(|i| Datum::UInt64(u64::cast_from(*i)));
                            packer.push_list(list)
                        }
                        _ => unreachable!("Addresses relation has three columns"),
                    });
                    session.give((data, time, diff));
                }
            },
        );

        let parks =
            consolidate_and_pack::<ColumnationChunker<_>, KB<_, _, _>, ColumnBuilder<_>, _, _, _>(
                parks,
                TimelyLog::Parks,
                move |data, packer, session| {
                    for ((datum, ()), time, diff) in data.iter() {
                        let data = packer.pack_slice(&[
                            Datum::UInt64(u64::cast_from(worker_id)),
                            Datum::UInt64(datum.duration_pow),
                            datum
                                .requested_pow
                                .map(Datum::UInt64)
                                .unwrap_or(Datum::Null),
                        ]);
                        session.give((data, time, diff));
                    }
                },
            );

        let batches_sent =
            consolidate_and_pack::<ColumnationChunker<_>, KB<_, _, _>, ColumnBuilder<_>, _, _, _>(
                batches_sent,
                TimelyLog::BatchesSent,
                move |data, packer, session| {
                    for ((datum, ()), time, diff) in data.iter() {
                        let data = packer.pack_slice(&[
                            Datum::UInt64(u64::cast_from(datum.channel)),
                            Datum::UInt64(u64::cast_from(worker_id)),
                            Datum::UInt64(u64::cast_from(datum.worker)),
                        ]);
                        session.give((data, time, diff));
                    }
                },
            );

        let batches_received =
            consolidate_and_pack::<ColumnationChunker<_>, KB<_, _, _>, ColumnBuilder<_>, _, _, _>(
                batches_received,
                TimelyLog::BatchesReceived,
                move |data, packer, session| {
                    for ((datum, ()), time, diff) in data.iter() {
                        let data = packer.pack_slice(&[
                            Datum::UInt64(u64::cast_from(datum.channel)),
                            Datum::UInt64(u64::cast_from(datum.worker)),
                            Datum::UInt64(u64::cast_from(worker_id)),
                        ]);
                        session.give((data, time, diff));
                    }
                },
            );

        let messages_sent =
            consolidate_and_pack::<ColumnationChunker<_>, KB<_, _, _>, ColumnBuilder<_>, _, _, _>(
                messages_sent,
                TimelyLog::MessagesSent,
                move |data, packer, session| {
                    for ((datum, ()), time, diff) in data.iter() {
                        let data = packer.pack_slice(&[
                            Datum::UInt64(u64::cast_from(datum.channel)),
                            Datum::UInt64(u64::cast_from(worker_id)),
                            Datum::UInt64(u64::cast_from(datum.worker)),
                        ]);
                        session.give((data, time, diff));
                    }
                },
            );

        let messages_received =
            consolidate_and_pack::<ColumnationChunker<_>, KB<_, _, _>, ColumnBuilder<_>, _, _, _>(
                messages_received,
                TimelyLog::MessagesReceived,
                move |data, packer, session| {
                    for ((datum, ()), time, diff) in data.iter() {
                        let data = packer.pack_slice(&[
                            Datum::UInt64(u64::cast_from(datum.channel)),
                            Datum::UInt64(u64::cast_from(datum.worker)),
                            Datum::UInt64(u64::cast_from(worker_id)),
                        ]);
                        session.give((data, time, diff));
                    }
                },
            );

        let elapsed =
            consolidate_and_pack::<ColumnationChunker<_>, KB<_, _, _>, ColumnBuilder<_>, _, _, _>(
                schedules_duration,
                TimelyLog::Elapsed,
                move |data, packer, session| {
                    for ((operator, ()), time, diff) in data.iter() {
                        let data = packer.pack_slice(&[
                            Datum::UInt64(u64::cast_from(*operator)),
                            Datum::UInt64(u64::cast_from(worker_id)),
                        ]);
                        session.give((data, time, diff));
                    }
                },
            );

        let histogram =
            consolidate_and_pack::<ColumnationChunker<_>, KB<_, _, _>, ColumnBuilder<_>, _, _, _>(
                schedules_histogram,
                TimelyLog::Histogram,
                move |data, packer, session| {
                    for ((datum, ()), time, diff) in data.iter() {
                        let data = packer.pack_slice(&[
                            Datum::UInt64(u64::cast_from(datum.operator)),
                            Datum::UInt64(u64::cast_from(worker_id)),
                            Datum::UInt64(datum.duration_pow),
                        ]);
                        session.give((data, time, diff));
                    }
                },
            );

        let logs = {
            use TimelyLog::*;
            [
                (Operates, operates),
                (Channels, channels),
                (Elapsed, elapsed),
                (Histogram, histogram),
                (Addresses, addresses),
                (Parks, parks),
                (MessagesSent, messages_sent),
                (MessagesReceived, messages_received),
                (BatchesSent, batches_sent),
                (BatchesReceived, batches_received),
            ]
        };

        // Build the output arrangements.
        let mut collections = BTreeMap::new();
        for (variant, collection) in logs {
            let variant = LogVariant::Timely(variant);
            if config.index_logs.contains_key(&variant) {
                // Extract types to make rustfmt happy.
                type Batcher<K, V, T, R> = Col2ValBatcher<K, V, T, R>;
                type Builder<T, R> = RowRowBuilder<T, R>;
                let trace = collection
                    .mz_arrange_core::<
                        _,
                        batcher::Chunker<_>,
                        Batcher<_, _, _, _>,
                        Builder<_, _>,
                        RowRowSpine<_, _>,
                    >(
                        ExchangeCore::<ColumnBuilder<_>, _>::new_core(
                            columnar_exchange::<mz_repr::Row, mz_repr::Row, Timestamp, Diff>,
                        ),
                        &format!("Arrange {variant:?}"),
                    )
                    .trace;
                let combined_token: Rc<dyn std::any::Any> =
                    Rc::new((Rc::clone(&token), Rc::clone(&storage_token)));
                let collection = LogCollection {
                    trace,
                    token: combined_token,
                };
                collections.insert(variant, collection);
            }
        }

        Return { collections }
    })
}

/// State maintained by the demux operator.
#[derive(Default)]
struct DemuxState {
    /// Information about live operators, indexed by operator ID.
    operators: BTreeMap<usize, OperatesEvent>,
    /// Maps dataflow IDs to channels in the dataflow.
    dataflow_channels: BTreeMap<usize, Vec<ChannelsEvent>>,
    /// Information about the last requested park.
    last_park: Option<Park>,
    /// Maps channel IDs to boxed slices counting the messages sent to each target worker.
    messages_sent: BTreeMap<usize, Box<[MessageCount]>>,
    /// Maps channel IDs to boxed slices counting the messages received from each source worker.
    messages_received: BTreeMap<usize, Box<[MessageCount]>>,
    /// Stores for scheduled operators the time when they were scheduled.
    schedule_starts: Vec<(usize, Duration)>,
    /// Maps operator IDs to a vector recording the (count, elapsed_ns) values in each histogram
    /// bucket.
    schedules_data: BTreeMap<usize, Vec<(isize, Diff)>>,
}

struct Park {
    /// Time when the park occurred.
    time: Duration,
    /// Requested park time.
    requested: Option<Duration>,
}

/// Organize message counts into number of batches and records.
#[derive(Default, Copy, Clone, Debug)]
struct MessageCount {
    /// The number of batches sent across a channel.
    batches: i64,
    /// The number of records sent across a channel.
    records: Diff,
}

/// Bundled output buffers used by the demux operator.
//
// We use tuples rather than dedicated `*Datum` structs for `operates` and `addresses` to avoid
// having to manually implement `Columnation`. If `Columnation` could be `#[derive]`ed, that
// wouldn't be an issue.
struct DemuxOutput<'a, 'b> {
    operates: OutputSessionVec<'a, 'b, Update<(usize, String)>>,
    channels: OutputSessionColumnar<'a, 'b, Update<(ChannelDatum, ())>>,
    addresses: OutputSessionVec<'a, 'b, Update<(usize, Vec<usize>)>>,
    parks: OutputSessionVec<'a, 'b, Update<(ParkDatum, ())>>,
    batches_sent: OutputSessionVec<'a, 'b, Update<(MessageDatum, ())>>,
    batches_received: OutputSessionVec<'a, 'b, Update<(MessageDatum, ())>>,
    messages_sent: OutputSessionVec<'a, 'b, Update<(MessageDatum, ())>>,
    messages_received: OutputSessionVec<'a, 'b, Update<(MessageDatum, ())>>,
    schedules_duration: OutputSessionVec<'a, 'b, Update<(usize, ())>>,
    schedules_histogram: OutputSessionVec<'a, 'b, Update<(ScheduleHistogramDatum, ())>>,
}

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Columnar)]
struct ChannelDatum {
    id: usize,
    source: (usize, usize),
    target: (usize, usize),
    typ: String,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct ParkDatum {
    duration_pow: u64,
    requested_pow: Option<u64>,
}

impl Columnation for ParkDatum {
    type InnerRegion = CopyRegion<Self>;
}

#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct MessageDatum {
    channel: usize,
    worker: usize,
}

impl Columnation for MessageDatum {
    type InnerRegion = CopyRegion<Self>;
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct ScheduleHistogramDatum {
    operator: usize,
    duration_pow: u64,
}

impl Columnation for ScheduleHistogramDatum {
    type InnerRegion = CopyRegion<Self>;
}

/// 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 buffers.
    output: &'a mut DemuxOutput<'b, 'c>,
    /// The logging interval specifying the time granularity for the updates.
    logging_interval_ms: u128,
    /// The number of timely workers.
    peers: usize,
    /// 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 timely event.
    fn handle(&mut self, event: TimelyEvent) {
        use TimelyEvent::*;

        match event {
            Operates(e) => self.handle_operates(e),
            Channels(e) => self.handle_channels(e),
            Shutdown(e) => self.handle_shutdown(e),
            Park(e) => self.handle_park(e),
            Messages(e) => self.handle_messages(e),
            Schedule(e) => self.handle_schedule(e),
            _ => (),
        }
    }

    fn handle_operates(&mut self, event: OperatesEvent) {
        let ts = self.ts();
        let datum = (event.id, event.name.clone());
        self.output.operates.give((datum, ts, Diff::ONE));

        let datum = (event.id, event.addr.clone());
        self.output.addresses.give((datum, ts, Diff::ONE));

        self.state.operators.insert(event.id, event);
    }

    fn handle_channels(&mut self, event: ChannelsEvent) {
        let ts = self.ts();
        let datum = ChannelDatumReference {
            id: event.id,
            source: event.source,
            target: event.target,
            typ: &event.typ,
        };
        self.output.channels.give(((datum, ()), ts, Diff::ONE));

        let datum = (event.id, event.scope_addr.clone());
        self.output.addresses.give((datum, ts, Diff::ONE));

        let dataflow_index = event.scope_addr[0];
        self.state
            .dataflow_channels
            .entry(dataflow_index)
            .or_default()
            .push(event);
    }

    fn handle_shutdown(&mut self, event: ShutdownEvent) {
        // Dropped operators should result in a negative record for
        // the `operates` collection, cancelling out the initial
        // operator announcement.
        // Remove logging for this operator.

        let Some(operator) = self.state.operators.remove(&event.id) else {
            error!(operator_id = ?event.id, "missing operator entry at time of shutdown");
            return;
        };

        // Retract operator information.
        let ts = self.ts();
        let datum = (operator.id, operator.name);
        self.output.operates.give((datum, ts, Diff::MINUS_ONE));

        // Retract schedules information for the operator
        if let Some(schedules) = self.state.schedules_data.remove(&event.id) {
            for (bucket, (count, elapsed_ns)) in IntoIterator::into_iter(schedules)
                .enumerate()
                .filter(|(_, (count, _))| *count != 0)
            {
                self.output
                    .schedules_duration
                    .give(((event.id, ()), ts, Diff::from(-elapsed_ns)));

                let datum = ScheduleHistogramDatum {
                    operator: event.id,
                    duration_pow: 1 << bucket,
                };
                let diff = Diff::cast_from(-count);
                self.output
                    .schedules_histogram
                    .give(((datum, ()), ts, diff));
            }
        }

        if operator.addr.len() == 1 {
            let dataflow_index = operator.addr[0];
            self.handle_dataflow_shutdown(dataflow_index);
        }

        let datum = (operator.id, operator.addr);
        self.output.addresses.give((datum, ts, Diff::MINUS_ONE));
    }

    fn handle_dataflow_shutdown(&mut self, dataflow_index: usize) {
        // Notify compute logging about the shutdown.
        self.shared_state.compute_logger.as_ref().map(|logger| {
            logger.log(&(ComputeEvent::DataflowShutdown(DataflowShutdown { dataflow_index })))
        });

        // When a dataflow shuts down, we need to retract all its channels.
        let Some(channels) = self.state.dataflow_channels.remove(&dataflow_index) else {
            return;
        };

        let ts = self.ts();
        for channel in channels {
            // Retract channel description.
            let datum = ChannelDatumReference {
                id: channel.id,
                source: channel.source,
                target: channel.target,
                typ: &channel.typ,
            };
            self.output
                .channels
                .give(((datum, ()), ts, Diff::MINUS_ONE));

            let datum = (channel.id, channel.scope_addr);
            self.output.addresses.give((datum, ts, Diff::MINUS_ONE));

            // Retract messages logged for this channel.
            if let Some(sent) = self.state.messages_sent.remove(&channel.id) {
                for (target_worker, count) in sent.iter().enumerate() {
                    let datum = MessageDatum {
                        channel: channel.id,
                        worker: target_worker,
                    };
                    self.output
                        .messages_sent
                        .give(((datum, ()), ts, Diff::from(-count.records)));
                    self.output
                        .batches_sent
                        .give(((datum, ()), ts, Diff::from(-count.batches)));
                }
            }
            if let Some(received) = self.state.messages_received.remove(&channel.id) {
                for (source_worker, count) in received.iter().enumerate() {
                    let datum = MessageDatum {
                        channel: channel.id,
                        worker: source_worker,
                    };
                    self.output.messages_received.give((
                        (datum, ()),
                        ts,
                        Diff::from(-count.records),
                    ));
                    self.output.batches_received.give((
                        (datum, ()),
                        ts,
                        Diff::from(-count.batches),
                    ));
                }
            }
        }
    }

    fn handle_park(&mut self, event: ParkEvent) {
        match event {
            ParkEvent::Park(requested) => {
                let park = Park {
                    time: self.time,
                    requested,
                };
                let existing = self.state.last_park.replace(park);
                if existing.is_some() {
                    error!("park without a succeeding unpark");
                }
            }
            ParkEvent::Unpark => {
                let Some(park) = self.state.last_park.take() else {
                    error!("unpark without a preceding park");
                    return;
                };

                let duration_ns = self.time.saturating_sub(park.time).as_nanos();
                let duration_pow =
                    u64::try_from(duration_ns.next_power_of_two()).expect("must fit");
                let requested_pow = park
                    .requested
                    .map(|r| u64::try_from(r.as_nanos().next_power_of_two()).expect("must fit"));

                let ts = self.ts();
                let datum = ParkDatum {
                    duration_pow,
                    requested_pow,
                };
                self.output.parks.give(((datum, ()), ts, Diff::ONE));
            }
        }
    }

    fn handle_messages(&mut self, event: MessagesEvent) {
        let ts = self.ts();
        let count = Diff::from(event.record_count);

        if event.is_send {
            let datum = MessageDatum {
                channel: event.channel,
                worker: event.target,
            };
            self.output.messages_sent.give(((datum, ()), ts, count));
            self.output.batches_sent.give(((datum, ()), ts, Diff::ONE));

            let sent_counts = self
                .state
                .messages_sent
                .entry(event.channel)
                .or_insert_with(|| vec![Default::default(); self.peers].into_boxed_slice());
            sent_counts[event.target].records += count;
            sent_counts[event.target].batches += 1;
        } else {
            let datum = MessageDatum {
                channel: event.channel,
                worker: event.source,
            };
            self.output.messages_received.give(((datum, ()), ts, count));
            self.output
                .batches_received
                .give(((datum, ()), ts, Diff::ONE));

            let received_counts = self
                .state
                .messages_received
                .entry(event.channel)
                .or_insert_with(|| vec![Default::default(); self.peers].into_boxed_slice());
            received_counts[event.source].records += count;
            received_counts[event.source].batches += 1;
        }
    }

    fn handle_schedule(&mut self, event: ScheduleEvent) {
        match event.start_stop {
            timely::logging::StartStop::Start => {
                self.state.schedule_starts.push((event.id, self.time));
            }
            timely::logging::StartStop::Stop => {
                let Some((old_id, start_time)) = self.state.schedule_starts.pop() else {
                    error!(operator_id = ?event.id, "schedule stop without preceding start");
                    return;
                };

                if old_id != event.id {
                    error!(start_id = ?old_id, stop_id = ?event.id, "schedule stop without preceding start");
                    return;
                }

                let elapsed_ns = self.time.saturating_sub(start_time).as_nanos();
                let elapsed_i64 = i64::try_from(elapsed_ns).expect("must fit");
                let elapsed_diff = Diff::from(elapsed_i64);
                let elapsed_pow = u64::try_from(elapsed_ns.next_power_of_two()).expect("must fit");

                let ts = self.ts();
                let datum = event.id;
                self.output
                    .schedules_duration
                    .give(((datum, ()), ts, elapsed_diff));

                let datum = ScheduleHistogramDatum {
                    operator: event.id,
                    duration_pow: elapsed_pow,
                };
                self.output
                    .schedules_histogram
                    .give(((datum, ()), ts, Diff::ONE));

                // Record count and elapsed time for later retraction.
                let index = usize::cast_from(elapsed_pow.trailing_zeros());
                let data = self.state.schedules_data.entry(event.id).or_default();
                grow_vec(data, index);
                let (count, duration) = &mut data[index];
                *count += 1;
                *duration += elapsed_diff;
            }
        }
    }
}

/// Grow the given vector so it fits the given index.
///
/// This does nothing if the vector is already large enough.
fn grow_vec<T>(vec: &mut Vec<T>, index: usize)
where
    T: Clone + Default,
{
    if vec.len() <= index {
        vec.resize(index + 1, Default::default());
    }
}