mz_compute/

logging.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 various subsystems.

pub mod compute;
mod differential;
pub(super) mod initialize;
mod prometheus;
mod reachability;
mod resource_usage;
mod timely;

use std::any::Any;
use std::collections::BTreeMap;
use std::marker::PhantomData;
use std::rc::Rc;
use std::time::{Duration, Instant};

use ::prometheus::IntCounter;
use ::timely::container::{CapacityContainerBuilder, PushInto};
use ::timely::dataflow::Stream;
use ::timely::dataflow::channels::pact::Pipeline;
use ::timely::dataflow::operators::capture::{Event, EventLink, EventPusher};
use ::timely::dataflow::operators::generic::Session;
use ::timely::dataflow::operators::{Capability, CapabilityTrait, InputCapability, Operator};
use ::timely::progress::Stamp;
use ::timely::progress::Timestamp as TimelyTimestamp;
use ::timely::scheduling::Activator;
use ::timely::{Container, ContainerBuilder};
use differential_dataflow::trace::Batcher;
use mz_compute_client::logging::{ComputeLog, DifferentialLog, LogVariant, TimelyLog};
use mz_expr::{MirScalarExpr, permutation_for_arrangement};
use mz_repr::{Datum, Diff, Row, RowPacker, RowRef, Timestamp};
use mz_timely_util::activator::RcActivator;
use mz_timely_util::columnar::builder::ColumnBuilder;
use mz_timely_util::operator::{BatcherNew, consolidate_pact};

use crate::logging::compute::Logger as ComputeLogger;
use crate::typedefs::RowRowAgent;

pub use crate::logging::initialize::initialize;

/// An update of value `D` at a time and with a diff.
pub(super) type Update<D> = (D, Timestamp, Diff);
/// A pusher for containers `C`.
/// An output session for the specified container builder.
pub(super) type OutputSession<'a, 'b, CB> =
    Session<'a, 'b, Timestamp, CB, InputCapability<Timestamp>>;
/// An output session for vector-based containers of updates `D`, using a capacity container builder.
pub(super) type OutputSessionVec<'a, 'b, D> =
    OutputSession<'a, 'b, CapacityContainerBuilder<Vec<D>>>;
/// An output session for columnar containers of updates `D`, using a column builder.
pub(super) type OutputSessionColumnar<'a, 'b, D> = OutputSession<'a, 'b, ColumnBuilder<D>>;

/// Logs events as a timely stream, with progress statements.
struct BatchLogger<C, P>
where
    P: EventPusher<Timestamp, C>,
{
    /// Time in milliseconds of the current expressed capability.
    time_ms: Timestamp,
    /// Pushes events to the logging dataflow.
    event_pusher: P,
    /// Each time is advanced to the strictly next millisecond that is a multiple of this interval.
    /// This means we should be able to perform the same action on timestamp capabilities, and only
    /// flush buffers when this timestamp advances.
    interval_ms: u128,
    /// Counts the records of published batches.
    records_total: IntCounter,
    _marker: PhantomData<C>,
}

impl<C, P> BatchLogger<C, P>
where
    P: EventPusher<Timestamp, C>,
{
    /// Creates a new batch logger.
    fn new(event_pusher: P, interval_ms: u128, records_total: IntCounter) -> Self {
        BatchLogger {
            time_ms: Timestamp::minimum(),
            event_pusher,
            interval_ms,
            records_total,
            _marker: PhantomData,
        }
    }
}

impl<C, P> BatchLogger<C, P>
where
    P: EventPusher<Timestamp, C>,
    C: Container,
{
    /// Publishes a batch of logged events.
    fn publish_batch(&mut self, data: C) {
        let records = u64::try_from(data.record_count()).expect("record count is non-negative");
        self.records_total.inc_by(records);
        // One capability per logged batch: the stamp is a singleton.
        self.event_pusher
            .push(Event::Messages(Stamp::from_elem(self.time_ms), data));
    }

    /// Indicate progress up to `time`, advances the capability.
    ///
    /// Returns `true` if the capability was advanced.
    fn report_progress(&mut self, time: Duration) -> bool {
        let time_ms = ((time.as_millis() / self.interval_ms) + 1) * self.interval_ms;
        let new_time_ms: Timestamp = time_ms.try_into().expect("must fit");
        if self.time_ms < new_time_ms {
            self.event_pusher
                .push(Event::Progress(vec![(new_time_ms, 1), (self.time_ms, -1)]));
            self.time_ms = new_time_ms;
            true
        } else {
            false
        }
    }
}

