<commit_before>//! Demosaicing using linear interpolation.
//!
//! ```text
//!   green_kernel = (1 / 4) *
//!       [ 0 1 0
//!       ; 1 4 1
//!       ; 0 1 0 ];
//!
//!   red/blue_kernel = (1 / 4) *
//!       [ 1 2 1
//!       ; 2 4 2
//!       ; 1 2 1 ];
//! ```

use std::io::Read;

use ::{BayerDepth,BayerError,BayerResult,CFA,RasterMut};
use bayer::{BayerRead8,BayerRead16};
use border_replicate::*;
use demosaic::check_depth;

const PADDING: usize = 1;

pub fn run(r: &mut Read,
        depth: BayerDepth, cfa: CFA, dst: &mut RasterMut)
        -> BayerResult<()> {
    if dst.w < 2 || dst.h < 2 {
        return Err(BayerError::WrongResolution);
    }
    if !check_depth(depth, dst.depth) {
        return Err(BayerError::WrongDepth);
    }

    match depth {
        BayerDepth::Depth8 => debayer_u8(r, cfa, dst),
        BayerDepth::Depth16BE => debayer_u16(r, true, cfa, dst),
        BayerDepth::Depth16LE => debayer_u16(r, false, cfa, dst),
    }
}

macro_rules! apply_kernel_row {
    ($T:ty; $row:ident, $prev:expr, $curr:expr, $next:expr, $cfa:expr, $w:expr) => {{
        let (mut i, cfa_c, cfa_g) =
            if $cfa == CFA::BGGR || $cfa == CFA::RGGB {
                (0, $cfa, $cfa.next_x())
            } else {
                apply_kernel_g!($T; $row, $prev, $curr, $next, $cfa, 0);
                (1, $cfa.next_x(), $cfa)
            };

        while i + 1 < $w {
            apply_kernel_c!($T; $row, $prev, $curr, $next, cfa_c, i);
            apply_kernel_g!($T; $row, $prev, $curr, $next, cfa_g, i + 1);
            i = i + 2;
        }

        if i < $w {
            apply_kernel_c!($T; $row, $prev, $curr, $next, cfa_c, i);
        }
    }}
}

macro_rules! apply_kernel_c {
    ($T:ty; $row:ident, $prev:expr, $curr:expr, $next:expr, $cfa:expr, $i:expr) => {{
        // current = B/R, diagonal = R/B.
        let (c, d) = if $cfa == CFA::BGGR { (2, 0) } else { (0, 2) };
        let j = $i + PADDING;

        $row[3 * $i + c] = $curr[j];
        $row[3 * $i + 1]
            = ((  $prev[j] as u32
                + $curr[j - 1] as u32 + $curr[j + 1] as u32
                + $next[j] as u32) / 4) as $T;
        $row[3 * $i + d]
            = ((  $prev[j - 1] as u32 + $prev[j + 1] as u32
                + $next[j - 1] as u32 + $next[j + 1] as u32) / 4) as $T;
    }}
}

macro_rules! apply_kernel_g {
    ($T:ty; $row:ident, $prev:expr, $curr:expr, $next:expr, $cfa:expr, $i:expr) => {{
        // horizontal = B/R, vertical = R/G.
        let (h, v) = if $cfa == CFA::GBRG { (2, 0) } else { (0, 2) };
        let j = $i + PADDING;

        $row[3 * $i + h]
            = (($curr[j - 1] as u32 + $curr[j + 1] as u32) / 2) as $T;
        $row[3 * $i + 1] = $curr[j];
        $row[3 * $i + v]
            = (($prev[j] as u32 + $next[j] as u32) / 2) as $T;
    }}
}

/*--------------------------------------------------------------*/

fn debayer_u8(r: &mut Read, cfa: CFA, dst: &mut RasterMut)
        -> BayerResult<()> {
    let (w, h) = (dst.w, dst.h);
    let mut prev = vec![0u8; 2 * PADDING + w];
    let mut curr = vec![0u8; 2 * PADDING + w];
    let mut next = vec![0u8; 2 * PADDING + w];
    let mut cfa = cfa;

    let rdr = BorderReplicate8::new(w, PADDING);
    rdr.read_line(r, &mut curr)?;
    rdr.read_line(r, &mut next)?;

    {   // y = 0.
        let row = dst.borrow_row_u8_mut(0);
        apply_kernel_row!(u8; row, next, curr, next, cfa, w);
        cfa = cfa.next_y();
    }

    for y in 1..(h - 1) {
        rotate!(prev <- curr <- next);
        rdr.read_line(r, &mut next)?;

        let row = dst.borrow_row_u8_mut(y);
        apply_kernel_row!(u8; row, prev, curr, next, cfa, w);
        cfa = cfa.next_y();
    }

    {   // y = h - 1.
        let row = dst.borrow_row_u8_mut(h - 1);
        apply_kernel_row!(u8; row, curr, next, curr, cfa, w);
    }

    Ok(())
}

