cggmp24_keygen/

progress.rs

//! Traces progress of protocol execution
//!
//! Provides [`Tracer`] trait that can be used to trace progress of ongoing MPC protocol execution.
//! For instance, it can be implemented to report progress to the end user.
//!
//! Out of box, there's [`PerfProfiler`] which can be used to bechmark a protocol.

/// Traces progress of protocol execution
///
/// See [module level documentation](self) for more details
pub trait Tracer: Send + Sync {
    /// Traces occurred event
    fn trace_event(&mut self, event: Event);

    /// Traces [`Event::ProtocolBegins`] event
    fn protocol_begins(&mut self) {
        self.trace_event(Event::ProtocolBegins)
    }
    /// Traces [`Event::RoundBegins`] event
    fn round_begins(&mut self) {
        self.trace_event(Event::RoundBegins { name: None })
    }
    /// Traces [`Event::RoundBegins`] event
    fn named_round_begins(&mut self, round_name: &'static str) {
        self.trace_event(Event::RoundBegins {
            name: Some(round_name),
        })
    }
    /// Traces [`Event::Stage`] event
    fn stage(&mut self, stage: &'static str) {
        self.trace_event(Event::Stage { name: stage })
    }
    /// Traces [`Event::ReceiveMsgs`] event
    fn receive_msgs(&mut self) {
        self.trace_event(Event::ReceiveMsgs)
    }
    /// Traces [`Event::MsgsReceived`] event
    fn msgs_received(&mut self) {
        self.trace_event(Event::MsgsReceived)
    }
    /// Traces [`Event::SendMsg`] event
    fn send_msg(&mut self) {
        self.trace_event(Event::SendMsg)
    }
    /// Traces [`Event::MsgSent`] event
    fn msg_sent(&mut self) {
        self.trace_event(Event::MsgSent)
    }
    /// Traces [`Event::ProtocolEnds`] event
    fn protocol_ends(&mut self) {
        self.trace_event(Event::ProtocolEnds)
    }
}

/// Event occurred during the protocol execution
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
pub enum Event {
    /// Protocol begins
    ///
    /// This event is always emitted before any other events
    ProtocolBegins,

    /// Round begins
    RoundBegins {
        /// Optional name of the round
        name: Option<&'static str>,
    },
    /// Stage begins
    Stage {
        /// Name of the stage
        name: &'static str,
    },

    /// Protocol waits for some messages to be received
    ReceiveMsgs,
    /// Protocol received messages, round continues
    MsgsReceived,

    /// Protocol starts sending a message
    SendMsg,
    /// Protocol sent a message, round continues
    MsgSent,

    /// Protocol completed
    ProtocolEnds,
}

impl Tracer for &mut dyn Tracer {
    fn trace_event(&mut self, event: Event) {
        (*self).trace_event(event)
    }
}

impl<T: Tracer> Tracer for &mut T {
    fn trace_event(&mut self, event: Event) {
        <T as Tracer>::trace_event(self, event)
    }
}

impl<T: Tracer> Tracer for Option<T> {
    fn trace_event(&mut self, event: Event) {
        match self {
            Some(tracer) => tracer.trace_event(event),
            None => {
                // no-op
            }
        }
    }
}

#[cfg(feature = "std")]
pub use requires_std::*;
#[cfg(feature = "std")]
mod requires_std {
    use alloc::{vec, vec::Vec};
    use core::fmt;
    use std::time::{Duration, Instant};

    use thiserror::Error;

    use super::*;

    /// Profiles performance of the protocol
    ///
    /// Implements [`Tracer`] trait so it can be embedded into protocol execution. `PerfProfiler` keeps track of time
    /// passed between each step of protocol. After protocol is completed, you can obtain a [`PerfReport`] via
    /// [`.get_report()`](PerfProfiler::get_report) method that contains all the measurements.
    pub struct PerfProfiler {
        last_timestamp: Option<Instant>,
        ongoing_stage: Option<usize>,
        protocol_began: Option<Instant>,
        report: PerfReport,
        error: Option<ProfileError>,
    }

    /// Performance report generated by [`PerfProfiler`]
    #[derive(Debug, Clone)]
    pub struct PerfReport {
        /// Duration of setup phase (time after protocol began and before first round started)
        pub setup: Duration,
        /// Stages of setup phase
        pub setup_stages: Vec<StageDuration>,
        /// Performance report for each round
        pub rounds: Vec<RoundDuration>,
        display_io: bool,
    }

    /// Performance of specific round (part of [`PerfReport`])
    #[derive(Debug, Clone)]
    pub struct RoundDuration {
        /// Round name (if provided)
        pub round_name: Option<&'static str>,
        /// Stages of the round
        pub stages: Vec<StageDuration>,
        /// Total duration of pure computation performed during the round
        pub computation: Duration,
        /// Total time we spent during this round on sending messages
        pub sending: Duration,
        /// Total time we spent during this round on receiving messages
        pub receiving: Duration,
    }

    /// Performance of specific stage (part of [`PerfReport`])
    #[derive(Debug, Clone)]
    pub struct StageDuration {
        /// Stage name
        pub name: &'static str,
        /// Duration of the stage
        pub duration: Duration,
    }

