<commit_before>use std::fmt;
use std::ops;
use std::num::Zero;
use std::cmp::{PartialEq, PartialOrd};

pub trait NumType: fmt::Debug+Clone+Zero+ops::Add<Output=Self>+ops::Sub<Output=Self>+ops::Mul<Output=Self>+ops::Div<Output=Self>+PartialEq+PartialOrd+Default {}

impl NumType for f32 {}
impl NumType for f64 {}
impl NumType for u32 {}
impl NumType for u64 {}

#[derive(Debug, Eq, PartialEq)]
pub enum ExprError {
    /// In case of division by zero.
    DivByZero,
    /// When the argument position `n` in Arg(n) is out of bounds.
    InvalidArg,
    /// In case of an invalid operation.
    InvalidOperation,
}

/// An expression evaluates to a numeric value of `NumType`.
#[derive(Clone, PartialEq, Eq)]
pub enum Expr<T: NumType> {
    Const(T),
    // References an actual Argument by position
    Arg(usize),
    Add(Box<Expr<T>>, Box<Expr<T>>),
    Sub(Box<Expr<T>>, Box<Expr<T>>),
    Mul(Box<Expr<T>>, Box<Expr<T>>),
    Div(Box<Expr<T>>, Box<Expr<T>>),

    // Safe division with x/0 = 0.0
    Divz(Box<Expr<T>>, Box<Expr<T>>),
}

impl<T:NumType> fmt::Debug for Expr<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Expr::Const(ref c) => write!(f, "{:?}", c),
            &Expr::Arg(num) => write!(f, "a{}", num),
            &Expr::Add(ref op1, ref op2) => write!(f, "({:?}+{:?})", op1, op2),
            &Expr::Sub(ref op1, ref op2) => write!(f, "({:?}-{:?})", op1, op2),
            &Expr::Mul(ref op1, ref op2) => write!(f, "{:?}*{:?}", op1, op2),
            &Expr::Div(ref op1, ref op2) => write!(f, "{:?}/{:?}", op1, op2),
            &Expr::Divz(ref op1, ref op2) => write!(f, "{:?}\\{:?}", op1, op2),
        }
    }
}

impl<T:NumType> Expr<T> {
    pub fn value(&self) -> Result<T, ExprError> {
        if let &Expr::Const(ref c) = self {
            Ok(c.clone())
        } else {
            Err(ExprError::InvalidArg)
        }
    }

    pub fn evaluate(&self, args: &[Expr<T>]) -> Result<T, ExprError> {
        Ok(match *self {
            Expr::Arg(n) => {
                try!(try!(args.get(n).ok_or(ExprError::InvalidArg)).value())
            }
            Expr::Const(ref c) => c.clone(),
            Expr::Add(ref e1, ref e2) => try!(e1.evaluate(args)) + try!(e2.evaluate(args)),
            Expr::Sub(ref e1, ref e2) => try!(e1.evaluate(args)) - try!(e2.evaluate(args)),
            Expr::Mul(ref e1, ref e2) => try!(e1.evaluate(args)) * try!(e2.evaluate(args)),
            Expr::Div(ref e1, ref e2) => {
                let a = try!(e1.evaluate(args));
                let b = try!(e2.evaluate(args));
                if b == T::zero() {
                    return Err(ExprError::DivByZero);
                }
                a / b
            }
            Expr::Divz(ref e1, ref e2) => {
                let a = try!(e1.evaluate(args));
                let b = try!(e2.evaluate(args));
                if b == T::zero() {
                    b
                } else {
                    a / b
                }
            }
        })
    }
}

/// A condition evaluates to either `true` or `false`.
#[derive(Debug, Clone)]
pub enum Condition<T: NumType> {
    True,
    False,
    Not(Box<Condition<T>>),
    And(Box<Condition<T>>, Box<Condition<T>>),
    Or(Box<Condition<T>>, Box<Condition<T>>),

    /// If two expressions are equal
    Equal(Box<Expr<T>>, Box<Expr<T>>),

    Less(Box<Expr<T>>, Box<Expr<T>>),
    Greater(Box<Expr<T>>, Box<Expr<T>>),

    LessEqual(Box<Expr<T>>, Box<Expr<T>>),
    GreaterEqual(Box<Expr<T>>, Box<Expr<T>>),
}

impl<T:NumType> Condition<T> {
    pub fn evaluate(&self, args: &[Expr<T>]) -> Result<bool, ExprError> {
        Ok(match *self {
            Condition::True => true,
            Condition::False => false,
            Condition::Not(ref c) => !try!(c.evaluate(args)),
            Condition::And(ref c1, ref c2) => try!(c1.evaluate(args)) && try!(c2.evaluate(args)),
            Condition::Or(ref c1, ref c2) => try!(c1.evaluate(args)) || try!(c2.evaluate(args)),
            Condition::Equal(ref e1, ref e2) => try!(e1.evaluate(args)) == try!(e2.evaluate(args)),
            Condition::Less(ref e1, ref e2) => try!(e1.evaluate(args)) < try!(e2.evaluate(args)),
            Condition::Greater(ref e1, ref e2) => try!(e1.evaluate(args)) > try!(e2.evaluate(args)),
            Condition::LessEqual(ref e1, ref e2) =>
                try!(e1.evaluate(args)) <= try!(e2.evaluate(args)),
            Condition::GreaterEqual(ref e1, ref e2) =>
                try!(e1.evaluate(args)) >= try!(e2.evaluate(args)),
        })
    }
}

#[test]
fn test_expr_divz() {
    let expr = Expr::Divz(box Expr::Const(1.0), box Expr::Const(0.0));
    assert_eq!(Ok(0.0), expr.evaluate(&[]));
}

#[test]
fn test_expr() {
    let expr = Expr::Sub(box Expr::Const(0.0),
                         box Expr::Div(box Expr::Mul(box Expr::Add(box Expr::Const(1.0),
                                                                   box Expr::Arg(0)),
                                                     box Expr::Arg(1)),
                                       box Expr::Const(2.0)));

    fn fun(a: f32, b: f32) -> f32 {
        0.0 - ((1.0 + a) * b) / 2.0
    }

    fn check(expr: &Expr<f32>, a: f32, b: f32) {
        assert_eq!(Ok(fun(a, b)),
                   expr.evaluate(&[Expr::Const(a), Expr::Const(b)]))
    }

    check(&expr, 123.0, 4444.0);
    check(&expr, 0.0, -12.0);
}

#[test]
fn test_condition() {
    let cond = Condition::Greater(box Expr::Arg(0), box Expr::Const(0.0));

    fn fun(a: f32) -> bool {
        a > 0.0
    }

    fn check(cond: &Condition<f32>, a: f32) {
        assert_eq!(Ok(fun(a)), cond.evaluate(&[Expr::Const(a)]))
    }

    check(&cond, 123.0);
    check(&cond, 0.0);
    check(&cond, -1.4);

    assert_eq!(Ok(true), (Condition::True::<f32>).evaluate(&[]));
}
<commit_msg>Add Expr::{Recip, Recipz} for reciprocal<commit_after>use std::fmt;
use std::ops;
use std::num::{Zero, One};
use std::cmp::{PartialEq, PartialOrd};

pub trait NumType: fmt::Debug+Clone+Zero+One+ops::Add<Output=Self>+ops::Sub<Output=Self>+ops::Mul<Output=Self>+ops::Div<Output=Self>+PartialEq+PartialOrd+Default {}

impl NumType for f32 {}
impl NumType for f64 {}
impl NumType for u32 {}
impl NumType for u64 {}

#[derive(Debug, Eq, PartialEq)]
pub enum ExprError {
    /// In case of division by zero.
    DivByZero,
    /// When the argument position `n` in Arg(n) is out of bounds.
    InvalidArg,
    /// In case of an invalid operation.
    InvalidOperation,
}

/// An expression evaluates to a numeric value of `NumType`.
#[derive(Clone, PartialEq, Eq)]
pub enum Expr<T: NumType> {
    Const(T),
    // References an actual Argument by position
    Arg(usize),
    Add(Box<Expr<T>>, Box<Expr<T>>),
    Sub(Box<Expr<T>>, Box<Expr<T>>),
    Mul(Box<Expr<T>>, Box<Expr<T>>),
    Div(Box<Expr<T>>, Box<Expr<T>>),

    // Safe division with x/0 = 0.0
    Divz(Box<Expr<T>>, Box<Expr<T>>),

    // Reciprocal (1 / x).
    Recip(Box<Expr<T>>),

    // Reciprocal using safe division
    Recipz(Box<Expr<T>>),
}

impl<T: NumType> fmt::Debug for Expr<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            &Expr::Const(ref c) => write!(f, "{:?}", c),
            &Expr::Arg(num) => write!(f, "a{}", num),
            &Expr::Add(ref op1, ref op2) => write!(f, "({:?}+{:?})", op1, op2),
            &Expr::Sub(ref op1, ref op2) => write!(f, "({:?}-{:?})", op1, op2),
            &Expr::Mul(ref op1, ref op2) => write!(f, "{:?}*{:?}", op1, op2),
            &Expr::Div(ref op1, ref op2) => write!(f, "{:?}/{:?}", op1, op2),
            &Expr::Divz(ref op1, ref op2) => write!(f, "{:?}\\{:?}", op1, op2),
            &Expr::Recip(ref op) => write!(f, "{:?}/{:?}", T::one(), op),
            &Expr::Recipz(ref op) => write!(f, "{:?}\\{:?}", T::one(), op),
        }
    }
}

impl<T: NumType> Expr<T> {
    pub fn value(&self) -> Result<T, ExprError> {
        if let &Expr::Const(ref c) = self {
            Ok(c.clone())
        } else {
            Err(ExprError::InvalidArg)
        }
    }

    pub fn evaluate(&self, args: &[Expr<T>]) -> Result<T, ExprError> {
        Ok(match *self {
            Expr::Arg(n) => try!(try!(args.get(n).ok_or(ExprError::InvalidArg)).value()),
            Expr::Const(ref c) => c.clone(),
            Expr::Add(ref e1, ref e2) => try!(e1.evaluate(args)) + try!(e2.evaluate(args)),
            Expr::Sub(ref e1, ref e2) => try!(e1.evaluate(args)) - try!(e2.evaluate(args)),
            Expr::Mul(ref e1, ref e2) => try!(e1.evaluate(args)) * try!(e2.evaluate(args)),
            Expr::Div(ref e1, ref e2) => {
                let a = try!(e1.evaluate(args));
                let div = try!(e2.evaluate(args));
                if div == T::zero() {
                    return Err(ExprError::DivByZero);
                }
                a / div
            }
            Expr::Divz(ref e1, ref e2) => {
                let a = try!(e1.evaluate(args));
                let div = try!(e2.evaluate(args));
                if div == T::zero() {
                    div
                } else {
                    a / div
                }
            }
            Expr::Recip(ref e1) => {
                let div = try!(e1.evaluate(args));
                if div == T::zero() {
                    return Err(ExprError::DivByZero);
                } else {
                    T::one() / div
                }
            }
            Expr::Recipz(ref e1) => {
                let div = try!(e1.evaluate(args));
                if div == T::zero() {
                    div
                } else {
                    T::one() / div
                }
            }
        })
    }
}