fn debayer_u16(r: &mut Read, be: bool, cfa: CFA, dst: &mut RasterMut)
        -> BayerResult<()> {
    let (w, h) = (dst.w, dst.h);
    let mut prev = vec![0u16; 2 * PADDING + w];
    let mut curr = vec![0u16; 2 * PADDING + w];
    let mut next = vec![0u16; 2 * PADDING + w];
    let mut cfa = cfa;

    let rdr: Box<BayerRead16> = if be {
        Box::new(BorderReplicate16BE::new(w, PADDING))
    } else {
        Box::new(BorderReplicate16LE::new(w, PADDING))
    };
    rdr.read_line(r, &mut curr)?;
    rdr.read_line(r, &mut next)?;

    {   // y = 0.
        let row = dst.borrow_row_u16_mut(0);
        apply_kernel_row!(u16; row, next, curr, next, cfa, w);
        cfa = cfa.next_y();
    }

    for y in 1..(h - 1) {
        rotate!(prev <- curr <- next);
        rdr.read_line(r, &mut next)?;

        let row = dst.borrow_row_u16_mut(y);
        apply_kernel_row!(u16; row, prev, curr, next, cfa, w);
        cfa = cfa.next_y();
    }

    {   // y = h - 1.
        let row = dst.borrow_row_u16_mut(h - 1);
        apply_kernel_row!(u16; row, curr, next, curr, cfa, w);
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;
    use ::{CFA,RasterDepth,RasterMut};
    use super::debayer_u8;

    #[test]
    fn test_even() {
        // R: set.seed(0); matrix(floor(runif(n=16, min=0, max=256)), nrow=4, byrow=TRUE)
        let src = [
            229, 67, 95,146,
            232, 51,229,241,
            169,161, 15, 52,
             45,175, 98,197 ];

        let expected = [
            229,149, 51,  162, 67, 51,   95,167,146,   95,146,241,
            199,232, 51,  127,172, 51,   55,229,146,   55,164,241,
            169,149,113,   92,161,113,   15,135,166,   15, 52,219,
            169, 45,175,   92,116,175,   15, 98,186,   15, 75,197 ];

        const IMG_W: usize = 4;
        const IMG_H: usize = 4;
        let mut dst = [0u8; 3 * IMG_W * IMG_H];

        let res = debayer_u8(&mut Cursor::new(&src[..]), CFA::RGGB,
                &mut RasterMut::new(IMG_W, IMG_H, RasterDepth::Depth8, &mut dst[..]));
        assert!(res.is_ok());
        assert_eq!(&dst[..], &expected[..]);
    }

    #[test]
    fn test_odd() {
        // R: set.seed(0); matrix(floor(runif(n=9, min=0, max=256)), nrow=3, byrow=TRUE)
        let src = [
            229, 67, 95,
            146,232, 51,
            229,241,169 ];

        let expected = [
            229,106,232,  162, 67,232,   95, 59,232,
            229,146,232,  180,126,232,  132, 51,232,
            229,193,232,  199,241,232,  169,146,232 ];

        const IMG_W: usize = 3;
        const IMG_H: usize = 3;
        let mut buf = [0u8; 3 * IMG_W * IMG_H];

        let res = debayer_u8(&mut Cursor::new(&src[..]), CFA::RGGB,
                &mut RasterMut::new(IMG_W, IMG_H, RasterDepth::Depth8, &mut buf));
        assert!(res.is_ok());
        assert_eq!(&buf[..], &expected[..]);
    }
}
<commit_msg>Parallel linear interpolation using Rayon.<commit_after>//! Demosaicing using linear interpolation.
//!
//! ```text
//!   green_kernel = (1 / 4) *
//!       [ 0 1 0
//!       ; 1 4 1
//!       ; 0 1 0 ];
//!
//!   red/blue_kernel = (1 / 4) *
//!       [ 1 2 1
//!       ; 2 4 2
//!       ; 1 2 1 ];
//! ```

use std::io::Read;

#[cfg(feature = "rayon")]
use std::slice;

#[cfg(feature = "rayon")]
use rayon::prelude::*;

use ::{BayerDepth,BayerError,BayerResult,CFA,RasterMut};
use bayer::{BayerRead8,BayerRead16};
use border_replicate::*;
use demosaic::check_depth;

const PADDING: usize = 1;

pub fn run(r: &mut Read,
        depth: BayerDepth, cfa: CFA, dst: &mut RasterMut)
        -> BayerResult<()> {
    if dst.w < 2 || dst.h < 2 {
        return Err(BayerError::WrongResolution);
    }
    if !check_depth(depth, dst.depth) {
        return Err(BayerError::WrongDepth);
    }

    match depth {
        BayerDepth::Depth8 => debayer_u8(r, cfa, dst),
        BayerDepth::Depth16BE => debayer_u16(r, true, cfa, dst),
        BayerDepth::Depth16LE => debayer_u16(r, false, cfa, dst),
    }
}

macro_rules! apply_kernel_row {
    ($T:ty; $row:ident, $prev:expr, $curr:expr, $next:expr, $cfa:expr, $w:expr) => {{
        let (mut i, cfa_c, cfa_g) =
            if $cfa == CFA::BGGR || $cfa == CFA::RGGB {
                (0, $cfa, $cfa.next_x())
            } else {
                apply_kernel_g!($T; $row, $prev, $curr, $next, $cfa, 0);
                (1, $cfa.next_x(), $cfa)
            };

        while i + 1 < $w {
            apply_kernel_c!($T; $row, $prev, $curr, $next, cfa_c, i);
            apply_kernel_g!($T; $row, $prev, $curr, $next, cfa_g, i + 1);
            i = i + 2;
        }

        if i < $w {
            apply_kernel_c!($T; $row, $prev, $curr, $next, cfa_c, i);
        }
    }}
}

macro_rules! apply_kernel_c {
    ($T:ty; $row:ident, $prev:expr, $curr:expr, $next:expr, $cfa:expr, $i:expr) => {{
        // current = B/R, diagonal = R/B.
        let (c, d) = if $cfa == CFA::BGGR { (2, 0) } else { (0, 2) };
        let j = $i + PADDING;

        $row[3 * $i + c] = $curr[j];
        $row[3 * $i + 1]
            = ((  $prev[j] as u32
                + $curr[j - 1] as u32 + $curr[j + 1] as u32
                + $next[j] as u32) / 4) as $T;
        $row[3 * $i + d]
            = ((  $prev[j - 1] as u32 + $prev[j + 1] as u32
                + $next[j - 1] as u32 + $next[j + 1] as u32) / 4) as $T;
    }}
}

macro_rules! apply_kernel_g {
    ($T:ty; $row:ident, $prev:expr, $curr:expr, $next:expr, $cfa:expr, $i:expr) => {{
        // horizontal = B/R, vertical = R/G.
        let (h, v) = if $cfa == CFA::GBRG { (2, 0) } else { (0, 2) };
        let j = $i + PADDING;

        $row[3 * $i + h]
            = (($curr[j - 1] as u32 + $curr[j + 1] as u32) / 2) as $T;
        $row[3 * $i + 1] = $curr[j];
        $row[3 * $i + v]
            = (($prev[j] as u32 + $next[j] as u32) / 2) as $T;
    }}
}

/*--------------------------------------------------------------*/
/* Rayon                                                        */
/*--------------------------------------------------------------*/

#[cfg(feature = "rayon")]
fn debayer_u8(r: &mut Read, cfa: CFA, dst: &mut RasterMut)
        -> BayerResult<()> {
    let (w, h) = (dst.w, dst.h);
    let mut data = vec![0u8; (2 * PADDING + w) * (2 * PADDING + h)];

    // Read all data.
    {
        let stride = 2 * PADDING + w;
        let rdr = BorderReplicate8::new(w, PADDING);

        for mut row in data.chunks_mut(stride).skip(PADDING).take(h) {
            rdr.read_line(r, &mut row)?;
        }

        {
            let (top, src) = data.split_at_mut(stride * PADDING);
            top[(stride * 0)..(stride * 1)].copy_from_slice(
                    &src[(stride * 1)..(stride * 2)]);
        }

        {
            let (src, bottom) = data.split_at_mut(stride * (h + PADDING));
            let yy = PADDING + h;
            bottom[(stride * 0)..(stride * 1)].copy_from_slice(
                    &src[(stride * (yy - 2))..(stride * (yy - 1))]);
        }
    }

    dst.buf.par_chunks_mut(dst.stride).enumerate()
            .for_each(|(y, mut row)| {
        let stride = 2 * PADDING + w;
        let prev = &data[(stride * (PADDING + y - 1)) .. (stride * (PADDING + y + 0))];
        let curr = &data[(stride * (PADDING + y + 0)) .. (stride * (PADDING + y + 1))];
        let next = &data[(stride * (PADDING + y + 1)) .. (stride * (PADDING + y + 2))];
        let cfa_y = if y % 2 == 0 { cfa } else { cfa.next_y() };

        apply_kernel_row!(u8; row, prev, curr, next, cfa_y, w);
    });

    Ok(())
}

#[cfg(feature = "rayon")]
fn debayer_u16(r: &mut Read, be: bool, cfa: CFA, dst: &mut RasterMut)
        -> BayerResult<()> {
    let (w, h) = (dst.w, dst.h);
    let mut data = vec![0u16; (2 * PADDING + w) * (2 * PADDING + h)];

    // Read all data.
    {
        let stride = 2 * PADDING + w;
        let rdr: Box<BayerRead16> = if be {
            Box::new(BorderReplicate16BE::new(w, PADDING))
        } else {
            Box::new(BorderReplicate16LE::new(w, PADDING))
        };

        for mut row in data.chunks_mut(stride).skip(PADDING).take(h) {
            rdr.read_line(r, &mut row)?;
        }

        {
            let (top, src) = data.split_at_mut(stride * PADDING);
            top[(stride * 0)..(stride * 1)].copy_from_slice(
                    &src[(stride * 1)..(stride * 2)]);
        }

        {
            let (src, bottom) = data.split_at_mut(stride * (h + PADDING));
            let yy = PADDING + h;
            bottom[(stride * 0)..(stride * 1)].copy_from_slice(
                    &src[(stride * (yy - 2))..(stride * (yy - 1))]);
        }
    }

    dst.buf.par_chunks_mut(dst.stride).enumerate()
            .for_each(|(y, mut row)| {
        let stride = 2 * PADDING + w;
        let prev = &data[(stride * (PADDING + y - 1)) .. (stride * (PADDING + y + 0))];
        let curr = &data[(stride * (PADDING + y + 0)) .. (stride * (PADDING + y + 1))];
        let next = &data[(stride * (PADDING + y + 1)) .. (stride * (PADDING + y + 2))];
        let cfa_y = if y % 2 == 0 { cfa } else { cfa.next_y() };

        let row16 = unsafe{ slice::from_raw_parts_mut(row.as_mut_ptr() as *mut u16, row.len() / 2) };
        apply_kernel_row!(u16; row16, prev, curr, next, cfa_y, w);
    });

    Ok(())
}

/*--------------------------------------------------------------*/
/* Naive                                                        */
/*--------------------------------------------------------------*/

#[cfg(not(feature = "rayon"))]
fn debayer_u8(r: &mut Read, cfa: CFA, dst: &mut RasterMut)
        -> BayerResult<()> {
    let (w, h) = (dst.w, dst.h);
    let mut prev = vec![0u8; 2 * PADDING + w];
    let mut curr = vec![0u8; 2 * PADDING + w];
    let mut next = vec![0u8; 2 * PADDING + w];
    let mut cfa = cfa;

    let rdr = BorderReplicate8::new(w, PADDING);
    rdr.read_line(r, &mut curr)?;
    rdr.read_line(r, &mut next)?;

    {   // y = 0.
        let row = dst.borrow_row_u8_mut(0);
        apply_kernel_row!(u8; row, next, curr, next, cfa, w);
        cfa = cfa.next_y();
    }

    for y in 1..(h - 1) {
        rotate!(prev <- curr <- next);
        rdr.read_line(r, &mut next)?;

        let row = dst.borrow_row_u8_mut(y);
        apply_kernel_row!(u8; row, prev, curr, next, cfa, w);
        cfa = cfa.next_y();
    }

    {   // y = h - 1.
        let row = dst.borrow_row_u8_mut(h - 1);
        apply_kernel_row!(u8; row, curr, next, curr, cfa, w);
    }

    Ok(())
}

#[cfg(not(feature = "rayon"))]
fn debayer_u16(r: &mut Read, be: bool, cfa: CFA, dst: &mut RasterMut)
        -> BayerResult<()> {
    let (w, h) = (dst.w, dst.h);
    let mut prev = vec![0u16; 2 * PADDING + w];
    let mut curr = vec![0u16; 2 * PADDING + w];
    let mut next = vec![0u16; 2 * PADDING + w];
    let mut cfa = cfa;

    let rdr: Box<BayerRead16> = if be {
        Box::new(BorderReplicate16BE::new(w, PADDING))
    } else {
        Box::new(BorderReplicate16LE::new(w, PADDING))
    };
    rdr.read_line(r, &mut curr)?;
    rdr.read_line(r, &mut next)?;

    {   // y = 0.
        let row = dst.borrow_row_u16_mut(0);
        apply_kernel_row!(u16; row, next, curr, next, cfa, w);
        cfa = cfa.next_y();
    }

    for y in 1..(h - 1) {
        rotate!(prev <- curr <- next);
        rdr.read_line(r, &mut next)?;

        let row = dst.borrow_row_u16_mut(y);
        apply_kernel_row!(u16; row, prev, curr, next, cfa, w);
        cfa = cfa.next_y();
    }

    {   // y = h - 1.
        let row = dst.borrow_row_u16_mut(h - 1);
        apply_kernel_row!(u16; row, curr, next, curr, cfa, w);
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;
    use ::{CFA,RasterDepth,RasterMut};
    use super::debayer_u8;

    #[test]
    fn test_even() {
        // R: set.seed(0); matrix(floor(runif(n=16, min=0, max=256)), nrow=4, byrow=TRUE)
        let src = [
            229, 67, 95,146,
            232, 51,229,241,
            169,161, 15, 52,
             45,175, 98,197 ];

        let expected = [
            229,149, 51,  162, 67, 51,   95,167,146,   95,146,241,
            199,232, 51,  127,172, 51,   55,229,146,   55,164,241,
            169,149,113,   92,161,113,   15,135,166,   15, 52,219,
            169, 45,175,   92,116,175,   15, 98,186,   15, 75,197 ];

        const IMG_W: usize = 4;
        const IMG_H: usize = 4;
        let mut dst = [0u8; 3 * IMG_W * IMG_H];

        let res = debayer_u8(&mut Cursor::new(&src[..]), CFA::RGGB,
                &mut RasterMut::new(IMG_W, IMG_H, RasterDepth::Depth8, &mut dst[..]));
        assert!(res.is_ok());
        assert_eq!(&dst[..], &expected[..]);
    }

    #[test]
    fn test_odd() {
        // R: set.seed(0); matrix(floor(runif(n=9, min=0, max=256)), nrow=3, byrow=TRUE)
        let src = [
            229, 67, 95,
            146,232, 51,
            229,241,169 ];

        let expected = [
            229,106,232,  162, 67,232,   95, 59,232,
            229,146,232,  180,126,232,  132, 51,232,
            229,193,232,  199,241,232,  169,146,232 ];

        const IMG_W: usize = 3;
        const IMG_H: usize = 3;
        let mut buf = [0u8; 3 * IMG_W * IMG_H];

        let res = debayer_u8(&mut Cursor::new(&src[..]), CFA::RGGB,
                &mut RasterMut::new(IMG_W, IMG_H, RasterDepth::Depth8, &mut buf));
        assert!(res.is_ok());
        assert_eq!(&buf[..], &expected[..]);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2017 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

extern crate futures;
extern crate tokio_core;
extern crate redis_async;

use std::sync::Arc;
use std::env;

use futures::{future, Future};

use tokio_core::reactor::Core;

use redis_async::client;

// An artificial "realistic" non-trivial example to demonstrate usage
fn main() {
    // Create some completely arbitrary "test data"
    let test_data_size = 10;
    let test_data: Vec<_> = (0..test_data_size).map(|x| (x, x.to_string())).collect();

    let mut core = Core::new().unwrap();

    let addr = env::args()
        .nth(1)
        .unwrap_or("127.0.0.1:6379".to_string())
        .parse()
        .unwrap();

    let test_f = client::paired_connect(&addr, &core.handle());

    let send_data = test_f.and_then(|connection| {
        let connection = Arc::new(connection);
        let futures: Vec<_> = test_data
            .into_iter()
            .map(move |data| {
                let connection_inner = connection.clone();
                connection
                    .send(["INCR", "realistic_test_ctr"].as_ref())
                    .and_then(move |ctr| {
                                  let key = format!("rt_{}", ctr.into_string().unwrap());
                                  let d_val = data.0.to_string();
                                  connection_inner.send(["SET", &key, &d_val].as_ref());
                                  connection_inner.send(["SET", &data.1, &key].as_ref())
                              })
            })
            .collect();
        future::join_all(futures)
    });

    let result = core.run(send_data).unwrap();
    assert_eq!(result.len(), test_data_size);
}<commit_msg>Unnecessary vector<commit_after>/*
 * Copyright 2017 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
 * http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
 * <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
 * option. This file may not be copied, modified, or distributed
 * except according to those terms.
 */

extern crate futures;
extern crate tokio_core;
extern crate redis_async;

use std::sync::Arc;
use std::env;

use futures::{future, Future};

use tokio_core::reactor::Core;

use redis_async::client;

// An artificial "realistic" non-trivial example to demonstrate usage
fn main() {
    // Create some completely arbitrary "test data"
    let test_data_size = 10;
    let test_data: Vec<_> = (0..test_data_size).map(|x| (x, x.to_string())).collect();

    let mut core = Core::new().unwrap();

    let addr = env::args()
        .nth(1)
        .unwrap_or("127.0.0.1:6379".to_string())
        .parse()
        .unwrap();

    let test_f = client::paired_connect(&addr, &core.handle());

    let send_data = test_f.and_then(|connection| {
        let connection = Arc::new(connection);
        let futures = test_data
            .into_iter()
            .map(move |data| {
                let connection_inner = connection.clone();
                connection
                    .send(["INCR", "realistic_test_ctr"].as_ref())
                    .and_then(move |ctr| {
                                  let key = format!("rt_{}", ctr.into_string().unwrap());
                                  let d_val = data.0.to_string();
                                  connection_inner.send(["SET", &key, &d_val].as_ref());
                                  connection_inner.send(["SET", &data.1, &key].as_ref())
                              })
            });
        future::join_all(futures)
    });

    let result = core.run(send_data).unwrap();
    assert_eq!(result.len(), test_data_size);
}<|endoftext|>


<commit_before>extern crate collections;
use std::io::{File, Open, Write, TempDir};
use std::os::args;
use std::cast;
use std::from_str::from_str;
use collections::priority_queue::PriorityQueue;
#[cfg(test)]
use std::io::{ReadWrite, SeekSet};

/*
 * an entry is the number that we read and the file it came from, so after we took
 * out one Entry from the priority queue, we need to know where to take the next number
 * from.
 */
struct Entry {
	value: u64,
	source: File,
}

/* The priority queue depends on the ordering relation and always returns the biggest */
impl Ord for Entry {
	// lesser than implementation
	fn lt(&self, other: &Entry) -> bool {
		// we invert: for us the smallest will have the highest priority
		!(self.value < other.value)
	}
}

/* ordering needs equality for some reason, so there */
impl Eq for Entry {
	fn eq(&self, other: &Entry) -> bool {
		self.value == other.value
	}
}

/* A struct to hold our queue */
struct PriorityFile {
	queue: PriorityQueue<Entry>,
}

/*
 * 'Class' that takes a list of files and does a k-way merge of the smallest values
 * from the files
 */
impl PriorityFile {
	pub fn new(files: Vec<File>) -> PriorityFile {
		let mut queue: PriorityQueue<Entry> = PriorityQueue::new();
		// now that we got the queue, populate it with initial values

		// move_iter is required to get the values themselves, not borrowed pointers to them
		for mut f in files.move_iter() {
			// read from file, append to queue
			let number = match f.read_le_u64() {
				Ok(num) => num,
				Err(e) => fail!("failed to read u64 from file: {}", e),
			};
			queue.push(Entry {value: number, source: f});
		};
		PriorityFile {queue: queue}
	}

	/* Returns the next smallest number from the list of files or None if exhausted */
	pub fn next(&mut self) -> Option<u64> {
		match self.queue.maybe_pop() {
			Some(mut res) => {
				let v = res.value;
				match res.source.read_le_u64() {
					Ok(num) => self.queue.push(Entry {value: num, source: res.source}),
					Err(_) => (),
				};
				Some(v)
			},
			None => None,
		}
	}
}

fn read_u64(from: &mut File, n: uint) -> Vec<u64> {
	// rust post 0.10 will directly return a Vec<u8> here
	let mut buf: ~[u8] = match from.read_exact(n*8) {
		Ok(b) => b,
		Err(e) => fail!("reading failed {}", e),
	};
	let mut buffer = Vec::from_slice(buf);

	let mut read: Vec<u64> = unsafe { cast::transmute(buffer) };
	let read_len = read.len();
	unsafe { read.set_len(read_len / 8); };
	read
}

fn write_u64(to: &mut File, items: Vec<u64>) {
	let mut written: Vec<u8> = unsafe { cast::transmute(items) };
	let written_len = written.len();
	unsafe { written.set_len(written_len * 8)};
	match to.write(written.as_slice()) {
		Ok(_) => (),
		Err(e) => fail!("writing failed: {}", e),
	};
}

fn externalSort(mut fdInput: File, size: u64, mut fdOutput: File, memSize: u64) {
	let mut runs = size / memSize;
	let items_per_run = (memSize / 8) as uint;
	let over = ((size / 8) - (items_per_run as u64 * runs)) as uint;

	println!("There will be {} runs with {} elements each.", runs, items_per_run);
	if over > 0 {
		println!("And one additional run with {} items.", over);
	}

	// let's create a temporary directory to store the sorted chunks
	// neat: TempDir deletes the directory and its contents when it goes out of
	// scope
	let overflow_dir = match TempDir::new("externalsort") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let overflow_path = overflow_dir.path();
	println!("temp dir path {}", overflow_path.display());

	// start the sorting runs
	for n in range(0, runs) {
		// iterate over our 'run' buffer and read from the input file
		// I suppose the values are little endian, which considering we are most
		// likely on x86, is a reasonable assumption
		println!("Starting read");
		let mut run = read_u64(&mut fdInput, items_per_run);
		println!("Read finished");
		// O(n log n) sort from the stdlib, hopefully more or less equivalent to
		// C++ std::sort
		run.sort();

		// now write out the results into numbered files
		let file_path = overflow_path.join(n.to_str());
		let mut file_file = match File::open_mode(&file_path, Open, Write) {
			Ok(f) => f,
			Err(e) => fail!("overflow file failed opening for write: {}", e),
		};

		println!("Starting write");
		write_u64(&mut file_file, run);
		println!("Write finished");
	};

	/* additional run to catch remaining objects */
	if over > 0 {
		let mut run = read_u64(&mut fdInput, over);
		run.sort();
		let file_path = overflow_path.join(runs.to_str());
		let mut file_file = match File::open_mode(&file_path, Open, Write) {
			Ok(f) => f,
			Err(e) => fail!("overflow file failed opening for write: {}", e),
		};
		write_u64(&mut file_file, run);

		runs += 1;
	};

	/* sorting done, now merging */

	// collect all file handles. Hopefully less than 4096
	let mut files = Vec::with_capacity(runs as uint);
	for n in range(0, runs) {
		let file_path = overflow_path.join(n.to_str());
		let fd = match File::open(&file_path) {
			Ok(f) => f,
			Err(e) => fail!("overflow file opening error {}", e),
		};
		files.push(fd);
	}

	// we got a nice class here that spits out sorted numbers until exhausted
	let mut pf = PriorityFile::new(files);
	loop {
		match pf.next() {
			// if we got a number, write it to the output
			Some (number) => match fdOutput.write_le_u64(number) {
				Ok(_) => (),
				Err(e) => fail!("writing output file failed: {}", e),
			},
			None => break,
		};
	};
}

fn main() {
	if args().len() < 4 {
		fail!("Argument error: <inputFile> <outputFile> <memoryBufferInMB>");
	}
	let input_file_path = &Path::new(args()[1]);
	let size = match input_file_path.stat() {
		Ok(stat) => stat.size,
		Err(e) => fail!("Couldn't read {}", e),
	};
	let fin = match File::open(input_file_path) {
		Ok(f) => f,
		Err(e) => fail!("input file error: {}", e),
	};
	let fout = match File::open_mode(&Path::new(args()[2]), Open, Write) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};
	let buffer_size:u64 = match from_str(args()[3]) {
		Some(num) => num,
		None => fail!("Not numeric input"),
	};
	println!("Input file size {}", size);

	//let values = read_u64(&mut fin, 2);
	//println!("Values: {}", values);
	externalSort(fin, size, fout, buffer_size * 1024 * 1024);
}

#[test]
fn generate_5gb_and_sort() {
	let mut random = match File::open(&Path::new("/dev/urandom")) {
		Ok(f) => f,
		Err(e) => fail!("Error opening urandom: {}", e),
	};

	let test_dir = match TempDir::new("externalsort-test") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let test_path = test_dir.path();

	let mut unordered_nums = match File::open_mode(&test_path.join("input"), Open, ReadWrite) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};

	let mut buffer = [0, .. 1024*1024];
	let mut left = 1 * 1024;
	while left > 0 {
		random.read(buffer);
		unordered_nums.write(buffer);
		left -= 1;
	}

	unordered_nums.seek(0, SeekSet);
	let ordered_nums = match File::open_mode(&test_path.join("output"), Open, Write) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};
	externalSort(unordered_nums, 1024 * 1024 * 1024 * 1, ordered_nums, 1024 * 1024 * 200);
	let mut ordered_nums = match File::open(&test_path.join("output")) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};

	let mut last = 0_u64;
	loop {
		let number = match ordered_nums.read_le_u64() {
			Ok(num) => num,
			Err(e) => fail!("failed to read u64 from file: {}", e),
		};
		assert!(number >= last);
		last = number;
	}
}
<commit_msg>Made writing stuff out actually work<commit_after>#[feature(phase)];
#[phase(syntax, link)] extern crate log;
extern crate collections;
use std::io::{File, Open, Write, TempDir};
use std::os::args;
use std::cast;
use std::from_str::from_str;
use std::raw::Slice;
use collections::priority_queue::PriorityQueue;
#[cfg(test)]
use std::io::{ReadWrite, SeekSet};

