<commit_before>use bitflags::bitflags;
#[cfg(feature = "chrono")]
use chrono::{DateTime, Local, TimeZone};
use colored::{Color, Colorize};
use std::fmt::{self, Write};
use {At, History, Meta, Record};

/// Configurable display formatting of structures.
#[derive(Copy, Clone, Debug)]
pub struct Display<'a, T: 'a> {
    data: &'a T,
    view: View,
}

impl<'a, T: 'a> Display<'a, T> {
    /// Show colored output (off by default).
    #[inline]
    pub fn colored(&mut self, on: bool) -> &mut Self {
        if on {
            self.view.insert(View::COLORED);
        } else {
            self.view.remove(View::COLORED);
        }
        self
    }

    /// Show detailed output (on by default).
    #[inline]
    pub fn detailed(&mut self, on: bool) -> &mut Self {
        if on {
            self.view.insert(View::DETAILED);
        } else {
            self.view.remove(View::DETAILED);
        }
        self
    }

    /// Show the position of the command (on by default).
    #[inline]
    pub fn position(&mut self, on: bool) -> &mut Self {
        if on {
            self.view.insert(View::POSITION);
        } else {
            self.view.remove(View::POSITION);
        }
        self
    }

    /// Show the saved command (on by default).
    #[inline]
    pub fn saved(&mut self, on: bool) -> &mut Self {
        if on {
            self.view.insert(View::SAVED);
        } else {
            self.view.remove(View::SAVED);
        }
        self
    }
}

impl<'a, R> Display<'a, History<R>> {
    /// Show the history as a graph (off by default).
    #[inline]
    pub fn graph(&mut self, on: bool) -> &mut Self {
        if on {
            self.view.insert(View::GRAPH);
        } else {
            self.view.remove(View::GRAPH);
        }
        self
    }
}

impl<'a, R> Display<'a, Record<R>> {
    #[inline]
    fn fmt_list(&self, f: &mut fmt::Formatter, at: At, meta: &Meta<R>) -> fmt::Result {
        self.view.mark(f)?;
        self.view.position(f, at, false)?;
        if self.view.contains(View::DETAILED) {
            #[cfg(feature = "chrono")]
            self.view.timestamp(f, &meta.timestamp)?;
        }
        self.view.current(
            f,
            at,
            At {
                branch: 0,
                cursor: self.data.cursor(),
            },
        )?;
        self.view.saved(
            f,
            at,
            self.data.saved.map(|saved| At {
                branch: 0,
                cursor: saved,
            }),
        )?;
        if self.view.contains(View::DETAILED) {
            writeln!(f)?;
            self.view.message(f, meta, 0)
        } else {
            f.write_char(' ')?;
            self.view.message(f, meta, 0)?;
            writeln!(f)
        }
    }
}

impl<'a, R> Display<'a, History<R>> {
    #[inline]
    fn fmt_list(
        &self,
        f: &mut fmt::Formatter,
        at: At,
        meta: &Meta<R>,
        level: usize,
    ) -> fmt::Result {
        self.view.mark(f)?;
        self.view.position(f, at, true)?;
        if self.view.contains(View::DETAILED) {
            #[cfg(feature = "chrono")]
            self.view.timestamp(f, &meta.timestamp)?;
        }
        self.view.current(
            f,
            at,
            At {
                branch: self.data.root(),
                cursor: self.data.cursor(),
            },
        )?;
        self.view.saved(
            f,
            at,
            self.data
                .record
                .saved
                .map(|saved| At {
                    branch: self.data.root(),
                    cursor: saved,
                })
                .or(self.data.saved),
        )?;
        if self.view.contains(View::DETAILED) {
            writeln!(f)?;
            self.view.message(f, meta, level)
        } else {
            f.write_char(' ')?;
            self.view.message(f, meta, level)?;
            writeln!(f)
        }
    }

    #[inline]
    fn fmt_graph(
        &self,
        f: &mut fmt::Formatter,
        at: At,
        meta: &Meta<R>,
        level: usize,
    ) -> fmt::Result {
        for (&i, branch) in self
            .data
            .branches
            .iter()
            .filter(|(_, branch)| branch.parent == at)
        {
            for (j, cmd) in branch.commands.iter().enumerate().rev() {
                let at = At {
                    branch: i,
                    cursor: j + branch.parent.cursor + 1,
                };
                self.fmt_graph(f, at, cmd, level + 1)?;
            }
            for j in 0..level {
                self.view.edge(f, j)?;
                f.write_char(' ')?;
            }
            self.view.split(f, level)?;
            writeln!(f)?;
        }
        for i in 0..level {
            self.view.edge(f, i)?;
            f.write_char(' ')?;
        }
        self.fmt_list(f, at, meta, level)
    }
}

impl<'a, T: 'a> From<&'a T> for Display<'a, T> {
    #[inline]
    fn from(data: &'a T) -> Self {
        Display {
            data,
            view: View::default(),
        }
    }
}

impl<'a, R> fmt::Display for Display<'a, Record<R>> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for (i, cmd) in self.data.commands.iter().enumerate().rev() {
            let at = At {
                branch: 0,
                cursor: i + 1,
            };
            self.fmt_list(f, at, cmd)?;
        }
        Ok(())
    }
}

impl<'a, R> fmt::Display for Display<'a, History<R>> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for (i, cmd) in self.data.record.commands.iter().enumerate().rev() {
            let at = At {
                branch: self.data.root(),
                cursor: i + 1,
            };
            if self.view.contains(View::GRAPH) {
                self.fmt_graph(f, at, cmd, 0)?;
            } else {
                self.fmt_list(f, at, cmd, 0)?;
            }
        }
        Ok(())
    }
}

bitflags! {
    struct View: u8 {
        const COLORED   = 0b_0000_0001;
        const CURRENT   = 0b_0000_0010;
        const DETAILED  = 0b_0000_0100;
        const GRAPH     = 0b_0000_1000;
        const LIGATURES = 0b_0001_0000;
        const MARK      = 0b_0010_0000;
        const POSITION  = 0b_0100_0000;
        const SAVED     = 0b_1000_0000;
    }
}

impl Default for View {
    #[inline]
    fn default() -> Self {
        View::CURRENT | View::DETAILED | View::MARK | View::POSITION | View::SAVED
    }
}

impl View {
    #[inline]
    fn message(self, f: &mut fmt::Formatter, msg: &impl ToString, level: usize) -> fmt::Result {
        let msg = msg.to_string();
        let lines = msg.lines();
        if self.contains(View::DETAILED) {
            for line in lines {
                for i in 0..=level {
                    self.edge(f, i)?;
                    f.write_char(' ')?;
                }
                writeln!(f, "{}", line.trim())?;
            }
        } else if let Some(line) = lines.map(|s| s.trim()).find(|s| !s.is_empty()) {
            f.write_str(&line)?;
        }
        Ok(())
    }

    #[inline]
    fn mark(self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.contains(View::MARK) {
            f.write_str("*")
        } else {
            Ok(())
        }
    }

    #[inline]
    fn edge(self, f: &mut fmt::Formatter, level: usize) -> fmt::Result {
        if self.contains(View::COLORED) {
            write!(f, "{}", "|".color(color(level)))
        } else {
            f.write_str("|")
        }
    }

    #[inline]
    fn split(self, f: &mut fmt::Formatter, level: usize) -> fmt::Result {
        if self.contains(View::COLORED) {
            write!(
                f,
                "{}{}",
                "|".color(color(level)),
                "/".color(color(level + 1))
            )
        } else {
            f.write_str("|/")
        }
    }

    #[inline]
    fn position(self, f: &mut fmt::Formatter, at: At, use_branch: bool) -> fmt::Result {
        if self.contains(View::POSITION) {
            if self.contains(View::COLORED) {
                let position = if use_branch {
                    format!("[{}:{}]", at.branch, at.cursor)
                } else {
                    format!("[{}]", at.cursor)
                };
                write!(f, " {}", position.yellow())
            } else if use_branch {
                write!(f, " [{}:{}]", at.branch, at.cursor)
            } else {
                write!(f, " [{}]", at.cursor)
            }
        } else {
            Ok(())
        }
    }

    #[inline]
    fn current(self, f: &mut fmt::Formatter, at: At, current: At) -> fmt::Result {
        if self.contains(View::CURRENT) && at == current {
            if self.contains(View::COLORED) {
                write!(f, " {}{}{}", "(".yellow(), "current".cyan(), ")".yellow())
            } else {
                f.write_str(" (current)")
            }
        } else {
            Ok(())
        }
    }

    #[inline]
    fn saved(self, f: &mut fmt::Formatter, at: At, saved: Option<At>) -> fmt::Result {
        if self.contains(View::SAVED) && saved.map_or(false, |saved| saved == at) {
            if self.contains(View::COLORED) {
                write!(
                    f,
                    " {}{}{}",
                    "(".yellow(),
                    "saved".bright_green(),
                    ")".yellow()
                )
            } else {
                f.write_str(" (saved)")
            }
        } else {
            Ok(())
        }
    }

    #[inline]
    #[cfg(feature = "chrono")]
    fn timestamp(self, f: &mut fmt::Formatter, timestamp: &DateTime<impl TimeZone>) -> fmt::Result {
        let local = Local.from_utc_datetime(&timestamp.naive_utc());
        if self.contains(View::COLORED) {
            let ts = format!("[{}]", local.format("%T %F"));
            write!(f, " {}", ts.yellow())
        } else {
            write!(f, " [{}]", local.format("%T %F"))
        }
    }
}

#[inline]
fn color(i: usize) -> Color {
    match i % 6 {
        0 => Color::Red,
        1 => Color::Blue,
        2 => Color::Magenta,
        3 => Color::Yellow,
        4 => Color::Green,
        5 => Color::Cyan,
        _ => unreachable!(),
    }
}
<commit_msg>Use set method instead of insert and remove<commit_after>use bitflags::bitflags;
#[cfg(feature = "chrono")]
use chrono::{DateTime, Local, TimeZone};
use colored::{Color, Colorize};
use std::fmt::{self, Write};
use {At, History, Meta, Record};

/// Configurable display formatting of structures.
#[derive(Copy, Clone, Debug)]
pub struct Display<'a, T: 'a> {
    data: &'a T,
    view: View,
}

impl<'a, T: 'a> Display<'a, T> {
    /// Show colored output (off by default).
    #[inline]
    pub fn colored(&mut self, on: bool) -> &mut Self {
        self.view.set(View::COLORED, on);
        self
    }

    /// Show detailed output (on by default).
    #[inline]
    pub fn detailed(&mut self, on: bool) -> &mut Self {
        self.view.set(View::DETAILED, on);
        self
    }

    /// Show the position of the command (on by default).
    #[inline]
    pub fn position(&mut self, on: bool) -> &mut Self {
        self.view.set(View::POSITION, on);
        self
    }

    /// Show the saved command (on by default).
    #[inline]
    pub fn saved(&mut self, on: bool) -> &mut Self {
        self.view.set(View::SAVED, on);
        self
    }
}