/// A condition evaluates to either `true` or `false`.
#[derive(Debug, Clone)]
pub enum Condition<T: NumType> {
    True,
    False,
    Not(Box<Condition<T>>),
    And(Box<Condition<T>>, Box<Condition<T>>),
    Or(Box<Condition<T>>, Box<Condition<T>>),

    /// If two expressions are equal
    Equal(Box<Expr<T>>, Box<Expr<T>>),

    Less(Box<Expr<T>>, Box<Expr<T>>),
    Greater(Box<Expr<T>>, Box<Expr<T>>),

    LessEqual(Box<Expr<T>>, Box<Expr<T>>),
    GreaterEqual(Box<Expr<T>>, Box<Expr<T>>),
}

impl<T: NumType> Condition<T> {
    pub fn evaluate(&self, args: &[Expr<T>]) -> Result<bool, ExprError> {
        Ok(match *self {
            Condition::True => true,
            Condition::False => false,
            Condition::Not(ref c) => !try!(c.evaluate(args)),
            Condition::And(ref c1, ref c2) => try!(c1.evaluate(args)) && try!(c2.evaluate(args)),
            Condition::Or(ref c1, ref c2) => try!(c1.evaluate(args)) || try!(c2.evaluate(args)),
            Condition::Equal(ref e1, ref e2) => try!(e1.evaluate(args)) == try!(e2.evaluate(args)),
            Condition::Less(ref e1, ref e2) => try!(e1.evaluate(args)) < try!(e2.evaluate(args)),
            Condition::Greater(ref e1, ref e2) => try!(e1.evaluate(args)) > try!(e2.evaluate(args)),
            Condition::LessEqual(ref e1, ref e2) => {
                try!(e1.evaluate(args)) <= try!(e2.evaluate(args))
            }
            Condition::GreaterEqual(ref e1, ref e2) => {
                try!(e1.evaluate(args)) >= try!(e2.evaluate(args))
            }
        })
    }
}

#[test]
fn test_expr_divz() {
    let expr = Expr::Divz(box Expr::Const(1.0), box Expr::Const(0.0));
    assert_eq!(Ok(0.0), expr.evaluate(&[]));
}

#[test]
fn test_expr_recipz() {
    let expr = Expr::Recipz(box Expr::Const(0.0));
    assert_eq!(Ok(0.0), expr.evaluate(&[]));

    let expr = Expr::Recipz(box Expr::Const(1.0));
    assert_eq!(Ok(1.0), expr.evaluate(&[]));

    let expr = Expr::Recipz(box Expr::Const(0.5));
    assert_eq!(Ok(2.0), expr.evaluate(&[]));
}

#[test]
fn test_expr() {
    let expr = Expr::Sub(box Expr::Const(0.0),
                         box Expr::Div(box Expr::Mul(box Expr::Add(box Expr::Const(1.0),
                                                                   box Expr::Arg(0)),
                                                     box Expr::Arg(1)),
                                       box Expr::Const(2.0)));

    fn fun(a: f32, b: f32) -> f32 {
        0.0 - ((1.0 + a) * b) / 2.0
    }

    fn check(expr: &Expr<f32>, a: f32, b: f32) {
        assert_eq!(Ok(fun(a, b)),
                   expr.evaluate(&[Expr::Const(a), Expr::Const(b)]))
    }

    check(&expr, 123.0, 4444.0);
    check(&expr, 0.0, -12.0);
}

#[test]
fn test_condition() {
    let cond = Condition::Greater(box Expr::Arg(0), box Expr::Const(0.0));

    fn fun(a: f32) -> bool {
        a > 0.0
    }

    fn check(cond: &Condition<f32>, a: f32) {
        assert_eq!(Ok(fun(a)), cond.evaluate(&[Expr::Const(a)]))
    }

    check(&cond, 123.0);
    check(&cond, 0.0);
    check(&cond, -1.4);

    assert_eq!(Ok(true), (Condition::True::<f32>).evaluate(&[]));
}
<|endoftext|>


<commit_before>use std::io::{fs, IoResult};
use std::io;
use std::str::CowString;

use ansi_term::{ANSIString, Colour, Style};
use ansi_term::Style::Plain;
use ansi_term::Colour::{Red, Green, Yellow, Blue, Purple, Cyan, Fixed};

use users::Users;

use number_prefix::{binary_prefix, decimal_prefix, Prefixed, Standalone, PrefixNames};

use column::{Column, SizeFormat};
use column::Column::*;
use dir::Dir;
use filetype::HasType;

pub static GREY: Colour = Fixed(244);

// Instead of working with Rust's Paths, we have our own File object
// that holds the Path and various cached information. Each file is
// definitely going to have its filename used at least once, its stat
// information queried at least once, and its file extension extracted
// at least once, so we may as well carry around that information with
// the actual path.

pub struct File<'a> {
    pub name:  String,
    pub dir:   Option<&'a Dir>,
    pub ext:   Option<String>,
    pub path:  Path,
    pub stat:  io::FileStat,
}