/*
 * an entry is the number that we read and the file it came from, so after we took
 * out one Entry from the priority queue, we need to know where to take the next number
 * from.
 */
struct Entry {
	value: u64,
	source: File,
}

/* The priority queue depends on the ordering relation and always returns the biggest */
impl Ord for Entry {
	// lesser than implementation
	fn lt(&self, other: &Entry) -> bool {
		// we invert: for us the smallest will have the highest priority
		!(self.value < other.value)
	}
}

/* ordering needs equality for some reason, so there */
impl Eq for Entry {
	fn eq(&self, other: &Entry) -> bool {
		self.value == other.value
	}
}

/* A struct to hold our queue */
struct PriorityFile {
	queue: PriorityQueue<Entry>,
}

/*
 * 'Class' that takes a list of files and does a k-way merge of the smallest values
 * from the files
 */
impl PriorityFile {
	pub fn new(files: Vec<File>) -> PriorityFile {
		let mut queue: PriorityQueue<Entry> = PriorityQueue::new();
		// now that we got the queue, populate it with initial values

		// move_iter is required to get the values themselves, not borrowed pointers to them
		for mut f in files.move_iter() {
			// read from file, append to queue
			let number = match f.read_le_u64() {
				Ok(num) => num,
				Err(e) => fail!("failed to read u64 from file: {}", e),
			};
			queue.push(Entry {value: number, source: f});
		};
		PriorityFile {queue: queue}
	}

	/* Returns the next smallest number from the list of files or None if exhausted */
	pub fn next(&mut self) -> Option<u64> {
		match self.queue.maybe_pop() {
			Some(mut res) => {
				let v = res.value;
				match res.source.read_le_u64() {
					Ok(num) => self.queue.push(Entry {value: num, source: res.source}),
					Err(_) => (),
				};
				Some(v)
			},
			None => None,
		}
	}
}

fn read_u64(from: &mut File, n: uint) -> Vec<u64> {
	// rust post 0.10 will directly return a Vec<u8> here
	let buf = match from.read_exact(n*8) {
		Ok(b) => b,
		Err(e) => fail!("reading failed {}", e),
	};
	//let u = buf.as_ptr();
	let u = buf.as_slice();

	let mut r: Slice<u64> = unsafe { cast::transmute(u) };
	assert_eq!(r.len % 8, 0);
	r.len /= 8;
	Vec::from_slice(unsafe { cast::transmute(r) })
}

fn write_u64(to: &mut File, items: &Vec<u64>) {
	let s = items.as_slice();
	let mut r: Slice<u8> = unsafe { cast::transmute(s) };
	r.len *= 8;
	let buf: &[u8] = unsafe { cast::transmute(r) };

	match to.write(buf) {
		Ok(_) => (),
		Err(e) => fail!("writing failed: {}", e),
	};
}

