mz_compute/logging/

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

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

use differential_dataflow::consolidation::ConsolidatingContainerBuilder;
use differential_dataflow::logging::{
    BatchEvent, BatcherEvent, DifferentialEvent, DropEvent, MergeEvent, TraceShare,
};
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::channels::pact::{ExchangeCore, Pipeline};
use timely::dataflow::operators::InputCapability;
use timely::dataflow::operators::generic::builder_rc::OperatorBuilder;
use timely::dataflow::operators::generic::operator::empty;
use timely::dataflow::operators::generic::{OutputBuilder, Session};
use timely::dataflow::{Scope, StreamVec};

use crate::extensions::arrange::MzArrangeCore;
use crate::logging::compute::{ArrangementHeapSizeOperatorDrop, ComputeEvent};
use crate::logging::{
    DifferentialLog, EventQueue, LogCollection, LogVariant, SharedLoggingState,
    consolidate_and_pack,
};
use crate::typedefs::{KeyBatcher, 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 differential logs.
///
/// Params
/// * `scope`: The Timely scope hosting the log analysis dataflow.
/// * `config`: Logging configuration
/// * `event_queue`: The source to read log events from.
/// * `shared_state`: Shared state across all logging dataflow fragments.
pub(super) fn construct(
    scope: Scope<'_, Timestamp>,
    config: &mz_compute_client::logging::LoggingConfig,
    event_queue: EventQueue<Vec<(Duration, DifferentialEvent)>>,
    shared_state: Rc<RefCell<SharedLoggingState>>,
) -> Return {
    let logging_interval_ms = std::cmp::max(1, config.interval.as_millis());

    scope.scoped("differential logging", move |scope| {
        let enable_logging = config.enable_logging;
        let (logs, token) = if enable_logging {
            event_queue.links.mz_replay(
                scope,
                "differential 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("Differential Logging Demux".to_string(), scope.clone());
        let mut input = demux.new_input(logs, Pipeline);
        let (batches_out, batches) = demux.new_output();
        let (records_out, records) = demux.new_output();
        let (sharing_out, sharing) = demux.new_output();
        let (batcher_records_out, batcher_records) = demux.new_output();
        let (batcher_size_out, batcher_size) = demux.new_output();
        let (batcher_capacity_out, batcher_capacity) = demux.new_output();
        let (batcher_allocations_out, batcher_allocations) = demux.new_output();

        let mut batches_out = OutputBuilder::from(batches_out);
        let mut records_out = OutputBuilder::from(records_out);
        let mut sharing_out = OutputBuilder::from(sharing_out);
        let mut batcher_records_out = OutputBuilder::from(batcher_records_out);
        let mut batcher_size_out = OutputBuilder::from(batcher_size_out);
        let mut batcher_capacity_out = OutputBuilder::from(batcher_capacity_out);
        let mut batcher_allocations_out = OutputBuilder::from(batcher_allocations_out);

        let mut demux_state = Default::default();
        demux.build(move |_capability| {
            move |_frontiers| {
                let mut batches = batches_out.activate();
                let mut records = records_out.activate();
                let mut sharing = sharing_out.activate();
                let mut batcher_records = batcher_records_out.activate();
                let mut batcher_size = batcher_size_out.activate();
                let mut batcher_capacity = batcher_capacity_out.activate();
                let mut batcher_allocations = batcher_allocations_out.activate();

                input.for_each_time(|cap, data| {
                    let mut output_buffers = DemuxOutput {
                        batches: batches.session_with_builder(&cap),
                        records: records.session_with_builder(&cap),
                        sharing: sharing.session_with_builder(&cap),
                        batcher_records: batcher_records.session_with_builder(&cap),
                        batcher_size: batcher_size.session_with_builder(&cap),
                        batcher_capacity: batcher_capacity.session_with_builder(&cap),
                        batcher_allocations: batcher_allocations.session_with_builder(&cap),
                    };

                    for (time, event) in data.flat_map(|data: &mut Vec<_>| data.drain(..)) {
                        DemuxHandler {
                            state: &mut demux_state,
                            output: &mut output_buffers,
                            logging_interval_ms,
                            time,
                            shared_state: &mut shared_state.borrow_mut(),
                        }
                        .handle(event);
                    }
                });
            }
        });

        // We're lucky and the differential logs all have the same stream format, so just implement
        // the call once.
        fn stream_to_collection<'scope>(
            input: StreamVec<'scope, Timestamp, ((usize, ()), Timestamp, Diff)>,
            log: DifferentialLog,
            worker_id: usize,
        ) -> timely::dataflow::Stream<
            'scope,
            Timestamp,
            mz_timely_util::columnar::Column<((mz_repr::Row, mz_repr::Row), Timestamp, Diff)>,
        > {
            consolidate_and_pack::<
                ColumnationChunker<_>,
                KeyBatcher<_, _, _>,
                ColumnBuilder<_>,
                _,
                _,
                _,
            >(input, log, move |data, packer, session| {
                for ((op, ()), time, diff) in data.iter() {
                    let data = packer.pack_slice(&[
                        Datum::UInt64(u64::cast_from(*op)),
                        Datum::UInt64(u64::cast_from(worker_id)),
                    ]);
                    session.give((data, *time, *diff))
                }
            })
        }
        let worker_id = scope.index();

        // Encode the contents of each logging stream into its expected `Row` format.
        let arrangement_batches = stream_to_collection(batches, ArrangementBatches, worker_id);
        let arrangement_records = stream_to_collection(records, ArrangementRecords, worker_id);
        let sharing = stream_to_collection(sharing, Sharing, worker_id);
        let batcher_records = stream_to_collection(batcher_records, BatcherRecords, worker_id);
        let batcher_size = stream_to_collection(batcher_size, BatcherSize, worker_id);
        let batcher_capacity = stream_to_collection(batcher_capacity, BatcherCapacity, worker_id);
        let batcher_allocations =
            stream_to_collection(batcher_allocations, BatcherAllocations, worker_id);

        use DifferentialLog::*;
        let logs = [
            (ArrangementBatches, arrangement_batches),
            (ArrangementRecords, arrangement_records),
            (Sharing, sharing),
            (BatcherRecords, batcher_records),
            (BatcherSize, batcher_size),
            (BatcherCapacity, batcher_capacity),
            (BatcherAllocations, batcher_allocations),
        ];

        // Build the output arrangements.
        let mut collections = BTreeMap::new();
        for (variant, collection) in logs {
            let variant = LogVariant::Differential(variant);
            if config.index_logs.contains_key(&variant) {
                let exchange = ExchangeCore::<ColumnBuilder<_>, _>::new_core(
                    columnar_exchange::<mz_repr::Row, mz_repr::Row, Timestamp, mz_repr::Diff>,
                );
                let trace = collection
                    .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 }
    })
}

type OutputSession<'a, 'b, D> = Session<
    'a,
    'b,
    Timestamp,
    ConsolidatingContainerBuilder<Vec<(D, Timestamp, Diff)>>,
    InputCapability<Timestamp>,
>;

/// Bundled output buffers used by the demux operator.
struct DemuxOutput<'a, 'b> {
    batches: OutputSession<'a, 'b, (usize, ())>,
    records: OutputSession<'a, 'b, (usize, ())>,
    sharing: OutputSession<'a, 'b, (usize, ())>,
    batcher_records: OutputSession<'a, 'b, (usize, ())>,
    batcher_size: OutputSession<'a, 'b, (usize, ())>,
    batcher_capacity: OutputSession<'a, 'b, (usize, ())>,
    batcher_allocations: OutputSession<'a, 'b, (usize, ())>,
}

/// State maintained by the demux operator.
#[derive(Default)]
struct DemuxState {
    /// Arrangement trace sharing
    sharing: BTreeMap<usize, usize>,
}

/// Event handler of the demux operator.
struct DemuxHandler<'a, 'b, 'c> {
    /// State kept by the demux operator
    state: &'a mut DemuxState,
    /// Demux output buffers.
    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,
    /// State shared across log receivers.
    shared_state: &'a mut SharedLoggingState,
}

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 differential event.
    fn handle(&mut self, event: DifferentialEvent) {
        use DifferentialEvent::*;

        match event {
            Batch(e) => self.handle_batch(e),
            Merge(e) => self.handle_merge(e),
            Drop(e) => self.handle_drop(e),
            TraceShare(e) => self.handle_trace_share(e),
            Batcher(e) => self.handle_batcher_event(e),
            MergeShortfall(_) => (),
        }
    }

    fn handle_batch(&mut self, event: BatchEvent) {
        let ts = self.ts();
        let operator_id = event.operator;
        self.output.batches.give(((operator_id, ()), ts, Diff::ONE));

        let diff = Diff::try_from(event.length).expect("must fit");
        self.output.records.give(((operator_id, ()), ts, diff));
        self.notify_arrangement_size(operator_id);
    }

    fn handle_merge(&mut self, event: MergeEvent) {
        let Some(done) = event.complete else { return };

        let ts = self.ts();
        let operator_id = event.operator;
        self.output
            .batches
            .give(((operator_id, ()), ts, Diff::MINUS_ONE));

        let diff = Diff::try_from(done).expect("must fit")
            - Diff::try_from(event.length1 + event.length2).expect("must fit");
        if diff != Diff::ZERO {
            self.output.records.give(((operator_id, ()), ts, diff));
        }
        self.notify_arrangement_size(operator_id);
    }

    fn handle_drop(&mut self, event: DropEvent) {
        let ts = self.ts();
        let operator_id = event.operator;
        self.output
            .batches
            .give(((operator_id, ()), ts, Diff::MINUS_ONE));

        let diff = -Diff::try_from(event.length).expect("must fit");
        if diff != Diff::ZERO {
            self.output.records.give(((operator_id, ()), ts, diff));
        }
        self.notify_arrangement_size(operator_id);
    }

    fn handle_trace_share(&mut self, event: TraceShare) {
        let ts = self.ts();
        let operator_id = event.operator;
        let diff = Diff::cast_from(event.diff);
        debug_assert_ne!(diff, Diff::ZERO);
        self.output.sharing.give(((operator_id, ()), ts, diff));

        let sharing = self.state.sharing.entry(operator_id).or_default();
        *sharing = (Diff::try_from(*sharing).expect("must fit") + diff)
            .into_inner()
            .try_into()
            .expect("under/overflow");
        if *sharing == 0 {
            self.state.sharing.remove(&operator_id);
            self.shared_state.compute_logger.as_ref().map(|logger| {
                logger.log(&ComputeEvent::ArrangementHeapSizeOperatorDrop(
                    ArrangementHeapSizeOperatorDrop { operator_id },
                ))
            });
        }
    }

    fn handle_batcher_event(&mut self, event: BatcherEvent) {
        let ts = self.ts();
        let operator_id = event.operator;
        let records_diff = Diff::cast_from(event.records_diff);
        let size_diff = Diff::cast_from(event.size_diff);
        let capacity_diff = Diff::cast_from(event.capacity_diff);
        let allocations_diff = Diff::cast_from(event.allocations_diff);
        self.output
            .batcher_records
            .give(((operator_id, ()), ts, records_diff));
        self.output
            .batcher_size
            .give(((operator_id, ()), ts, size_diff));
        self.output
            .batcher_capacity
            .give(((operator_id, ()), ts, capacity_diff));
        self.output
            .batcher_allocations
            .give(((operator_id, ()), ts, allocations_diff));
    }

    fn notify_arrangement_size(&self, operator: usize) {
        // While every arrangement should have a corresponding arrangement size operator,
        // we have no guarantee that it already/still exists. Otherwise we could print a warning
        // here, but it's difficult to implement without maintaining state for a longer period than
        // while the arrangement actually exists.
        if let Some(activator) = self.shared_state.arrangement_size_activators.get(&operator) {
            activator.activate();
        }
    }
}