impl<'a, R> Display<'a, History<R>> {
    /// Show the history as a graph (off by default).
    #[inline]
    pub fn graph(&mut self, on: bool) -> &mut Self {
        self.view.set(View::GRAPH, on);
        self
    }
}

impl<'a, R> Display<'a, Record<R>> {
    #[inline]
    fn fmt_list(&self, f: &mut fmt::Formatter, at: At, meta: &Meta<R>) -> fmt::Result {
        self.view.mark(f)?;
        self.view.position(f, at, false)?;
        if self.view.contains(View::DETAILED) {
            #[cfg(feature = "chrono")]
            self.view.timestamp(f, &meta.timestamp)?;
        }
        self.view.current(
            f,
            at,
            At {
                branch: 0,
                cursor: self.data.cursor(),
            },
        )?;
        self.view.saved(
            f,
            at,
            self.data.saved.map(|saved| At {
                branch: 0,
                cursor: saved,
            }),
        )?;
        if self.view.contains(View::DETAILED) {
            writeln!(f)?;
            self.view.message(f, meta, 0)
        } else {
            f.write_char(' ')?;
            self.view.message(f, meta, 0)?;
            writeln!(f)
        }
    }
}

impl<'a, R> Display<'a, History<R>> {
    #[inline]
    fn fmt_list(
        &self,
        f: &mut fmt::Formatter,
        at: At,
        meta: &Meta<R>,
        level: usize,
    ) -> fmt::Result {
        self.view.mark(f)?;
        self.view.position(f, at, true)?;
        if self.view.contains(View::DETAILED) {
            #[cfg(feature = "chrono")]
            self.view.timestamp(f, &meta.timestamp)?;
        }
        self.view.current(
            f,
            at,
            At {
                branch: self.data.root(),
                cursor: self.data.cursor(),
            },
        )?;
        self.view.saved(
            f,
            at,
            self.data
                .record
                .saved
                .map(|saved| At {
                    branch: self.data.root(),
                    cursor: saved,
                })
                .or(self.data.saved),
        )?;
        if self.view.contains(View::DETAILED) {
            writeln!(f)?;
            self.view.message(f, meta, level)
        } else {
            f.write_char(' ')?;
            self.view.message(f, meta, level)?;
            writeln!(f)
        }
    }

    #[inline]
    fn fmt_graph(
        &self,
        f: &mut fmt::Formatter,
        at: At,
        meta: &Meta<R>,
        level: usize,
    ) -> fmt::Result {
        for (&i, branch) in self
            .data
            .branches
            .iter()
            .filter(|(_, branch)| branch.parent == at)
        {
            for (j, cmd) in branch.commands.iter().enumerate().rev() {
                let at = At {
                    branch: i,
                    cursor: j + branch.parent.cursor + 1,
                };
                self.fmt_graph(f, at, cmd, level + 1)?;
            }
            for j in 0..level {
                self.view.edge(f, j)?;
                f.write_char(' ')?;
            }
            self.view.split(f, level)?;
            writeln!(f)?;
        }
        for i in 0..level {
            self.view.edge(f, i)?;
            f.write_char(' ')?;
        }
        self.fmt_list(f, at, meta, level)
    }
}

impl<'a, T: 'a> From<&'a T> for Display<'a, T> {
    #[inline]
    fn from(data: &'a T) -> Self {
        Display {
            data,
            view: View::default(),
        }
    }
}

impl<'a, R> fmt::Display for Display<'a, Record<R>> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for (i, cmd) in self.data.commands.iter().enumerate().rev() {
            let at = At {
                branch: 0,
                cursor: i + 1,
            };
            self.fmt_list(f, at, cmd)?;
        }
        Ok(())
    }
}

impl<'a, R> fmt::Display for Display<'a, History<R>> {
    #[inline]
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        for (i, cmd) in self.data.record.commands.iter().enumerate().rev() {
            let at = At {
                branch: self.data.root(),
                cursor: i + 1,
            };
            if self.view.contains(View::GRAPH) {
                self.fmt_graph(f, at, cmd, 0)?;
            } else {
                self.fmt_list(f, at, cmd, 0)?;
            }
        }
        Ok(())
    }
}

bitflags! {
    struct View: u8 {
        const COLORED   = 0b_0000_0001;
        const CURRENT   = 0b_0000_0010;
        const DETAILED  = 0b_0000_0100;
        const GRAPH     = 0b_0000_1000;
        const LIGATURES = 0b_0001_0000;
        const MARK      = 0b_0010_0000;
        const POSITION  = 0b_0100_0000;
        const SAVED     = 0b_1000_0000;
    }
}

impl Default for View {
    #[inline]
    fn default() -> Self {
        View::CURRENT | View::DETAILED | View::MARK | View::POSITION | View::SAVED
    }
}

impl View {
    #[inline]
    fn message(self, f: &mut fmt::Formatter, msg: &impl ToString, level: usize) -> fmt::Result {
        let msg = msg.to_string();
        let lines = msg.lines();
        if self.contains(View::DETAILED) {
            for line in lines {
                for i in 0..=level {
                    self.edge(f, i)?;
                    f.write_char(' ')?;
                }
                writeln!(f, "{}", line.trim())?;
            }
        } else if let Some(line) = lines.map(|s| s.trim()).find(|s| !s.is_empty()) {
            f.write_str(&line)?;
        }
        Ok(())
    }

    #[inline]
    fn mark(self, f: &mut fmt::Formatter) -> fmt::Result {
        if self.contains(View::MARK) {
            f.write_str("*")
        } else {
            Ok(())
        }
    }

    #[inline]
    fn edge(self, f: &mut fmt::Formatter, level: usize) -> fmt::Result {
        if self.contains(View::COLORED) {
            write!(f, "{}", "|".color(color(level)))
        } else {
            f.write_str("|")
        }
    }

    #[inline]
    fn split(self, f: &mut fmt::Formatter, level: usize) -> fmt::Result {
        if self.contains(View::COLORED) {
            write!(
                f,
                "{}{}",
                "|".color(color(level)),
                "/".color(color(level + 1))
            )
        } else {
            f.write_str("|/")
        }
    }

    #[inline]
    fn position(self, f: &mut fmt::Formatter, at: At, use_branch: bool) -> fmt::Result {
        if self.contains(View::POSITION) {
            if self.contains(View::COLORED) {
                let position = if use_branch {
                    format!("[{}:{}]", at.branch, at.cursor)
                } else {
                    format!("[{}]", at.cursor)
                };
                write!(f, " {}", position.yellow())
            } else if use_branch {
                write!(f, " [{}:{}]", at.branch, at.cursor)
            } else {
                write!(f, " [{}]", at.cursor)
            }
        } else {
            Ok(())
        }
    }

    #[inline]
    fn current(self, f: &mut fmt::Formatter, at: At, current: At) -> fmt::Result {
        if self.contains(View::CURRENT) && at == current {
            if self.contains(View::COLORED) {
                write!(f, " {}{}{}", "(".yellow(), "current".cyan(), ")".yellow())
            } else {
                f.write_str(" (current)")
            }
        } else {
            Ok(())
        }
    }

    #[inline]
    fn saved(self, f: &mut fmt::Formatter, at: At, saved: Option<At>) -> fmt::Result {
        if self.contains(View::SAVED) && saved.map_or(false, |saved| saved == at) {
            if self.contains(View::COLORED) {
                write!(
                    f,
                    " {}{}{}",
                    "(".yellow(),
                    "saved".bright_green(),
                    ")".yellow()
                )
            } else {
                f.write_str(" (saved)")
            }
        } else {
            Ok(())
        }
    }

    #[inline]
    #[cfg(feature = "chrono")]
    fn timestamp(self, f: &mut fmt::Formatter, timestamp: &DateTime<impl TimeZone>) -> fmt::Result {
        let local = Local.from_utc_datetime(&timestamp.naive_utc());
        if self.contains(View::COLORED) {
            let ts = format!("[{}]", local.format("%T %F"));
            write!(f, " {}", ts.yellow())
        } else {
            write!(f, " [{}]", local.format("%T %F"))
        }
    }
}

#[inline]
fn color(i: usize) -> Color {
    match i % 6 {
        0 => Color::Red,
        1 => Color::Blue,
        2 => Color::Magenta,
        3 => Color::Yellow,
        4 => Color::Green,
        5 => Color::Cyan,
        _ => unreachable!(),
    }
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::experiments::ExperimentParam;
use crate::tx_emitter::EmitJobRequest;
use crate::{
    cluster::Cluster,
    effects::{Effect, StopContainer},
    experiments::Context,
    experiments::Experiment,
    instance,
    instance::Instance,
    stats,
    util::unix_timestamp_now,
};
use async_trait::async_trait;
use futures::future::join_all;
use slog_scope::info;
use std::sync::atomic::Ordering;
use std::{
    collections::HashSet,
    fmt::{Display, Error, Formatter},
    time::Duration,
};
use structopt::StructOpt;

#[derive(StructOpt, Debug)]
pub struct PerformanceBenchmarkNodesDownParams {
    #[structopt(
        long,
        default_value = "0",
        help = "Number of nodes which should be down"
    )]
    pub num_nodes_down: usize,
    #[structopt(
        long,
        help = "Whether cluster test should run against validators or full nodes"
    )]
    pub is_fullnode: bool,
    #[structopt(
    long,
    default_value = Box::leak(format!("{}", DEFAULT_BENCH_DURATION).into_boxed_str()),
    help = "Duration of an experiment in seconds"
    )]
    pub duration: u64,
}

pub struct PerformanceBenchmarkNodesDown {
    down_instances: Vec<Instance>,
    up_instances: Vec<Instance>,
    num_nodes_down: usize,
    duration: Duration,
}

pub const DEFAULT_BENCH_DURATION: u64 = 180;

impl PerformanceBenchmarkNodesDownParams {
    pub fn new_nodes_down(num_nodes_down: usize) -> Self {
        Self {
            num_nodes_down,
            is_fullnode: false,
            duration: DEFAULT_BENCH_DURATION,
        }
    }
}