fn externalSort(mut fdInput: File, size: u64, mut fdOutput: File, memSize: u64) {
	let mut runs = size / memSize;
	let items_per_run = (memSize / 8) as uint;
	let over = ((size / 8) - (items_per_run as u64 * runs)) as uint;

	info!("There will be {} runs with {} elements each.", runs, items_per_run);
	if over > 0 {
		info!("And one additional run with {} items.", over);
	}

	// let's create a temporary directory to store the sorted chunks
	// neat: TempDir deletes the directory and its contents when it goes out of
	// scope
	let overflow_dir = match TempDir::new("externalsort") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let overflow_path = overflow_dir.path();
	info!("temp dir:\ncd {}", overflow_path.display());

	// start the sorting runs
	for n in range(0, runs) {
		// iterate over our 'run' buffer and read from the input file
		// I suppose the values are little endian, which considering we are most
		// likely on x86, is a reasonable assumption
		info!("Run {}, starting read", n);
		let mut run = read_u64(&mut fdInput, items_per_run);
		info!("Run {}, read finished", n);
		// O(n log n) sort from the stdlib, hopefully more or less equivalent to
		// C++ std::sort
		run.sort();

		// now write out the results into numbered files
		let file_path = overflow_path.join(n.to_str());
		let mut file_file = match File::open_mode(&file_path, Open, Write) {
			Ok(f) => f,
			Err(e) => fail!("overflow file failed opening for write: {}", e),
		};

		println!("Starting write");
		write_u64(&mut file_file, &run);
		println!("Write finished");
	};

	/* additional run to catch remaining objects */
	if over > 0 {
		println!("overrun read start");
		let mut run = read_u64(&mut fdInput, over);
		println!("overrun read done");
		run.sort();

		let file_path = overflow_path.join(runs.to_str());
		let mut file_file = match File::open_mode(&file_path, Open, Write) {
			Ok(f) => f,
			Err(e) => fail!("overflow file failed opening for write: {}", e),
		};
		println!("overrun write start");
		write_u64(&mut file_file, &run);
		println!("overrun write done");

		runs += 1;
	};
	println!("remaining run and sort done")

	/* sorting done, now merging */

	// collect all file handles. Hopefully less than 4096
	let mut files = Vec::with_capacity(runs as uint);
	for n in range(0, runs) {
		let file_path = overflow_path.join(n.to_str());
		let fd = match File::open(&file_path) {
			Ok(f) => f,
			Err(e) => fail!("overflow file opening error {}", e),
		};
		files.push(fd);
	}

	// we got a nice class here that spits out sorted numbers until exhausted
	let mut pf = PriorityFile::new(files);
	// put the sorted numbers in a cache first, so they can all be written in bulk
	let mut write_cache = Vec::with_capacity(items_per_run);
	loop {
		match pf.next() {
			Some (number) => {
				// a new number was found, add it to the cache
				// and if the cache is full, write that to disk
				write_cache.push(number);
				if write_cache.len() >= items_per_run {
					write_u64(&mut fdOutput, &write_cache);
					write_cache.clear();
				};
			},
			None => {
				// we finished, write remaining cache to disk
				write_u64(&mut fdOutput, &write_cache);
				break;
			},
		};
	};
}

fn main() {
	if args().len() < 4 {
		fail!("Argument error: <inputFile> <outputFile> <memoryBufferInMB>");
	}
	let input_file_path = &Path::new(args()[1]);
	let size = match input_file_path.stat() {
		Ok(stat) => stat.size,
		Err(e) => fail!("Couldn't read {}", e),
	};
	let fin = match File::open(input_file_path) {
		Ok(f) => f,
		Err(e) => fail!("input file error: {}", e),
	};
	let fout = match File::open_mode(&Path::new(args()[2]), Open, Write) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};
	let buffer_size:u64 = match from_str(args()[3]) {
		Some(num) => num,
		None => fail!("Not numeric input"),
	};
	info!("Input file size {}", size);

	//let values = read_u64(&mut fin, 2);
	//println!("Values: {}", values);
	externalSort(fin, size, fout, buffer_size * 1024 * 1024);
}

#[test]
fn generate_5gb_and_sort() {
	let mut random = match File::open(&Path::new("/dev/urandom")) {
		Ok(f) => f,
		Err(e) => fail!("Error opening urandom: {}", e),
	};

	let test_dir = match TempDir::new("externalsort-test") {
		Some(temp_dir) => temp_dir,
		None => fail!("creation of temporary directory"),
	};
	let test_path = test_dir.path();

	let mut unordered_nums = match File::open_mode(&test_path.join("input"), Open, ReadWrite) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};

	let mut buffer = [0, .. 1024*1024];
	let mut left = 1 * 1024;
	while left > 0 {
		random.read(buffer);
		unordered_nums.write(buffer);
		left -= 1;
	}

	unordered_nums.seek(0, SeekSet);
	let ordered_nums = match File::open_mode(&test_path.join("output"), Open, Write) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};
	externalSort(unordered_nums, 1024 * 1024 * 1024 * 1, ordered_nums, 1024 * 1024 * 200);
	let mut ordered_nums = match File::open(&test_path.join("output")) {
		Ok(f) => f,
		Err(e) => fail!("output file error: {}", e),
	};

	let mut last = 0_u64;
	loop {
		let number = match ordered_nums.read_le_u64() {
			Ok(num) => num,
			Err(e) => fail!("failed to read u64 from file: {}", e),
		};
		assert!(number >= last);
		last = number;
	}
}
<|endoftext|>


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

//! Protocol used to ensure peer liveness
//!
//! The HealthChecker is responsible for ensuring liveness of all peers of a node.
//! It does so by periodically selecting a random connected peer and sending a Ping probe. A
//! healthy peer is expected to respond with a corresponding Pong message.
//!
//! If a certain number of successive liveness probes for a peer fail, the HealthChecker initiates a
//! disconnect from the peer. It relies on ConnectivityManager or the remote peer to re-establish
//! the connection.
//!
//! Future Work
//! -----------
//! We can make a few other improvements to the health checker. These are:
//! - Make the policy for interpreting ping failures pluggable
//! - Use successful inbound pings as a sign of remote note being healthy
//! - Ping a peer only in periods of no application-level communication with the peer
use crate::{
    counters,
    error::NetworkError,
    peer_manager::{ConnectionRequestSender, PeerManagerRequestSender},
    protocols::{
        network::{Event, NetworkEvents, NetworkSender},
        rpc::error::RpcError,
    },
    validator_network::network_builder::NetworkBuilder,
    ProtocolId,
};
use bytes::Bytes;
use channel::message_queues::QueueStyle;
use futures::{
    channel::oneshot,
    stream::{FusedStream, FuturesUnordered, Stream, StreamExt},
};
use libra_logger::prelude::*;
use libra_security_logger::{security_log, SecurityEvent};
use libra_types::PeerId;
use rand::{rngs::SmallRng, seq::SliceRandom, FromEntropy, Rng};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, time::Duration};

#[cfg(test)]
mod test;

/// The interface from Network to HealthChecker layer.
///
/// `HealthCheckerNetworkEvents` is a `Stream` of `PeerManagerNotification` where the
/// raw `Bytes` rpc messages are deserialized into
/// `HealthCheckerMsg` types. `HealthCheckerNetworkEvents` is a thin wrapper
/// around an `channel::Receiver<PeerManagerNotification>`.
pub type HealthCheckerNetworkEvents = NetworkEvents<HealthCheckerMsg>;

/// The interface from HealthChecker to Networking layer.
///
/// This is a thin wrapper around a `NetworkSender<HealthCheckerMsg>`, so it is
/// easy to clone and send off to a separate task. For example, the rpc requests
/// return Futures that encapsulate the whole flow, from sending the request to
/// remote, to finally receiving the response and deserializing. It therefore
/// makes the most sense to make the rpc call on a separate async task, which
/// requires the `HealthCheckerNetworkSender` to be `Clone` and `Send`.
#[derive(Clone)]
pub struct HealthCheckerNetworkSender {
    inner: NetworkSender<HealthCheckerMsg>,
}

pub fn add_to_network(
    network: &mut NetworkBuilder,
) -> (HealthCheckerNetworkSender, HealthCheckerNetworkEvents) {
    let (sender, receiver, connection_reqs_tx, connection_notifs_rx) = network
        .add_protocol_handler(
            vec![ProtocolId::HealthCheckerRpc],
            vec![],
            QueueStyle::LIFO,
            Some(&counters::PENDING_HEALTH_CHECKER_NETWORK_EVENTS),
        );
    (
        HealthCheckerNetworkSender::new(sender, connection_reqs_tx),
        HealthCheckerNetworkEvents::new(receiver, connection_notifs_rx),
    )
}

impl HealthCheckerNetworkSender {
    pub fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self {
        Self {
            inner: NetworkSender::new(peer_mgr_reqs_tx, connection_reqs_tx),
        }
    }

    /// Send a HealthChecker Ping RPC request to remote peer `recipient`. Returns
    /// the remote peer's future `Pong` reply.
    ///
    /// The rpc request can be canceled at any point by dropping the returned
    /// future.
    pub async fn send_rpc(
        &mut self,
        recipient: PeerId,
        req_msg: HealthCheckerMsg,
        timeout: Duration,
    ) -> Result<HealthCheckerMsg, RpcError> {
        let protocol = ProtocolId::HealthCheckerRpc;
        self.inner
            .send_rpc(recipient, protocol, req_msg, timeout)
            .await
    }