impl<'a> File<'a> {
    pub fn from_path(path: &Path, parent: Option<&'a Dir>) -> IoResult<File<'a>> {
        // Use lstat here instead of file.stat(), as it doesn't follow
        // symbolic links. Otherwise, the stat() call will fail if it
        // encounters a link that's target is non-existent.
        fs::lstat(path).map(|stat| File::with_stat(stat, path, parent))
    }

    pub fn with_stat(stat: io::FileStat, path: &Path, parent: Option<&'a Dir>) -> File<'a> {
        let v = path.filename().unwrap();  // fails if / or . or ..
        let filename = String::from_utf8_lossy(v);

        File {
            path:  path.clone(),
            dir:   parent,
            stat:  stat,
            name:  filename.to_string(),
            ext:   File::ext(filename),
        }
    }

    fn ext(name: CowString) -> Option<String> {
        // The extension is the series of characters after a dot at
        // the end of a filename. This deliberately also counts
        // dotfiles - the ".git" folder has the extension "git".
        let re = regex!(r"\.([^.]+)$");
        re.captures(name.as_slice()).map(|caps| caps.at(1).to_string())
    }

    pub fn is_dotfile(&self) -> bool {
        self.name.as_slice().starts_with(".")
    }

    pub fn is_tmpfile(&self) -> bool {
        let name = self.name.as_slice();
        name.ends_with("~") || (name.starts_with("#") && name.ends_with("#"))
    }

    // Highlight the compiled versions of files. Some of them, like .o,
    // get special highlighting when they're alone because there's no
    // point in existing without their source. Others can be perfectly
    // content without their source files, such as how .js is valid
    // without a .coffee.

    pub fn get_source_files(&self) -> Vec<Path> {
        if let Some(ref ext) = self.ext {
            let ext = ext.as_slice();
            match ext {
                "class" => vec![self.path.with_extension("java")],  // Java
                "css"   => vec![self.path.with_extension("sass"),   self.path.with_extension("less")],  // SASS, Less
                "elc"   => vec![self.path.with_extension("el")],    // Emacs Lisp
                "hi"    => vec![self.path.with_extension("hs")],    // Haskell
                "js"    => vec![self.path.with_extension("coffee"), self.path.with_extension("ts")],  // CoffeeScript, TypeScript
                "o"     => vec![self.path.with_extension("c"),      self.path.with_extension("cpp")], // C, C++
                "pyc"   => vec![self.path.with_extension("py")],    // Python

                "aux" => vec![self.path.with_extension("tex")],  // TeX: auxiliary file
                "bbl" => vec![self.path.with_extension("tex")],  // BibTeX bibliography file
                "blg" => vec![self.path.with_extension("tex")],  // BibTeX log file
                "lof" => vec![self.path.with_extension("tex")],  // TeX list of figures
                "log" => vec![self.path.with_extension("tex")],  // TeX log file
                "lot" => vec![self.path.with_extension("tex")],  // TeX list of tables
                "toc" => vec![self.path.with_extension("tex")],  // TeX table of contents

                _ => vec![],  // No source files if none of the above
            }
        }
        else {
            vec![]  // No source files if there's no extension, either!
        }
    }

    pub fn display<U: Users>(&self, column: &Column, users_cache: &mut U) -> String {
        match *column {
            Permissions => {
                self.permissions_string()
            },

            FileName => {
                self.file_name()
            },

            FileSize(use_iec) => {
                self.file_size(use_iec)
            },

            // A file with multiple links is interesting, but
            // directories and suchlike can have multiple links all
            // the time.
            HardLinks => {
                let style = if self.has_multiple_links() { Red.on(Yellow) } else { Red.normal() };
                style.paint(self.stat.unstable.nlink.to_string().as_slice()).to_string()
            },

            Inode => {
                Purple.paint(self.stat.unstable.inode.to_string().as_slice()).to_string()
            },

            Blocks => {
                if self.stat.kind == io::FileType::RegularFile || self.stat.kind == io::FileType::Symlink {
                    Cyan.paint(self.stat.unstable.blocks.to_string().as_slice()).to_string()
                }
                else {
                    GREY.paint("-").to_string()
                }
            },

            // Display the ID if the user/group doesn't exist, which
            // usually means it was deleted but its files weren't.
            User => {
                let uid = self.stat.unstable.uid as i32;

                let user_name = match users_cache.get_user_by_uid(uid) {
                    Some(user) => user.name,
                    None => uid.to_string(),
                };

                let style = if users_cache.get_current_uid() == uid { Yellow.bold() } else { Plain };
                style.paint(user_name.as_slice()).to_string()
            },

            Group => {
                let gid = self.stat.unstable.gid as u32;
                let mut style = Plain;

                let group_name = match users_cache.get_group_by_gid(gid) {
                    Some(group) => {
                        let current_uid = users_cache.get_current_uid();
                        if let Some(current_user) = users_cache.get_user_by_uid(current_uid) {
                            if current_user.primary_group == group.gid || group.members.contains(&current_user.name) {
                                style = Yellow.bold();
                            }
                        }
                        group.name
                    },
                    None => gid.to_string(),
                };

                style.paint(group_name.as_slice()).to_string()
            },
        }
    }

    pub fn file_name(&self) -> String {
        let name = self.name.as_slice();
        let displayed_name = self.file_colour().paint(name);
        if self.stat.kind == io::FileType::Symlink {
            match fs::readlink(&self.path) {
                Ok(path) => {
                    let target_path = match self.dir {
                        Some(dir) => dir.path.join(path),
                        None => path,
                    };
                    format!("{} {}", displayed_name, self.target_file_name_and_arrow(&target_path))
                }
                Err(_) => displayed_name.to_string(),
            }
        }
        else {
            displayed_name.to_string()
        }
    }

    pub fn file_name_width(&self) -> uint {
        self.name.as_slice().width(false)
    }

    fn target_file_name_and_arrow(&self, target_path: &Path) -> String {
        let v = target_path.filename().unwrap();
        let filename = String::from_utf8_lossy(v);

        // Use stat instead of lstat - we *want* to follow links.
        let link_target = fs::stat(target_path).map(|stat| File {
            path:  target_path.clone(),
            dir:   self.dir,
            stat:  stat,
            name:  filename.to_string(),
            ext:   File::ext(filename.clone()),
        });

        // Statting a path usually fails because the file at the
        // other end doesn't exist. Show this by highlighting the
        // target file in red instead of displaying an error, because
        // the error would be shown out of context (before the
        // results, not right by the file) and it's almost always for
        // that reason anyway.

        match link_target {
            Ok(file) => format!("{} {}{}{}", GREY.paint("=>"), Cyan.paint(target_path.dirname_str().unwrap()), Cyan.paint("/"), file.file_colour().paint(filename.as_slice())),
            Err(_)   => format!("{} {}",     Red.paint("=>"),  Red.underline().paint(filename.as_slice())),
        }
    }

    fn file_size(&self, size_format: SizeFormat) -> String {
        // Don't report file sizes for directories. I've never looked
        // at one of those numbers and gained any information from it.
        if self.stat.kind == io::FileType::Directory {
            GREY.paint("-").to_string()
        }
        else {
        	let result = match size_format {
        		SizeFormat::DecimalBytes => decimal_prefix(self.stat.size as f64),
        		SizeFormat::BinaryBytes => binary_prefix(self.stat.size as f64),
        	};

			match result {
				Standalone(bytes) => Green.bold().paint(bytes.to_string().as_slice()).to_string(),
				Prefixed(prefix, n) => {
					let number = if n < 10f64 { format!("{:.1}", n) } else { format!("{:.0}", n) };
					format!("{}{}", Green.bold().paint(number.as_slice()), Green.paint(prefix.symbol()))
				}
			}
        }
    }

    fn type_char(&self) -> ANSIString {
        return match self.stat.kind {
            io::FileType::RegularFile  => Plain.paint("."),
            io::FileType::Directory    => Blue.paint("d"),
            io::FileType::NamedPipe    => Yellow.paint("|"),
            io::FileType::BlockSpecial => Purple.paint("s"),
            io::FileType::Symlink      => Cyan.paint("l"),
            io::FileType::Unknown      => Plain.paint("?"),
        }
    }

    pub fn file_colour(&self) -> Style {
        self.get_type().style()
    }

    fn has_multiple_links(&self) -> bool {
        self.stat.kind == io::FileType::RegularFile && self.stat.unstable.nlink > 1
    }

    fn permissions_string(&self) -> String {
        let bits = self.stat.perm;
        return format!("{}{}{}{}{}{}{}{}{}{}",
            self.type_char(),

            // The first three are bold because they're the ones used
            // most often.
            File::permission_bit(bits, io::USER_READ,     "r", Yellow.bold()),
            File::permission_bit(bits, io::USER_WRITE,    "w", Red.bold()),
            File::permission_bit(bits, io::USER_EXECUTE,  "x", Green.bold().underline()),
            File::permission_bit(bits, io::GROUP_READ,    "r", Yellow.normal()),
            File::permission_bit(bits, io::GROUP_WRITE,   "w", Red.normal()),
            File::permission_bit(bits, io::GROUP_EXECUTE, "x", Green.normal()),
            File::permission_bit(bits, io::OTHER_READ,    "r", Yellow.normal()),
            File::permission_bit(bits, io::OTHER_WRITE,   "w", Red.normal()),
            File::permission_bit(bits, io::OTHER_EXECUTE, "x", Green.normal()),
       );
    }

    fn permission_bit(bits: io::FilePermission, bit: io::FilePermission, character: &'static str, style: Style) -> ANSIString {
        if bits.contains(bit) {
            style.paint(character)
        }
        else {
            GREY.paint("-")
        }
    }
}
<commit_msg>Only underline executable *files* in perms view<commit_after>use std::io::{fs, IoResult};
use std::io;
use std::str::CowString;

use ansi_term::{ANSIString, Colour, Style};
use ansi_term::Style::Plain;
use ansi_term::Colour::{Red, Green, Yellow, Blue, Purple, Cyan, Fixed};

use users::Users;

use number_prefix::{binary_prefix, decimal_prefix, Prefixed, Standalone, PrefixNames};

use column::{Column, SizeFormat};
use column::Column::*;
use dir::Dir;
use filetype::HasType;

pub static GREY: Colour = Fixed(244);

// Instead of working with Rust's Paths, we have our own File object
// that holds the Path and various cached information. Each file is
// definitely going to have its filename used at least once, its stat
// information queried at least once, and its file extension extracted
// at least once, so we may as well carry around that information with
// the actual path.

pub struct File<'a> {
    pub name:  String,
    pub dir:   Option<&'a Dir>,
    pub ext:   Option<String>,
    pub path:  Path,
    pub stat:  io::FileStat,
}

impl<'a> File<'a> {
    pub fn from_path(path: &Path, parent: Option<&'a Dir>) -> IoResult<File<'a>> {
        // Use lstat here instead of file.stat(), as it doesn't follow
        // symbolic links. Otherwise, the stat() call will fail if it
        // encounters a link that's target is non-existent.
        fs::lstat(path).map(|stat| File::with_stat(stat, path, parent))
    }

    pub fn with_stat(stat: io::FileStat, path: &Path, parent: Option<&'a Dir>) -> File<'a> {
        let v = path.filename().unwrap();  // fails if / or . or ..
        let filename = String::from_utf8_lossy(v);

        File {
            path:  path.clone(),
            dir:   parent,
            stat:  stat,
            name:  filename.to_string(),
            ext:   File::ext(filename),
        }
    }

    fn ext(name: CowString) -> Option<String> {
        // The extension is the series of characters after a dot at
        // the end of a filename. This deliberately also counts
        // dotfiles - the ".git" folder has the extension "git".
        let re = regex!(r"\.([^.]+)$");
        re.captures(name.as_slice()).map(|caps| caps.at(1).to_string())
    }

    pub fn is_dotfile(&self) -> bool {
        self.name.as_slice().starts_with(".")
    }

    pub fn is_tmpfile(&self) -> bool {
        let name = self.name.as_slice();
        name.ends_with("~") || (name.starts_with("#") && name.ends_with("#"))
    }

    // Highlight the compiled versions of files. Some of them, like .o,
    // get special highlighting when they're alone because there's no
    // point in existing without their source. Others can be perfectly
    // content without their source files, such as how .js is valid
    // without a .coffee.

    pub fn get_source_files(&self) -> Vec<Path> {
        if let Some(ref ext) = self.ext {
            let ext = ext.as_slice();
            match ext {
                "class" => vec![self.path.with_extension("java")],  // Java
                "css"   => vec![self.path.with_extension("sass"),   self.path.with_extension("less")],  // SASS, Less
                "elc"   => vec![self.path.with_extension("el")],    // Emacs Lisp
                "hi"    => vec![self.path.with_extension("hs")],    // Haskell
                "js"    => vec![self.path.with_extension("coffee"), self.path.with_extension("ts")],  // CoffeeScript, TypeScript
                "o"     => vec![self.path.with_extension("c"),      self.path.with_extension("cpp")], // C, C++
                "pyc"   => vec![self.path.with_extension("py")],    // Python

                "aux" => vec![self.path.with_extension("tex")],  // TeX: auxiliary file
                "bbl" => vec![self.path.with_extension("tex")],  // BibTeX bibliography file
                "blg" => vec![self.path.with_extension("tex")],  // BibTeX log file
                "lof" => vec![self.path.with_extension("tex")],  // TeX list of figures
                "log" => vec![self.path.with_extension("tex")],  // TeX log file
                "lot" => vec![self.path.with_extension("tex")],  // TeX list of tables
                "toc" => vec![self.path.with_extension("tex")],  // TeX table of contents

                _ => vec![],  // No source files if none of the above
            }
        }
        else {
            vec![]  // No source files if there's no extension, either!
        }
    }

    pub fn display<U: Users>(&self, column: &Column, users_cache: &mut U) -> String {
        match *column {
            Permissions => {
                self.permissions_string()
            },

            FileName => {
                self.file_name()
            },

            FileSize(use_iec) => {
                self.file_size(use_iec)
            },

            // A file with multiple links is interesting, but
            // directories and suchlike can have multiple links all
            // the time.
            HardLinks => {
                let style = if self.has_multiple_links() { Red.on(Yellow) } else { Red.normal() };
                style.paint(self.stat.unstable.nlink.to_string().as_slice()).to_string()
            },

            Inode => {
                Purple.paint(self.stat.unstable.inode.to_string().as_slice()).to_string()
            },

            Blocks => {
                if self.stat.kind == io::FileType::RegularFile || self.stat.kind == io::FileType::Symlink {
                    Cyan.paint(self.stat.unstable.blocks.to_string().as_slice()).to_string()
                }
                else {
                    GREY.paint("-").to_string()
                }
            },

            // Display the ID if the user/group doesn't exist, which
            // usually means it was deleted but its files weren't.
            User => {
                let uid = self.stat.unstable.uid as i32;

                let user_name = match users_cache.get_user_by_uid(uid) {
                    Some(user) => user.name,
                    None => uid.to_string(),
                };

                let style = if users_cache.get_current_uid() == uid { Yellow.bold() } else { Plain };
                style.paint(user_name.as_slice()).to_string()
            },

            Group => {
                let gid = self.stat.unstable.gid as u32;
                let mut style = Plain;

                let group_name = match users_cache.get_group_by_gid(gid) {
                    Some(group) => {
                        let current_uid = users_cache.get_current_uid();
                        if let Some(current_user) = users_cache.get_user_by_uid(current_uid) {
                            if current_user.primary_group == group.gid || group.members.contains(&current_user.name) {
                                style = Yellow.bold();
                            }
                        }
                        group.name
                    },
                    None => gid.to_string(),
                };

                style.paint(group_name.as_slice()).to_string()
            },
        }
    }

    pub fn file_name(&self) -> String {
        let name = self.name.as_slice();
        let displayed_name = self.file_colour().paint(name);
        if self.stat.kind == io::FileType::Symlink {
            match fs::readlink(&self.path) {
                Ok(path) => {
                    let target_path = match self.dir {
                        Some(dir) => dir.path.join(path),
                        None => path,
                    };
                    format!("{} {}", displayed_name, self.target_file_name_and_arrow(&target_path))
                }
                Err(_) => displayed_name.to_string(),
            }
        }
        else {
            displayed_name.to_string()
        }
    }

    pub fn file_name_width(&self) -> uint {
        self.name.as_slice().width(false)
    }

    fn target_file_name_and_arrow(&self, target_path: &Path) -> String {
        let v = target_path.filename().unwrap();
        let filename = String::from_utf8_lossy(v);

        // Use stat instead of lstat - we *want* to follow links.
        let link_target = fs::stat(target_path).map(|stat| File {
            path:  target_path.clone(),
            dir:   self.dir,
            stat:  stat,
            name:  filename.to_string(),
            ext:   File::ext(filename.clone()),
        });

        // Statting a path usually fails because the file at the
        // other end doesn't exist. Show this by highlighting the
        // target file in red instead of displaying an error, because
        // the error would be shown out of context (before the
        // results, not right by the file) and it's almost always for
        // that reason anyway.

        match link_target {
            Ok(file) => format!("{} {}{}{}", GREY.paint("=>"), Cyan.paint(target_path.dirname_str().unwrap()), Cyan.paint("/"), file.file_colour().paint(filename.as_slice())),
            Err(_)   => format!("{} {}",     Red.paint("=>"),  Red.underline().paint(filename.as_slice())),
        }
    }

    fn file_size(&self, size_format: SizeFormat) -> String {
        // Don't report file sizes for directories. I've never looked
        // at one of those numbers and gained any information from it.
        if self.stat.kind == io::FileType::Directory {
            GREY.paint("-").to_string()
        }
        else {
        	let result = match size_format {
        		SizeFormat::DecimalBytes => decimal_prefix(self.stat.size as f64),
        		SizeFormat::BinaryBytes => binary_prefix(self.stat.size as f64),
        	};

			match result {
				Standalone(bytes) => Green.bold().paint(bytes.to_string().as_slice()).to_string(),
				Prefixed(prefix, n) => {
					let number = if n < 10f64 { format!("{:.1}", n) } else { format!("{:.0}", n) };
					format!("{}{}", Green.bold().paint(number.as_slice()), Green.paint(prefix.symbol()))
				}
			}
        }
    }

    fn type_char(&self) -> ANSIString {
        return match self.stat.kind {
            io::FileType::RegularFile  => Plain.paint("."),
            io::FileType::Directory    => Blue.paint("d"),
            io::FileType::NamedPipe    => Yellow.paint("|"),
            io::FileType::BlockSpecial => Purple.paint("s"),
            io::FileType::Symlink      => Cyan.paint("l"),
            io::FileType::Unknown      => Plain.paint("?"),
        }
    }

    pub fn file_colour(&self) -> Style {
        self.get_type().style()
    }

    fn has_multiple_links(&self) -> bool {
        self.stat.kind == io::FileType::RegularFile && self.stat.unstable.nlink > 1
    }

    fn permissions_string(&self) -> String {
        let bits = self.stat.perm;
        let executable_colour = match self.stat.kind {
        	io::FileType::RegularFile => Green.bold().underline(),
        	_ => Green.bold(),
        };

        return format!("{}{}{}{}{}{}{}{}{}{}",
            self.type_char(),

            // The first three are bold because they're the ones used
            // most often.
            File::permission_bit(bits, io::USER_READ,     "r", Yellow.bold()),
            File::permission_bit(bits, io::USER_WRITE,    "w", Red.bold()),
            File::permission_bit(bits, io::USER_EXECUTE,  "x", executable_colour),
            File::permission_bit(bits, io::GROUP_READ,    "r", Yellow.normal()),
            File::permission_bit(bits, io::GROUP_WRITE,   "w", Red.normal()),
            File::permission_bit(bits, io::GROUP_EXECUTE, "x", Green.normal()),
            File::permission_bit(bits, io::OTHER_READ,    "r", Yellow.normal()),
            File::permission_bit(bits, io::OTHER_WRITE,   "w", Red.normal()),
            File::permission_bit(bits, io::OTHER_EXECUTE, "x", Green.normal()),
       );
    }

    fn permission_bit(bits: io::FilePermission, bit: io::FilePermission, character: &'static str, style: Style) -> ANSIString {
        if bits.contains(bit) {
            style.paint(character)
        }
        else {
            GREY.paint("-")
        }
    }
}
<|endoftext|>


<commit_before>use futures::{Future, future};
use futures::sync::oneshot;
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, Mutex, RwLock, Weak};
use uuid::Uuid;

#[derive(Debug, PartialEq)]
pub enum Error {
    Invalid,
    Dropped,
    NotReady,
    Eof,
    Other,
}

pub const MAX_SIZE: usize = 1048576;
pub const MAX_CAPACITY: u64 = 16777216;

#[derive(Debug)]
pub enum Chunk {
    Data(u64, Vec<u8>),
    Eof(u64),
}

impl Chunk {
    fn data(data: Vec<u8>) -> Self {
        Chunk::Data(0, data)
    }