    /// Protocol profiling resulted into error
    #[derive(Debug, Error, Clone)]
    #[error("profiler failed to trace protocol: it behaved unexpectedly")]
    pub struct ProfileError(
        #[source]
        #[from]
        ErrorReason,
    );

    #[derive(Debug, Error, Clone)]
    enum ErrorReason {
        #[error("protocol has never began")]
        ProtocolNeverBegan,
        #[error("tracing stage or sending/receiving message but round never began")]
        RoundNeverBegan,
        #[error("stage is ongoing, but it can't be finished with that event: {event:?}")]
        CantFinishStage { event: Event },
    }

    impl Tracer for PerfProfiler {
        fn trace_event(&mut self, event: Event) {
            if self.error.is_none() {
                if let Err(err) = self.try_trace_event(event) {
                    self.error = Some(err)
                }
            }
        }
    }

    impl PerfProfiler {
        /// Constructs new [`PerfProfiler`]
        pub fn new() -> Self {
            Self {
                last_timestamp: None,
                ongoing_stage: None,
                protocol_began: None,
                report: PerfReport {
                    setup: Duration::ZERO,
                    setup_stages: vec![],
                    rounds: vec![],
                    display_io: true,
                },
                error: None,
            }
        }

        /// Obtains a report
        ///
        /// Returns error if protocol behaved unexpectedly
        pub fn get_report(&self) -> Result<PerfReport, ProfileError> {
            if let Some(err) = self.error.clone() {
                Err(err)
            } else {
                Ok(self.report.clone())
            }
        }

        fn try_trace_event(&mut self, event: Event) -> Result<(), ProfileError> {
            let now = Instant::now();

            if Self::event_can_finish_ongoing_stage(&event) {
                if let Some(stage_i) = self.ongoing_stage.take() {
                    let last_timestamp = self.last_timestamp()?;

                    if !self.report.rounds.is_empty() {
                        let last_round = self.last_round_mut()?;
                        last_round.stages[stage_i].duration += now - last_timestamp;
                    } else {
                        self.report.setup_stages[stage_i].duration += now - last_timestamp;
                    }
                }
            } else if self.ongoing_stage.is_some() {
                return Err(ErrorReason::CantFinishStage { event }.into());
            }
            match event {
                Event::ProtocolBegins => {
                    self.protocol_began = Some(now);
                }
                Event::RoundBegins { name } => {
                    let last_timestamp = self.last_timestamp()?;
                    match self.report.rounds.last_mut() {
                        None => self.report.setup += now - last_timestamp,
                        Some(last_round) => last_round.computation += now - last_timestamp,
                    }
                    self.report.rounds.push(RoundDuration {
                        round_name: name,
                        stages: vec![],
                        computation: Duration::ZERO,
                        sending: Duration::ZERO,
                        receiving: Duration::ZERO,
                    })
                }
                Event::Stage { name } => {
                    let last_timestamp = self.last_timestamp()?;

                    let stages = if !self.report.rounds.is_empty() {
                        let last_round = self.last_round_mut()?;
                        last_round.computation += now - last_timestamp;

                        &mut last_round.stages
                    } else {
                        self.report.setup += now - last_timestamp;
                        &mut self.report.setup_stages
                    };

                    let stage_i = stages.iter().position(|s| s.name == name);
                    let stage_i = match stage_i {
                        Some(i) => i,
                        None => {
                            stages.push(StageDuration {
                                name,
                                duration: Duration::ZERO,
                            });
                            stages.len() - 1
                        }
                    };
                    self.ongoing_stage = Some(stage_i);
                }
                Event::ReceiveMsgs => {
                    let last_timestamp = self.last_timestamp()?;
                    let last_round = self.last_round_mut()?;
                    last_round.computation += now - last_timestamp;
                }
                Event::MsgsReceived => {
                    let last_timestamp = self.last_timestamp()?;
                    let last_round = self.last_round_mut()?;
                    last_round.receiving += now - last_timestamp;
                }
                Event::SendMsg => {
                    let last_timestamp = self.last_timestamp()?;
                    let last_round = self.last_round_mut()?;
                    last_round.computation += now - last_timestamp;
                }
                Event::MsgSent => {
                    let last_timestamp = self.last_timestamp()?;
                    let last_round = self.last_round_mut()?;
                    last_round.sending += now - last_timestamp;
                }
                Event::ProtocolEnds => {
                    let last_timestamp = self.last_timestamp()?;
                    let last_round = self.last_round_mut()?;
                    last_round.computation += now - last_timestamp;
                }
            }

            self.last_timestamp = Some(now);
            Ok(())
        }

        fn last_timestamp(&self) -> Result<Instant, ProfileError> {
            let last_timestamp = self.last_timestamp.ok_or(ErrorReason::ProtocolNeverBegan)?;
            Ok(last_timestamp)
        }
        fn last_round_mut(&mut self) -> Result<&mut RoundDuration, ProfileError> {
            let last_round = self
                .report
                .rounds
                .last_mut()
                .ok_or(ErrorReason::RoundNeverBegan)?;
            Ok(last_round)
        }
        fn event_can_finish_ongoing_stage(event: &Event) -> bool {
            matches!(
                event,
                Event::RoundBegins { .. }
                    | Event::Stage { .. }
                    | Event::ReceiveMsgs
                    | Event::SendMsg
                    | Event::ProtocolEnds
            )
        }
    }