    pub async fn disconnect_peer(&mut self, peer_id: PeerId) -> Result<(), NetworkError> {
        self.inner.disconnect_peer(peer_id).await
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub enum HealthCheckerMsg {
    Ping(Ping),
    Pong(Pong),
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Ping(u32);

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Pong(u32);

/// The actor performing health checks by running the Ping protocol
pub struct HealthChecker<TTicker> {
    /// Ticker to trigger ping to a random peer. In production, the ticker is likely to be
    /// fixed duration interval timer.
    ticker: TTicker,
    /// Channel to send requests to Network layer.
    network_tx: HealthCheckerNetworkSender,
    /// Channel to receive notifications from Network layer about new/lost connections.
    network_rx: HealthCheckerNetworkEvents,
    /// Map from connected peer to last round of successful ping, and number of failures since
    /// then.
    connected: HashMap<PeerId, (u64, u64)>,
    /// Random-number generator.
    rng: SmallRng,
    /// Ping timmeout duration.
    ping_timeout: Duration,
    /// Number of successive ping failures we tolerate before declaring a node as unhealthy and
    /// disconnecting from it. In the future, this can be replaced with a more general failure
    /// detection policy.
    ping_failures_tolerated: u64,
    /// Counter incremented in each round of health checks
    round: u64,
}

impl<TTicker> HealthChecker<TTicker>
where
    TTicker: Stream + FusedStream + Unpin,
{
    /// Create new instance of the [`HealthChecker`] actor.
    pub fn new(
        ticker: TTicker,
        network_tx: HealthCheckerNetworkSender,
        network_rx: HealthCheckerNetworkEvents,
        ping_timeout: Duration,
        ping_failures_tolerated: u64,
    ) -> Self {
        HealthChecker {
            ticker,
            network_tx,
            network_rx,
            connected: HashMap::new(),
            rng: SmallRng::from_entropy(),
            ping_timeout,
            ping_failures_tolerated,
            round: 0,
        }
    }

    pub async fn start(mut self) {
        let mut tick_handlers = FuturesUnordered::new();
        loop {
            futures::select! {
                event = self.network_rx.select_next_some() => {
                    match event {
                        Ok(Event::NewPeer(peer_id)) => {
                            self.connected.insert(peer_id, (self.round, 0));
                        },
                        Ok(Event::LostPeer(peer_id)) => {
                            self.connected.remove(&peer_id);
                        },
                        Ok(Event::RpcRequest((peer_id, msg, res_tx))) => {
                            match msg {
                            HealthCheckerMsg::Ping(ping) => self.handle_ping_request(peer_id, ping, res_tx),
                            _ => security_log(SecurityEvent::InvalidHealthCheckerMsg)
                                .error("Unexpected rpc message")
                                    .data(&msg)
                                    .data(&peer_id)
                                    .log(),
                            };
                        }
                        Ok(Event::Message(_)) => {
                            security_log(SecurityEvent::InvalidNetworkEventHC)
                                .error("Unexpected network event")
                                .data(&event)
                                .log();
                            debug_assert!(false, "Unexpected network event");
                        },
                        Err(err) => {
                            security_log(SecurityEvent::InvalidNetworkEventHC)
                                .error(&err)
                                .log();
                            debug_assert!(false, "Unexpected network error");
                        }
                    }
                }
                _ = self.ticker.select_next_some() => {
                    self.round += 1;
                    debug!("Tick: Round number: {}", self.round);
                    match self.sample_random_peer() {
                        Some(peer_id) => {
                            debug!("Will ping: {}", peer_id.short_str());

                            let nonce = self.sample_nonce();

                            tick_handlers.push(
                                Self::ping_peer(
                                    self.network_tx.clone(),
                                    peer_id,
                                    self.round,
                                    nonce,
                                    self.ping_timeout.clone()));
                        }
                        None => {
                            debug!("No connected peer to ping");
                        }
                    }
                }
                res = tick_handlers.select_next_some() => {
                    let (peer_id, round, nonce, ping_result) = res;
                    self.handle_ping_response(peer_id, round, nonce, ping_result).await;
                }
                complete => {
                    break;
                }
            }
        }
        crit!("Health checker actor terminated");
    }

    fn handle_ping_request(
        &mut self,
        peer_id: PeerId,
        ping: Ping,
        res_tx: oneshot::Sender<Result<Bytes, RpcError>>,
    ) {
        let message = match lcs::to_bytes(&HealthCheckerMsg::Pong(Pong(ping.0))) {
            Ok(msg) => msg,
            Err(e) => {
                warn!("Unable to serialize pong response: {}", e);
                return;
            }
        };
        debug!(
            "Sending Pong response to peer: {} with nonce: {}",
            peer_id.short_str(),
            ping.0,
        );
        let _ = res_tx.send(Ok(message.into()));
    }

    async fn handle_ping_response(
        &mut self,
        peer_id: PeerId,
        round: u64,
        req_nonce: u32,
        ping_result: Result<Pong, RpcError>,
    ) {
        debug!("Got result for ping round: {}", round);
        match ping_result {
            Ok(pong) => {
                if pong.0 == req_nonce {
                    debug!("Ping successful for peer: {}", peer_id.short_str());
                    // Update last successful ping to current round.
                    self.connected
                        .entry(peer_id)
                        .and_modify(|(ref mut r, ref mut count)| {
                            if round > *r {
                                *r = round;
                                *count = 0;
                            }
                        });
                } else {
                    security_log(SecurityEvent::InvalidHealthCheckerMsg)
                        .error("Pong nonce doesn't match our challenge Ping nonce")
                        .data(&peer_id)
                        .data(req_nonce)
                        .data(pong.0)
                        .log();
                    debug_assert!(false, "Pong nonce doesn't match our challenge Ping nonce");
                }
            }
            Err(err) => {
                warn!(
                    "Ping failed for peer: {} with error: {:?}",
                    peer_id.short_str(),
                    err
                );
                match self.connected.get_mut(&peer_id) {
                    None => {
                        // If we are no longer connected to the peer, we ignore ping
                        // failure.
                    }
                    Some((ref mut prev, ref mut failures)) => {
                        // If this is the result of an older ping, we ignore it.
                        if *prev > round {
                            return;
                        }
                        // Increment num of failures. If the ping failures are now more than
                        // `self.ping_failures_tolerated`, we disconnect from the node.
                        // The HealthChecker only performs the disconnect. It relies on
                        // ConnectivityManager or the remote peer to re-establish the connection.
                        *failures += 1;
                        if *failures > self.ping_failures_tolerated {
                            info!("Disonnecting from peer: {}", peer_id.short_str());
                            if let Err(err) = self.network_tx.disconnect_peer(peer_id).await {
                                warn!(
                                    "Failed to disconnect from peer: {} with error: {:?}",
                                    peer_id.short_str(),
                                    err
                                );
                            }
                        }
                    }
                }
            }
        }
    }

    async fn ping_peer(
        mut network_tx: HealthCheckerNetworkSender,
        peer_id: PeerId,
        round: u64,
        nonce: u32,
        ping_timeout: Duration,
    ) -> (PeerId, u64, u32, Result<Pong, RpcError>) {
        debug!(
            "Sending Ping request to peer: {} with nonce: {}",
            peer_id.short_str(),
            nonce
        );
        let res_pong_msg = network_tx
            .send_rpc(peer_id, HealthCheckerMsg::Ping(Ping(nonce)), ping_timeout)
            .await
            .and_then(|msg| match msg {
                HealthCheckerMsg::Pong(res) => Ok(res),
                _ => Err(RpcError::InvalidRpcResponse),
            });
        (peer_id, round, nonce, res_pong_msg)
    }

    fn sample_random_peer(&mut self) -> Option<PeerId> {
        let peers: Vec<_> = self.connected.keys().cloned().collect();
        peers.choose(&mut self.rng).cloned()
    }

    fn sample_nonce(&mut self) -> u32 {
        self.rng.gen::<u32>()
    }
}
<commit_msg>[network] Fix typo Disonnecting -> Disconnecting<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Protocol used to ensure peer liveness
//!
//! The HealthChecker is responsible for ensuring liveness of all peers of a node.
//! It does so by periodically selecting a random connected peer and sending a Ping probe. A
//! healthy peer is expected to respond with a corresponding Pong message.
//!
//! If a certain number of successive liveness probes for a peer fail, the HealthChecker initiates a
//! disconnect from the peer. It relies on ConnectivityManager or the remote peer to re-establish
//! the connection.
//!
//! Future Work
//! -----------
//! We can make a few other improvements to the health checker. These are:
//! - Make the policy for interpreting ping failures pluggable
//! - Use successful inbound pings as a sign of remote note being healthy
//! - Ping a peer only in periods of no application-level communication with the peer
use crate::{
    counters,
    error::NetworkError,
    peer_manager::{ConnectionRequestSender, PeerManagerRequestSender},
    protocols::{
        network::{Event, NetworkEvents, NetworkSender},
        rpc::error::RpcError,
    },
    validator_network::network_builder::NetworkBuilder,
    ProtocolId,
};
use bytes::Bytes;
use channel::message_queues::QueueStyle;
use futures::{
    channel::oneshot,
    stream::{FusedStream, FuturesUnordered, Stream, StreamExt},
};
use libra_logger::prelude::*;
use libra_security_logger::{security_log, SecurityEvent};
use libra_types::PeerId;
use rand::{rngs::SmallRng, seq::SliceRandom, FromEntropy, Rng};
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, time::Duration};

#[cfg(test)]
mod test;

/// The interface from Network to HealthChecker layer.
///
/// `HealthCheckerNetworkEvents` is a `Stream` of `PeerManagerNotification` where the
/// raw `Bytes` rpc messages are deserialized into
/// `HealthCheckerMsg` types. `HealthCheckerNetworkEvents` is a thin wrapper
/// around an `channel::Receiver<PeerManagerNotification>`.
pub type HealthCheckerNetworkEvents = NetworkEvents<HealthCheckerMsg>;

/// The interface from HealthChecker to Networking layer.
///
/// This is a thin wrapper around a `NetworkSender<HealthCheckerMsg>`, so it is
/// easy to clone and send off to a separate task. For example, the rpc requests
/// return Futures that encapsulate the whole flow, from sending the request to
/// remote, to finally receiving the response and deserializing. It therefore
/// makes the most sense to make the rpc call on a separate async task, which
/// requires the `HealthCheckerNetworkSender` to be `Clone` and `Send`.
#[derive(Clone)]
pub struct HealthCheckerNetworkSender {
    inner: NetworkSender<HealthCheckerMsg>,
}

pub fn add_to_network(
    network: &mut NetworkBuilder,
) -> (HealthCheckerNetworkSender, HealthCheckerNetworkEvents) {
    let (sender, receiver, connection_reqs_tx, connection_notifs_rx) = network
        .add_protocol_handler(
            vec![ProtocolId::HealthCheckerRpc],
            vec![],
            QueueStyle::LIFO,
            Some(&counters::PENDING_HEALTH_CHECKER_NETWORK_EVENTS),
        );
    (
        HealthCheckerNetworkSender::new(sender, connection_reqs_tx),
        HealthCheckerNetworkEvents::new(receiver, connection_notifs_rx),
    )
}

impl HealthCheckerNetworkSender {
    pub fn new(
        peer_mgr_reqs_tx: PeerManagerRequestSender,
        connection_reqs_tx: ConnectionRequestSender,
    ) -> Self {
        Self {
            inner: NetworkSender::new(peer_mgr_reqs_tx, connection_reqs_tx),
        }
    }

    /// Send a HealthChecker Ping RPC request to remote peer `recipient`. Returns
    /// the remote peer's future `Pong` reply.
    ///
    /// The rpc request can be canceled at any point by dropping the returned
    /// future.
    pub async fn send_rpc(
        &mut self,
        recipient: PeerId,
        req_msg: HealthCheckerMsg,
        timeout: Duration,
    ) -> Result<HealthCheckerMsg, RpcError> {
        let protocol = ProtocolId::HealthCheckerRpc;
        self.inner
            .send_rpc(recipient, protocol, req_msg, timeout)
            .await
    }

    pub async fn disconnect_peer(&mut self, peer_id: PeerId) -> Result<(), NetworkError> {
        self.inner.disconnect_peer(peer_id).await
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub enum HealthCheckerMsg {
    Ping(Ping),
    Pong(Pong),
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Ping(u32);

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Pong(u32);

/// The actor performing health checks by running the Ping protocol
pub struct HealthChecker<TTicker> {
    /// Ticker to trigger ping to a random peer. In production, the ticker is likely to be
    /// fixed duration interval timer.
    ticker: TTicker,
    /// Channel to send requests to Network layer.
    network_tx: HealthCheckerNetworkSender,
    /// Channel to receive notifications from Network layer about new/lost connections.
    network_rx: HealthCheckerNetworkEvents,
    /// Map from connected peer to last round of successful ping, and number of failures since
    /// then.
    connected: HashMap<PeerId, (u64, u64)>,
    /// Random-number generator.
    rng: SmallRng,
    /// Ping timmeout duration.
    ping_timeout: Duration,
    /// Number of successive ping failures we tolerate before declaring a node as unhealthy and
    /// disconnecting from it. In the future, this can be replaced with a more general failure
    /// detection policy.
    ping_failures_tolerated: u64,
    /// Counter incremented in each round of health checks
    round: u64,
}

impl<TTicker> HealthChecker<TTicker>
where
    TTicker: Stream + FusedStream + Unpin,
{
    /// Create new instance of the [`HealthChecker`] actor.
    pub fn new(
        ticker: TTicker,
        network_tx: HealthCheckerNetworkSender,
        network_rx: HealthCheckerNetworkEvents,
        ping_timeout: Duration,
        ping_failures_tolerated: u64,
    ) -> Self {
        HealthChecker {
            ticker,
            network_tx,
            network_rx,
            connected: HashMap::new(),
            rng: SmallRng::from_entropy(),
            ping_timeout,
            ping_failures_tolerated,
            round: 0,
        }
    }

    pub async fn start(mut self) {
        let mut tick_handlers = FuturesUnordered::new();
        loop {
            futures::select! {
                event = self.network_rx.select_next_some() => {
                    match event {
                        Ok(Event::NewPeer(peer_id)) => {
                            self.connected.insert(peer_id, (self.round, 0));
                        },
                        Ok(Event::LostPeer(peer_id)) => {
                            self.connected.remove(&peer_id);
                        },
                        Ok(Event::RpcRequest((peer_id, msg, res_tx))) => {
                            match msg {
                            HealthCheckerMsg::Ping(ping) => self.handle_ping_request(peer_id, ping, res_tx),
                            _ => security_log(SecurityEvent::InvalidHealthCheckerMsg)
                                .error("Unexpected rpc message")
                                    .data(&msg)
                                    .data(&peer_id)
                                    .log(),
                            };
                        }
                        Ok(Event::Message(_)) => {
                            security_log(SecurityEvent::InvalidNetworkEventHC)
                                .error("Unexpected network event")
                                .data(&event)
                                .log();
                            debug_assert!(false, "Unexpected network event");
                        },
                        Err(err) => {
                            security_log(SecurityEvent::InvalidNetworkEventHC)
                                .error(&err)
                                .log();
                            debug_assert!(false, "Unexpected network error");
                        }
                    }
                }
                _ = self.ticker.select_next_some() => {
                    self.round += 1;
                    debug!("Tick: Round number: {}", self.round);
                    match self.sample_random_peer() {
                        Some(peer_id) => {
                            debug!("Will ping: {}", peer_id.short_str());

                            let nonce = self.sample_nonce();

                            tick_handlers.push(
                                Self::ping_peer(
                                    self.network_tx.clone(),
                                    peer_id,
                                    self.round,
                                    nonce,
                                    self.ping_timeout.clone()));
                        }
                        None => {
                            debug!("No connected peer to ping");
                        }
                    }
                }
                res = tick_handlers.select_next_some() => {
                    let (peer_id, round, nonce, ping_result) = res;
                    self.handle_ping_response(peer_id, round, nonce, ping_result).await;
                }
                complete => {
                    break;
                }
            }
        }
        crit!("Health checker actor terminated");
    }

    fn handle_ping_request(
        &mut self,
        peer_id: PeerId,
        ping: Ping,
        res_tx: oneshot::Sender<Result<Bytes, RpcError>>,
    ) {
        let message = match lcs::to_bytes(&HealthCheckerMsg::Pong(Pong(ping.0))) {
            Ok(msg) => msg,
            Err(e) => {
                warn!("Unable to serialize pong response: {}", e);
                return;
            }
        };
        debug!(
            "Sending Pong response to peer: {} with nonce: {}",
            peer_id.short_str(),
            ping.0,
        );
        let _ = res_tx.send(Ok(message.into()));
    }

    async fn handle_ping_response(
        &mut self,
        peer_id: PeerId,
        round: u64,
        req_nonce: u32,
        ping_result: Result<Pong, RpcError>,
    ) {
        debug!("Got result for ping round: {}", round);
        match ping_result {
            Ok(pong) => {
                if pong.0 == req_nonce {
                    debug!("Ping successful for peer: {}", peer_id.short_str());
                    // Update last successful ping to current round.
                    self.connected
                        .entry(peer_id)
                        .and_modify(|(ref mut r, ref mut count)| {
                            if round > *r {
                                *r = round;
                                *count = 0;
                            }
                        });
                } else {
                    security_log(SecurityEvent::InvalidHealthCheckerMsg)
                        .error("Pong nonce doesn't match our challenge Ping nonce")
                        .data(&peer_id)
                        .data(req_nonce)
                        .data(pong.0)
                        .log();
                    debug_assert!(false, "Pong nonce doesn't match our challenge Ping nonce");
                }
            }
            Err(err) => {
                warn!(
                    "Ping failed for peer: {} with error: {:?}",
                    peer_id.short_str(),
                    err
                );
                match self.connected.get_mut(&peer_id) {
                    None => {
                        // If we are no longer connected to the peer, we ignore ping
                        // failure.
                    }
                    Some((ref mut prev, ref mut failures)) => {
                        // If this is the result of an older ping, we ignore it.
                        if *prev > round {
                            return;
                        }
                        // Increment num of failures. If the ping failures are now more than
                        // `self.ping_failures_tolerated`, we disconnect from the node.
                        // The HealthChecker only performs the disconnect. It relies on
                        // ConnectivityManager or the remote peer to re-establish the connection.
                        *failures += 1;
                        if *failures > self.ping_failures_tolerated {
                            info!("Disconnecting from peer: {}", peer_id.short_str());
                            if let Err(err) = self.network_tx.disconnect_peer(peer_id).await {
                                warn!(
                                    "Failed to disconnect from peer: {} with error: {:?}",
                                    peer_id.short_str(),
                                    err
                                );
                            }
                        }
                    }
                }
            }
        }
    }

    async fn ping_peer(
        mut network_tx: HealthCheckerNetworkSender,
        peer_id: PeerId,
        round: u64,
        nonce: u32,
        ping_timeout: Duration,
    ) -> (PeerId, u64, u32, Result<Pong, RpcError>) {
        debug!(
            "Sending Ping request to peer: {} with nonce: {}",
            peer_id.short_str(),
            nonce
        );
        let res_pong_msg = network_tx
            .send_rpc(peer_id, HealthCheckerMsg::Ping(Ping(nonce)), ping_timeout)
            .await
            .and_then(|msg| match msg {
                HealthCheckerMsg::Pong(res) => Ok(res),
                _ => Err(RpcError::InvalidRpcResponse),
            });
        (peer_id, round, nonce, res_pong_msg)
    }

    fn sample_random_peer(&mut self) -> Option<PeerId> {
        let peers: Vec<_> = self.connected.keys().cloned().collect();
        peers.choose(&mut self.rng).cloned()
    }

    fn sample_nonce(&mut self) -> u32 {
        self.rng.gen::<u32>()
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;
use std::fmt;
use std::fs::File;
use std::io::BufReader;
use std::io::BufRead;
use std::path::PathBuf;

use itertools::Itertools;
use serde_json;

use errors::*;
use telemetry::{LogMessage, TelemetryEvent};

pub struct TelemetryAnalysis {
    telemetry_dir: PathBuf,
    rustc_statistics: RustcStatistics,
    rustc_success_statistics: RustcStatistics,
    rustc_error_statistics: RustcStatistics,
}

#[derive(Default)]
pub struct RustcStatistics {
    rustc_execution_count: u32,
    compile_time_ms_total: u64,
    compile_time_ms_mean: u64,
    compile_time_ms_ntile_75: u64,
    compile_time_ms_ntile_90: u64,
    compile_time_ms_ntile_95: u64,
    compile_time_ms_ntile_99: u64,
    compile_time_ms_stdev: f64,
    exit_codes_with_count: HashMap<i32, i32>,
    error_codes_with_counts: HashMap<String, i32>,
}

impl RustcStatistics {
    pub fn new() -> RustcStatistics {
        Default::default()
    }
}

impl fmt::Display for RustcStatistics {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut errors: String = String::new();

        if !self.error_codes_with_counts.is_empty() {
            errors = "  rustc errors\n".to_owned();
            for (error, count) in &self.error_codes_with_counts {
                errors = errors + &format!("    '{}': {}\n", error, count);
            }
        }

        let mut exits: String = String::new();

        if !self.exit_codes_with_count.is_empty() {
            exits = "  rustc exit codes\n".to_owned();

            for (exit, count) in &self.exit_codes_with_count {
                exits = exits + &format!("    {}: {}\n", exit, count);
            }
        }

        write!(
            f,
            r"
  Total compiles: {}
  Compile Time (ms)
    Total : {}
    Mean  : {}
    STDEV : {}
    75th  : {}
    90th  : {}
    95th  : {}
    99th  : {}

{}

{}",
            self.rustc_execution_count,
            self.compile_time_ms_total,
            self.compile_time_ms_mean,
            self.compile_time_ms_stdev,
            self.compile_time_ms_ntile_75,
            self.compile_time_ms_ntile_90,
            self.compile_time_ms_ntile_95,
            self.compile_time_ms_ntile_99,
            errors,
            exits
        )
    }
}

impl fmt::Display for TelemetryAnalysis {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            r"
Overall rustc statistics:
{}

rustc successful execution statistics
{}

rustc error statistics
{}",
            self.rustc_statistics, self.rustc_success_statistics, self.rustc_error_statistics
        )
    }
}

impl TelemetryAnalysis {
    pub fn new(telemetry_dir: PathBuf) -> TelemetryAnalysis {
        TelemetryAnalysis {
            telemetry_dir: telemetry_dir,
            rustc_statistics: RustcStatistics::new(),
            rustc_success_statistics: RustcStatistics::new(),
            rustc_error_statistics: RustcStatistics::new(),
        }
    }

    pub fn import_telemery(&mut self) -> Result<Vec<TelemetryEvent>> {
        let mut events: Vec<TelemetryEvent> = Vec::new();
        let contents = self.telemetry_dir
            .read_dir()
            .chain_err(|| ErrorKind::TelemetryAnalysisError)?;

        let mut telemetry_files: Vec<PathBuf> = Vec::new();

        for c in contents {
            let x = c.unwrap();
            let filename = x.path().file_name().unwrap().to_str().unwrap().to_owned();

            if filename.starts_with("log") && filename.ends_with("json") {
                telemetry_files.push(x.path());
                match self.read_telemetry_file(x.path()) {
                    Ok(y) => events.extend(y),
                    Err(e) => return Err(e),
                };
            }
        }

        Ok(events)
    }

    fn read_telemetry_file(&self, path: PathBuf) -> Result<Vec<TelemetryEvent>> {
        let mut events: Vec<TelemetryEvent> = Vec::new();

        let f = File::open(&path).chain_err(|| ErrorKind::TelemetryAnalysisError)?;

        let file = BufReader::new(&f);

        for line in file.lines() {
            let l = line.unwrap();
            let log_message_result = serde_json::from_str(&l);

            if log_message_result.is_ok() {
                let log_message: LogMessage = log_message_result.unwrap();
                let event: TelemetryEvent = log_message.get_event();
                events.push(event);
            }
        }

        Ok(events)
    }

    pub fn analyze_telemetry_events(&mut self, events: &[TelemetryEvent]) -> Result<()> {
        let mut rustc_durations = Vec::new();
        let mut rustc_exit_codes = Vec::new();

        let mut rustc_successful_durations = Vec::new();

        let mut rustc_error_durations = Vec::new();
        let mut error_list: Vec<Vec<String>> = Vec::new();
        let mut error_codes_with_counts: HashMap<String, i32> = HashMap::new();

        let mut toolchains = Vec::new();
        let mut toolchains_with_errors = Vec::new();
        let mut targets = Vec::new();

        let mut updated_toolchains = Vec::new();
        let mut updated_toolchains_with_errors = Vec::new();

        for event in events {
            match *event {
                TelemetryEvent::RustcRun {
                    duration_ms,
                    ref exit_code,
                    ref errors,
                } => {
                    self.rustc_statistics.rustc_execution_count += 1;
                    rustc_durations.push(duration_ms);

                    let exit_count = self.rustc_statistics
                        .exit_codes_with_count
                        .entry(*exit_code)
                        .or_insert(0);
                    *exit_count += 1;

                    rustc_exit_codes.push(exit_code);

                    if errors.is_some() {
                        let errors = errors.clone().unwrap();

                        for e in &errors {
                            let error_count =
                                error_codes_with_counts.entry(e.to_owned()).or_insert(0);
                            *error_count += 1;
                        }

                        error_list.push(errors);
                        rustc_error_durations.push(duration_ms);
                    } else {
                        rustc_successful_durations.push(duration_ms);
                    }
                }
                TelemetryEvent::TargetAdd {
                    ref toolchain,
                    ref target,
                    success,
                } => {
                    toolchains.push(toolchain.to_owned());
                    targets.push(target.to_owned());
                    if !success {
                        toolchains_with_errors.push(toolchain.to_owned());
                    }
                }
                TelemetryEvent::ToolchainUpdate {
                    ref toolchain,
                    success,
                } => {
                    updated_toolchains.push(toolchain.to_owned());
                    if !success {
                        updated_toolchains_with_errors.push(toolchain.to_owned());
                    }
                }
            }
        }

        self.rustc_statistics = compute_rustc_percentiles(&rustc_durations);
        self.rustc_error_statistics = compute_rustc_percentiles(&rustc_error_durations);
        self.rustc_error_statistics.error_codes_with_counts = error_codes_with_counts;
        self.rustc_success_statistics = compute_rustc_percentiles(&rustc_successful_durations);

        let error_list = error_list.into_iter().flatten();

        for e in error_list {
            let error_count = self.rustc_statistics
                .error_codes_with_counts
                .entry(e)
                .or_insert(0);
            *error_count += 1;
        }

        Ok(())
    }
}

pub fn compute_rustc_percentiles(values: &[u64]) -> RustcStatistics {
    RustcStatistics {
        rustc_execution_count: (values.len() as u32),
        compile_time_ms_total: values.iter().fold(0, |sum, val| sum + val),
        compile_time_ms_mean: mean(values),
        compile_time_ms_ntile_75: ntile(75, values),
        compile_time_ms_ntile_90: ntile(90, values),
        compile_time_ms_ntile_95: ntile(95, values),
        compile_time_ms_ntile_99: ntile(99, values),
        compile_time_ms_stdev: stdev(values),
        exit_codes_with_count: HashMap::new(),
        error_codes_with_counts: HashMap::new(),
    }
}

pub fn ntile(percentile: i32, values: &[u64]) -> u64 {
    if values.is_empty() {
        return 0u64;
    }

    let mut values = values.to_owned();
    values.sort();

    let count = values.len() as f32;
    let percentile = (percentile as f32) / 100f32;

    let n = (count * percentile).ceil() - 1f32;
    let n = n as usize;

    values[n]
}

pub fn mean(values: &[u64]) -> u64 {
    if values.is_empty() {
        return 0;
    }

    let count = values.len() as f64;

    let sum = values.iter().fold(0, |sum, val| sum + val) as f64;

    (sum / count) as u64
}

pub fn variance(values: &[u64]) -> f64 {
    if values.is_empty() {
        return 0f64;
    }

    let mean = mean(values);

    let mut deviations: Vec<i64> = Vec::new();

    for v in values.iter() {
        let x = (*v as i64) - (mean as i64);

        deviations.push(x * x);
    }

    let sum = deviations.iter().fold(0, |sum, val| sum + val) as f64;

    sum / (values.len() as f64)
}

pub fn stdev(values: &[u64]) -> f64 {
    if values.is_empty() {
        return 0f64;
    }

    let variance = variance(values);

    variance.sqrt()
}
<commit_msg>Fix Itertools::flatten() warning<commit_after>use std::collections::HashMap;
use std::fmt;
use std::fs::File;
use std::io::BufReader;
use std::io::BufRead;
use std::path::PathBuf;

use itertools::Itertools;
use serde_json;

use errors::*;
use telemetry::{LogMessage, TelemetryEvent};

pub struct TelemetryAnalysis {
    telemetry_dir: PathBuf,
    rustc_statistics: RustcStatistics,
    rustc_success_statistics: RustcStatistics,
    rustc_error_statistics: RustcStatistics,
}

#[derive(Default)]
pub struct RustcStatistics {
    rustc_execution_count: u32,
    compile_time_ms_total: u64,
    compile_time_ms_mean: u64,
    compile_time_ms_ntile_75: u64,
    compile_time_ms_ntile_90: u64,
    compile_time_ms_ntile_95: u64,
    compile_time_ms_ntile_99: u64,
    compile_time_ms_stdev: f64,
    exit_codes_with_count: HashMap<i32, i32>,
    error_codes_with_counts: HashMap<String, i32>,
}

impl RustcStatistics {
    pub fn new() -> RustcStatistics {
        Default::default()
    }
}

impl fmt::Display for RustcStatistics {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut errors: String = String::new();

        if !self.error_codes_with_counts.is_empty() {
            errors = "  rustc errors\n".to_owned();
            for (error, count) in &self.error_codes_with_counts {
                errors = errors + &format!("    '{}': {}\n", error, count);
            }
        }

        let mut exits: String = String::new();

        if !self.exit_codes_with_count.is_empty() {
            exits = "  rustc exit codes\n".to_owned();

            for (exit, count) in &self.exit_codes_with_count {
                exits = exits + &format!("    {}: {}\n", exit, count);
            }
        }

        write!(
            f,
            r"
  Total compiles: {}
  Compile Time (ms)
    Total : {}
    Mean  : {}
    STDEV : {}
    75th  : {}
    90th  : {}
    95th  : {}
    99th  : {}

{}

{}",
            self.rustc_execution_count,
            self.compile_time_ms_total,
            self.compile_time_ms_mean,
            self.compile_time_ms_stdev,
            self.compile_time_ms_ntile_75,
            self.compile_time_ms_ntile_90,
            self.compile_time_ms_ntile_95,
            self.compile_time_ms_ntile_99,
            errors,
            exits
        )
    }
}

impl fmt::Display for TelemetryAnalysis {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            r"
Overall rustc statistics:
{}

rustc successful execution statistics
{}

rustc error statistics
{}",
            self.rustc_statistics, self.rustc_success_statistics, self.rustc_error_statistics
        )
    }
}

impl TelemetryAnalysis {
    pub fn new(telemetry_dir: PathBuf) -> TelemetryAnalysis {
        TelemetryAnalysis {
            telemetry_dir: telemetry_dir,
            rustc_statistics: RustcStatistics::new(),
            rustc_success_statistics: RustcStatistics::new(),
            rustc_error_statistics: RustcStatistics::new(),
        }
    }