    fn eof() -> Self {
        Chunk::Eof(0)
    }

    fn count(&self) -> u64 {
        match *self {
            Chunk::Data(count, ..) => count,
            Chunk::Eof(count, ..) => count,
        }
    }

    fn len(&self) -> u64 {
        match *self {
            Chunk::Data(_, ref data) => data.len() as u64,
            Chunk::Eof(..) => 0,
        }
    }
}

type SharedChunk = Arc<Mutex<Chunk>>;

#[derive(Clone, Debug)]
enum State {
    Streaming,
    Stop,
}

struct Config {
    pub length: Option<u64>,
    pub capacity: u64,
    pub lifecount: u64,
}

struct Statistic {
    pub pushed: u64,
    pub dropped: u64,
    pub buffered: u64,
}

pub struct Flow {
    reference: Weak<RwLock<Flow>>,
    pub id: String,
    config: Config,
    statistic: Statistic,
    state: State,
    next_index: u64,
    tail_index: u64,
    bucket: HashMap<u64, SharedChunk>,
    waiting_push: VecDeque<(u64, oneshot::Sender<()>)>,
    waiting_pull: Arc<Mutex<HashMap<u64, Vec<oneshot::Sender<SharedChunk>>>>>,
}

type FlowFuture<T> = future::BoxFuture<T, Error>;

impl Flow {
    pub fn new(length: Option<u64>) -> Arc<RwLock<Self>> {
        let flow = Flow {
            reference: Weak::new(),
            id: Uuid::new_v4().simple().to_string(),
            config: Config {
                length,
                capacity: MAX_CAPACITY,
                lifecount: 1,
            },
            statistic: Statistic {
                pushed: 0,
                dropped: 0,
                buffered: 0,
            },
            state: State::Streaming,
            next_index: 0,
            tail_index: 0,
            bucket: HashMap::new(),
            waiting_push: VecDeque::new(),
            waiting_pull: Arc::new(Mutex::new(HashMap::new())),
        };
        let flow_ptr = Arc::new(RwLock::new(flow));
        flow_ptr.write().unwrap().reference = Arc::downgrade(&flow_ptr);
        flow_ptr
    }

    pub fn get_range(&self) -> (u64, u64) {
        (self.tail_index, self.next_index)
    }

    fn is_streaming(&self) -> bool {
        match self.state {
            State::Streaming => true,
            _ => false,
        }
    }

    fn acquire_chunk(&mut self, chunk: Chunk) -> FlowFuture<u64> {
        let chunk_len = chunk.len() as u64;

        // Check if the flow is EOF.
        if !self.is_streaming() {
            return future::err(Error::Invalid).boxed();
        }
        // Check if the flow is already overflow.
        if self.statistic.buffered > self.config.capacity {
            return future::err(Error::Invalid).boxed();
        }

        // Update statistic.
        let new_pushed = match self.statistic.pushed.checked_add(chunk_len) {
            Some(new_pushed) => {
                if let Some(length) = self.config.length {
                    // Check if the length is over.
                    if new_pushed > length {
                        return future::err(Error::Invalid).boxed();
                    }
                }
                new_pushed
            }
            None => return future::err(Error::Invalid).boxed(),
        };
        let new_buffered = match self.statistic.buffered.checked_add(chunk_len) {
            Some(new_buffered) => new_buffered,
            None => return future::err(Error::Invalid).boxed(),
        };
        self.statistic.pushed = new_pushed;
        self.statistic.buffered = new_buffered;

        // Update EOF state.
        if let Chunk::Eof(..) = chunk {
            self.state = State::Stop;
        }

        // Acquire the chunk index.
        let chunk_index = self.next_index;
        self.next_index += 1;

        let shared_chunk = Arc::new(Mutex::new(chunk));
        // Insert the chunk.
        self.bucket.insert(chunk_index, shared_chunk.clone());

        // Check if we need to block for overflow.
        let fut = if self.statistic.buffered > self.config.capacity {
            let (tx, rx) = oneshot::channel();
            self.waiting_push.push_back((chunk_len, tx));
            rx.map_err(|_| Error::Other).boxed()
        } else {
            future::ok(()).boxed()
        };

        // Wake up the waiting pulls.
        let mut waiting_pull = self.waiting_pull.lock().unwrap();
        if let Some(waits) = waiting_pull.remove(&chunk_index) {
            for wait in waits {
                // Don't unwrap, since the receiver can early quit.
                wait.send(shared_chunk.clone()).is_ok();
            }
        }

        fut.map(move |_| chunk_index).boxed()
    }

    fn sanitize_buffer(&mut self) {
        let mut tail_index = self.tail_index;
        let next_index = self.next_index;
        let lifecount = self.config.lifecount;

        while tail_index < next_index {
            {
                // Get should always success.
                let chunk = self.bucket.get(&tail_index).unwrap().lock().unwrap();
                if chunk.count() < lifecount {
                    break;
                }

                self.statistic.buffered -= chunk.len();
                self.statistic.dropped += chunk.len();

                let mut remain = chunk.len();
                while remain > 0 {
                    let new_len = if let Some(wait) = self.waiting_push.front_mut() {
                        let (new_remain, new_len) = if remain >= wait.0 {
                            (remain - wait.0, 0)
                        } else {
                            (0, wait.0 - remain)
                        };
                        remain = new_remain;
                        wait.0 = new_len;
                        wait.0
                    } else {
                        break;
                    };
                    if new_len == 0 {
                        // Remove should always success.
                        let wait = self.waiting_push.pop_front().unwrap();
                        // Don't unwrap, since the receiver can early quit.
                        wait.1.send(()).is_ok();
                    }
                }
            }

            // Remove should always success.
            self.bucket.remove(&tail_index).unwrap();

            tail_index += 1;
        }

        self.tail_index = tail_index;
    }

    pub fn push(&mut self, data: &[u8]) -> FlowFuture<u64> {
        if data.len() > MAX_SIZE {
            return future::err(Error::Invalid).boxed();
        }
        self.acquire_chunk(Chunk::data(data.to_vec()))
    }

    pub fn close(&mut self) -> FlowFuture<()> {
        if self.is_streaming() {
            self.acquire_chunk(Chunk::eof()).map(|_| ()).boxed()
        } else {
            future::err(Error::Invalid).boxed()
        }
    }