impl ExperimentParam for PerformanceBenchmarkNodesDownParams {
    type E = PerformanceBenchmarkNodesDown;
    fn build(self, cluster: &Cluster) -> Self::E {
        if self.is_fullnode {
            let (down_instances, up_instances) =
                cluster.split_n_fullnodes_random(self.num_nodes_down);
            Self::E {
                down_instances: down_instances.into_fullnode_instances(),
                up_instances: up_instances.into_fullnode_instances(),
                num_nodes_down: self.num_nodes_down,
                duration: Duration::from_secs(self.duration),
            }
        } else {
            let (down_instances, up_instances) =
                cluster.split_n_validators_random(self.num_nodes_down);
            Self::E {
                down_instances: down_instances.into_validator_instances(),
                up_instances: up_instances.into_validator_instances(),
                num_nodes_down: self.num_nodes_down,
                duration: Duration::from_secs(self.duration),
            }
        }
    }
}

#[async_trait]
impl Experiment for PerformanceBenchmarkNodesDown {
    fn affected_validators(&self) -> HashSet<String> {
        instance::instancelist_to_set(&self.down_instances)
    }

    async fn run(&mut self, context: &mut Context<'_>) -> anyhow::Result<()> {
        let stop_effects: Vec<_> = self
            .down_instances
            .clone()
            .into_iter()
            .map(StopContainer::new)
            .collect();
        let futures = stop_effects.iter().map(|e| e.activate());
        join_all(futures).await;
        let buffer = Duration::from_secs(60);
        let window = self.duration + buffer * 2;
        let emit_job_request = EmitJobRequest::for_instances(
            self.up_instances.clone(),
            context.global_emit_job_request,
        );
        let stats = context
            .tx_emitter
            .emit_txn_for(window, emit_job_request)
            .await?;
        let end = unix_timestamp_now() - buffer;
        let start = end - window + 2 * buffer;
        let (avg_tps, avg_latency) = stats::txn_stats(&context.prometheus, start, end)?;
        let avg_txns_per_block = stats::avg_txns_per_block(&context.prometheus, start, end)?;
        info!(
            "Link to dashboard : {}",
            context.prometheus.link_to_dashboard(start, end)
        );
        let futures = stop_effects.iter().map(|e| e.deactivate());
        join_all(futures).await;
        let submitted_txn = stats.submitted.load(Ordering::Relaxed);
        let expired_txn = stats.expired.load(Ordering::Relaxed);
        context
            .report
            .report_metric(&self, "submitted_txn", submitted_txn as f64);
        context
            .report
            .report_metric(&self, "expired_txn", expired_txn as f64);
        context
            .report
            .report_metric(&self, "avg_txns_per_block", avg_txns_per_block as f64);
        context.report.report_metric(&self, "avg_tps", avg_tps);
        context
            .report
            .report_metric(&self, "avg_latency", avg_latency);
        info!("avg_txns_per_block: {}", avg_txns_per_block);
        let expired_text = if expired_txn == 0 {
            "no expired txns".to_string()
        } else {
            format!("(!) expired {} out of {} txns", expired_txn, submitted_txn)
        };
        context.report.report_text(format!(
            "{} : {:.0} TPS, {:.1} ms latency, {}",
            self, avg_tps, avg_latency, expired_text
        ));
        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(480)
    }
}

impl Display for PerformanceBenchmarkNodesDown {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        if self.num_nodes_down == 0 {
            write!(f, "all up")
        } else {
            write!(f, "{}% down", self.num_nodes_down)
        }
    }
}
<commit_msg>[cluster-test] Increase deadline to 600 and reduce test dur to 120<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::experiments::ExperimentParam;
use crate::tx_emitter::EmitJobRequest;
use crate::{
    cluster::Cluster,
    effects::{Effect, StopContainer},
    experiments::Context,
    experiments::Experiment,
    instance,
    instance::Instance,
    stats,
    util::unix_timestamp_now,
};
use async_trait::async_trait;
use futures::future::join_all;
use slog_scope::info;
use std::sync::atomic::Ordering;
use std::{
    collections::HashSet,
    fmt::{Display, Error, Formatter},
    time::Duration,
};
use structopt::StructOpt;

#[derive(StructOpt, Debug)]
pub struct PerformanceBenchmarkNodesDownParams {
    #[structopt(
        long,
        default_value = "0",
        help = "Number of nodes which should be down"
    )]
    pub num_nodes_down: usize,
    #[structopt(
        long,
        help = "Whether cluster test should run against validators or full nodes"
    )]
    pub is_fullnode: bool,
    #[structopt(
    long,
    default_value = Box::leak(format!("{}", DEFAULT_BENCH_DURATION).into_boxed_str()),
    help = "Duration of an experiment in seconds"
    )]
    pub duration: u64,
}

pub struct PerformanceBenchmarkNodesDown {
    down_instances: Vec<Instance>,
    up_instances: Vec<Instance>,
    num_nodes_down: usize,
    duration: Duration,
}

pub const DEFAULT_BENCH_DURATION: u64 = 120;

impl PerformanceBenchmarkNodesDownParams {
    pub fn new_nodes_down(num_nodes_down: usize) -> Self {
        Self {
            num_nodes_down,
            is_fullnode: false,
            duration: DEFAULT_BENCH_DURATION,
        }
    }
}

impl ExperimentParam for PerformanceBenchmarkNodesDownParams {
    type E = PerformanceBenchmarkNodesDown;
    fn build(self, cluster: &Cluster) -> Self::E {
        if self.is_fullnode {
            let (down_instances, up_instances) =
                cluster.split_n_fullnodes_random(self.num_nodes_down);
            Self::E {
                down_instances: down_instances.into_fullnode_instances(),
                up_instances: up_instances.into_fullnode_instances(),
                num_nodes_down: self.num_nodes_down,
                duration: Duration::from_secs(self.duration),
            }
        } else {
            let (down_instances, up_instances) =
                cluster.split_n_validators_random(self.num_nodes_down);
            Self::E {
                down_instances: down_instances.into_validator_instances(),
                up_instances: up_instances.into_validator_instances(),
                num_nodes_down: self.num_nodes_down,
                duration: Duration::from_secs(self.duration),
            }
        }
    }
}

#[async_trait]
impl Experiment for PerformanceBenchmarkNodesDown {
    fn affected_validators(&self) -> HashSet<String> {
        instance::instancelist_to_set(&self.down_instances)
    }

    async fn run(&mut self, context: &mut Context<'_>) -> anyhow::Result<()> {
        let stop_effects: Vec<_> = self
            .down_instances
            .clone()
            .into_iter()
            .map(StopContainer::new)
            .collect();
        let futures = stop_effects.iter().map(|e| e.activate());
        join_all(futures).await;
        let buffer = Duration::from_secs(60);
        let window = self.duration + buffer * 2;
        let emit_job_request = EmitJobRequest::for_instances(
            self.up_instances.clone(),
            context.global_emit_job_request,
        );
        let stats = context
            .tx_emitter
            .emit_txn_for(window, emit_job_request)
            .await?;
        let end = unix_timestamp_now() - buffer;
        let start = end - window + 2 * buffer;
        let (avg_tps, avg_latency) = stats::txn_stats(&context.prometheus, start, end)?;
        let avg_txns_per_block = stats::avg_txns_per_block(&context.prometheus, start, end)?;
        info!(
            "Link to dashboard : {}",
            context.prometheus.link_to_dashboard(start, end)
        );
        let futures = stop_effects.iter().map(|e| e.deactivate());
        join_all(futures).await;
        let submitted_txn = stats.submitted.load(Ordering::Relaxed);
        let expired_txn = stats.expired.load(Ordering::Relaxed);
        context
            .report
            .report_metric(&self, "submitted_txn", submitted_txn as f64);
        context
            .report
            .report_metric(&self, "expired_txn", expired_txn as f64);
        context
            .report
            .report_metric(&self, "avg_txns_per_block", avg_txns_per_block as f64);
        context.report.report_metric(&self, "avg_tps", avg_tps);
        context
            .report
            .report_metric(&self, "avg_latency", avg_latency);
        info!("avg_txns_per_block: {}", avg_txns_per_block);
        let expired_text = if expired_txn == 0 {
            "no expired txns".to_string()
        } else {
            format!("(!) expired {} out of {} txns", expired_txn, submitted_txn)
        };
        context.report.report_text(format!(
            "{} : {:.0} TPS, {:.1} ms latency, {}",
            self, avg_tps, avg_latency, expired_text
        ));
        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(600)
    }
}

impl Display for PerformanceBenchmarkNodesDown {
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error> {
        if self.num_nodes_down == 0 {
            write!(f, "all up")
        } else {
            write!(f, "{}% down", self.num_nodes_down)
        }
    }
}
<|endoftext|>


<commit_before>extern crate kafka;
extern crate getopts;
extern crate env_logger;
extern crate time;

use std::ascii::AsciiExt;
use std::env;
use std::fmt;
use std::io::{self, stdout, Write, BufWriter};
use std::process;
use std::thread;
use std::time as stdtime;

use kafka::client::{KafkaClient, FetchOffset, GroupOffsetStorage};

/// A very simple offset monitor for a particular topic able to show
/// the lag for a particular consumer group. Dumps the offset/lag of
/// the monitored topic/group to stdout every few seconds.
fn main() {
    env_logger::init().unwrap();

    let cfg = match Config::from_cmdline() {
        Ok(cfg) => cfg,
        Err(e) => {
            println!("error: {}", e);
            process::exit(1);
        }
    };

    if let Err(e) = run(cfg) {
        println!("error: {}", e);
        process::exit(1);
    }
}

#[derive(Copy, Clone)]
struct Offsets {
    latest: i64,
    group: i64,
}

impl Default for Offsets {
    fn default() -> Self {
        Offsets { latest: -1, group: -1 }
    }
}