impl<C, P> Drop for BatchLogger<C, P>
where
    P: EventPusher<Timestamp, C>,
{
    fn drop(&mut self) {
        self.event_pusher
            .push(Event::Progress(vec![(self.time_ms, -1)]));
    }
}

/// Parts to connect a logging dataflows the timely runtime.
///
/// This is just a bundle-type intended to make passing around its contents in the logging
/// initialization code more convenient.
///
/// The `N` type parameter specifies the number of links to create for the event queue. We need
/// separate links for queues that feed from multiple loggers because the `EventLink` type is not
/// multi-producer safe (it is a linked-list, and multiple writers would blindly append, replacing
/// existing new data, and cutting off other writers).
#[derive(Clone)]
struct EventQueue<C, const N: usize = 1> {
    links: [Rc<EventLink<Timestamp, C>>; N],
    activator: RcActivator,
}

impl<C, const N: usize> EventQueue<C, N> {
    fn new(name: &str) -> Self {
        let activator_name = format!("{name}_activator");
        let activate_after = 128;
        Self {
            links: [(); N].map(|_| Rc::new(EventLink::new())),
            activator: RcActivator::new(activator_name, activate_after),
        }
    }
}

/// State shared between different logging dataflow fragments.
#[derive(Default)]
struct SharedLoggingState {
    /// Activators for arrangement heap size operators.
    arrangement_size_activators: BTreeMap<usize, Activator>,
    /// Shared compute logger.
    compute_logger: Option<ComputeLogger>,
}

/// Helper to pack collections of [`Datum`]s into key and value row.
pub(crate) struct PermutedRowPacker {
    key: Vec<usize>,
    value: Vec<usize>,
    key_row: Row,
    value_row: Row,
}

impl PermutedRowPacker {
    /// Construct based on the information within the log variant.
    pub(crate) fn new<V: Into<LogVariant>>(variant: V) -> Self {
        let variant = variant.into();
        let key = variant.index_by();
        let (_, value) = permutation_for_arrangement(
            &key.iter()
                .cloned()
                .map(MirScalarExpr::column)
                .collect::<Vec<_>>(),
            variant.desc().arity(),
        );
        Self {
            key,
            value,
            key_row: Row::default(),
            value_row: Row::default(),
        }
    }

    /// Pack a slice of datums suitable for the key columns in the log variant.
    pub(crate) fn pack_slice(&mut self, datums: &[Datum]) -> (&RowRef, &RowRef) {
        self.pack_by_index(|packer, index| packer.push(datums[index]))
    }

    /// Pack using a callback suitable for the key columns in the log variant.
    pub(crate) fn pack_by_index<F: Fn(&mut RowPacker, usize)>(
        &mut self,
        logic: F,
    ) -> (&RowRef, &RowRef) {
        let mut packer = self.key_row.packer();
        for index in &self.key {
            logic(&mut packer, *index);
        }

        let mut packer = self.value_row.packer();
        for index in &self.value {
            logic(&mut packer, *index);
        }

        (&self.key_row, &self.value_row)
    }
}

/// Downgrade `cap` to the next logging-interval boundary and schedule the operator's next
/// activation there. Returns the time the capability now holds.
///
/// `now` and `start_offset` must be the ones the logging dataflow was constructed with, so that
/// every collection in it reports on the same boundaries. Scheduling off the boundary rather than
/// off a fixed delay keeps the output frontier progressing at the logging rate without drifting
/// from wall-clock elapsed time.
///
/// NOTE: downgrading the capability asserts the collection is complete up to the new time, so an
/// operator that samples less often than the logging interval publishes a stale value rather than
/// withholding it.
pub(super) fn downgrade_to_interval_boundary(
    cap: &mut Capability<Timestamp>,
    activator: &Activator,
    now: Instant,
    start_offset: Duration,
    interval_ms: u128,
) -> Timestamp {
    let elapsed = now.elapsed().as_millis();
    let time_ms: u128 = ((elapsed + start_offset.as_millis()) / interval_ms + 1) * interval_ms;
    let ts: Timestamp = time_ms.try_into().expect("must fit");
    cap.downgrade(&ts);

    let next_boundary_ms = time_ms - start_offset.as_millis();
    let next_activation =
        now + Duration::from_millis(next_boundary_ms.try_into().expect("must fit"));
    activator.activate_after(next_activation.saturating_duration_since(Instant::now()));

    ts
}