    pub fn pull(&self, chunk_index: u64, timeout: Option<u64>) -> FlowFuture<Vec<u8>> {
        // Clone the chunk if exists.
        let chunk = match self.bucket.get(&chunk_index) {
            Some(chunk) => Some(chunk.clone()),
            None => None,
        };
        // Try to get the chunk.
        let fut = if let Some(chunk) = chunk {
            future::ok(chunk).boxed()
        } else {
            if chunk_index < self.next_index {
                future::err(Error::Dropped).boxed()
            } else if let Some(0) = timeout {
                future::err(Error::NotReady).boxed()
            } else {
                let (tx, rx) = oneshot::channel();
                let mut waiting_pull = self.waiting_pull.lock().unwrap();
                let waits = waiting_pull.entry(chunk_index).or_insert(Vec::new());
                waits.push(tx);
                rx.map_err(|_| Error::Other).boxed()
            }
        };

        let flow_ref = self.reference.clone();
        let lifecount = self.config.lifecount;
        fut.and_then(move |chunk| {
                let flow_ptr = match flow_ref.upgrade() {
                    Some(ptr) => ptr.clone(),
                    None => return future::err(Error::Other),
                };

                let (count, result) = {
                    let mut chunk = chunk.lock().unwrap();
                    let (count, result) = match *chunk {
                        Chunk::Data(ref mut count, ref data) => (count, Ok(data.to_vec())),
                        Chunk::Eof(ref mut count) => (count, Err(Error::Eof)),
                    };
                    *count += 1;
                    (*count, result)
                };

                // Fast check if we need to sanitize.
                if count >= lifecount {
                    let mut flow = flow_ptr.write().unwrap();
                    flow.sanitize_buffer();
                }

                future::result(result)
            })
            .boxed()
    }
}

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

    macro_rules! sync_assert_eq {
        ($a:expr, $b:expr) => {
            let fut = $a;
            assert_eq!(fut.wait(), $b);
        }
    }