    pub fn import_telemery(&mut self) -> Result<Vec<TelemetryEvent>> {
        let mut events: Vec<TelemetryEvent> = Vec::new();
        let contents = self.telemetry_dir
            .read_dir()
            .chain_err(|| ErrorKind::TelemetryAnalysisError)?;

        let mut telemetry_files: Vec<PathBuf> = Vec::new();

        for c in contents {
            let x = c.unwrap();
            let filename = x.path().file_name().unwrap().to_str().unwrap().to_owned();

            if filename.starts_with("log") && filename.ends_with("json") {
                telemetry_files.push(x.path());
                match self.read_telemetry_file(x.path()) {
                    Ok(y) => events.extend(y),
                    Err(e) => return Err(e),
                };
            }
        }

        Ok(events)
    }

    fn read_telemetry_file(&self, path: PathBuf) -> Result<Vec<TelemetryEvent>> {
        let mut events: Vec<TelemetryEvent> = Vec::new();

        let f = File::open(&path).chain_err(|| ErrorKind::TelemetryAnalysisError)?;

        let file = BufReader::new(&f);

        for line in file.lines() {
            let l = line.unwrap();
            let log_message_result = serde_json::from_str(&l);

            if log_message_result.is_ok() {
                let log_message: LogMessage = log_message_result.unwrap();
                let event: TelemetryEvent = log_message.get_event();
                events.push(event);
            }
        }

        Ok(events)
    }

