Rudimentary profiler.

- `Circuit::map_node_recursive` - apply a function to all nodes in a
  circuit.

- `Profiler` - ties together several existing components into a
  rudimentary profiler:
  - `TraceMonitor` - to capture circuit topology and dump it in graphviz
    format.
  - `map_node_recursive` - to extract node metadata from the circuit
  - `CPUProfiler` - to measure CPU usage

- Modify `DBSPHandle` to create a `Profiler` isntance on startup.
  `DBSPHandle::enable_cpu_profiler` API to enable CPU profiling in all
  worker threads.  `DBSPHandler::dump_profile` API to extract profiles
  from all worker threads and store them in a directory.

- Add `--profile-path` switch to the Nexmark benchmark to enable
  profiling and dump profile after running each query.

Author: ryzhyk

benches/nexmark/main.rs

@@ -28,6 +28,7 @@ use mimalloc::{AllocStats, MiMalloc};
 use num_format::{Locale, ToFormattedString};
 use size_of::HumanBytes;
 use std::{
+    path::Path,
     sync::mpsc,
     thread::{self, JoinHandle},
     time::{Duration, Instant},
@@ -56,17 +57,29 @@ enum StepCompleted {
 }
 
 fn spawn_dbsp_consumer(
+    query: &str,
+    profile_path: Option<&str>,
     mut dbsp: DBSPHandle,
     step_do_rx: mpsc::Receiver<()>,
     step_done_tx: mpsc::SyncSender<StepCompleted>,
 ) -> JoinHandle<()> {
+    let query = query.to_string();
+    let profile_path = profile_path.map(ToString::to_string);
+
     thread::Builder::new()
         .name("benchmark_consumer".into())
         .spawn(move || {
+            if profile_path.is_some() {
+                dbsp.enable_cpu_profiler().unwrap();
+            }
             while let Ok(()) = step_do_rx.recv() {
                 dbsp.step().unwrap();
                 step_done_tx.send(StepCompleted::Dbsp).unwrap();
             }
+            if let Some(profile_path) = profile_path {
+                dbsp.dump_profile(<String as AsRef<Path>>::as_ref(&profile_path).join(query))
+                    .unwrap();
+            }
         })
         .unwrap()
 }

benches/nexmark/run_queries.rs

@@ -15,7 +15,7 @@ macro_rules! run_queries {
         // so. The DBSP processing happens in its own thread where the resource usage
         // calculation can also happen.
         let (dbsp_step_tx, dbsp_step_rx) = mpsc::sync_channel(1);
-        let dbsp_join_handle = spawn_dbsp_consumer(dbsp, dbsp_step_rx, step_done_tx.clone());
+        let dbsp_join_handle = spawn_dbsp_consumer(stringify!($query), $nexmark_config.profile_path.as_deref(), dbsp, dbsp_step_rx, step_done_tx.clone());
 
         // Start the generator inputting the specified number of batches to the circuit
         // whenever it receives a message.

src/circuit/circuit_builder.rs

@@ -325,6 +325,8 @@ pub trait Node {
     fn metadata(&self, output: &mut OperatorMeta);
 
     fn fixedpoint(&self, scope: Scope) -> bool;
+
+    fn map_nodes_recursive(&self, _f: &mut dyn FnMut(&dyn Node)) {}
 }
 
 /// Id of an operator, guaranteed to be unique within a circuit.
@@ -957,6 +959,14 @@ impl<P> Circuit<P> {
             .collect()
     }
 
+    /// Recursively apply `f` to all nodes in `self` and its children.
+    pub(crate) fn map_nodes_recursive(&self, f: &mut dyn FnMut(&dyn Node)) {
+        for node in self.inner().nodes.iter() {
+            f(node.as_ref());
+            node.map_nodes_recursive(f);
+        }
+    }
+
     /// Deliver `clock_start` notification to all nodes in the circuit.
     pub(super) fn clock_start(&self, scope: Scope) {
         for node in self.inner_mut().nodes.iter_mut() {
@@ -3065,6 +3075,10 @@ where
     fn fixedpoint(&self, scope: Scope) -> bool {
         self.circuit.inner().fixedpoint(scope + 1)
     }
+
+    fn map_nodes_recursive(&self, f: &mut dyn FnMut(&dyn Node)) {
+        self.circuit.map_nodes_recursive(f);
+    }
 }
 
 /// Top-level circuit with executor.

src/circuit/dbsp_handle.rs

@@ -1,9 +1,9 @@
 use crate::{
-    circuit::runtime::RuntimeHandle, Circuit, Error as DBSPError, Runtime, RuntimeError,
-    SchedulerError,
+    circuit::runtime::RuntimeHandle, profile::Profiler, Circuit, Error as DBSPError, Runtime,
+    RuntimeError, SchedulerError,
 };
 use crossbeam::channel::{bounded, Receiver, Sender, TryRecvError};
-use std::thread::Result as ThreadResult;
+use std::{fs, fs::create_dir_all, path::Path, thread::Result as ThreadResult, time::Instant};
 
 impl Runtime {
     /// Instantiate a circuit in a multithreaded runtime.
@@ -55,12 +55,16 @@ impl Runtime {
             let status_sender = status_senders.into_iter().nth(worker_index).unwrap();
             let command_receiver = command_receivers.into_iter().nth(worker_index).unwrap();
 
-            let circuit = match Circuit::build(constructor) {
-                Ok((circuit, res)) => {
+            let (circuit, profiler) = match Circuit::build(|circuit| {
+                let profiler = Profiler::new(circuit);
+                let res = constructor(circuit);
+                (res, profiler)
+            }) {
+                Ok((circuit, (res, profiler))) => {
                     if init_sender.send(Ok(res)).is_err() {
                         return;
                     }
-                    circuit
+                    (circuit, profiler)
                 }
                 Err(e) => {
                     let _ = init_sender.send(Err(e));
@@ -74,14 +78,29 @@ impl Runtime {
             while !Runtime::kill_in_progress() {
                 // Wait for command.
                 match command_receiver.try_recv() {
-                    Ok(()) => {
+                    Ok(Command::Step) => {
                         //moregc = true;
-                        let status = circuit.step();
+                        let status = circuit.step().map(|_| Response::Unit);
                         // Send response.
                         if status_sender.send(status).is_err() {
                             return;
                         };
                     }
+                    Ok(Command::EnableProfiler) => {
+                        profiler.enable_cpu_profiler();
+                        // Send response.
+                        if status_sender.send(Ok(Response::Unit)).is_err() {
+                            return;
+                        }
+                    }
+                    Ok(Command::DumpProfile) => {
+                        if status_sender
+                            .send(Ok(Response::Profile(profiler.dump_profile())))
+                            .is_err()
+                        {
+                            return;
+                        }
+                    }
                     // Nothing to do: do some housekeeping and relinquish the CPU if there's none
                     // left.
                     Err(TryRecvError::Empty) => {
@@ -116,9 +135,13 @@ impl Runtime {
         }
 
         // On error, kill the runtime.
-        if let Some(error) = init_status.iter().find_map(|status| status.as_ref().err()) {
+        if init_status.iter().any(Result::is_err) {
+            let error = init_status
+                .into_iter()
+                .find_map(|status| status.err())
+                .unwrap();
             let _ = runtime.kill();
-            return Err(error.clone());
+            return Err(error);
         }
 
         let dbsp = DBSPHandle::new(runtime, command_senders, status_receivers);
@@ -129,25 +152,39 @@ impl Runtime {
     }
 }
 
+#[derive(Clone)]
+enum Command {
+    Step,
+    EnableProfiler,
+    DumpProfile,
+}
+
+enum Response {
+    Unit,
+    Profile(String),
+}
+
 /// A handle to control the execution of a circuit in a multithreaded runtime.
 #[derive(Debug)]
 pub struct DBSPHandle {
+    // Time when the handle was created.
+    start_time: Instant,
     runtime: Option<RuntimeHandle>,
-    // Channels used to send commands to workers.  Currently the only supported
-    // command is 'step', so we can use `()` to represent commands.
-    command_senders: Vec<Sender<()>>,
+    // Channels used to send commands to workers.
+    command_senders: Vec<Sender<Command>>,
     // Channels used to receive command completion status from
     // workers.
-    status_receivers: Vec<Receiver<Result<(), SchedulerError>>>,
+    status_receivers: Vec<Receiver<Result<Response, SchedulerError>>>,
 }
 
 impl DBSPHandle {
     fn new(
         runtime: RuntimeHandle,
-        command_senders: Vec<Sender<()>>,
-        status_receivers: Vec<Receiver<Result<(), SchedulerError>>>,
+        command_senders: Vec<Sender<Command>>,
+        status_receivers: Vec<Receiver<Result<Response, SchedulerError>>>,
     ) -> Self {
         Self {
+            start_time: Instant::now(),
             runtime: Some(runtime),
             command_senders,
             status_receivers,
@@ -160,15 +197,17 @@ impl DBSPHandle {
         self.runtime.take().unwrap().kill()
     }
 
-    /// Evaluate the circuit for one clock cycle.
-    pub fn step(&mut self) -> Result<(), DBSPError> {
+    fn broadcast_command<F>(&mut self, command: Command, mut handler: F) -> Result<(), DBSPError>
+    where
+        F: FnMut(Response),
+    {
         if self.runtime.is_none() {
             return Err(DBSPError::Runtime(RuntimeError::Killed));
         }
 
         // Send command.
         for (worker, sender) in self.command_senders.iter().enumerate() {
-            if matches!(sender.send(()), Err(_)) {
+            if matches!(sender.send(command.clone()), Err(_)) {
                 let _ = self.kill_inner();
                 return Err(DBSPError::Runtime(RuntimeError::WorkerPanic(worker)));
             }
@@ -186,8 +225,55 @@ impl DBSPHandle {
                     let _ = self.kill_inner();
                     return Err(DBSPError::Scheduler(e));
                 }
-                Ok(Ok(_)) => {}
+                Ok(Ok(resp)) => handler(resp),
+            }
+        }
+
+        Ok(())
+    }
+
+    pub fn num_workers(&self) -> usize {
+        self.status_receivers.len()
+    }
+
+    /// Evaluate the circuit for one clock cycle.
+    pub fn step(&mut self) -> Result<(), DBSPError> {
+        self.broadcast_command(Command::Step, |_| {})
+    }
+
+    /// Enable CPU profiler.
+    ///
+    /// Enable recording of CPU usage info.  When CPU profiling is enabled,
+    /// [`Self::dump_profile`] outputs CPU usage info along with memory
+    /// usage and other circuit metadata.  CPU profiling introduces small
+    /// runtime overhead.
+    pub fn enable_cpu_profiler(&mut self) -> Result<(), DBSPError> {
+        self.broadcast_command(Command::EnableProfiler, |_| {})
+    }
+
+    /// Dump profiling information to the specified directory.
+    ///
+    /// Creates `dir_path` if it doesn't exist.  For each worker thread, creates
+    /// `dir_path/<timestamp>/<worker>.dot` file containing worker profile in
+    /// the graphviz format.  If CPU profiling was enabled (see
+    /// [`Self::enable_cpu_profiler`]), the profile will contain both CPU and
+    /// memory usage information; otherwise only memory usage details are
+    /// reported.
+    pub fn dump_profile<P: AsRef<Path>>(&mut self, dir_path: P) -> Result<(), DBSPError> {
+        let elapsed = self.start_time.elapsed().as_micros();
+        let mut profiles = Vec::with_capacity(self.num_workers());
+
+        let dir_path = dir_path.as_ref().join(elapsed.to_string());
+        create_dir_all(&dir_path)?;
+
+        self.broadcast_command(Command::DumpProfile, |resp| {
+            if let Response::Profile(prof) = resp {
+                profiles.push(prof);
             }
+        })?;
+
+        for (worker, profile) in profiles.into_iter().enumerate() {
+            fs::write(dir_path.join(format!("{worker}.dot")), profile)?;
         }
 
         Ok(())
@@ -242,10 +328,11 @@ mod tests {
             circuit.add_source(Generator::new(|| 5usize));
         });
 
-        assert_eq!(
-            res.unwrap_err(),
-            DBSPError::Runtime(RuntimeError::WorkerPanic(0))
-        );
+        if let DBSPError::Runtime(err) = res.unwrap_err() {
+            assert_eq!(err, RuntimeError::WorkerPanic(0));
+        } else {
+            panic!();
+        }
     }
 
     // TODO: initialization error in worker thread (the `constructor` closure
@@ -275,10 +362,11 @@ mod tests {
         })
         .unwrap();
 
-        assert_eq!(
-            handle.step().unwrap_err(),
-            DBSPError::Runtime(RuntimeError::WorkerPanic(0))
-        );
+        if let DBSPError::Runtime(err) = handle.step().unwrap_err() {
+            assert_eq!(err, RuntimeError::WorkerPanic(0));
+        } else {
+            panic!();
+        }
     }
 
     // Kill the runtime.
@@ -298,7 +386,11 @@ mod tests {
         })
         .unwrap();
 
+        handle.enable_cpu_profiler().unwrap();
         handle.step().unwrap();
+        handle
+            .dump_profile(std::env::temp_dir().join("test_kill"))
+            .unwrap();
         handle.kill().unwrap();
     }
 

src/error.rs

@@ -1,8 +1,34 @@
 use crate::{RuntimeError, SchedulerError};
+use std::io::Error as IOError;
 
-#[derive(Clone, Debug, Eq, PartialEq)]
+#[derive(Debug)]
 pub enum Error {
     Scheduler(SchedulerError),
     Runtime(RuntimeError),
+    IO(IOError),
     Custom(String),
 }
+
+impl From<IOError> for Error {
+    fn from(error: IOError) -> Self {
+        Self::IO(error)
+    }
+}
+
+impl From<SchedulerError> for Error {
+    fn from(error: SchedulerError) -> Self {
+        Self::Scheduler(error)
+    }
+}
+
+impl From<RuntimeError> for Error {
+    fn from(error: RuntimeError) -> Self {
+        Self::Runtime(error)
+    }
+}
+
+impl From<String> for Error {
+    fn from(error: String) -> Self {
+        Self::Custom(error)
+    }
+}

src/nexmark/config.rs

@@ -91,6 +91,10 @@ pub struct Config {
     #[clap(long, default_value = "1", env = "NEXMARK_PERSON_PROPORTION")]
     pub person_proportion: usize,
 
+    /// Dump DBSP profiles for all executed queries to the specified directory.
+    #[clap(long, env = "NEXMARK_PROFILE_PATH")]
+    pub profile_path: Option<String>,
+
     /// Queries to run, all by default.
     #[clap(long, env = "NEXMARK_QUERIES", value_enum)]
     pub query: Vec<Query>,
@@ -133,6 +137,7 @@ impl Default for Config {
             num_in_flight_auctions: 100,
             out_of_order_group_size: 1,
             person_proportion: 1,
+            profile_path: None,
             query: Vec::new(),
             source_buffer_size: 10_000,
             input_batch_size: 40_000,

src/nexmark/queries/q16.rs

@@ -793,9 +793,11 @@ mod tests {
         })
         .unwrap();
 
+        dbsp.enable_cpu_profiler().unwrap();
         for mut vec in input_vecs {
             input_handle.append(&mut vec);
             dbsp.step().unwrap();
         }
+        dbsp.dump_profile(std::env::temp_dir().join("q16")).unwrap();
     }
 }