fn run(cfg: Config) -> Result<(), Error> {
    let mut client = KafkaClient::new(cfg.brokers);
    client.set_group_offset_storage(cfg.offset_storage);
    try!(client.load_metadata_all());

    let mut stdout = stdout();

    if cfg.topic.is_empty() {
        let ts = client.topics();
        let num_topics = ts.len();
        if num_topics == 0 {
            return Err(Error::Literal("no topics available".to_owned()));
        }
        let mut names: Vec<&str> = Vec::with_capacity(ts.len());
        names.extend(ts.names());
        names.sort();

        let mut buf = BufWriter::with_capacity(4096, stdout);
        for name in names {
            try!(write!(buf, "topic: {}\n", name));
        }
        return Err(Error::Literal("choose a topic".to_owned()));
    }

    let num_partitions = match client.topics().partitions(&cfg.topic) {
        None => return Err(Error::Literal(format!("no such topic: {}", &cfg.topic))),
        Some(partitions) => partitions.len(),
    };
    let mut offs = vec![Offsets::default(); num_partitions];
    let mut offs = offs.as_mut_slice();

    let mut fmt_buf = String::with_capacity(16);
    let mut out_buf = String::with_capacity(128);

    // ~ shall we produce the lag of the specified group?
    let show_lag = !cfg.group.is_empty();

    // ~ print a header
    {
        use std::fmt::Write;

        let _ = write!(out_buf, "{:>8} ", "time");
        for i in 0..num_partitions {
            fmt_buf.clear();
            let _ = write!(fmt_buf, "p-{}", i);
            let _ = write!(out_buf, " {:>10}", fmt_buf);
            if show_lag {
                let _ = write!(out_buf, " {:6}", "(lag)");
            }
        }
        if cfg.summary {
            let _ = write!(out_buf, " {:>12}", "summary");
            if show_lag {
                let _ = write!(out_buf, " {:8}", "(lag)");
            }
        }
        out_buf.push('\n');
        let _ = stdout.write_all(out_buf.as_bytes());
    }
    // ~ now periodically fetch the offsets and dump them to stdout
    loop {
        let t = time::now();
        // ~ latest topic offsets
        let latests = try!(client.fetch_topic_offsets(&cfg.topic, FetchOffset::Latest));
        for l in latests {
            let v = l.offset.unwrap_or(-1);
            offs.get_mut(l.partition as usize).expect("available partition").latest = v;
        }
        if show_lag {
            // ~ group offsets
            let groups = try!(client.fetch_group_topic_offsets(&cfg.group, &cfg.topic));
            for g in groups {
                let v = g.offset.unwrap_or(-1);
                offs.get_mut(g.partition as usize).expect("available partition").group = v;
            }
        }
        // ~ format the output
        {
            use std::fmt::Write;

            let mut total_latest = 0;
            let mut total_lag = 0;

            out_buf.clear();
            let _ = write!(out_buf, "{} ", t.strftime("%H:%M:%S").unwrap());
            for o in &*offs {
                total_latest += o.latest;
                let _ = write!(out_buf, " {:>10}", o.latest);
                if show_lag {
                    fmt_buf.clear();
                    if o.group < 0 {
                        let _ = write!(fmt_buf, "()");
                        total_lag = -1;
                    } else {
                        let mut lag = if cfg.commited_not_consumed {
                            o.latest - o.group
                        } else {
                            o.latest - o.group - 1
                        };
                        if lag < 0 {
                            // ~ it's quite likely that we fetch group offsets which
                            // are a bit ahead of the topic's latest offset since we
                            // issued the fetch-latest-offset request earlier than
                            // the request for the group-offsets
                            lag = 0;
                        }
                        if total_lag >= 0 {
                            total_lag += lag;
                        }
                        let _ = write!(fmt_buf, "({:<})", lag);
                    }
                    let _ = write!(out_buf, " {:6}", fmt_buf);
                }
            }
            if cfg.summary {
                let _ = write!(out_buf, " {:12}", total_latest);
                if show_lag {
                    fmt_buf.clear();
                    if total_lag >= 0 {
                        let _ = write!(fmt_buf, "({:<})", total_lag);
                    } else {
                        let _ = write!(fmt_buf, "()");
                    }
                    let _ = write!(out_buf, " {:8}", fmt_buf);
                }
            }
        }
        // ~ write the output
        out_buf.push('\n');
        let _ = stdout.write_all(out_buf.as_bytes());

        thread::sleep(cfg.period);
    }
}

// --------------------------------------------------------------------

struct Config {
    brokers: Vec<String>,
    topic: String,
    group: String,
    offset_storage: GroupOffsetStorage,
    period: stdtime::Duration,
    commited_not_consumed: bool,
    summary: bool,
}

impl Config {
    fn from_cmdline() -> Result<Config, Error> {
        let args: Vec<String> = env::args().collect();
        let mut opts = getopts::Options::new();
        opts.optflag("h", "help", "Print this help screen");
        opts.optopt("", "brokers", "Specify kafka bootstrap brokers (comma separated)", "HOSTS");
        opts.optopt("", "topic", "Specify the topic to monitor", "TOPIC");
        opts.optopt("", "group", "Specify the group to monitor", "GROUP");
        opts.optopt("", "offsets", "Specify offset store [zookeeper, kafka]", "STORE");
        opts.optopt("", "sleep", "Specify the sleep time", "SECS");
        opts.optflag("", "summary", "Print a summary for the topic");
        opts.optflag("", "committed-not-yet-consumed",
                     "Assume commited group offsets specify \
                      messages the group will start consuming \
                      (including those at these offsets)");

        let m = match opts.parse(&args[1..]) {
            Ok(m) => m,
            Err(e) => return Err(Error::Literal(e.to_string())),
        };
        if m.opt_present("help") {
            let brief = format!("{} [options]", args[0]);
            return Err(Error::Literal(opts.usage(&brief)));
        }
        let mut offset_storage = GroupOffsetStorage::Zookeeper;
        if let Some(s) = m.opt_str("offsets") {
            if s.eq_ignore_ascii_case("zookeeper") {
                offset_storage = GroupOffsetStorage::Zookeeper;
            } else if s.eq_ignore_ascii_case("kafka") {
                offset_storage = GroupOffsetStorage::Kafka;
            } else {
                return Err(Error::Literal(format!("unknown offset store: {}", s)));
            }
        }
        let mut period = stdtime::Duration::from_secs(5);
        if let Some(s) = m.opt_str("sleep") {
            match s.parse::<u64>() {
                Ok(n) if n != 0 => period = stdtime::Duration::from_secs(n),
                _ => return Err(Error::Literal(format!("not a number greater than zero: {}", s))),
            }
        }
        Ok(Config {
            brokers: m.opt_str("brokers")
                .unwrap_or_else(|| "localhost:9092".to_owned())
                .split(',')
                .map(|s| s.trim().to_owned())
                .collect(),
            topic: m.opt_str("topic")
                .unwrap_or_else(|| String::new()),
            group: m.opt_str("group")
                .unwrap_or_else(|| String::new()),
            offset_storage: offset_storage,
            period: period,
            commited_not_consumed: m.opt_present("committed-not-yet-consumed"),
            summary: m.opt_present("summary"),
        })
    }
}

// --------------------------------------------------------------------

enum Error {
    Kafka(kafka::error::Error),
    Io(io::Error),
    Literal(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::Kafka(ref e) => write!(f, "kafka-error: {}", e),
            &Error::Io(ref e) => write!(f, "io-error: {}", e),
            &Error::Literal(ref s) => write!(f, "{}", s),
        }
    }
}

impl From<kafka::error::Error> for Error {
    fn from(e: kafka::error::Error) -> Self { Error::Kafka(e) }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { Error::Io(e) }
}
<commit_msg>Wording<commit_after>extern crate kafka;
extern crate getopts;
extern crate env_logger;
extern crate time;

use std::ascii::AsciiExt;
use std::env;
use std::fmt;
use std::io::{self, stdout, Write, BufWriter};
use std::process;
use std::thread;
use std::time as stdtime;

use kafka::client::{KafkaClient, FetchOffset, GroupOffsetStorage};

/// A very simple offset monitor for a particular topic able to show
/// the lag for a particular consumer group. Dumps the offset/lag of
/// the monitored topic/group to stdout every few seconds.
fn main() {
    env_logger::init().unwrap();

    let cfg = match Config::from_cmdline() {
        Ok(cfg) => cfg,
        Err(e) => {
            println!("error: {}", e);
            process::exit(1);
        }
    };

    if let Err(e) = run(cfg) {
        println!("error: {}", e);
        process::exit(1);
    }
}

#[derive(Copy, Clone)]
struct Offsets {
    latest: i64,
    group: i64,
}

impl Default for Offsets {
    fn default() -> Self {
        Offsets { latest: -1, group: -1 }
    }
}

fn run(cfg: Config) -> Result<(), Error> {
    let mut client = KafkaClient::new(cfg.brokers);
    client.set_group_offset_storage(cfg.offset_storage);
    try!(client.load_metadata_all());

    let mut stdout = stdout();

    if cfg.topic.is_empty() {
        let ts = client.topics();
        let num_topics = ts.len();
        if num_topics == 0 {
            return Err(Error::Literal("no topics available".to_owned()));
        }
        let mut names: Vec<&str> = Vec::with_capacity(ts.len());
        names.extend(ts.names());
        names.sort();

        let mut buf = BufWriter::with_capacity(4096, stdout);
        for name in names {
            try!(write!(buf, "topic: {}\n", name));
        }
        return Err(Error::Literal("choose a topic".to_owned()));
    }

    let num_partitions = match client.topics().partitions(&cfg.topic) {
        None => return Err(Error::Literal(format!("no such topic: {}", &cfg.topic))),
        Some(partitions) => partitions.len(),
    };
    let mut offs = vec![Offsets::default(); num_partitions];
    let mut offs = offs.as_mut_slice();

    let mut fmt_buf = String::with_capacity(16);
    let mut out_buf = String::with_capacity(128);

    // ~ shall we produce the lag of the specified group?
    let show_lag = !cfg.group.is_empty();

    // ~ print a header
    {
        use std::fmt::Write;

        let _ = write!(out_buf, "{:>8} ", "time");
        for i in 0..num_partitions {
            fmt_buf.clear();
            let _ = write!(fmt_buf, "p-{}", i);
            let _ = write!(out_buf, " {:>10}", fmt_buf);
            if show_lag {
                let _ = write!(out_buf, " {:6}", "(lag)");
            }
        }
        if cfg.summary {
            let _ = write!(out_buf, " {:>12}", "summary");
            if show_lag {
                let _ = write!(out_buf, " {:8}", "(lag)");
            }
        }
        out_buf.push('\n');
        let _ = stdout.write_all(out_buf.as_bytes());
    }
    // ~ now periodically fetch the offsets and dump them to stdout
    loop {
        let t = time::now();
        // ~ latest topic offsets
        let latests = try!(client.fetch_topic_offsets(&cfg.topic, FetchOffset::Latest));
        for l in latests {
            let v = l.offset.unwrap_or(-1);
            offs.get_mut(l.partition as usize).expect("available partition").latest = v;
        }
        if show_lag {
            // ~ group offsets
            let groups = try!(client.fetch_group_topic_offsets(&cfg.group, &cfg.topic));
            for g in groups {
                let v = g.offset.unwrap_or(-1);
                offs.get_mut(g.partition as usize).expect("available partition").group = v;
            }
        }
        // ~ format the output
        {
            use std::fmt::Write;

            let mut total_latest = 0;
            let mut total_lag = 0;

            out_buf.clear();
            let _ = write!(out_buf, "{} ", t.strftime("%H:%M:%S").unwrap());
            for o in &*offs {
                total_latest += o.latest;
                let _ = write!(out_buf, " {:>10}", o.latest);
                if show_lag {
                    fmt_buf.clear();
                    if o.group < 0 {
                        let _ = write!(fmt_buf, "()");
                        total_lag = -1;
                    } else {
                        let mut lag = if cfg.commited_not_consumed {
                            o.latest - o.group
                        } else {
                            o.latest - o.group - 1
                        };
                        if lag < 0 {
                            // ~ it's quite likely that we fetched group offsets
                            // which are a bit ahead of the topic's latest offset
                            // since we issued the fetch-latest-offset request
                            // earlier than the request for the group-offsets
                            lag = 0;
                        }
                        if total_lag >= 0 {
                            total_lag += lag;
                        }
                        let _ = write!(fmt_buf, "({:<})", lag);
                    }
                    let _ = write!(out_buf, " {:6}", fmt_buf);
                }
            }
            if cfg.summary {
                let _ = write!(out_buf, " {:12}", total_latest);
                if show_lag {
                    fmt_buf.clear();
                    if total_lag >= 0 {
                        let _ = write!(fmt_buf, "({:<})", total_lag);
                    } else {
                        let _ = write!(fmt_buf, "()");
                    }
                    let _ = write!(out_buf, " {:8}", fmt_buf);
                }
            }
        }
        // ~ write the output
        out_buf.push('\n');
        let _ = stdout.write_all(out_buf.as_bytes());

        thread::sleep(cfg.period);
    }
}

