compute*: log dataflow shutdown durations

This commit adds a new compute logging variant, `ShutdownDuration`, that
provides an histogram over the time elapsed between dropping a dataflow
in the `ComputeState` and shutting down all dataflow operators.

The new logging variant requires some cross-talk between the timely and
the compute logging dataflow, since only the timely logging dataflow
gets to know when all operators of a dataflow have shut down, while the
compute logging dataflow knows when a dataflow was dropped by the
`ComputeState`. The timely dataflow communicates shutdown info to the
compute dataflow through the compute logger and a new
`ComputeEvent::DataflowShutdown` event.

Author: teskje

src/compute-client/src/logging.proto

@@ -55,6 +55,7 @@ message ProtoComputeLog {
         google.protobuf.Empty arrangement_heap_size = 8;
         google.protobuf.Empty arrangement_heap_capacity = 9;
         google.protobuf.Empty arrangement_heap_allocations = 10;
+        google.protobuf.Empty shutdown_duration = 11;
     }
 }
 message ProtoLogVariant {

src/compute-client/src/logging.rs

@@ -223,6 +223,7 @@ pub enum ComputeLog {
     ArrangementHeapSize,
     ArrangementHeapCapacity,
     ArrangementHeapAllocations,
+    ShutdownDuration,
 }
 
 impl RustType<ProtoComputeLog> for ComputeLog {
@@ -240,6 +241,7 @@ impl RustType<ProtoComputeLog> for ComputeLog {
                 ComputeLog::ArrangementHeapSize => ArrangementHeapSize(()),
                 ComputeLog::ArrangementHeapCapacity => ArrangementHeapCapacity(()),
                 ComputeLog::ArrangementHeapAllocations => ArrangementHeapAllocations(()),
+                ComputeLog::ShutdownDuration => ShutdownDuration(()),
             }),
         }
     }
@@ -257,6 +259,7 @@ impl RustType<ProtoComputeLog> for ComputeLog {
             Some(ArrangementHeapSize(())) => Ok(ComputeLog::ArrangementHeapSize),
             Some(ArrangementHeapCapacity(())) => Ok(ComputeLog::ArrangementHeapCapacity),
             Some(ArrangementHeapAllocations(())) => Ok(ComputeLog::ArrangementHeapAllocations),
+            Some(ShutdownDuration(())) => Ok(ComputeLog::ShutdownDuration),
             None => Err(TryFromProtoError::missing_field("ProtoComputeLog::kind")),
         }
     }
@@ -419,6 +422,10 @@ impl LogVariant {
             LogVariant::Compute(ComputeLog::PeekDuration) => RelationDesc::empty()
                 .with_column("worker_id", ScalarType::UInt64.nullable(false))
                 .with_column("duration_ns", ScalarType::UInt64.nullable(false)),
+
+            LogVariant::Compute(ComputeLog::ShutdownDuration) => RelationDesc::empty()
+                .with_column("worker_id", ScalarType::UInt64.nullable(false))
+                .with_column("duration_ns", ScalarType::UInt64.nullable(false)),
         }
     }
 
@@ -471,6 +478,7 @@ impl LogVariant {
             LogVariant::Compute(ComputeLog::FrontierDelay) => vec![],
             LogVariant::Compute(ComputeLog::PeekCurrent) => vec![],
             LogVariant::Compute(ComputeLog::PeekDuration) => vec![],
+            LogVariant::Compute(ComputeLog::ShutdownDuration) => vec![],
         }
     }
 }

src/compute/src/logging/compute.rs

@@ -11,7 +11,7 @@
 
 use std::any::Any;
 use std::cell::RefCell;
-use std::collections::{BTreeMap, VecDeque};
+use std::collections::{BTreeMap, BTreeSet, VecDeque};
 use std::rc::Rc;
 use std::time::Duration;
 
@@ -113,6 +113,11 @@ pub enum ComputeEvent {
         /// Operator index
         operator: usize,
     },