/// Emit the difference between two snapshots of a sampled source as updates at `ts`.
///
/// A sampled source reports its whole state on every read, so a changed value has to be expressed
/// as a retraction of the previous one paired with an insertion of the new one. A key absent from
/// `current` is retracted without a replacement, so a source that stops being readable drops out of
/// the collection rather than lingering at its last value.
///
/// `pack` takes the packer as an argument instead of closing over it, because it hands back rows
/// borrowed from it.
pub(super) fn emit_snapshot_diff<K, V, CB, P, F>(
    session: &mut Session<'_, '_, Timestamp, CB, P>,
    packer: &mut PermutedRowPacker,
    prev: &BTreeMap<K, V>,
    current: &BTreeMap<K, V>,
    ts: Timestamp,
    pack: F,
) where
    K: Ord,
    V: PartialEq,
    CB: ContainerBuilder + for<'a> PushInto<((&'a RowRef, &'a RowRef), Timestamp, Diff)>,
    P: CapabilityTrait<Timestamp>,
    F: for<'a> Fn(&'a mut PermutedRowPacker, &K, &V) -> (&'a RowRef, &'a RowRef),
{
    for (key, value) in prev {
        if current.get(key) != Some(value) {
            let row = pack(packer, key, value);
            session.give((row, ts, Diff::MINUS_ONE));
        }
    }
    for (key, value) in current {
        if prev.get(key) != Some(value) {
            let row = pack(packer, key, value);
            session.give((row, ts, Diff::ONE));
        }
    }
}

/// Information about a collection exported from a logging dataflow.
struct LogCollection {
    /// Trace handle providing access to the logged records.
    trace: RowRowAgent<Timestamp, Diff>,
    /// Token that should be dropped to drop this collection.
    token: Rc<dyn Any>,
}

/// A single-purpose function to consolidate and pack updates for log collection.
///
/// The function first consolidates worker-local updates using the [`Pipeline`] pact, then converts
/// the updates into `(Row, Row)` pairs using the provided logic function. It is crucial that the
/// data is not exchanged between workers, as the consolidation would not function as desired
/// otherwise.
pub(super) fn consolidate_and_pack<'scope, B, CB, L, F, C, Ch>(
    input: Stream<'scope, Timestamp, C>,
    log: L,
    mut logic: F,
) -> Stream<'scope, Timestamp, CB::Container>
where
    B: Batcher<C, Time = Timestamp, Output = Vec<Ch>> + BatcherNew + 'static,
    C: Container + Clone + 'static,
    Ch: Clone + 'static,
    CB: ContainerBuilder,
    L: Into<LogVariant>,
    F: FnMut(Ch, &mut PermutedRowPacker, &mut OutputSession<CB>) + 'static,
{
    let log = log.into();
    // TODO: Use something other than the debug representation of the log variant as a name.
    let c_name = &format!("Consolidate {log:?}");
    let u_name = &format!("ToRow {log:?}");
    let mut packer = PermutedRowPacker::new(log);
    let consolidated = consolidate_pact::<B, _, _>(input, Pipeline, c_name, B::new_batcher);
    consolidated.unary::<CB, _, _, _>(Pipeline, u_name, |_, _| {
        move |input, output| {
            input.for_each_stamp(|time, data| {
                let mut session = output.session_with_builder(&time);
                for item in data.flatten().flat_map(|data| data.drain(..)) {
                    logic(item, &mut packer, &mut session);
                }
            });
        }
    })
}