    #[test]
    fn basic_flow_operations() {
        let ptr = Flow::new(None);
        sync_assert_eq!(ptr.write().unwrap().push(&[1u8; 1234]), Ok(0));
        sync_assert_eq!(ptr.write().unwrap().push(&[2u8; MAX_SIZE]), Ok(1));
        sync_assert_eq!(ptr.write().unwrap().push(b"hello"), Ok(2));
        sync_assert_eq!(ptr.write().unwrap().push(&[2u8; MAX_SIZE + 1]), Err(Error::Invalid));
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Ok(Vec::from(&[1u8; 1234] as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(2, Some(0)), Ok(Vec::from(b"hello" as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(100, Some(0)), Err(Error::NotReady));
    }

    #[test]
    fn fixed_size_flow() {
        let ptr = Flow::new(Some(10));
        sync_assert_eq!(ptr.write().unwrap().push(b"hello"), Ok(0));
        sync_assert_eq!(ptr.write().unwrap().push(b"world"), Ok(1));
        sync_assert_eq!(ptr.write().unwrap().push(b"!"), Err(Error::Invalid));
    }

    #[test]
    fn close_flow() {
        let ptr = Flow::new(None);
        sync_assert_eq!(ptr.write().unwrap().push(b"hello"), Ok(0));
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Ok(Vec::from(b"hello" as &[u8])));
        sync_assert_eq!(ptr.write().unwrap().close(), Ok(()));
        sync_assert_eq!(ptr.write().unwrap().push(b"hello"), Err(Error::Invalid));
        sync_assert_eq!(ptr.read().unwrap().pull(1, Some(0)), Err(Error::Eof));
        sync_assert_eq!(ptr.write().unwrap().close(), Err(Error::Invalid));
    }

    #[test]
    fn dropped_chunk() {
        let ptr = Flow::new(None);
        sync_assert_eq!(ptr.write().unwrap().push(b"A"), Ok(0));
        sync_assert_eq!(ptr.write().unwrap().push(b"B"), Ok(1));
        sync_assert_eq!(ptr.write().unwrap().push(b"C"), Ok(2));
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Ok(Vec::from(b"A" as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Err(Error::Dropped));
        sync_assert_eq!(ptr.read().unwrap().pull(1, Some(0)), Ok(Vec::from(b"B" as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(1, Some(0)), Err(Error::Dropped));
        assert_eq!(ptr.read().unwrap().get_range(), (2, 3));
    }

    #[test]
    fn waiting_pull() {
        let ptr = Flow::new(None);
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Err(Error::NotReady));
        let fut = ptr.read().unwrap().pull(1, None);
        let fut1 = ptr.write().unwrap().push(b"ello");
        let fut2 = ptr.write().unwrap().push(b"hello");
        sync_assert_eq!(fut.join3(fut1, fut2), Ok((Vec::from(b"hello" as &[u8]), 0, 1)));
    }

    #[test]
    fn waiting_push() {
        let ptr = Flow::new(None);

        sync_assert_eq!(ptr.write().unwrap().push(b"A"), Ok(0));
        for idx in 1..(MAX_CAPACITY / MAX_SIZE as u64) {
            sync_assert_eq!(ptr.write().unwrap().push(&[0u8; MAX_SIZE]), Ok(idx));
        }
        let base_idx = MAX_CAPACITY / MAX_SIZE as u64;
        sync_assert_eq!(ptr.write().unwrap().push(b"D"), Ok(base_idx));

        let fut = ptr.write().unwrap().push(&[0u8; MAX_SIZE]);
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Ok(Vec::from(b"A" as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(1, Some(0)), Ok(vec![0u8; MAX_SIZE]));
        sync_assert_eq!(fut, Ok(base_idx + 1));

        ptr.write().unwrap().push(&[0u8; MAX_SIZE]);
        sync_assert_eq!(ptr.write().unwrap().push(b"C"), Err(Error::Invalid));
    }

    #[test]
    fn outlive() {
        let fut = {
            let ptr = Flow::new(None);
            sync_assert_eq!(ptr.write().unwrap().push(b"A"), Ok(0));
            let flow = ptr.read().unwrap();
            flow.pull(0, Some(0))
        };
        sync_assert_eq!(fut, Err(Error::Other));
    }
}
<commit_msg>Improve tests<commit_after>use futures::{Future, future};
use futures::sync::oneshot;
use std::collections::{HashMap, VecDeque};
use std::sync::{Arc, Mutex, RwLock, Weak};
use uuid::Uuid;

#[derive(Debug, PartialEq)]
pub enum Error {
    Invalid,
    Dropped,
    NotReady,
    Eof,
    Other,
}

pub const MAX_SIZE: usize = 1048576;
pub const MAX_CAPACITY: u64 = 16777216;

#[derive(Debug)]
pub enum Chunk {
    Data(u64, Vec<u8>),
    Eof(u64),
}

impl Chunk {
    fn data(data: Vec<u8>) -> Self {
        Chunk::Data(0, data)
    }

    fn eof() -> Self {
        Chunk::Eof(0)
    }

    fn count(&self) -> u64 {
        match *self {
            Chunk::Data(count, ..) => count,
            Chunk::Eof(count, ..) => count,
        }
    }

    fn len(&self) -> u64 {
        match *self {
            Chunk::Data(_, ref data) => data.len() as u64,
            Chunk::Eof(..) => 0,
        }
    }
}

type SharedChunk = Arc<Mutex<Chunk>>;

#[derive(Clone, Debug)]
enum State {
    Streaming,
    Stop,
}

struct Config {
    pub length: Option<u64>,
    pub capacity: u64,
    pub lifecount: u64,
}

struct Statistic {
    pub pushed: u64,
    pub dropped: u64,
    pub buffered: u64,
}

pub struct Flow {
    reference: Weak<RwLock<Flow>>,
    pub id: String,
    config: Config,
    statistic: Statistic,
    state: State,
    next_index: u64,
    tail_index: u64,
    bucket: HashMap<u64, SharedChunk>,
    waiting_push: VecDeque<(u64, oneshot::Sender<()>)>,
    waiting_pull: Arc<Mutex<HashMap<u64, Vec<oneshot::Sender<SharedChunk>>>>>,
}

type FlowFuture<T> = future::BoxFuture<T, Error>;

impl Flow {
    pub fn new(length: Option<u64>) -> Arc<RwLock<Self>> {
        let flow = Flow {
            reference: Weak::new(),
            id: Uuid::new_v4().simple().to_string(),
            config: Config {
                length,
                capacity: MAX_CAPACITY,
                lifecount: 1,
            },
            statistic: Statistic {
                pushed: 0,
                dropped: 0,
                buffered: 0,
            },
            state: State::Streaming,
            next_index: 0,
            tail_index: 0,
            bucket: HashMap::new(),
            waiting_push: VecDeque::new(),
            waiting_pull: Arc::new(Mutex::new(HashMap::new())),
        };
        let flow_ptr = Arc::new(RwLock::new(flow));
        flow_ptr.write().unwrap().reference = Arc::downgrade(&flow_ptr);
        flow_ptr
    }

    pub fn get_range(&self) -> (u64, u64) {
        (self.tail_index, self.next_index)
    }

    fn is_streaming(&self) -> bool {
        match self.state {
            State::Streaming => true,
            _ => false,
        }
    }

    fn acquire_chunk(&mut self, chunk: Chunk) -> FlowFuture<u64> {
        let chunk_len = chunk.len() as u64;

        // Check if the flow is EOF.
        if !self.is_streaming() {
            return future::err(Error::Invalid).boxed();
        }
        // Check if the flow is already overflow.
        if self.statistic.buffered > self.config.capacity {
            return future::err(Error::Invalid).boxed();
        }

        // Update statistic.
        let new_pushed = match self.statistic.pushed.checked_add(chunk_len) {
            Some(new_pushed) => {
                if let Some(length) = self.config.length {
                    // Check if the length is over.
                    if new_pushed > length {
                        return future::err(Error::Invalid).boxed();
                    }
                }
                new_pushed
            }
            None => return future::err(Error::Invalid).boxed(),
        };
        let new_buffered = match self.statistic.buffered.checked_add(chunk_len) {
            Some(new_buffered) => new_buffered,
            None => return future::err(Error::Invalid).boxed(),
        };
        self.statistic.pushed = new_pushed;
        self.statistic.buffered = new_buffered;

        // Update EOF state.
        if let Chunk::Eof(..) = chunk {
            self.state = State::Stop;
        }

        // Acquire the chunk index.
        let chunk_index = self.next_index;
        self.next_index += 1;

        let shared_chunk = Arc::new(Mutex::new(chunk));
        // Insert the chunk.
        self.bucket.insert(chunk_index, shared_chunk.clone());

        // Check if we need to block for overflow.
        let fut = if self.statistic.buffered > self.config.capacity {
            let (tx, rx) = oneshot::channel();
            self.waiting_push.push_back((chunk_len, tx));
            rx.map_err(|_| Error::Other).boxed()
        } else {
            future::ok(()).boxed()
        };

        // Wake up the waiting pulls.
        let mut waiting_pull = self.waiting_pull.lock().unwrap();
        if let Some(waits) = waiting_pull.remove(&chunk_index) {
            for wait in waits {
                // Don't unwrap, since the receiver can early quit.
                wait.send(shared_chunk.clone()).is_ok();
            }
        }

        fut.map(move |_| chunk_index).boxed()
    }

    fn sanitize_buffer(&mut self) {
        let mut tail_index = self.tail_index;
        let next_index = self.next_index;
        let lifecount = self.config.lifecount;

        while tail_index < next_index {
            {
                // Get should always success.
                let chunk = self.bucket.get(&tail_index).unwrap().lock().unwrap();
                if chunk.count() < lifecount {
                    break;
                }

                self.statistic.buffered -= chunk.len();
                self.statistic.dropped += chunk.len();

                let mut remain = chunk.len();
                while remain > 0 {
                    let new_len = if let Some(wait) = self.waiting_push.front_mut() {
                        let (new_remain, new_len) = if remain >= wait.0 {
                            (remain - wait.0, 0)
                        } else {
                            (0, wait.0 - remain)
                        };
                        remain = new_remain;
                        wait.0 = new_len;
                        wait.0
                    } else {
                        break;
                    };
                    if new_len == 0 {
                        // Remove should always success.
                        let wait = self.waiting_push.pop_front().unwrap();
                        // Don't unwrap, since the receiver can early quit.
                        wait.1.send(()).is_ok();
                    }
                }
            }

            // Remove should always success.
            self.bucket.remove(&tail_index).unwrap();

            tail_index += 1;
        }

        self.tail_index = tail_index;
    }

    pub fn push(&mut self, data: &[u8]) -> FlowFuture<u64> {
        if data.len() > MAX_SIZE {
            return future::err(Error::Invalid).boxed();
        }
        self.acquire_chunk(Chunk::data(data.to_vec()))
    }

    pub fn close(&mut self) -> FlowFuture<()> {
        if self.is_streaming() {
            self.acquire_chunk(Chunk::eof()).map(|_| ()).boxed()
        } else {
            future::err(Error::Invalid).boxed()
        }
    }

    pub fn pull(&self, chunk_index: u64, timeout: Option<u64>) -> FlowFuture<Vec<u8>> {
        // Clone the chunk if exists.
        let chunk = match self.bucket.get(&chunk_index) {
            Some(chunk) => Some(chunk.clone()),
            None => None,
        };
        // Try to get the chunk.
        let fut = if let Some(chunk) = chunk {
            future::ok(chunk).boxed()
        } else {
            if chunk_index < self.next_index {
                future::err(Error::Dropped).boxed()
            } else if let Some(0) = timeout {
                future::err(Error::NotReady).boxed()
            } else {
                let (tx, rx) = oneshot::channel();
                let mut waiting_pull = self.waiting_pull.lock().unwrap();
                let waits = waiting_pull.entry(chunk_index).or_insert(Vec::new());
                waits.push(tx);
                rx.map_err(|_| Error::Other).boxed()
            }
        };

        let flow_ref = self.reference.clone();
        let lifecount = self.config.lifecount;
        fut.and_then(move |chunk| {
                let flow_ptr = match flow_ref.upgrade() {
                    Some(ptr) => ptr.clone(),
                    None => return future::err(Error::Other),
                };

                let (count, result) = {
                    let mut chunk = chunk.lock().unwrap();
                    let (count, result) = match *chunk {
                        Chunk::Data(ref mut count, ref data) => (count, Ok(data.to_vec())),
                        Chunk::Eof(ref mut count) => (count, Err(Error::Eof)),
                    };
                    *count += 1;
                    (*count, result)
                };

                // Fast check if we need to sanitize.
                if count >= lifecount {
                    let mut flow = flow_ptr.write().unwrap();
                    flow.sanitize_buffer();
                }

                future::result(result)
            })
            .boxed()
    }
}

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

    macro_rules! sync_assert_eq {
        ($a:expr, $b:expr) => {
            let fut = $a;
            assert_eq!(fut.wait(), $b);
        }
    }

    #[test]
    fn basic_flow_operations() {
        let ptr = Flow::new(None);
        sync_assert_eq!(ptr.write().unwrap().push(&[1u8; 1234]), Ok(0));
        sync_assert_eq!(ptr.write().unwrap().push(&[2u8; MAX_SIZE]), Ok(1));
        sync_assert_eq!(ptr.write().unwrap().push(b"hello"), Ok(2));
        sync_assert_eq!(ptr.write().unwrap().push(&[2u8; MAX_SIZE + 1]), Err(Error::Invalid));
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Ok(Vec::from(&[1u8; 1234] as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(2, Some(0)), Ok(Vec::from(b"hello" as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(100, Some(0)), Err(Error::NotReady));
    }

    #[test]
    fn fixed_size_flow() {
        let ptr = Flow::new(Some(10));
        sync_assert_eq!(ptr.write().unwrap().push(b"hello"), Ok(0));
        sync_assert_eq!(ptr.write().unwrap().push(b"world"), Ok(1));
        sync_assert_eq!(ptr.write().unwrap().push(b"!"), Err(Error::Invalid));
    }

    #[test]
    fn close_flow() {
        let ptr = Flow::new(None);
        sync_assert_eq!(ptr.write().unwrap().push(b"hello"), Ok(0));
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Ok(Vec::from(b"hello" as &[u8])));
        sync_assert_eq!(ptr.write().unwrap().close(), Ok(()));
        sync_assert_eq!(ptr.write().unwrap().push(b"hello"), Err(Error::Invalid));
        sync_assert_eq!(ptr.read().unwrap().pull(1, Some(0)), Err(Error::Eof));
        sync_assert_eq!(ptr.write().unwrap().close(), Err(Error::Invalid));
    }

    #[test]
    fn dropped_chunk() {
        let ptr = Flow::new(None);
        sync_assert_eq!(ptr.write().unwrap().push(b"A"), Ok(0));
        sync_assert_eq!(ptr.write().unwrap().push(b"B"), Ok(1));
        sync_assert_eq!(ptr.write().unwrap().push(b"C"), Ok(2));
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Ok(Vec::from(b"A" as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Err(Error::Dropped));
        sync_assert_eq!(ptr.read().unwrap().pull(1, Some(0)), Ok(Vec::from(b"B" as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(1, Some(0)), Err(Error::Dropped));
        assert_eq!(ptr.read().unwrap().get_range(), (2, 3));
    }

    #[test]
    fn waiting_pull() {
        let ptr = Flow::new(None);
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Err(Error::NotReady));
        let fut = ptr.read().unwrap().pull(1, None);
        let fut1 = ptr.write().unwrap().push(b"ello");
        let fut2 = ptr.write().unwrap().push(b"hello");
        sync_assert_eq!(fut.join3(fut1, fut2), Ok((Vec::from(b"hello" as &[u8]), 0, 1)));
    }

    #[test]
    fn waiting_push() {
        let ptr = Flow::new(None);

        sync_assert_eq!(ptr.write().unwrap().push(b"A"), Ok(0));
        for idx in 1..(MAX_CAPACITY / MAX_SIZE as u64) {
            sync_assert_eq!(ptr.write().unwrap().push(&[0u8; MAX_SIZE]), Ok(idx));
        }
        let base_idx = MAX_CAPACITY / MAX_SIZE as u64;
        sync_assert_eq!(ptr.write().unwrap().push(b"D"), Ok(base_idx));

        let fut = ptr.write().unwrap().push(&[0u8; MAX_SIZE]);
        sync_assert_eq!(ptr.write().unwrap().push(b"C"), Err(Error::Invalid));
        sync_assert_eq!(ptr.read().unwrap().pull(0, Some(0)), Ok(Vec::from(b"A" as &[u8])));
        sync_assert_eq!(ptr.read().unwrap().pull(1, Some(0)), Ok(vec![0u8; MAX_SIZE]));
        sync_assert_eq!(fut, Ok(base_idx + 1));
    }

    #[test]
    fn outlive() {
        let fut = {
            let ptr = Flow::new(None);
            sync_assert_eq!(ptr.write().unwrap().push(b"A"), Ok(0));
            let flow = ptr.read().unwrap();
            flow.pull(0, Some(0))
        };
        sync_assert_eq!(fut, Err(Error::Other));
    }
}
<|endoftext|>


<commit_before>use regex::{Regex};
use std::str::{FromStr};

use super::super::metrics::*;

/// Reader that takes a log line string and returns any metrics found in it.
pub trait LogLineReader: Send + Sync {
    fn read(&self, &str) -> Vec<Metric>;
}

/// Reads metrics from log lines in the standard formats:
///
/// - Measures: `measure#metric=1.2`
/// - Counts: `count#metric=3`
pub struct StandardLogLineReader;

lazy_static! {
    static ref LOG_MEASURE_REGEX: Regex =
        Regex::new(r"measure#([a-zA-Z0-9._]+)=(\d+(?:\.\d+)?)").unwrap();

    static ref LOG_SAMPLE_REGEX: Regex =
        Regex::new(r"sample#([a-zA-Z0-9._]+)=(\d+(?:\.\d+)?)").unwrap();

    static ref LOG_COUNT_REGEX: Regex =
        Regex::new(r"count#([a-zA-Z0-9._]+)=(\d+)").unwrap();
}

impl LogLineReader for StandardLogLineReader {
    fn read(&self, line: &str) -> Vec<Metric> {
        let mut metrics = vec![];

        // Look for counts
        for cap in LOG_COUNT_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = u64::from_str(cap.at(2).unwrap()) {
                metrics.push(Count(Dimension::with_name(name), value))
            }
        }

        // Look for measures
        for cap in LOG_MEASURE_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = f64::from_str(cap.at(2).unwrap()) {
                metrics.push(Measure(Dimension::with_name(name), value))
            }
        }

        // Look for samples
        for cap in LOG_SAMPLE_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = f64::from_str(cap.at(2).unwrap()) {
                metrics.push(Sample(Dimension::with_name(name), value))
            }
        }

        metrics
    } // fn read
}

/// Reads Heroku's logging metrics.
pub struct HerokuLogLineReader;

lazy_static! {
    static ref DYNO_TYPE_REGEX: Regex =
        Regex::new(r"dyno=(\w+)").unwrap();

    static ref SERVICE_REGEX: Regex =
        Regex::new(r"service=(\d+)ms").unwrap();

    static ref STATUS_REGEX: Regex =
        Regex::new(r"status=(\d+)").unwrap();

    static ref HEROKU_HTTP_ERROR_CODE_REGEX: Regex =
        Regex::new(r"code=(H\d+)").unwrap();

    static ref HEROKU_RUNTIME_ERROR_CODE_REGEX: Regex =
        Regex::new(r"Error (R\d+)").unwrap();

    static ref LOAD_AVG_1M_REGEX: Regex =
        Regex::new(r"sample#load_avg_1m=([0-9.]+)").unwrap();

    static ref SOURCE_REGEX: Regex =
        Regex::new(r"source=([\w.]+)").unwrap();
}

impl HerokuLogLineReader {
    /// Parses Heroku router status lines for the response service time (how
    /// long it took) and the HTTP response status.
    pub fn parse_status(line: &str) -> Option<Vec<Metric>> {
        let mut metrics: Vec<Metric> = vec![];

        let status = match STATUS_REGEX.captures(line)
            .and_then(|c| c.at(1))
            .and_then(|s| u16::from_str(s).ok()) {
                Some(s) => s,
                None => return None,
            };

        let service = match SERVICE_REGEX.captures(line)
            .and_then(|c| c.at(1))
            .and_then(|s| u32::from_str(s).ok()) {
                Some(s) => s,
                None => return None,
            };

        let dyno_type = match DYNO_TYPE_REGEX.captures(line)
            .and_then(|c| c.at(1)) {
                Some(d) => d,
                None => return None,
            };

        let base = format!("dyno.{}", dyno_type);

        // Counting a 499 as a 500
        let is_500 = status >= 499 && status < 600;

        // Don't record timing for 499 and 5xx errors
        if !is_500 {
            let service_time_name = format!("{}.service_time", base);
            metrics.push(Measure(Dimension::with_name(service_time_name), service as f64));
        }

        // Count the status
        let status_name = format!("{}.status.{}", base, status);
        metrics.push(Count(Dimension::with_name(status_name), 1));

        Some(metrics)
    }

    /// Parses Heroku warning and error codes like "Hxx" and "Rxx" where "xx" is a pair
    /// of numbers. See the [Heroku][] site for more details.
    ///
    /// [Heroku]: https://devcenter.heroku.com/articles/error-codes
    pub fn parse_heroku_code(line: &str) -> Option<Metric> {
        let code: &str;

        if let Some(http_code) = HEROKU_HTTP_ERROR_CODE_REGEX.captures(line).and_then(|c| c.at(1)) {
            code = http_code;

        } else if let Some(runtime_code) = HEROKU_RUNTIME_ERROR_CODE_REGEX.captures(line).and_then(|c| c.at(1)) {
            code = runtime_code;

        } else {
            return None
        }

        let name = format!("heroku.error.{}", code);
        Some(Count(Dimension::with_name(name), 1))
    }

    /// Parses the `sample#load_avg_1m=` metrics from Heroku logs.
    pub fn parse_load(line: &str) -> Option<Metric> {
        let load_avg_1m = match LOAD_AVG_1M_REGEX.captures(line)
                                                 .and_then(|c| c.at(1))
                                                 .and_then(|c| f64::from_str(c).ok()) {
            Some(l) => l,
            None => return None,
        };

        let source = match SOURCE_REGEX.captures(line)
                                       .and_then(|c| c.at(1)) {
            Some(t) => t,
            None => return None,
        };

        let dyno_type = match source.split('.').nth(0) {
            Some(s) => s,
            None => return None,
        };

        let name = format!("dyno.{}.load_avg_1m", dyno_type);
        Some(Measure(Dimension::with_name(name), load_avg_1m))
    }
}

impl LogLineReader for HerokuLogLineReader {
    fn read(&self, line: &str) -> Vec<Metric> {
        if !line.contains("heroku") { return vec![] }

        let mut metrics: Vec<Metric> = vec![];

        if let Some(statuses) = HerokuLogLineReader::parse_status(line)      { metrics.extend(statuses) }
        if let Some(code)     = HerokuLogLineReader::parse_heroku_code(line) { metrics.push(code) }
        if let Some(load)     = HerokuLogLineReader::parse_load(line)        { metrics.push(load) }

        metrics
    }
}

#[cfg(test)]
mod tests {
    use super::{
        LogLineReader,
        StandardLogLineReader,
        HerokuLogLineReader
    };
    use super::super::super::metrics::*;

    #[test]
    fn standard_reader_reads_measure() {
        let reader = StandardLogLineReader;
        let line = "measure#foo=1.2\n";

        assert_eq!(
            reader.read(line),
            vec![ Measure(Dimension::with_name("foo"), 1.2) ]
        )
    }

    #[test]
    fn standard_reader_reads_count() {
        let reader = StandardLogLineReader;
        let line = "count#foo=1\n";

        assert_eq!(
            reader.read(line),
            vec![ Count(Dimension::with_name("foo"), 1) ]
        )
    }

    #[test]
    fn standard_reader_reads_sample() {
        let reader = StandardLogLineReader;
        let line = "sample#bar=3.4\n";

        assert_eq!(
            reader.read(line),
            vec![ Sample(Dimension::with_name("bar"), 3.4) ]
        )
    }

    #[test]
    fn standard_reader_returns_nothing_on_failed_read() {
        let reader = StandardLogLineReader;
        let line = "metric#bar=3.4\n";

        assert_eq!(
            reader.read(line),
            vec![]
        )
    }

    #[test]
    fn heroku_reader_reads_loads() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:34:59.429615+00:00 heroku web.2 - - source=web.2 dyno=heroku.123.XYZ sample#load_avg_1m=0.56 sample#load_avg_5m=0.26 sample#load_avg_15m=0.17\n";

        assert_eq!(
            reader.read(line),
            vec![ Measure(Dimension::with_name("dyno.web.load_avg_1m"), 0.56) ]
        )
    }

    #[test]
    fn heroku_reader_reads_http_errors() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:50:36.352129+00:00 heroku router - - sock=backend at=error code=H18 desc=\"Server Request Interrupted\" method=GET path=\"/\" host=www.example.com request_id=XYZ fwd=\"1.2.3.4\" dyno=web.5 connect=0ms service=495ms status=503 bytes=1648\n";

        assert_eq!(
            reader.read(line),
            vec![
                Count(Dimension::with_name("dyno.web.status.503"), 1),
                Count(Dimension::with_name("heroku.error.H18"), 1),
            ]
        )
    }

    #[test]
    fn heroku_reader_reads_runtime_errors() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-25 21:35:34.990292+00:00 heroku scheduler.5451 - - Error R14 (Memory quota exceeded)\n";

        assert_eq!(
            reader.read(line),
            vec![
                Count(Dimension::with_name("heroku.error.R14"), 1),
            ]
        )
    }

    #[test]
    fn heroku_reader_reads_service_times() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:34:59.370813+00:00 heroku router - - at=info method=PUT path=\"/\" host=www.example.com request_id=XYZ fwd=\"1.2.3.4\" dyno=web.1 connect=0ms service=39ms status=200 bytes=1627\n";

        assert_eq!(
            reader.read(line),
            vec![
                Measure(Dimension::with_name("dyno.web.service_time"), 39.0),
                Count(Dimension::with_name("dyno.web.status.200"), 1),
            ]
        )
    }
}
<commit_msg>Start parsing sources in log line reader<commit_after>use regex::{Regex};
use std::str::{FromStr};

use super::super::metrics::*;

/// Reader that takes a log line string and returns any metrics found in it.
pub trait LogLineReader: Send + Sync {
    fn read(&self, &str) -> Vec<Metric>;
}

/// Reads metrics from log lines in the standard formats:
///
/// - Measures: `measure#metric=1.2`
/// - Counts: `count#metric=3`
pub struct StandardLogLineReader;

lazy_static! {
    static ref LOG_MEASURE_REGEX: Regex =
        Regex::new(r"measure#([a-zA-Z0-9._]+)=(\d+(?:\.\d+)?)").unwrap();

    static ref LOG_SAMPLE_REGEX: Regex =
        Regex::new(r"sample#([a-zA-Z0-9._]+)=(\d+(?:\.\d+)?)").unwrap();

    static ref LOG_COUNT_REGEX: Regex =
        Regex::new(r"count#([a-zA-Z0-9._]+)=(\d+)").unwrap();

    static ref SOURCE_REGEX: Regex =
        Regex::new(r"source=([\w.]+)").unwrap();
}

impl StandardLogLineReader {
    fn parse_source<'a>(line: &'a str) -> Option<&'a str> {
        SOURCE_REGEX.captures(line)
                    .and_then(|c| c.at(1))
    }
}