+    /// All operators of a dataflow have shut down.
+    DataflowShutdown {
+        /// Timely worker index of the dataflow.
+        dataflow_index: usize,
+    },
 }
 
 /// A logged peek event.
@@ -176,6 +181,7 @@ pub(super) fn construct<A: Allocate + 'static>(
         let (mut frontier_delay_out, frontier_delay) = 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) =
@@ -192,6 +198,7 @@ pub(super) fn construct<A: Allocate + 'static>(
                 let mut frontier_delay = frontier_delay_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();
@@ -207,6 +214,7 @@ pub(super) fn construct<A: Allocate + 'static>(
                         frontier_delay: frontier_delay.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),
@@ -280,7 +288,13 @@ pub(super) fn construct<A: Allocate + 'static>(
         let peek_duration = peek_duration.as_collection().map(move |bucket| {
             Row::pack_slice(&[
                 Datum::UInt64(u64::cast_from(worker_id)),
-                Datum::UInt64(bucket.try_into().expect("pow too big")),
+                Datum::UInt64(bucket.try_into().expect("bucket too big")),
+            ])
+        });
+        let shutdown_duration = shutdown_duration.as_collection().map(move |bucket| {
+            Row::pack_slice(&[
+                Datum::UInt64(u64::cast_from(worker_id)),
+                Datum::UInt64(bucket.try_into().expect("bucket too big")),
             ])
         });
 
@@ -312,6 +326,7 @@ pub(super) fn construct<A: Allocate + 'static>(
             (FrontierDelay, frontier_delay),
             (PeekCurrent, peek_current),
             (PeekDuration, peek_duration),
+            (ShutdownDuration, shutdown_duration),
             (ArrangementHeapSize, arrangement_heap_size),
             (ArrangementHeapCapacity, arrangement_heap_capacity),
             (ArrangementHeapAllocations, arrangement_heap_allocations),
@@ -361,6 +376,12 @@ struct DemuxState<A: Allocate> {
     export_dataflows: BTreeMap<GlobalId, usize>,
     /// Maps dataflow exports to their imports and frontier delay tracking state.
     export_imports: BTreeMap<GlobalId, BTreeMap<GlobalId, FrontierDelayState>>,
+    /// Maps live dataflows to counts of their exports.
+    dataflow_export_counts: BTreeMap<usize, u32>,
+    /// Maps dropped dataflows to their drop time (in ns).
+    dataflow_drop_times: BTreeMap<usize, u128>,
+    /// Contains dataflows that have shut down but not yet been dropped.
+    shutdown_dataflows: BTreeSet<usize>,
     /// Maps pending peeks to their installation time (in ns).
     peek_stash: BTreeMap<Uuid, u128>,
     /// Arrangement size stash
@@ -373,6 +394,9 @@ impl<A: Allocate> DemuxState<A> {
             worker,
             export_dataflows: Default::default(),
             export_imports: Default::default(),
+            dataflow_export_counts: Default::default(),
+            dataflow_drop_times: Default::default(),
+            shutdown_dataflows: Default::default(),
             peek_stash: Default::default(),
             arrangement_size: Default::default(),
         }
@@ -403,6 +427,7 @@ struct DemuxOutput<'a> {
     frontier_delay: OutputSession<'a, FrontierDelayDatum>,
     peek: OutputSession<'a, Peek>,
     peek_duration: OutputSession<'a, u128>,
+    shutdown_duration: OutputSession<'a, u128>,
     arrangement_heap_size: OutputSession<'a, ArrangementHeapDatum>,
     arrangement_heap_capacity: OutputSession<'a, ArrangementHeapDatum>,
     arrangement_heap_allocations: OutputSession<'a, ArrangementHeapDatum>,
@@ -514,6 +539,7 @@ impl<A: Allocate> DemuxHandler<'_, '_, A> {
             ArrangementHeapSizeOperatorDrop { operator } => {
                 self.handle_arrangement_heap_size_operator_dropped(operator)
             }
+            DataflowShutdown { dataflow_index } => self.handle_dataflow_shutdown(dataflow_index),
         }
     }
 