// --------------------------------------------------------------------

struct Config {
    brokers: Vec<String>,
    topic: String,
    group: String,
    offset_storage: GroupOffsetStorage,
    period: stdtime::Duration,
    commited_not_consumed: bool,
    summary: bool,
}

impl Config {
    fn from_cmdline() -> Result<Config, Error> {
        let args: Vec<String> = env::args().collect();
        let mut opts = getopts::Options::new();
        opts.optflag("h", "help", "Print this help screen");
        opts.optopt("", "brokers", "Specify kafka bootstrap brokers (comma separated)", "HOSTS");
        opts.optopt("", "topic", "Specify the topic to monitor", "TOPIC");
        opts.optopt("", "group", "Specify the group to monitor", "GROUP");
        opts.optopt("", "offsets", "Specify offset store [zookeeper, kafka]", "STORE");
        opts.optopt("", "sleep", "Specify the sleep time", "SECS");
        opts.optflag("", "summary", "Print a summary for the topic");
        opts.optflag("", "committed-not-yet-consumed",
                     "Assume commited group offsets specify \
                      messages the group will start consuming \
                      (including those at these offsets)");

        let m = match opts.parse(&args[1..]) {
            Ok(m) => m,
            Err(e) => return Err(Error::Literal(e.to_string())),
        };
        if m.opt_present("help") {
            let brief = format!("{} [options]", args[0]);
            return Err(Error::Literal(opts.usage(&brief)));
        }
        let mut offset_storage = GroupOffsetStorage::Zookeeper;
        if let Some(s) = m.opt_str("offsets") {
            if s.eq_ignore_ascii_case("zookeeper") {
                offset_storage = GroupOffsetStorage::Zookeeper;
            } else if s.eq_ignore_ascii_case("kafka") {
                offset_storage = GroupOffsetStorage::Kafka;
            } else {
                return Err(Error::Literal(format!("unknown offset store: {}", s)));
            }
        }
        let mut period = stdtime::Duration::from_secs(5);
        if let Some(s) = m.opt_str("sleep") {
            match s.parse::<u64>() {
                Ok(n) if n != 0 => period = stdtime::Duration::from_secs(n),
                _ => return Err(Error::Literal(format!("not a number greater than zero: {}", s))),
            }
        }
        Ok(Config {
            brokers: m.opt_str("brokers")
                .unwrap_or_else(|| "localhost:9092".to_owned())
                .split(',')
                .map(|s| s.trim().to_owned())
                .collect(),
            topic: m.opt_str("topic")
                .unwrap_or_else(|| String::new()),
            group: m.opt_str("group")
                .unwrap_or_else(|| String::new()),
            offset_storage: offset_storage,
            period: period,
            commited_not_consumed: m.opt_present("committed-not-yet-consumed"),
            summary: m.opt_present("summary"),
        })
    }
}

// --------------------------------------------------------------------

enum Error {
    Kafka(kafka::error::Error),
    Io(io::Error),
    Literal(String),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Error::Kafka(ref e) => write!(f, "kafka-error: {}", e),
            &Error::Io(ref e) => write!(f, "io-error: {}", e),
            &Error::Literal(ref s) => write!(f, "{}", s),
        }
    }
}

impl From<kafka::error::Error> for Error {
    fn from(e: kafka::error::Error) -> Self { Error::Kafka(e) }
}

impl From<io::Error> for Error {
    fn from(e: io::Error) -> Self { Error::Io(e) }
}
<|endoftext|>


<commit_before>use args;
use analyzer;
use hyper;
use prometheus::Encoder;
use render;
use render::Renderer;
use run;

struct HttpHandler {
    args: args::RequestLogAnalyzerArgs,
    run: fn(&args::RequestLogAnalyzerArgs) -> Option<analyzer::RequestLogAnalyzerResult>,
}

impl hyper::server::Handler for HttpHandler {
    fn handle(&self, _: hyper::server::Request, mut res: hyper::server::Response) {
        let result = (self.run)(&self.args);

        let mut renderer = render::prometheus::PrometheusRenderer::new();
        renderer.render(result);
        res.headers_mut()
            .set(hyper::header::ContentType(renderer.encoder
                .format_type()
                .parse::<hyper::mime::Mime>()
                .unwrap()));
        res.send(&renderer.buffer).unwrap();
    }
}

pub fn listen_http(args: args::RequestLogAnalyzerArgs, binding_address: &str) {
    let handler = HttpHandler {
        args: args,
        run: run,
    };

    info!("listening addr {:?}", binding_address);
    hyper::server::Server::http(binding_address).unwrap().handle(handler).unwrap();
}

#[cfg(test)]
mod tests {
    use std::io;
    use std::io::prelude::*;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4};
    use std::time::Duration;
    use hyper::server::Handler;
    use hyper;
    use filter;
    use analyzer;
    use super::*;

    #[test]
    fn test_handle() {
        let args = args::RequestLogAnalyzerArgs {
            filename: String::from("src/test/simple-1.log"),
            conditions: filter::FilterConditions {
                include_terms: None,
                exclude_terms: None,
                latest_time: None,
            },
            graphite_server: None,
            graphite_port: Some(2003),
            graphite_prefix: None,
            prometheus_listen: None,
        };

        fn run_fn(args: &args::RequestLogAnalyzerArgs)
                  -> Option<analyzer::RequestLogAnalyzerResult> {
            Some(analyzer::RequestLogAnalyzerResult {
                count: 3,
                max: 100,
                min: 1,
                avg: 37,
                median: 10,
                percentile90: 100,
            })
        };

        let handler = HttpHandler {
            args: args,
            run: run_fn,
        };

        struct MockNetworkStream {}

        impl Read for MockNetworkStream {
            fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
                Ok(1)
            }
        }

        impl Write for MockNetworkStream {
            fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
                Ok(1)
            }

            fn flush(&mut self) -> io::Result<()> {
                Ok(())
            }
        }

        impl hyper::net::NetworkStream for MockNetworkStream {
            fn peer_addr(&mut self) -> Result<SocketAddr, io::Error> {
                Ok(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::new(127, 0, 0, 1), 8080)))
            }

            fn set_read_timeout(&self, dur: Option<Duration>) -> Result<(), io::Error> {
                Ok(())
            }

            fn set_write_timeout(&self, dur: Option<Duration>) -> Result<(), io::Error> {
                Ok(())
            }
        }

        let mut request_mock_network_stream = MockNetworkStream {};

        let mut reader = hyper::buffer::BufReader::new(&mut request_mock_network_stream as
                                                       &mut hyper::net::NetworkStream);

        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080);
        let request = hyper::server::Request::new(&mut reader, socket).unwrap();

        let mut headers = hyper::header::Headers::new();
        let mut response_mock_network_stream = MockNetworkStream {};
        let response = hyper::server::Response::new(&mut response_mock_network_stream,
                                                    &mut headers);
        handler.handle(request, response);
    }
}
<commit_msg>Use mock from Hyper to test handler<commit_after>use args;
use analyzer;
use hyper;
use prometheus::Encoder;
use render;
use render::Renderer;
use run;

struct HttpHandler {
    args: args::RequestLogAnalyzerArgs,
    run: fn(&args::RequestLogAnalyzerArgs) -> Option<analyzer::RequestLogAnalyzerResult>,
}

impl hyper::server::Handler for HttpHandler {
    fn handle(&self, _: hyper::server::Request, mut res: hyper::server::Response) {
        let result = (self.run)(&self.args);

        let mut renderer = render::prometheus::PrometheusRenderer::new();
        renderer.render(result);
        res.headers_mut()
            .set(hyper::header::ContentType(renderer.encoder
                .format_type()
                .parse::<hyper::mime::Mime>()
                .unwrap()));
        res.send(&renderer.buffer).unwrap();
    }
}

pub fn listen_http(args: args::RequestLogAnalyzerArgs, binding_address: &str) {
    let handler = HttpHandler {
        args: args,
        run: run,
    };

    info!("listening addr {:?}", binding_address);
    hyper::server::Server::http(binding_address).unwrap().handle(handler).unwrap();
}

#[cfg(test)]
mod tests {
    use std::io::{self, Read, Write, Cursor};
    use std::str;
    use std::net::Shutdown;
    use std::time::Duration;
    use std::cell::Cell;
    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
    use hyper::net::NetworkStream;
    use hyper::server::Handler;
    use hyper;

    use filter;
    use analyzer;
    use super::*;

    // MockStream is copied from Hyper tests
    // https://github.com/hyperium/hyper/blob/0.10.x/src/mock.rs
    // MIT licensed, Copyright (c) 2014 Sean McArthur
    #[derive(Clone, Debug)]
    pub struct MockStream {
        pub read: Cursor<Vec<u8>>,
        next_reads: Vec<Vec<u8>>,
        pub write: Vec<u8>,
        pub is_closed: bool,
        pub error_on_write: bool,
        pub error_on_read: bool,
        pub read_timeout: Cell<Option<Duration>>,
        pub write_timeout: Cell<Option<Duration>>,
    }

    impl PartialEq for MockStream {
        fn eq(&self, other: &MockStream) -> bool {
            self.read.get_ref() == other.read.get_ref() && self.write == other.write
        }
    }

    impl MockStream {
        pub fn new() -> MockStream {
            MockStream::with_input(b"")
        }