    pub fn analyze_telemetry_events(&mut self, events: &[TelemetryEvent]) -> Result<()> {
        let mut rustc_durations = Vec::new();
        let mut rustc_exit_codes = Vec::new();

        let mut rustc_successful_durations = Vec::new();

        let mut rustc_error_durations = Vec::new();
        let mut error_list: Vec<Vec<String>> = Vec::new();
        let mut error_codes_with_counts: HashMap<String, i32> = HashMap::new();

        let mut toolchains = Vec::new();
        let mut toolchains_with_errors = Vec::new();
        let mut targets = Vec::new();

        let mut updated_toolchains = Vec::new();
        let mut updated_toolchains_with_errors = Vec::new();

        for event in events {
            match *event {
                TelemetryEvent::RustcRun {
                    duration_ms,
                    ref exit_code,
                    ref errors,
                } => {
                    self.rustc_statistics.rustc_execution_count += 1;
                    rustc_durations.push(duration_ms);

                    let exit_count = self.rustc_statistics
                        .exit_codes_with_count
                        .entry(*exit_code)
                        .or_insert(0);
                    *exit_count += 1;

                    rustc_exit_codes.push(exit_code);

                    if errors.is_some() {
                        let errors = errors.clone().unwrap();

                        for e in &errors {
                            let error_count =
                                error_codes_with_counts.entry(e.to_owned()).or_insert(0);
                            *error_count += 1;
                        }

                        error_list.push(errors);
                        rustc_error_durations.push(duration_ms);
                    } else {
                        rustc_successful_durations.push(duration_ms);
                    }
                }
                TelemetryEvent::TargetAdd {
                    ref toolchain,
                    ref target,
                    success,
                } => {
                    toolchains.push(toolchain.to_owned());
                    targets.push(target.to_owned());
                    if !success {
                        toolchains_with_errors.push(toolchain.to_owned());
                    }
                }
                TelemetryEvent::ToolchainUpdate {
                    ref toolchain,
                    success,
                } => {
                    updated_toolchains.push(toolchain.to_owned());
                    if !success {
                        updated_toolchains_with_errors.push(toolchain.to_owned());
                    }
                }
            }
        }

        self.rustc_statistics = compute_rustc_percentiles(&rustc_durations);
        self.rustc_error_statistics = compute_rustc_percentiles(&rustc_error_durations);
        self.rustc_error_statistics.error_codes_with_counts = error_codes_with_counts;
        self.rustc_success_statistics = compute_rustc_percentiles(&rustc_successful_durations);

        let error_list = Itertools::flatten(error_list.into_iter());

        for e in error_list {
            let error_count = self.rustc_statistics
                .error_codes_with_counts
                .entry(e)
                .or_insert(0);
            *error_count += 1;
        }

        Ok(())
    }
}

pub fn compute_rustc_percentiles(values: &[u64]) -> RustcStatistics {
    RustcStatistics {
        rustc_execution_count: (values.len() as u32),
        compile_time_ms_total: values.iter().fold(0, |sum, val| sum + val),
        compile_time_ms_mean: mean(values),
        compile_time_ms_ntile_75: ntile(75, values),
        compile_time_ms_ntile_90: ntile(90, values),
        compile_time_ms_ntile_95: ntile(95, values),
        compile_time_ms_ntile_99: ntile(99, values),
        compile_time_ms_stdev: stdev(values),
        exit_codes_with_count: HashMap::new(),
        error_codes_with_counts: HashMap::new(),
    }
}

pub fn ntile(percentile: i32, values: &[u64]) -> u64 {
    if values.is_empty() {
        return 0u64;
    }

    let mut values = values.to_owned();
    values.sort();

    let count = values.len() as f32;
    let percentile = (percentile as f32) / 100f32;

    let n = (count * percentile).ceil() - 1f32;
    let n = n as usize;

    values[n]
}

pub fn mean(values: &[u64]) -> u64 {
    if values.is_empty() {
        return 0;
    }

    let count = values.len() as f64;

    let sum = values.iter().fold(0, |sum, val| sum + val) as f64;

    (sum / count) as u64
}

pub fn variance(values: &[u64]) -> f64 {
    if values.is_empty() {
        return 0f64;
    }

    let mean = mean(values);

    let mut deviations: Vec<i64> = Vec::new();

    for v in values.iter() {
        let x = (*v as i64) - (mean as i64);

        deviations.push(x * x);
    }

    let sum = deviations.iter().fold(0, |sum, val| sum + val) as f64;

    sum / (values.len() as f64)
}

pub fn stdev(values: &[u64]) -> f64 {
    if values.is_empty() {
        return 0f64;
    }

    let variance = variance(values);

    variance.sqrt()
}
<|endoftext|>


<commit_before>//!
//! The script interpreter runs scripts directly from their parse tree. This is
//! a fairly slow way to run scripts but comparatively simple to implement.
//!

use std::result::Result;
use serde_json::*;

use silkthread_base::*;

use super::script::*;

///
/// A tool representing a script that will be interepreted
///
pub struct InterpretedScriptTool {
    statements: Vec<Script>
}

///
/// Creates an unquoted version of a string
///
fn unquote_string(string: &str) -> String {
    let chars: Vec<char>    = string.chars().collect();
    let mut result          = String::new();
    let mut index           = 1;
    while index < chars.len()-1 {
        // Push character
        let chr = chars[index];

        match chr {
            '\\' => {
                let quoted = chars[index+1];
                match quoted {
                    'n' => result.push('\n'),
                    'r' => result.push('\r'),
                    't' => result.push('\t'),
                    quoted => result.push(quoted)
                }
            },
            chr => result.push(chr)
        }

        // Next character
        index += 1;
    }

    result
}

impl InterpretedScriptTool {
    ///
    /// Creates a new interpreted script tool from a set of statements
    ///
    pub fn from_statements(statements: Vec<Script>) -> InterpretedScriptTool {
        InterpretedScriptTool { statements: statements }
    }