    impl Default for PerfProfiler {
        fn default() -> Self {
            Self::new()
        }
    }

    impl PerfReport {
        /// Specifies whether time spent on i/o should be rendered in the final report
        ///
        /// Time spent on i/o is the time when signer was sending messages or waiting other
        /// parties to send messages
        pub fn display_io(mut self, display: bool) -> Self {
            self.display_io = display;
            self
        }
    }

    impl fmt::Display for PerfReport {
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
            let total_computation =
                self.setup + self.rounds.iter().map(|r| r.computation).sum::<Duration>();
            let total_send = if self.display_io {
                self.rounds.iter().map(|r| r.sending).sum::<Duration>()
            } else {
                Duration::ZERO
            };
            let total_recv = if self.display_io {
                self.rounds.iter().map(|r| r.receiving).sum::<Duration>()
            } else {
                Duration::ZERO
            };
            let total_io = total_send + total_recv;
            let total = total_computation + total_io;

            writeln!(f, "Protocol Performance:")?;
            writeln!(f, "  - Protocol took {total:.2?} to complete")?;
            if self.display_io {
                writeln!(
                    f,
                    "    - Computation: {total_computation:.2?} ({})",
                    percent(total_computation, total)
                )?;
                writeln!(
                    f,
                    "    - I/O: {total_io:.2?} ({})",
                    percent(total_io, total)
                )?;
                writeln!(f, "      - Send: {total_send:.2?}")?;
                writeln!(f, "      - Recv: {total_recv:.2?}")?;
            }

            writeln!(f, "In particular:")?;
            Self::fmt_round(f, 0, Some("Stage"), &self.setup_stages, self.setup, None)?;

            for (i, round) in self.rounds.iter().enumerate() {
                Self::fmt_round(
                    f,
                    i + 1,
                    round.round_name,
                    &round.stages,
                    round.computation,
                    if self.display_io {
                        Some((round.sending, round.receiving))
                    } else {
                        None
                    },
                )?;
            }

            Ok(())
        }
    }

    impl PerfReport {
        fn fmt_round(
            f: &mut fmt::Formatter,
            i: usize,
            round_name: Option<&str>,
            stages: &[StageDuration],
            computation: Duration,
            io: Option<(Duration, Duration)>, // (sending, receiving)
        ) -> fmt::Result {
            let total_duration = computation + io.map(|(s, r)| s + r).unwrap_or_default();
            if let Some(round_name) = round_name {
                writeln!(f, "  - {round_name}: {total_duration:.2?}")?
            } else {
                writeln!(f, "  - Round {i}: {total_duration:.2?}")?
            }

            Self::fmt_stages(f, total_duration, stages)?;

            if let Some((sending, receiving)) = io {
                let total_io = sending + receiving;
                writeln!(
                    f,
                    "    - I/O: {:.2?} ({})",
                    total_io,
                    percent(total_io, total_duration)
                )?;
                writeln!(f, "      - Send: {sending:.2?}")?;
                writeln!(f, "      - Recv: {receiving:.2?}")?;
            }

            if !stages.is_empty() || io.is_some() {
                let stages_total = stages.iter().map(|s| s.duration).sum::<Duration>();
                let unstaged = computation - stages_total;
                let percent = percent(unstaged, total_duration);
                writeln!(f, "    - Unstaged: {unstaged:.2?} ({percent})")?;
            }

            Ok(())
        }

        fn fmt_stages(
            f: &mut fmt::Formatter,
            total: Duration,
            stages: &[StageDuration],
        ) -> fmt::Result {
            for stage in stages {
                writeln!(
                    f,
                    "    - {}: {:.2?} ({})",
                    stage.name,
                    stage.duration,
                    percent(stage.duration, total),
                )?;
            }
            Ok(())
        }
    }

    fn percent(part: Duration, total: Duration) -> impl fmt::Display {
        struct Percentage(Duration, Duration);

        impl fmt::Display for Percentage {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                let percent = self.0.as_secs_f64() / self.1.as_secs_f64() * 100.;
                write!(f, "{percent:.1}%")
            }
        }

        Percentage(part, total)
    }

    /// Prints progress of the protocol to stderr
    #[derive(Default)]
    pub struct Stderr {
        prefix: Option<std::string::String>,
    }

    impl Stderr {
        /// Constructs an stderr tracer
        pub fn new() -> Self {
            Self::default()
        }

        /// Sets a prefix to be printed for each event
        pub fn with_prefix(mut self, prefix: impl std::string::ToString) -> Self {
            self.prefix = Some(prefix.to_string());
            self
        }
    }

    impl Tracer for Stderr {
        fn trace_event(&mut self, event: Event) {
            if let Some(prefix) = &self.prefix {
                std::eprintln!("{prefix}: {event:?}")
            } else {
                std::eprintln!("{event:?}")
            }
        }
    }
}