        pub fn with_input(input: &[u8]) -> MockStream {
            MockStream::with_responses(vec![input])
        }

        pub fn with_responses(mut responses: Vec<&[u8]>) -> MockStream {
            MockStream {
                read: Cursor::new(responses.remove(0).to_vec()),
                next_reads: responses.into_iter().map(|arr| arr.to_vec()).collect(),
                write: vec![],
                is_closed: false,
                error_on_write: false,
                error_on_read: false,
                read_timeout: Cell::new(None),
                write_timeout: Cell::new(None),
            }
        }
    }

    impl Read for MockStream {
        fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
            if self.error_on_read {
                Err(io::Error::new(io::ErrorKind::Other, "mock error"))
            } else {
                match self.read.read(buf) {
                    Ok(n) => {
                        if self.read.position() as usize == self.read.get_ref().len() {
                            if self.next_reads.len() > 0 {
                                self.read = Cursor::new(self.next_reads.remove(0));
                            }
                        }
                        Ok(n)
                    }
                    r => r,
                }
            }
        }
    }

    impl Write for MockStream {
        fn write(&mut self, msg: &[u8]) -> io::Result<usize> {
            if self.error_on_write {
                Err(io::Error::new(io::ErrorKind::Other, "mock error"))
            } else {
                Write::write(&mut self.write, msg)
            }
        }

        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    impl NetworkStream for MockStream {
        fn peer_addr(&mut self) -> io::Result<SocketAddr> {
            Ok("127.0.0.1:1337".parse().unwrap())
        }

        fn set_read_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.read_timeout.set(dur);
            Ok(())
        }

        fn set_write_timeout(&self, dur: Option<Duration>) -> io::Result<()> {
            self.write_timeout.set(dur);
            Ok(())
        }

        fn close(&mut self, _how: Shutdown) -> io::Result<()> {
            self.is_closed = true;
            Ok(())
        }
    }

    #[test]
    fn test_handle() {
        let args = args::RequestLogAnalyzerArgs {
            filename: String::from("src/test/simple-1.log"),
            conditions: filter::FilterConditions {
                include_terms: None,
                exclude_terms: None,
                latest_time: None,
            },
            graphite_server: None,
            graphite_port: Some(2003),
            graphite_prefix: None,
            prometheus_listen: None,
        };

        fn run_fn(_: &args::RequestLogAnalyzerArgs) -> Option<analyzer::RequestLogAnalyzerResult> {
            Some(analyzer::RequestLogAnalyzerResult {
                count: 3,
                max: 100,
                min: 1,
                avg: 37,
                median: 10,
                percentile90: 100,
            })
        };

        let handler = HttpHandler {
            args: args,
            run: run_fn,
        };

        // Create a minimal HTTP request
        let mut request_mock_network_stream = MockStream::with_input(b"GET / HTTP/1.0\r\n\r\n");

        let mut reader = hyper::buffer::BufReader::new(&mut request_mock_network_stream as
                                                       &mut hyper::net::NetworkStream);

        let socket = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080);

        let request = hyper::server::Request::new(&mut reader, socket).unwrap();

        let mut headers = hyper::header::Headers::new();
        let mut response_mock_network_stream = MockStream::new();

        {
            let response = hyper::server::Response::new(&mut response_mock_network_stream,
                                                        &mut headers);

            handler.handle(request, response);
        }

        let result = str::from_utf8(&response_mock_network_stream.write).unwrap();

        assert!(result.contains("request_count 3"));
        assert!(result.contains("request_time_max 100"));
        assert!(result.contains("request_time_min 1"));
        assert!(result.contains("request_time_avg 37"));
        assert!(result.contains("request_time_median 10"));
        assert!(result.contains("request_time_percentile90 100"));

    }
}
<|endoftext|>


<commit_before>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

pub(crate) use self::ffi::is_nand;
use super::{Error, Result};
use std::{
    fs,
    io::{BufRead, BufReader},
    path::PathBuf,
};

pub(crate) fn target_device_from_ubi_volume_name(volume: &str) -> Result<PathBuf> {
    let re = regex::Regex::new(r"^Volume ID:   (?P<volume>\d+) \(on ubi(\d+)\)$").unwrap();
    walkdir::WalkDir::new("/dev")
        .min_depth(1)
        .into_iter()
        .filter_entry(|p| {
            p.file_name()
                .to_str()
                .map(|n| n.starts_with("ubi") && !n.contains('_'))
                .unwrap_or(false)
        })
        .filter_map(std::result::Result::ok)
        .find_map(|entry| {
            let path = entry.path();
            let output =
                easy_process::run(&format!("ubinfo {} -N {}", path.display(), volume)).ok()?;

            let line = output.stdout.lines().next()?;
            let re_match = re.captures(line)?;

            Some(PathBuf::from(format!(
                "{}_{}",
                path.display(),
                &re_match.name("volume").unwrap().as_str()
            )))
        })
        .ok_or_else(|| Error::NoUbiVolume(volume.to_owned()))
}

pub(crate) fn target_device_from_mtd_name(name: &str) -> Result<PathBuf> {
    let re =
        regex::Regex::new(r#"^(?P<dev>mtd\d): ([[:xdigit:]]+) ([[:xdigit:]]+) "(?P<name>.*)"$"#)
            .unwrap();
    let proc = fs::File::open("/proc/mtd")?;

    BufReader::new(proc)
        .lines()
        .filter_map(std::result::Result::ok)
        .find_map(|line| {
            re.captures(&line).and_then(|re_match| {
                let re_dev = re_match.name("dev").unwrap().as_str();
                let re_name = re_match.name("name").unwrap().as_str();
                if re_name == name { Some(PathBuf::from(format!("/dev/{}", re_dev))) } else { None }
            })
        })
        .ok_or_else(|| Error::NoMtdDevice(name.to_owned()))
}

mod ffi {
    use crate::utils::Result;
    use nix::ioctl_read;
    use std::{mem::MaybeUninit, os::unix::io::AsRawFd, path::Path};

    // From https://github.com/torvalds/linux/blob/master/include/uapi/mtd/mtd-abi.h
    const MTD_NANDFLASH: u8 = 4;
    const MTD_MLCNANDFLASH: u8 = 8;
    const MEMGETINFO: u8 = b'M';
    const MEMGETINFO_MODE: u8 = 1;

    #[repr(C)]
    pub struct mtd_info_user {
        kind: u8,
        flags: u32,
        size: u32,
        erasesize: u32,
        writesize: u32,
        oobsize: u32,
        padding: u64,
    }

    ioctl_read!(mtd_get_info, MEMGETINFO, MEMGETINFO_MODE, mtd_info_user);

    pub fn is_nand(device: &Path) -> Result<bool> {
        let device = std::fs::File::open(device)?;
        let info = unsafe {
            let mut info = MaybeUninit::<mtd_info_user>::uninit();
            mtd_get_info(device.as_raw_fd(), info.as_mut_ptr())?;
            info.assume_init()
        };

        Ok(info.kind == MTD_NANDFLASH || info.kind == MTD_MLCNANDFLASH)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use lazy_static::lazy_static;
    use pretty_assertions::assert_eq;
    use std::sync::{Arc, Mutex};

    pub(crate) struct FakeUbi {
        #[allow(dead_code)]
        pub(crate) mtd_guard: FakeMtd,
    }

    impl FakeUbi {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeUbi> {
            let mtd_guard = FakeMtd::new(&["system"], kind)?;
            easy_process::run("modprobe ubi mtd=0")?;

            // Ubi created here so if anything fails the Drop will still be executed
            let ubi = FakeUbi { mtd_guard };

            for name in names {
                easy_process::run(&format!("ubimkvol /dev/ubi0 -N {} -s 1MiB", name))?;
            }

            Ok(ubi)
        }
    }

    impl Drop for FakeUbi {
        fn drop(&mut self) {
            if let Err(e) = easy_process::run(&format!("rmmod ubi")) {
                eprintln!("Failed to cleanup FakeUbi, Error: {}", e);
            }
        }
    }

    pub(crate) struct FakeMtd {
        pub(crate) devices: Vec<PathBuf>,
        pub(crate) kind: MtdKind,
    }

    pub(crate) enum MtdKind {
        Nand,
        Nor,
    }

    impl FakeMtd {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeMtd> {
            match kind {
                MtdKind::Nand => easy_process::run("modprobe nandsim second_id_byte=0x36"),
                MtdKind::Nor => easy_process::run("modprobe mtdram total_size=20000"),
            }?;

            // FakeMtd created here so if any subsequent command fails the drop will still
            // be called to cleanup mtd devices
            let mut mtd = FakeMtd { devices: vec![], kind };
            names.iter().enumerate().try_for_each(|(i, name)| {
                easy_process::run(&format!("mtdpart add /dev/mtd0 {} {} {}", name, i * 100, 100))
                    .map(|_| mtd.devices.push(PathBuf::from(format!("/dev/mtd{}", i + 1))))
            })?;

            Ok(mtd)
        }
    }

    impl Drop for FakeMtd {
        fn drop(&mut self) {
            let module = match self.kind {
                MtdKind::Nand => "nandsim",
                MtdKind::Nor => "mtdram",
            };

            // Sleep time for nandsim to sync and avoid errors
            std::thread::sleep(std::time::Duration::from_millis(500));

            if let Err(e) = easy_process::run(&format!("rmmod {}", module)) {
                eprintln!("Failed to cleanup FakeMtd, Error: {}", e);
            }
        }
    }

    // Used to serialize access to MTD devices
    lazy_static! {
        pub static ref SERIALIZE: Arc<Mutex<()>> = Arc::new(Mutex::default());
    }

    #[test]
    #[ignore]
    fn device_from_mtd_name() {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0", "system1"];

        let mtd = FakeMtd::new(&dev_names, MtdKind::Nor).unwrap();

        assert_eq!(
            dev_names
                .into_iter()
                .map(target_device_from_mtd_name)
                .map(Result::unwrap)
                .collect::<Vec<_>>(),
            mtd.devices,
        );
        assert!(target_device_from_mtd_name("some_inexistent_device").is_err());
    }

    #[test]
    #[ignore]
    fn test_is_nand() {
        let _lock = SERIALIZE.lock();

        {
            let _mtd = FakeMtd::new(&[], MtdKind::Nand).unwrap();
            assert_eq!(is_nand(&PathBuf::from("/dev/mtd0")).unwrap(), true);
        }
        {
            let _mtd = FakeMtd::new(&[], MtdKind::Nor).unwrap();
            assert_eq!(is_nand(&PathBuf::from("/dev/mtd0")).unwrap(), false);
        }
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec!["some_ui_volume", "another_ubi_volume"];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[1]).unwrap(),
            PathBuf::from("/dev/ubi0_1")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[0]).unwrap(),
            PathBuf::from("/dev/ubi0_0")
        );
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name_multiple_volumes() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec![
            "volume0", "volume1", "volume2", "volume3", "volume4", "volume5", "volume6", "volume7",
            "volume8", "volume9", "volume10", "volume11", "volume12", "volume13",
        ];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[8]).unwrap(),
            PathBuf::from("/dev/ubi0_8")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[12]).unwrap(),
            PathBuf::from("/dev/ubi0_12")
        );
    }
}
<commit_msg>updatehub: Use pattern match for loop on FakeMtd::new<commit_after>// Copyright (C) 2019 O.S. Systems Sofware LTDA
//
// SPDX-License-Identifier: Apache-2.0