@@ -524,13 +550,30 @@ impl<A: Allocate> DemuxHandler<'_, '_, A> {
 
         self.state.export_dataflows.insert(id, dataflow_id);
         self.state.export_imports.insert(id, BTreeMap::new());
+        *self
+            .state
+            .dataflow_export_counts
+            .entry(dataflow_id)
+            .or_default() += 1;
     }
 
     fn handle_export_dropped(&mut self, id: GlobalId) {
         let ts = self.ts();
         if let Some(dataflow_id) = self.state.export_dataflows.remove(&id) {
             let datum = ExportDatum { id, dataflow_id };
             self.output.export.give((datum, ts, -1));
+
+            match self.state.dataflow_export_counts.get_mut(&dataflow_id) {
+                entry @ Some(0) | entry @ None => {
+                    error!(
+                        export = ?id,
+                        dataflow = ?dataflow_id,
+                        "invalid dataflow_export_counts entry at time of export drop: {entry:?}",
+                    );
+                }
+                Some(1) => self.handle_dataflow_dropped(dataflow_id),
+                Some(count) => *count -= 1,
+            }
         } else {
             error!(
                 export = ?id,
@@ -564,6 +607,41 @@ impl<A: Allocate> DemuxHandler<'_, '_, A> {
         }
     }
 
+    fn handle_dataflow_dropped(&mut self, id: usize) {
+        self.state.dataflow_export_counts.remove(&id);
+
+        if self.state.shutdown_dataflows.remove(&id) {
+            // Dataflow has already shut down before it was dropped.
+            self.output.shutdown_duration.give((0, self.ts(), 1));
+        } else {
+            // Dataflow has not yet shut down.
+            let existing = self
+                .state
+                .dataflow_drop_times
+                .insert(id, self.time.as_nanos());
+            if existing.is_some() {
+                error!(dataflow = ?id, "dataflow already dropped");
+            }
+        }
+    }
+
+    fn handle_dataflow_shutdown(&mut self, id: usize) {
+        if let Some(start) = self.state.dataflow_drop_times.remove(&id) {
+            // Dataflow has alredy been dropped.
+            let elapsed_ns = self.time.as_nanos() - start;
+            let elapsed_pow = elapsed_ns.next_power_of_two();
+            self.output
+                .shutdown_duration
+                .give((elapsed_pow, self.ts(), 1));
+        } else {
+            // Dataflow has not yet been dropped.
+            let was_new = self.state.shutdown_dataflows.insert(id);
+            if !was_new {
+                error!(dataflow = ?id, "dataflow already shutdown");
+            }
+        }
+    }
+
     fn handle_export_dependency(&mut self, export_id: GlobalId, import_id: GlobalId) {
         let ts = self.ts();
         let datum = DependencyDatum {

src/compute/src/logging/initialize.rs

@@ -94,6 +94,7 @@ impl<A: Allocate + 'static> LoggingContext<'_, A> {
             self.worker,
             self.config,
             self.t_event_queue.clone(),
+            Rc::clone(&self.shared_state),
         ));
         traces.extend(super::reachability::construct(
             self.worker,
@@ -131,21 +132,18 @@ impl<A: Allocate + 'static> LoggingContext<'_, A> {
     }
 
     fn register_loggers(&self) {
-        self.worker
-            .log_register()
-            .insert_logger("timely", self.simple_logger(self.t_event_queue.clone()));
-        self.worker
-            .log_register()
-            .insert_logger("timely/reachability", self.reachability_logger());
-        self.worker.log_register().insert_logger(
-            "differential/arrange",
-            self.simple_logger(self.d_event_queue.clone()),
-        );
-        let compute_logger = self.simple_logger(self.c_event_queue.clone());
-        self.worker
-            .log_register()
-            .insert_logger("materialize/compute", compute_logger.clone());
-        self.shared_state.borrow_mut().compute_logger = Some(compute_logger);
+        let t_logger = self.simple_logger(self.t_event_queue.clone());
+        let r_logger = self.reachability_logger();
+        let d_logger = self.simple_logger(self.d_event_queue.clone());
+        let c_logger = self.simple_logger(self.c_event_queue.clone());
+
+        let mut register = self.worker.log_register();
+        register.insert_logger("timely", t_logger);
+        register.insert_logger("timely/reachability", r_logger);
+        register.insert_logger("differential/arrange", d_logger);
+        register.insert_logger("materialize/compute", c_logger.clone());
+
+        self.shared_state.borrow_mut().compute_logger = Some(c_logger);
     }
 
     fn simple_logger<E: 'static>(&self, event_queue: EventQueue<E>) -> Logger<E> {

src/compute/src/logging/timely.rs

@@ -10,6 +10,7 @@
 //! Logging dataflows for events generated by timely dataflow.
 
 use std::any::Any;
+use std::cell::RefCell;
 use std::collections::BTreeMap;
 use std::rc::Rc;
 use std::time::Duration;
@@ -35,7 +36,8 @@ use timely::logging::{
 use tracing::error;
 
 use crate::extensions::arrange::MzArrange;
-use crate::logging::{EventQueue, LogVariant, TimelyLog};
+use crate::logging::compute::ComputeEvent;
+use crate::logging::{EventQueue, LogVariant, SharedLoggingState, TimelyLog};
 use crate::typedefs::{KeysValsHandle, RowSpine};
 
 /// Constructs the logging dataflow for timely logs.
@@ -50,6 +52,7 @@ pub(super) fn construct<A: Allocate>(
     worker: &mut timely::worker::Worker<A>,
     config: &LoggingConfig,
     event_queue: EventQueue<TimelyEvent>,
+    shared_state: Rc<RefCell<SharedLoggingState>>,
 ) -> BTreeMap<LogVariant, (KeysValsHandle, Rc<dyn Any>)> {
     let logging_interval_ms = std::cmp::max(1, config.interval.as_millis());
     let worker_id = worker.index();
@@ -123,6 +126,7 @@ pub(super) fn construct<A: Allocate>(
 
                         DemuxHandler {
                             state: &mut demux_state,
+                            shared_state: &mut shared_state.borrow_mut(),
                             output: &mut output_buffers,
                             logging_interval_ms,
                             peers,
@@ -407,6 +411,8 @@ impl Columnation for ScheduleHistogramDatum {
 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.
@@ -467,10 +473,10 @@ impl DemuxHandler<'_, '_, '_> {
         let datum = (event.id, event.scope_addr.clone());
         self.output.addresses.give(self.cap, (datum, ts, 1));
 
-        let dataflow_id = event.scope_addr[0];
+        let dataflow_index = event.scope_addr[0];
         self.state
             .dataflow_channels
-            .entry(dataflow_id)
+            .entry(dataflow_index)
             .or_default()
             .push(event);
     }
@@ -514,18 +520,22 @@ impl DemuxHandler<'_, '_, '_> {
         }
 
         if operator.addr.len() == 1 {
-            let dataflow_id = operator.addr[0];
-            self.handle_dataflow_shutdown(dataflow_id);
+            let dataflow_index = operator.addr[0];
+            self.handle_dataflow_shutdown(dataflow_index);
         }
 
         let datum = (operator.id, operator.addr);
         self.output.addresses.give(self.cap, (datum, ts, -1));
     }
 
-    fn handle_dataflow_shutdown(&mut self, dataflow_id: usize) {
-        // When a dataflow shuts down, we need to retract all its channels.
+    fn handle_dataflow_shutdown(&mut self, dataflow_index: usize) {
+        // Notify compute logging about the shutdown.
+        if let Some(logger) = &self.shared_state.compute_logger {
+            logger.log(ComputeEvent::DataflowShutdown { dataflow_index });
+        }
 
-        let Some(channels) = self.state.dataflow_channels.remove(&dataflow_id) else {
+        // When a dataflow shuts down, we need to retract all its channels.
+        let Some(channels) = self.state.dataflow_channels.remove(&dataflow_index) else {
             return;
         };