    ///
    /// Attempts to evaluate an expression to a tool
    ///
    pub fn evaluate_expression_to_tool(expression: &Expression, environment: &mut ScriptExecutionEnvironment) -> Result<Box<Tool>, Value> {
        match expression {
            &Expression::Identifier(ref name)   => environment.parent_environment.get_json_tool(&name.matched).map_err(|_| Value::Null),
            _                                   => Err(Value::Null)
        }
    }

    ///
    /// Calls an expression representing a tool and calls it with the specified parameters
    ///
    pub fn call_tool(tool_name: &Expression, parameters: Value, environment: &mut ScriptExecutionEnvironment) -> Result<Value, Value> {
        InterpretedScriptTool::evaluate_expression_to_tool(tool_name, environment).and_then(|tool| {
            // TODO: create environment for the tool to run in

            // Invoke the tool and generate the final result
            tool.invoke_json(parameters, environment.parent_environment)
        })
    }

    ///
    /// Evaluates a single expression
    ///
    pub fn evaluate_expression(expression: &Expression, environment: &mut ScriptExecutionEnvironment) -> Result<Value, Value> {
        match expression {
            &Expression::String(ref s)          => Ok(Value::String(unquote_string(&s.matched))),

            &Expression::Identifier(_)          => InterpretedScriptTool::call_tool(expression, Value::Null, environment),

            _                                   => Err(Value::Null)
        }
    }

    ///
    /// Evaluates the result of executing a single statement
    ///
    pub fn evaluate_statement(statement: &Script, environment: &mut ScriptExecutionEnvironment) -> Result<Value, Value> {
        unimplemented!()
    }
}

impl Tool for InterpretedScriptTool {
    fn invoke_json(&self, input: Value, environment: &Environment) -> Result<Value, Value> {
        // Make the environment that this script will run in
        let mut script_environment = ScriptExecutionEnvironment::new(environment);

        // Execute the script
        let mut result = vec![];
        for statement in self.statements.iter() {
            // Evaluate the next statement
            let next_result = match Self::evaluate_statement(statement, &mut script_environment) {
                Ok(result) => result,

                // Fail immediately if any statement generates an error
                Err(fail) => return Err(fail)
            };

            // The script result is built up from the result of each statement
            // TODO: unless there's something like a return statement?
            result.push(next_result);
        }

        // Script is done
        Ok(Value::Array(result))
    }
}

///
/// Represents an execution environment for a running script
///
pub struct ScriptExecutionEnvironment<'a> {
    /// The environment where tools are drawn from
    parent_environment: &'a Environment
}

impl<'a> ScriptExecutionEnvironment<'a> {
    ///
    /// Creates a new script execution environment
    ///
    pub fn new(parent_environment: &'a Environment) -> ScriptExecutionEnvironment<'a> {
        ScriptExecutionEnvironment { parent_environment: parent_environment }
    }
}

#[cfg(test)]
mod test {
    use silkthread_base::basic::*;
    use super::*;

    #[test]
    fn can_evaluate_string() {
        let string_expr         = Expression::string("\"Foo\"");
        let empty_environment   = EmptyEnvironment::new();
        let mut env             = ScriptExecutionEnvironment::new(&empty_environment);
        let result              = InterpretedScriptTool::evaluate_expression(&string_expr, &mut env);

        assert!(result == Ok(Value::String(String::from("Foo"))));
    }

    #[test]
    fn can_evaluate_tool_call() {
        let tool_expr           = Expression::identifier("test");
        let tool_environment    = DynamicEnvironment::new();

        tool_environment.define("test", Box::new(make_pure_tool(|_: ()| "Success")));

        let mut env             = ScriptExecutionEnvironment::new(&tool_environment);
        let result              = InterpretedScriptTool::evaluate_expression(&tool_expr, &mut env);

        println!("{:?}", make_pure_tool(|_: ()| "Success").invoke_json(Value::Null, &tool_environment));
        println!("{:?}", result);
        assert!(result == Ok(Value::String(String::from("Success"))));
    }
}
<commit_msg>Improve error handling<commit_after>//!
//! The script interpreter runs scripts directly from their parse tree. This is
//! a fairly slow way to run scripts but comparatively simple to implement.
//!

use std::result::Result;
use serde_json::*;

use silkthread_base::*;

use super::script::*;

///
/// A tool representing a script that will be interepreted
///
pub struct InterpretedScriptTool {
    statements: Vec<Script>
}

///
/// Creates an unquoted version of a string
///
fn unquote_string(string: &str) -> String {
    let chars: Vec<char>    = string.chars().collect();
    let mut result          = String::new();
    let mut index           = 1;
    while index < chars.len()-1 {
        // Push character
        let chr = chars[index];

        match chr {
            '\\' => {
                let quoted = chars[index+1];
                match quoted {
                    'n' => result.push('\n'),
                    'r' => result.push('\r'),
                    't' => result.push('\t'),
                    quoted => result.push(quoted)
                }
            },
            chr => result.push(chr)
        }

        // Next character
        index += 1;
    }

    result
}

///
/// Script evaluation error
///
#[derive(Serialize, Deserialize)]
pub enum ScriptEvaluationError {
    /// Tried to evaluate an expression type that's not implemented yet
    ExpressionNotImplemented,

    /// Tried to look up a tool and it couldn't be found
    ToolNameNotFound,

    /// Found an expression that can't be treated as a tool where a tool name was expected
    ExpressionDoesNotEvaluateToTool
}

///
/// Creates an execution error
///
fn generate_expression_error(error: ScriptEvaluationError, expr: &Expression) -> Value {
    json![{
        "error":                error,
        "failed-expression":    expr
    }]
}

impl InterpretedScriptTool {
    ///
    /// Creates a new interpreted script tool from a set of statements
    ///
    pub fn from_statements(statements: Vec<Script>) -> InterpretedScriptTool {
        InterpretedScriptTool { statements: statements }
    }

    ///
    /// Attempts to evaluate an expression to a tool
    ///
    pub fn evaluate_expression_to_tool(expression: &Expression, environment: &mut ScriptExecutionEnvironment) -> Result<Box<Tool>, Value> {
        match expression {
            &Expression::Identifier(ref name)   => environment.parent_environment.get_json_tool(&name.matched).map_err(|_| generate_expression_error(ScriptEvaluationError::ToolNameNotFound, expression)),
            _                                   => Err(generate_expression_error(ScriptEvaluationError::ExpressionDoesNotEvaluateToTool, expression))
        }
    }

    ///
    /// Calls an expression representing a tool and calls it with the specified parameters
    ///
    pub fn call_tool(tool_name: &Expression, parameters: Value, environment: &mut ScriptExecutionEnvironment) -> Result<Value, Value> {
        InterpretedScriptTool::evaluate_expression_to_tool(tool_name, environment).and_then(|tool| {
            // TODO: create environment for the tool to run in

            // Invoke the tool and generate the final result
            tool.invoke_json(parameters, environment.parent_environment)
        })
    }

    ///
    /// Evaluates a single expression
    ///
    pub fn evaluate_expression(expression: &Expression, environment: &mut ScriptExecutionEnvironment) -> Result<Value, Value> {
        match expression {
            &Expression::String(ref s)          => Ok(Value::String(unquote_string(&s.matched))),

            &Expression::Identifier(_)          => InterpretedScriptTool::call_tool(expression, Value::Null, environment),

            _                                   => Err(generate_expression_error(ScriptEvaluationError::ExpressionNotImplemented, expression))
        }
    }

    ///
    /// Evaluates the result of executing a single statement
    ///
    pub fn evaluate_statement(statement: &Script, environment: &mut ScriptExecutionEnvironment) -> Result<Value, Value> {
        unimplemented!()
    }
}

impl Tool for InterpretedScriptTool {
    fn invoke_json(&self, input: Value, environment: &Environment) -> Result<Value, Value> {
        // Make the environment that this script will run in
        let mut script_environment = ScriptExecutionEnvironment::new(environment);

        // Execute the script
        let mut result = vec![];
        for statement in self.statements.iter() {
            // Evaluate the next statement
            let next_result = match Self::evaluate_statement(statement, &mut script_environment) {
                Ok(result) => result,

                // Fail immediately if any statement generates an error
                Err(fail) => return Err(fail)
            };

            // The script result is built up from the result of each statement
            // TODO: unless there's something like a return statement?
            result.push(next_result);
        }

        // Script is done
        Ok(Value::Array(result))
    }
}

///
/// Represents an execution environment for a running script
///
pub struct ScriptExecutionEnvironment<'a> {
    /// The environment where tools are drawn from
    parent_environment: &'a Environment
}

impl<'a> ScriptExecutionEnvironment<'a> {
    ///
    /// Creates a new script execution environment
    ///
    pub fn new(parent_environment: &'a Environment) -> ScriptExecutionEnvironment<'a> {
        ScriptExecutionEnvironment { parent_environment: parent_environment }
    }
}

#[cfg(test)]
mod test {
    use silkthread_base::basic::*;
    use super::*;

    #[test]
    fn can_evaluate_string() {
        let string_expr         = Expression::string("\"Foo\"");
        let empty_environment   = EmptyEnvironment::new();
        let mut env             = ScriptExecutionEnvironment::new(&empty_environment);
        let result              = InterpretedScriptTool::evaluate_expression(&string_expr, &mut env);

        assert!(result == Ok(Value::String(String::from("Foo"))));
    }

    #[test]
    fn can_evaluate_tool_call() {
        let tool_expr           = Expression::identifier("test");
        let tool_environment    = DynamicEnvironment::new();

        tool_environment.define("test", Box::new(make_pure_tool(|_: ()| "Success")));

        let mut env             = ScriptExecutionEnvironment::new(&tool_environment);
        let result              = InterpretedScriptTool::evaluate_expression(&tool_expr, &mut env);

        println!("{:?}", tool_environment.get_json_tool("test").unwrap().invoke_json(Value::Null, &tool_environment));
        println!("{:?}", InterpretedScriptTool::call_tool(&tool_expr, Value::Null, &mut env));
        println!("{:?}", result);
        assert!(result == Ok(Value::String(String::from("Success"))));
    }
}
<|endoftext|>