pub(crate) use self::ffi::is_nand;
use super::{Error, Result};
use std::{
    fs,
    io::{BufRead, BufReader},
    path::PathBuf,
};

pub(crate) fn target_device_from_ubi_volume_name(volume: &str) -> Result<PathBuf> {
    let re = regex::Regex::new(r"^Volume ID:   (?P<volume>\d+) \(on ubi(\d+)\)$").unwrap();
    walkdir::WalkDir::new("/dev")
        .min_depth(1)
        .into_iter()
        .filter_entry(|p| {
            p.file_name()
                .to_str()
                .map(|n| n.starts_with("ubi") && !n.contains('_'))
                .unwrap_or(false)
        })
        .filter_map(std::result::Result::ok)
        .find_map(|entry| {
            let path = entry.path();
            let output =
                easy_process::run(&format!("ubinfo {} -N {}", path.display(), volume)).ok()?;

            let line = output.stdout.lines().next()?;
            let re_match = re.captures(line)?;

            Some(PathBuf::from(format!(
                "{}_{}",
                path.display(),
                &re_match.name("volume").unwrap().as_str()
            )))
        })
        .ok_or_else(|| Error::NoUbiVolume(volume.to_owned()))
}

pub(crate) fn target_device_from_mtd_name(name: &str) -> Result<PathBuf> {
    let re =
        regex::Regex::new(r#"^(?P<dev>mtd\d): ([[:xdigit:]]+) ([[:xdigit:]]+) "(?P<name>.*)"$"#)
            .unwrap();
    let proc = fs::File::open("/proc/mtd")?;

    BufReader::new(proc)
        .lines()
        .filter_map(std::result::Result::ok)
        .find_map(|line| {
            re.captures(&line).and_then(|re_match| {
                let re_dev = re_match.name("dev").unwrap().as_str();
                let re_name = re_match.name("name").unwrap().as_str();
                if re_name == name { Some(PathBuf::from(format!("/dev/{}", re_dev))) } else { None }
            })
        })
        .ok_or_else(|| Error::NoMtdDevice(name.to_owned()))
}

mod ffi {
    use crate::utils::Result;
    use nix::ioctl_read;
    use std::{mem::MaybeUninit, os::unix::io::AsRawFd, path::Path};

    // From https://github.com/torvalds/linux/blob/master/include/uapi/mtd/mtd-abi.h
    const MTD_NANDFLASH: u8 = 4;
    const MTD_MLCNANDFLASH: u8 = 8;
    const MEMGETINFO: u8 = b'M';
    const MEMGETINFO_MODE: u8 = 1;

    #[repr(C)]
    pub struct mtd_info_user {
        kind: u8,
        flags: u32,
        size: u32,
        erasesize: u32,
        writesize: u32,
        oobsize: u32,
        padding: u64,
    }

    ioctl_read!(mtd_get_info, MEMGETINFO, MEMGETINFO_MODE, mtd_info_user);

    pub fn is_nand(device: &Path) -> Result<bool> {
        let device = std::fs::File::open(device)?;
        let info = unsafe {
            let mut info = MaybeUninit::<mtd_info_user>::uninit();
            mtd_get_info(device.as_raw_fd(), info.as_mut_ptr())?;
            info.assume_init()
        };

        Ok(info.kind == MTD_NANDFLASH || info.kind == MTD_MLCNANDFLASH)
    }
}

#[cfg(test)]
pub(crate) mod tests {
    use super::*;
    use lazy_static::lazy_static;
    use pretty_assertions::assert_eq;
    use std::sync::{Arc, Mutex};

    pub(crate) struct FakeUbi {
        #[allow(dead_code)]
        pub(crate) mtd_guard: FakeMtd,
    }

    impl FakeUbi {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeUbi> {
            let mtd_guard = FakeMtd::new(&["system"], kind)?;
            easy_process::run("modprobe ubi mtd=0")?;

            // Ubi created here so if anything fails the Drop will still be executed
            let ubi = FakeUbi { mtd_guard };

            for name in names {
                easy_process::run(&format!("ubimkvol /dev/ubi0 -N {} -s 1MiB", name))?;
            }

            Ok(ubi)
        }
    }

    impl Drop for FakeUbi {
        fn drop(&mut self) {
            if let Err(e) = easy_process::run(&format!("rmmod ubi")) {
                eprintln!("Failed to cleanup FakeUbi, Error: {}", e);
            }
        }
    }

    pub(crate) struct FakeMtd {
        pub(crate) devices: Vec<PathBuf>,
        pub(crate) kind: MtdKind,
    }

    pub(crate) enum MtdKind {
        Nand,
        Nor,
    }

    impl FakeMtd {
        pub(crate) fn new(names: &[&str], kind: MtdKind) -> Result<FakeMtd> {
            match kind {
                MtdKind::Nand => easy_process::run("modprobe nandsim second_id_byte=0x36"),
                MtdKind::Nor => easy_process::run("modprobe mtdram total_size=20000"),
            }?;

            // FakeMtd created here so if any subsequent command fails the drop will still
            // be called to cleanup mtd devices
            let mut mtd = FakeMtd { devices: vec![], kind };
            for (i, name) in names.iter().enumerate() {
                easy_process::run(&format!("mtdpart add /dev/mtd0 {} {} {}", name, i * 100, 100))?;
                mtd.devices.push(PathBuf::from(format!("/dev/mtd{}", i + 1)));
            }

            Ok(mtd)
        }
    }

    impl Drop for FakeMtd {
        fn drop(&mut self) {
            let module = match self.kind {
                MtdKind::Nand => "nandsim",
                MtdKind::Nor => "mtdram",
            };

            // Sleep time for nandsim to sync and avoid errors
            std::thread::sleep(std::time::Duration::from_millis(500));

            if let Err(e) = easy_process::run(&format!("rmmod {}", module)) {
                eprintln!("Failed to cleanup FakeMtd, Error: {}", e);
            }
        }
    }

    // Used to serialize access to MTD devices
    lazy_static! {
        pub static ref SERIALIZE: Arc<Mutex<()>> = Arc::new(Mutex::default());
    }

    #[test]
    #[ignore]
    fn device_from_mtd_name() {
        let _lock = SERIALIZE.lock();
        let dev_names = vec!["system0", "system1"];

        let mtd = FakeMtd::new(&dev_names, MtdKind::Nor).unwrap();

        assert_eq!(
            dev_names
                .into_iter()
                .map(target_device_from_mtd_name)
                .map(Result::unwrap)
                .collect::<Vec<_>>(),
            mtd.devices,
        );
        assert!(target_device_from_mtd_name("some_inexistent_device").is_err());
    }

    #[test]
    #[ignore]
    fn test_is_nand() {
        let _lock = SERIALIZE.lock();

        {
            let _mtd = FakeMtd::new(&[], MtdKind::Nand).unwrap();
            assert_eq!(is_nand(&PathBuf::from("/dev/mtd0")).unwrap(), true);
        }
        {
            let _mtd = FakeMtd::new(&[], MtdKind::Nor).unwrap();
            assert_eq!(is_nand(&PathBuf::from("/dev/mtd0")).unwrap(), false);
        }
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec!["some_ui_volume", "another_ubi_volume"];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[1]).unwrap(),
            PathBuf::from("/dev/ubi0_1")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[0]).unwrap(),
            PathBuf::from("/dev/ubi0_0")
        );
    }

    #[test]
    #[ignore]
    fn device_from_ubi_volume_name_multiple_volumes() {
        let _lock = SERIALIZE.lock();
        let volume_names = vec![
            "volume0", "volume1", "volume2", "volume3", "volume4", "volume5", "volume6", "volume7",
            "volume8", "volume9", "volume10", "volume11", "volume12", "volume13",
        ];

        let _ubi = FakeUbi::new(&volume_names, MtdKind::Nor).unwrap();
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[8]).unwrap(),
            PathBuf::from("/dev/ubi0_8")
        );
        assert_eq!(
            target_device_from_ubi_volume_name(volume_names[12]).unwrap(),
            PathBuf::from("/dev/ubi0_12")
        );
    }
}
<|endoftext|>


<commit_before>use std::time::Duration;
use std::collections::HashMap;

use telebot;
use telebot::functions::*;
use tokio_core::reactor::{Interval, Handle, Timeout};
use futures::{self, Future, Stream, IntoFuture};
use tokio_curl::Session;
use regex::Regex;

use data;
use feed;
use utlis::{Escape, EscapeUrl, send_multiple_messages, format_and_split_msgs, to_chinese_error_msg,
            truncate_message, chat_is_unavailable, TELEGRAM_MAX_MSG_LEN};

// 5 minute
const FREQUENCY_SECOND: u64 = 300;

lazy_static!{
    // it's different from `feed::HOST`, so maybe need a better name?
    static ref HOST: Regex = Regex::new(r"^(?:https?://)?([^/]+)").unwrap();
}

pub fn spawn_fetcher(bot: telebot::RcBot, db: data::Database, handle: Handle) {
    handle.clone().spawn(Interval::new(Duration::from_secs(FREQUENCY_SECOND), &handle)
                             .expect("failed to start feed loop")
                             .map_err(|e| error!("feed loop error: {}", e))
                             .for_each(move |_| {
        let feeds = db.get_all_feeds();
        let grouped_feeds = grouping_by_host(feeds);
        let handle2 = handle.clone();
        let bot = bot.clone();
        let db = db.clone();
        let fetcher = futures::stream::iter(grouped_feeds.into_iter().map(Ok))
            .for_each(move |group| {
                let session = Session::new(handle2.clone());
                let bot = bot.clone();
                let db = db.clone();
                let group_fetcher = futures::stream::iter(group.into_iter().map(Ok))
                    .for_each(move |feed| {
                                  fetch_feed_updates(bot.clone(), db.clone(), &session, feed)
                              });
                handle2.spawn(group_fetcher);
                Timeout::new(Duration::from_secs(1), &handle2)
                    .expect("failed to start sleep")
                    .map_err(|e| error!("feed loop sleep error: {}", e))
            });
        handle.spawn(fetcher);
        Ok(())
    }))
}