impl LogLineReader for StandardLogLineReader {
    fn read(&self, line: &str) -> Vec<Metric> {
        let mut metrics = vec![];

        // Look for counts
        for cap in LOG_COUNT_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = u64::from_str(cap.at(2).unwrap()) {
                metrics.push(Count(Dimension::with_name(name), value))
            }
        }

        // Look for measures
        for cap in LOG_MEASURE_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = f64::from_str(cap.at(2).unwrap()) {
                metrics.push(Measure(Dimension::with_name(name), value))
            }
        }

        // Look for samples
        for cap in LOG_SAMPLE_REGEX.captures_iter(line) {
            let name = cap.at(1).unwrap();

            if let Ok(value) = f64::from_str(cap.at(2).unwrap()) {
                metrics.push(Sample(Dimension::with_name(name), value))
            }
        }

        metrics
    } // fn read
}

/// Reads Heroku's logging metrics.
pub struct HerokuLogLineReader;

lazy_static! {
    static ref DYNO_TYPE_REGEX: Regex =
        Regex::new(r"dyno=(\w+)").unwrap();

    static ref SERVICE_REGEX: Regex =
        Regex::new(r"service=(\d+)ms").unwrap();

    static ref STATUS_REGEX: Regex =
        Regex::new(r"status=(\d+)").unwrap();

    static ref HEROKU_HTTP_ERROR_CODE_REGEX: Regex =
        Regex::new(r"code=(H\d+)").unwrap();

    static ref HEROKU_RUNTIME_ERROR_CODE_REGEX: Regex =
        Regex::new(r"Error (R\d+)").unwrap();

    static ref LOAD_AVG_1M_REGEX: Regex =
        Regex::new(r"sample#load_avg_1m=([0-9.]+)").unwrap();
}

impl HerokuLogLineReader {
    /// Parses Heroku router status lines for the response service time (how
    /// long it took) and the HTTP response status.
    pub fn parse_status(line: &str) -> Option<Vec<Metric>> {
        let mut metrics: Vec<Metric> = vec![];

        let status = match STATUS_REGEX.captures(line)
            .and_then(|c| c.at(1))
            .and_then(|s| u16::from_str(s).ok()) {
                Some(s) => s,
                None => return None,
            };

        let service = match SERVICE_REGEX.captures(line)
            .and_then(|c| c.at(1))
            .and_then(|s| u32::from_str(s).ok()) {
                Some(s) => s,
                None => return None,
            };

        let dyno_type = match DYNO_TYPE_REGEX.captures(line)
            .and_then(|c| c.at(1)) {
                Some(d) => d,
                None => return None,
            };

        let base = format!("dyno.{}", dyno_type);

        // Counting a 499 as a 500
        let is_500 = status >= 499 && status < 600;

        // Don't record timing for 499 and 5xx errors
        if !is_500 {
            let service_time_name = format!("{}.service_time", base);
            metrics.push(Measure(Dimension::with_name(service_time_name), service as f64));
        }

        // Count the status
        let status_name = format!("{}.status.{}", base, status);
        metrics.push(Count(Dimension::with_name(status_name), 1));

        Some(metrics)
    }

    /// Parses Heroku warning and error codes like "Hxx" and "Rxx" where "xx" is a pair
    /// of numbers. See the [Heroku][] site for more details.
    ///
    /// [Heroku]: https://devcenter.heroku.com/articles/error-codes
    pub fn parse_heroku_code(line: &str) -> Option<Metric> {
        let code: &str;

        if let Some(http_code) = HEROKU_HTTP_ERROR_CODE_REGEX.captures(line).and_then(|c| c.at(1)) {
            code = http_code;

        } else if let Some(runtime_code) = HEROKU_RUNTIME_ERROR_CODE_REGEX.captures(line).and_then(|c| c.at(1)) {
            code = runtime_code;

        } else {
            return None
        }

        let name = format!("heroku.error.{}", code);
        Some(Count(Dimension::with_name(name), 1))
    }

    /// Parses the `sample#load_avg_1m=` metrics from Heroku logs.
    pub fn parse_load(line: &str) -> Option<Metric> {
        let load_avg_1m = match LOAD_AVG_1M_REGEX.captures(line)
                                                 .and_then(|c| c.at(1))
                                                 .and_then(|c| f64::from_str(c).ok()) {
            Some(l) => l,
            None => return None,
        };

        let source = match StandardLogLineReader::parse_source(line) {
            Some(s) => s,
            None => return None,
        };

        let dyno_type = match source.split('.').nth(0) {
            Some(s) => s,
            None => return None,
        };

        let dim = Dimension::with_name_and_source(format!("dyno.{}.load_avg_1m", dyno_type), source);
        Some(Measure(dim, load_avg_1m))
    }
}

impl LogLineReader for HerokuLogLineReader {
    fn read(&self, line: &str) -> Vec<Metric> {
        if !line.contains("heroku") { return vec![] }

        let mut metrics: Vec<Metric> = vec![];

        if let Some(statuses) = HerokuLogLineReader::parse_status(line)      { metrics.extend(statuses) }
        if let Some(code)     = HerokuLogLineReader::parse_heroku_code(line) { metrics.push(code) }
        if let Some(load)     = HerokuLogLineReader::parse_load(line)        { metrics.push(load) }

        metrics
    }
}

#[cfg(test)]
mod tests {
    use super::{
        LogLineReader,
        StandardLogLineReader,
        HerokuLogLineReader
    };
    use super::super::super::metrics::*;

    #[test]
    fn standard_reader_reads_measure() {
        let reader = StandardLogLineReader;
        let line = "measure#foo=1.2\n";

        assert_eq!(
            reader.read(line),
            vec![ Measure(Dimension::with_name("foo"), 1.2) ]
        )
    }

    #[test]
    fn standard_reader_reads_count() {
        let reader = StandardLogLineReader;
        let line = "count#foo=1\n";

        assert_eq!(
            reader.read(line),
            vec![ Count(Dimension::with_name("foo"), 1) ]
        )
    }

    #[test]
    fn standard_reader_reads_sample() {
        let reader = StandardLogLineReader;
        let line = "sample#bar=3.4\n";

        assert_eq!(
            reader.read(line),
            vec![ Sample(Dimension::with_name("bar"), 3.4) ]
        )
    }

    #[test]
    fn standard_reader_returns_nothing_on_failed_read() {
        let reader = StandardLogLineReader;
        let line = "metric#bar=3.4\n";

        assert_eq!(
            reader.read(line),
            vec![]
        )
    }

    #[test]
    fn standard_reader_parses_source() {
        let line = "source=something metric#other=5.6\n";

        assert_eq!(
            StandardLogLineReader::parse_source(line),
            Some("something")
        )
    }

    #[test]
    fn heroku_reader_reads_loads() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:34:59.429615+00:00 heroku web.2 - - source=web.2 dyno=heroku.123.XYZ sample#load_avg_1m=0.56 sample#load_avg_5m=0.26 sample#load_avg_15m=0.17\n";

        assert_eq!(
            reader.read(line),
            vec![ Measure(Dimension::with_name_and_source("dyno.web.load_avg_1m", "web.2"), 0.56) ]
        )
    }

    #[test]
    fn heroku_reader_reads_http_errors() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:50:36.352129+00:00 heroku router - - sock=backend at=error code=H18 desc=\"Server Request Interrupted\" method=GET path=\"/\" host=www.example.com request_id=XYZ fwd=\"1.2.3.4\" dyno=web.5 connect=0ms service=495ms status=503 bytes=1648\n";

        assert_eq!(
            reader.read(line),
            vec![
                Count(Dimension::with_name("dyno.web.status.503"), 1),
                Count(Dimension::with_name("heroku.error.H18"), 1),
            ]
        )
    }

    #[test]
    fn heroku_reader_reads_runtime_errors() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-25 21:35:34.990292+00:00 heroku scheduler.5451 - - Error R14 (Memory quota exceeded)\n";

        assert_eq!(
            reader.read(line),
            vec![
                Count(Dimension::with_name("heroku.error.R14"), 1),
            ]
        )
    }

    #[test]
    fn heroku_reader_reads_service_times() {
        let reader = HerokuLogLineReader;
        let line = "2016-02-26 21:34:59.370813+00:00 heroku router - - at=info method=PUT path=\"/\" host=www.example.com request_id=XYZ fwd=\"1.2.3.4\" dyno=web.1 connect=0ms service=39ms status=200 bytes=1627\n";

        assert_eq!(
            reader.read(line),
            vec![
                Measure(Dimension::with_name("dyno.web.service_time"), 39.0),
                Count(Dimension::with_name("dyno.web.status.200"), 1),
            ]
        )
    }
}
<|endoftext|>


<commit_before>use core::pin::Pin;
use futures_core::stream::{Stream, FusedStream};
use futures_core::task::{Context, Poll};
use futures_sink::{Sink};
use pin_utils::{unsafe_pinned, unsafe_unpinned};

/// Sink for the [`sink_map_err`](super::SinkExt::sink_map_err) method.
#[derive(Debug)]
#[must_use = "sinks do nothing unless polled"]
pub struct SinkMapErr<Si, F> {
    sink: Si,
    f: Option<F>,
}

impl<Si: Unpin, F> Unpin for SinkMapErr<Si, F> {}

impl<Si, F> SinkMapErr<Si, F> {
    unsafe_pinned!(sink: Si);
    unsafe_unpinned!(f: Option<F>);

    pub(super) fn new(sink: Si, f: F) -> SinkMapErr<Si, F> {
        SinkMapErr { sink, f: Some(f) }
    }

    /// Get a shared reference to the inner sink.
    pub fn get_ref(&self) -> &Si {
        &self.sink
    }

    /// Get a mutable reference to the inner sink.
    pub fn get_mut(&mut self) -> &mut Si {
        &mut self.sink
    }

    /// Get a pinned mutable reference to the inner sink.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Si> {
        self.sink()
    }

    /// Consumes this combinator, returning the underlying sink.
    ///
    /// Note that this may discard intermediate state of this combinator, so
    /// care should be taken to avoid losing resources when this is called.
    pub fn into_inner(self) -> Si {
        self.sink
    }

    fn take_f(self: Pin<&mut Self>) -> F {
        self.f().take().expect("polled MapErr after completion")
    }
}

impl<Si, F, E, Item> Sink<Item> for SinkMapErr<Si, F>
    where Si: Sink<Item>,
          F: FnOnce(Si::Error) -> E,
{
    type Error = E;

    fn poll_ready(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        self.as_mut().sink().poll_ready(cx).map_err(|e| self.as_mut().take_f()(e))
    }

    fn start_send(
        mut self: Pin<&mut Self>,
        item: Item,
    ) -> Result<(), Self::Error> {
        self.as_mut().sink().start_send(item).map_err(|e| self.as_mut().take_f()(e))
    }

    fn poll_flush(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        self.as_mut().sink().poll_flush(cx).map_err(|e| self.as_mut().take_f()(e))
    }

    fn poll_close(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        self.as_mut().sink().poll_close(cx).map_err(|e| self.as_mut().take_f()(e))
    }
}

// Forwarding impl of Stream from the underlying sink
impl<S: Stream, F> Stream for SinkMapErr<S, F> {
    type Item = S::Item;

    fn poll_next(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<S::Item>> {
        self.sink().poll_next(cx)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.sink.size_hint()
    }
}

impl<S: FusedStream, F> FusedStream for SinkMapErr<S, F> {
    fn is_terminated(&self) -> bool {
        self.sink.is_terminated()
    }
}
<commit_msg>Derive Clone on SinkMapErr<commit_after>use core::pin::Pin;
use futures_core::stream::{Stream, FusedStream};
use futures_core::task::{Context, Poll};
use futures_sink::{Sink};
use pin_utils::{unsafe_pinned, unsafe_unpinned};

/// Sink for the [`sink_map_err`](super::SinkExt::sink_map_err) method.
#[derive(Debug, Clone)]
#[must_use = "sinks do nothing unless polled"]
pub struct SinkMapErr<Si, F> {
    sink: Si,
    f: Option<F>,
}

impl<Si: Unpin, F> Unpin for SinkMapErr<Si, F> {}

impl<Si, F> SinkMapErr<Si, F> {
    unsafe_pinned!(sink: Si);
    unsafe_unpinned!(f: Option<F>);

    pub(super) fn new(sink: Si, f: F) -> SinkMapErr<Si, F> {
        SinkMapErr { sink, f: Some(f) }
    }

    /// Get a shared reference to the inner sink.
    pub fn get_ref(&self) -> &Si {
        &self.sink
    }

    /// Get a mutable reference to the inner sink.
    pub fn get_mut(&mut self) -> &mut Si {
        &mut self.sink
    }

    /// Get a pinned mutable reference to the inner sink.
    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Si> {
        self.sink()
    }

    /// Consumes this combinator, returning the underlying sink.
    ///
    /// Note that this may discard intermediate state of this combinator, so
    /// care should be taken to avoid losing resources when this is called.
    pub fn into_inner(self) -> Si {
        self.sink
    }

    fn take_f(self: Pin<&mut Self>) -> F {
        self.f().take().expect("polled MapErr after completion")
    }
}

impl<Si, F, E, Item> Sink<Item> for SinkMapErr<Si, F>
    where Si: Sink<Item>,
          F: FnOnce(Si::Error) -> E,
{
    type Error = E;

    fn poll_ready(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        self.as_mut().sink().poll_ready(cx).map_err(|e| self.as_mut().take_f()(e))
    }

    fn start_send(
        mut self: Pin<&mut Self>,
        item: Item,
    ) -> Result<(), Self::Error> {
        self.as_mut().sink().start_send(item).map_err(|e| self.as_mut().take_f()(e))
    }

    fn poll_flush(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        self.as_mut().sink().poll_flush(cx).map_err(|e| self.as_mut().take_f()(e))
    }

    fn poll_close(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), Self::Error>> {
        self.as_mut().sink().poll_close(cx).map_err(|e| self.as_mut().take_f()(e))
    }
}

// Forwarding impl of Stream from the underlying sink
impl<S: Stream, F> Stream for SinkMapErr<S, F> {
    type Item = S::Item;

    fn poll_next(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<S::Item>> {
        self.sink().poll_next(cx)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.sink.size_hint()
    }
}

impl<S: FusedStream, F> FusedStream for SinkMapErr<S, F> {
    fn is_terminated(&self) -> bool {
        self.sink.is_terminated()
    }
}
<|endoftext|>