fn grouping_by_host(feeds: Vec<data::Feed>) -> Vec<Vec<data::Feed>> {
    let mut result = HashMap::new();
    for feed in feeds {
        let host = get_host(&feed.link).to_owned();
        let group = result.entry(host).or_insert_with(Vec::new);
        group.push(feed);
    }
    result.into_iter().map(|(_, v)| v).collect()
}

fn get_host(url: &str) -> &str {
    HOST.captures(url).map_or(url, |r| r.get(0).unwrap().as_str())
}

fn fetch_feed_updates<'a>(bot: telebot::RcBot,
                          db: data::Database,
                          session: &Session,
                          feed: data::Feed)
                          -> impl Future<Item = (), Error = ()> + 'a {
    let bot_ = bot.clone();
    let db_ = db.clone();
    let feed_ = feed.clone();
    feed::fetch_feed(session, feed.link.to_owned())
        .map(move |rss| (bot_, db_, rss, feed_))
        .or_else(move |e| {
            // 1440 * 5 minute = 5 days
            if db.inc_error_count(&feed.link) > 1440 {
                    Err((bot, db, feed))
                } else {
                    Ok(())
                }
                .into_future()
                .or_else(|(bot, db, feed)| {
                    db.reset_error_count(&feed.link);
                    let err_msg = to_chinese_error_msg(e);
                    let mut msgs = Vec::with_capacity(feed.subscribers.len());
                    for &subscriber in &feed.subscribers {
                        let m = bot.message(subscriber,
                                            format!("《{}》\
                                                     已经连续 5 天拉取出错 ({}),\
                                                     可能已经关闭, 请取消订阅",
                                                    EscapeUrl(&feed.link),
                                                    Escape(&feed.title),
                                                    Escape(&err_msg)))
                            .parse_mode("HTML")
                            .disable_web_page_preview(true)
                            .send();
                        let db = db.clone();
                        let r = m.map_err(move |e| {
                            match e {
                                telebot::error::Error::Telegram(ref s)
                                    if chat_is_unavailable(s) => {
                                    db.delete_subscriber(subscriber);
                                }
                                _ => {
                                    warn!("failed to send error to {}, {:?}", subscriber, e);
                                }
                            };
                        });
                        // if not use Box, rustc will panic
                        msgs.push(Box::new(r) as Box<Future<Item = _, Error = _>>);
                    }
                    futures::future::join_all(msgs).then(|_| Err(()))
                })
                .and_then(|_| Err(()))
        })
        .and_then(|(bot, db, rss, feed)| {
            if rss.title != feed.title {
                db.update_title(&feed.link, &rss.title);
            }
            let updates = db.update(&feed.link, rss.items);
            if updates.is_empty() {
                futures::future::err(())
            } else {
                futures::future::ok((bot, db, feed, rss.title, rss.link, updates))
            }
        })
        .and_then(|(bot, db, feed, rss_title, rss_link, updates)| {
            let msgs =
                format_and_split_msgs(format!("{}", Escape(&rss_title)), &updates, |item| {
                    let title = item.title
                        .as_ref()
                        .map(|s| s.as_str())
                        .unwrap_or_else(|| &rss_title);
                    let link = item.link
                        .as_ref()
                        .map(|s| s.as_str())
                        .unwrap_or_else(|| &rss_link);
                    format!("{}",
                            EscapeUrl(link),
                            Escape(&truncate_message(title, TELEGRAM_MAX_MSG_LEN - 500)))
                });

            let mut msg_futures = Vec::with_capacity(feed.subscribers.len());
            for &subscriber in &feed.subscribers {
                let db = db.clone();
                let bot = bot.clone();
                let r = send_multiple_messages(&bot, subscriber, msgs.clone())
                    .map_err(move |e| {
                        match e {
                            telebot::error::Error::Telegram(ref s) if chat_is_unavailable(s) => {
                                db.delete_subscriber(subscriber);
                            }
                            _ => {
                                warn!("failed to send updates to {}, {:?}", subscriber, e);
                            }
                        };
                    });
                msg_futures.push(Box::new(r) as Box<Future<Item = _, Error = _>>);
            }
            futures::future::join_all(msg_futures).then(|_| Ok(()))
        })
}
<commit_msg>ignore error in `group_fetcher`<commit_after>use std::time::Duration;
use std::collections::HashMap;

use telebot;
use telebot::functions::*;
use tokio_core::reactor::{Interval, Handle, Timeout};
use futures::{self, Future, Stream, IntoFuture};
use tokio_curl::Session;
use regex::Regex;

use data;
use feed;
use utlis::{Escape, EscapeUrl, send_multiple_messages, format_and_split_msgs, to_chinese_error_msg,
            truncate_message, chat_is_unavailable, TELEGRAM_MAX_MSG_LEN};

// 5 minute
const FREQUENCY_SECOND: u64 = 300;

lazy_static!{
    // it's different from `feed::HOST`, so maybe need a better name?
    static ref HOST: Regex = Regex::new(r"^(?:https?://)?([^/]+)").unwrap();
}

pub fn spawn_fetcher(bot: telebot::RcBot, db: data::Database, handle: Handle) {
    handle.clone().spawn(Interval::new(Duration::from_secs(FREQUENCY_SECOND), &handle)
                             .expect("failed to start feed loop")
                             .map_err(|e| error!("feed loop error: {}", e))
                             .for_each(move |_| {
        let feeds = db.get_all_feeds();
        let grouped_feeds = grouping_by_host(feeds);
        let handle2 = handle.clone();
        let bot = bot.clone();
        let db = db.clone();
        let fetcher = futures::stream::iter(grouped_feeds.into_iter().map(Ok))
            .for_each(move |group| {
                let session = Session::new(handle2.clone());
                let bot = bot.clone();
                let db = db.clone();
                let group_fetcher = futures::stream::iter(group.into_iter().map(Ok))
                    .for_each(move |feed| {
                                  fetch_feed_updates(bot.clone(), db.clone(), &session, feed)
                                      .then(|_| Ok(()))
                              });
                handle2.spawn(group_fetcher);
                Timeout::new(Duration::from_secs(1), &handle2)
                    .expect("failed to start sleep")
                    .map_err(|e| error!("feed loop sleep error: {}", e))
            });
        handle.spawn(fetcher);
        Ok(())
    }))
}

fn grouping_by_host(feeds: Vec<data::Feed>) -> Vec<Vec<data::Feed>> {
    let mut result = HashMap::new();
    for feed in feeds {
        let host = get_host(&feed.link).to_owned();
        let group = result.entry(host).or_insert_with(Vec::new);
        group.push(feed);
    }
    result.into_iter().map(|(_, v)| v).collect()
}

fn get_host(url: &str) -> &str {
    HOST.captures(url).map_or(url, |r| r.get(0).unwrap().as_str())
}

fn fetch_feed_updates<'a>(bot: telebot::RcBot,
                          db: data::Database,
                          session: &Session,
                          feed: data::Feed)
                          -> impl Future<Item = (), Error = ()> + 'a {
    let bot_ = bot.clone();
    let db_ = db.clone();
    let feed_ = feed.clone();
    feed::fetch_feed(session, feed.link.to_owned())
        .map(move |rss| (bot_, db_, rss, feed_))
        .or_else(move |e| {
            // 1440 * 5 minute = 5 days
            if db.inc_error_count(&feed.link) > 1440 {
                    Err((bot, db, feed))
                } else {
                    Ok(())
                }
                .into_future()
                .or_else(|(bot, db, feed)| {
                    db.reset_error_count(&feed.link);
                    let err_msg = to_chinese_error_msg(e);
                    let mut msgs = Vec::with_capacity(feed.subscribers.len());
                    for &subscriber in &feed.subscribers {
                        let m = bot.message(subscriber,
                                            format!("《{}》\
                                                     已经连续 5 天拉取出错 ({}),\
                                                     可能已经关闭, 请取消订阅",
                                                    EscapeUrl(&feed.link),
                                                    Escape(&feed.title),
                                                    Escape(&err_msg)))
                            .parse_mode("HTML")
                            .disable_web_page_preview(true)
                            .send();
                        let db = db.clone();
                        let r = m.map_err(move |e| {
                            match e {
                                telebot::error::Error::Telegram(ref s)
                                    if chat_is_unavailable(s) => {
                                    db.delete_subscriber(subscriber);
                                }
                                _ => {
                                    warn!("failed to send error to {}, {:?}", subscriber, e);
                                }
                            };
                        });
                        // if not use Box, rustc will panic
                        msgs.push(Box::new(r) as Box<Future<Item = _, Error = _>>);
                    }
                    futures::future::join_all(msgs).then(|_| Err(()))
                })
                .and_then(|_| Err(()))
        })
        .and_then(|(bot, db, rss, feed)| {
            if rss.title != feed.title {
                db.update_title(&feed.link, &rss.title);
            }
            let updates = db.update(&feed.link, rss.items);
            if updates.is_empty() {
                futures::future::err(())
            } else {
                futures::future::ok((bot, db, feed, rss.title, rss.link, updates))
            }
        })
        .and_then(|(bot, db, feed, rss_title, rss_link, updates)| {
            let msgs =
                format_and_split_msgs(format!("{}", Escape(&rss_title)), &updates, |item| {
                    let title = item.title
                        .as_ref()
                        .map(|s| s.as_str())
                        .unwrap_or_else(|| &rss_title);
                    let link = item.link
                        .as_ref()
                        .map(|s| s.as_str())
                        .unwrap_or_else(|| &rss_link);
                    format!("{}",
                            EscapeUrl(link),
                            Escape(&truncate_message(title, TELEGRAM_MAX_MSG_LEN - 500)))
                });

            let mut msg_futures = Vec::with_capacity(feed.subscribers.len());
            for &subscriber in &feed.subscribers {
                let db = db.clone();
                let bot = bot.clone();
                let r = send_multiple_messages(&bot, subscriber, msgs.clone())
                    .map_err(move |e| {
                        match e {
                            telebot::error::Error::Telegram(ref s) if chat_is_unavailable(s) => {
                                db.delete_subscriber(subscriber);
                            }
                            _ => {
                                warn!("failed to send updates to {}, {:?}", subscriber, e);
                            }
                        };
                    });
                msg_futures.push(Box::new(r) as Box<Future<Item = _, Error = _>>);
            }
            futures::future::join_all(msg_futures).then(|_| Ok(()))
        })
}
<|endoftext|>