<commit_before>//! Support for ANSI terminal colors via the colored crate.
//!
//! To enable support for colors, add the `"colored"` feature in your
//! `Cargo.toml`:
//!
//! ```toml
//! [dependencies]
//! fern = { version = "0.5", features = ["colored"] }
//! ```
//!
//! ---
//!
//! Colors are mainly supported via coloring the log level itself, but it's
//! also possible to color each entire log line based off of the log level.
//!
//! First, here's an example which colors the log levels themselves ("INFO" /
//! "WARN" text will have color, but won't color the rest of the line).
//! [`ColoredLevelConfig`] lets us configure the colors per-level, but also has
//! sane defaults.
//!
//! ```
//! use fern::colors::{Color, ColoredLevelConfig};
//!
//! let mut colors = ColoredLevelConfig::new()
//!     // use builder methods
//!     .info(Color::Green);
//! // or access raw fields
//! colors.warn = Color::Magenta;
//! ```
//!
//! It can then be used within any regular fern formatting closure:
//!
//! ```
//! # let colors = fern::colors::ColoredLevelConfig::new();
//! fern::Dispatch::new()
//!     // ...
//!     .format(move |out, message, record| {
//!         out.finish(format_args!(
//!             "[{}] {}",
//!             // just use 'colors.color(..)' instead of the level
//!             // itself to insert ANSI colors.
//!             colors.color(record.level()),
//!             message,
//!         ))
//!     })
//!     # .into_log();
//! ```
//!
//! ---
//!
//! Coloring levels is nice, but the alternative is good to. For an example
//! application coloring each entire log line with the right color, see
//! [examples/pretty-colored.rs][ex].
//!
//! [`ColoredLevelConfig`]: struct.ColoredLevelConfig.html
//! [ex]: https://github.com/daboross/fern/blob/master/examples/pretty-colored.rs
use std::fmt;

pub use colored::Color;
use log::Level;

/// Extension crate allowing the use of `.colored` on Levels.
trait ColoredLogLevel {
    /// Colors this log level with the given color.
    fn colored(&self, color: Color) -> WithFgColor<Level>;
}

/// Opaque structure which represents some text data and a color to display it
/// with.
///
/// This implements [`fmt::Display`] to displaying the inner text (usually a
/// log level) with ANSI color markers before to set the color and after to
/// reset the color.
///
/// `WithFgColor` instances can be created and displayed without any allocation.
// this is necessary in order to avoid using colored::ColorString, which has a
// Display implementation involving many allocations, and would involve two
// more string allocations even to create it.
//
// [`fmt::Display`]: https://doc.rust-lang.org/std/fmt/trait.Display.html
pub struct WithFgColor<T>
where
    T: fmt::Display,
{
    text: T,
    color: Color,
}

impl<T> fmt::Display for WithFgColor<T>
where
    T: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "\x1B[{}m{}\x1B[0m", self.color.to_fg_str(), self.text)
    }
}

/// Configuration specifying colors a log level can be colored as.
///
/// Example usage setting custom 'info' and 'debug' colors:
///
/// ```
/// use fern::colors::{Color, ColoredLevelConfig};
///
/// let colors = ColoredLevelConfig::new()
///     .info(Color::Green)
///     .debug(Color::Magenta);
///
/// fern::Dispatch::new()
///     .format(move |out, message, record| {
///         out.finish(format_args!(
///             "[{}] {}",
///             colors.color(record.level()),
///             message
///         ))
///     })
///     .chain(std::io::stdout())
/// # /*
///     .apply()?;
/// # */
/// #   .into_log();
/// ```
#[derive(Copy, Clone)]
#[must_use = "builder methods take config by value and thus must be reassigned to variable"]
pub struct ColoredLevelConfig {
    /// The color to color logs with the [`Error`] level.
    ///
    /// [`Error`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Error
    pub error: Color,
    /// The color to color logs with the [`Warn`] level.
    ///
    /// [`Warn`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Warn
    pub warn: Color,
    /// The color to color logs with the [`Info`] level.
    ///
    /// [`Info`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Info
    pub info: Color,
    /// The color to color logs with the [`Debug`] level.
    ///
    /// [`Debug`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Debug
    pub debug: Color,
    /// The color to color logs with the [`Trace`] level.
    ///
    /// [`Trace`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Trace
    pub trace: Color,
}

impl ColoredLevelConfig {
    /// Creates a new ColoredLevelConfig with the default colors.
    ///
    /// This matches the behavior of [`ColoredLevelConfig::default`].
    ///
    /// [`ColoredLevelConfig::default`]: #method.default
    #[inline]
    pub fn new() -> Self {
        Self::default()
    }

    /// Overrides the [`Error`] level color with the given color.
    ///
    /// The default color is [`Color::Red`].
    ///
    /// [`Error`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Error
    /// [`Color::Red`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.Red
    pub fn error(mut self, error: Color) -> Self {
        self.error = error;
        self
    }

    /// Overrides the [`Warn`] level color with the given color.
    ///
    /// The default color is [`Color::Yellow`].
    ///
    /// [`Warn`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Warn
    /// [`Color::Yellow`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.Yellow
    pub fn warn(mut self, warn: Color) -> Self {
        self.warn = warn;
        self
    }

    /// Overrides the [`Info`] level color with the given color.
    ///
    /// The default color is [`Color::White`].
    ///
    /// [`Info`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Info
    /// [`Color::White`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.White
    pub fn info(mut self, info: Color) -> Self {
        self.info = info;
        self
    }

    /// Overrides the [`Debug`] level color with the given color.
    ///
    /// The default color is [`Color::White`].
    ///
    /// [`Debug`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Debug
    /// [`Color::White`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.White
    pub fn debug(mut self, debug: Color) -> Self {
        self.debug = debug;
        self
    }

    /// Overrides the [`Trace`] level color with the given color.
    ///
    /// The default color is [`Color::White`].
    ///
    /// [`Trace`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Trace
    /// [`Color::White`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.White
    pub fn trace(mut self, trace: Color) -> Self {
        self.trace = trace;
        self
    }

    /// Colors the given log level with the color in this configuration
    /// corresponding to it's level.
    ///
    /// The structure returned is opaque, but will print the Level surrounded
    /// by ANSI color codes when displayed. This will work correctly for
    /// UNIX terminals, but due to a lack of support from the [`colored`]
    /// crate, this will not function in Windows.
    ///
    /// [`colored`]: https://github.com/mackwic/colored
    pub fn color(&self, level: Level) -> WithFgColor<Level> {
        level.colored(self.get_color(&level))
    }

    /// Retrieves the color that a log level should be colored as.
    pub fn get_color(&self, level: &Level) -> Color {
        match *level {
            Level::Error => self.error,
            Level::Warn => self.warn,
            Level::Info => self.info,
            Level::Debug => self.debug,
            Level::Trace => self.trace,
        }
    }
}

impl Default for ColoredLevelConfig {
    /// Retrieves the default configuration. This has:
    ///
    /// - [`Error`] as [`Color::Red`]
    /// - [`Warn`] as [`Color::Yellow`]
    /// - [`Info`] as [`Color::White`]
    /// - [`Debug`] as [`Color::White`]
    /// - [`Trace`] as [`Color::White`]
    ///
    /// [`Error`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Error
    /// [`Warn`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Warn
    /// [`Info`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Info
    /// [`Debug`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Debug
    /// [`Trace`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Trace
    /// [`Color::White`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.White
    /// [`Color::Yellow`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.Yellow
    /// [`Color::Red`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.Red
    fn default() -> Self {
        ColoredLevelConfig {
            error: Color::Red,
            warn: Color::Yellow,
            debug: Color::White,
            info: Color::White,
            trace: Color::White,
        }
    }
}

impl ColoredLogLevel for Level {
    fn colored(&self, color: Color) -> WithFgColor<Level> {
        WithFgColor {
            text: *self,
            color: color,
        }
    }
}

#[cfg(test)]
mod test {
    use colored::Color::*;
    use colored::Colorize;

    use super::WithFgColor;

    #[test]
    fn fg_color_matches_colored_behavior() {
        for &color in &[
            Black,
            Red,
            Green,
            Yellow,
            Blue,
            Magenta,
            Cyan,
            White,
            BrightBlack,
            BrightRed,
            BrightGreen,
            BrightYellow,
            BrightBlue,
            BrightMagenta,
            BrightCyan,
            BrightWhite,
        ] {
            assert_eq!(
                format!("{}", "test".color(color)),
                format!(
                    "{}",
                    WithFgColor {
                        text: "test",
                        color: color,
                    }
                )
            );
        }
    }
}
<commit_msg>honor formatting flags in colored log levels<commit_after>//! Support for ANSI terminal colors via the colored crate.
//!
//! To enable support for colors, add the `"colored"` feature in your
//! `Cargo.toml`:
//!
//! ```toml
//! [dependencies]
//! fern = { version = "0.5", features = ["colored"] }
//! ```
//!
//! ---
//!
//! Colors are mainly supported via coloring the log level itself, but it's
//! also possible to color each entire log line based off of the log level.
//!
//! First, here's an example which colors the log levels themselves ("INFO" /
//! "WARN" text will have color, but won't color the rest of the line).
//! [`ColoredLevelConfig`] lets us configure the colors per-level, but also has
//! sane defaults.
//!
//! ```
//! use fern::colors::{Color, ColoredLevelConfig};
//!
//! let mut colors = ColoredLevelConfig::new()
//!     // use builder methods
//!     .info(Color::Green);
//! // or access raw fields
//! colors.warn = Color::Magenta;
//! ```
//!
//! It can then be used within any regular fern formatting closure:
//!
//! ```
//! # let colors = fern::colors::ColoredLevelConfig::new();
//! fern::Dispatch::new()
//!     // ...
//!     .format(move |out, message, record| {
//!         out.finish(format_args!(
//!             "[{}] {}",
//!             // just use 'colors.color(..)' instead of the level
//!             // itself to insert ANSI colors.
//!             colors.color(record.level()),
//!             message,
//!         ))
//!     })
//!     # .into_log();
//! ```
//!
//! ---
//!
//! Coloring levels is nice, but the alternative is good to. For an example
//! application coloring each entire log line with the right color, see
//! [examples/pretty-colored.rs][ex].
//!
//! [`ColoredLevelConfig`]: struct.ColoredLevelConfig.html
//! [ex]: https://github.com/daboross/fern/blob/master/examples/pretty-colored.rs
use std::fmt;

pub use colored::Color;
use log::Level;

/// Extension crate allowing the use of `.colored` on Levels.
trait ColoredLogLevel {
    /// Colors this log level with the given color.
    fn colored(&self, color: Color) -> WithFgColor<Level>;
}

/// Opaque structure which represents some text data and a color to display it
/// with.
///
/// This implements [`fmt::Display`] to displaying the inner text (usually a
/// log level) with ANSI color markers before to set the color and after to
/// reset the color.
///
/// `WithFgColor` instances can be created and displayed without any allocation.
// this is necessary in order to avoid using colored::ColorString, which has a
// Display implementation involving many allocations, and would involve two
// more string allocations even to create it.
//
// [`fmt::Display`]: https://doc.rust-lang.org/std/fmt/trait.Display.html
pub struct WithFgColor<T>
where
    T: fmt::Display,
{
    text: T,
    color: Color,
}

impl<T> fmt::Display for WithFgColor<T>
where
    T: fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "\x1B[{}m", self.color.to_fg_str())?;
        fmt::Display::fmt(&self.text, f)?;
        write!(f, "\x1B[0m")?;
        Ok(())
    }
}

/// Configuration specifying colors a log level can be colored as.
///
/// Example usage setting custom 'info' and 'debug' colors:
///
/// ```
/// use fern::colors::{Color, ColoredLevelConfig};
///
/// let colors = ColoredLevelConfig::new()
///     .info(Color::Green)
///     .debug(Color::Magenta);
///
/// fern::Dispatch::new()
///     .format(move |out, message, record| {
///         out.finish(format_args!(
///             "[{}] {}",
///             colors.color(record.level()),
///             message
///         ))
///     })
///     .chain(std::io::stdout())
/// # /*
///     .apply()?;
/// # */
/// #   .into_log();
/// ```
#[derive(Copy, Clone)]
#[must_use = "builder methods take config by value and thus must be reassigned to variable"]
pub struct ColoredLevelConfig {
    /// The color to color logs with the [`Error`] level.
    ///
    /// [`Error`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Error
    pub error: Color,
    /// The color to color logs with the [`Warn`] level.
    ///
    /// [`Warn`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Warn
    pub warn: Color,
    /// The color to color logs with the [`Info`] level.
    ///
    /// [`Info`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Info
    pub info: Color,
    /// The color to color logs with the [`Debug`] level.
    ///
    /// [`Debug`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Debug
    pub debug: Color,
    /// The color to color logs with the [`Trace`] level.
    ///
    /// [`Trace`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Trace
    pub trace: Color,
}

impl ColoredLevelConfig {
    /// Creates a new ColoredLevelConfig with the default colors.
    ///
    /// This matches the behavior of [`ColoredLevelConfig::default`].
    ///
    /// [`ColoredLevelConfig::default`]: #method.default
    #[inline]
    pub fn new() -> Self {
        Self::default()
    }

    /// Overrides the [`Error`] level color with the given color.
    ///
    /// The default color is [`Color::Red`].
    ///
    /// [`Error`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Error
    /// [`Color::Red`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.Red
    pub fn error(mut self, error: Color) -> Self {
        self.error = error;
        self
    }

    /// Overrides the [`Warn`] level color with the given color.
    ///
    /// The default color is [`Color::Yellow`].
    ///
    /// [`Warn`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Warn
    /// [`Color::Yellow`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.Yellow
    pub fn warn(mut self, warn: Color) -> Self {
        self.warn = warn;
        self
    }

    /// Overrides the [`Info`] level color with the given color.
    ///
    /// The default color is [`Color::White`].
    ///
    /// [`Info`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Info
    /// [`Color::White`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.White
    pub fn info(mut self, info: Color) -> Self {
        self.info = info;
        self
    }

    /// Overrides the [`Debug`] level color with the given color.
    ///
    /// The default color is [`Color::White`].
    ///
    /// [`Debug`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Debug
    /// [`Color::White`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.White
    pub fn debug(mut self, debug: Color) -> Self {
        self.debug = debug;
        self
    }

    /// Overrides the [`Trace`] level color with the given color.
    ///
    /// The default color is [`Color::White`].
    ///
    /// [`Trace`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Trace
    /// [`Color::White`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.White
    pub fn trace(mut self, trace: Color) -> Self {
        self.trace = trace;
        self
    }

    /// Colors the given log level with the color in this configuration
    /// corresponding to it's level.
    ///
    /// The structure returned is opaque, but will print the Level surrounded
    /// by ANSI color codes when displayed. This will work correctly for
    /// UNIX terminals, but due to a lack of support from the [`colored`]
    /// crate, this will not function in Windows.
    ///
    /// [`colored`]: https://github.com/mackwic/colored
    pub fn color(&self, level: Level) -> WithFgColor<Level> {
        level.colored(self.get_color(&level))
    }

    /// Retrieves the color that a log level should be colored as.
    pub fn get_color(&self, level: &Level) -> Color {
        match *level {
            Level::Error => self.error,
            Level::Warn => self.warn,
            Level::Info => self.info,
            Level::Debug => self.debug,
            Level::Trace => self.trace,
        }
    }
}

impl Default for ColoredLevelConfig {
    /// Retrieves the default configuration. This has:
    ///
    /// - [`Error`] as [`Color::Red`]
    /// - [`Warn`] as [`Color::Yellow`]
    /// - [`Info`] as [`Color::White`]
    /// - [`Debug`] as [`Color::White`]
    /// - [`Trace`] as [`Color::White`]
    ///
    /// [`Error`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Error
    /// [`Warn`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Warn
    /// [`Info`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Info
    /// [`Debug`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Debug
    /// [`Trace`]: https://docs.rs/log/0.4/log/enum.Level.html#variant.Trace
    /// [`Color::White`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.White
    /// [`Color::Yellow`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.Yellow
    /// [`Color::Red`]: https://docs.rs/colored/1/colored/enum.Color.html#variant.Red
    fn default() -> Self {
        ColoredLevelConfig {
            error: Color::Red,
            warn: Color::Yellow,
            debug: Color::White,
            info: Color::White,
            trace: Color::White,
        }
    }
}

impl ColoredLogLevel for Level {
    fn colored(&self, color: Color) -> WithFgColor<Level> {
        WithFgColor {
            text: *self,
            color: color,
        }
    }
}

#[cfg(test)]
mod test {
    use colored::Color::*;
    use colored::Colorize;

    use super::WithFgColor;

    #[test]
    fn fg_color_matches_colored_behavior() {
        for &color in &[
            Black,
            Red,
            Green,
            Yellow,
            Blue,
            Magenta,
            Cyan,
            White,
            BrightBlack,
            BrightRed,
            BrightGreen,
            BrightYellow,
            BrightBlue,
            BrightMagenta,
            BrightCyan,
            BrightWhite,
        ] {
            assert_eq!(
                format!("{}", "test".color(color)),
                format!(
                    "{}",
                    WithFgColor {
                        text: "test",
                        color: color,
                    }
                )
            );
            assert_eq!(
                format!("{:^7}", "test".color(color)),
                format!(
                    "{:^7}",
                    WithFgColor {
                        text: "test",
                        color: color,
                    }
                )
            );
        }
    }
}
<|endoftext|>


<commit_before>extern crate url;
use self::url::percent_encoding::from_hex;

extern crate hyper;
use hyper::uri::RequestUri;

use std::collections::HashMap;

/// Percent-decode the input string
///
/// Not using url::percent_decode because it does not throw an error for malformed percent encoding
/// such as %AK
pub fn percent_decode(input: &str) -> Result<Vec<u8>, String> {
    let mut output: Vec<u8> = Vec::new();

    let mut input_iterator = input.as_bytes().into_iter();
    while let Some(i) = input_iterator.next() {
        match i {
            &b'%' => {
                let hexdigit1 = input_iterator.next()
                    .and_then(|h| from_hex(*h));
                let hexdigit2 = input_iterator.next()
                    .and_then(|h| from_hex(*h));
                match (hexdigit1, hexdigit2) {
                    (Some(h1), Some(h2)) => {
                        output.push(h1 * 0x10 + h2);
                    },
                    _ => {
                        return Err("Invalid percent encoding".to_owned());
                    },
                }
            }
            _ => {
                output.push(*i);
            }
        }
    }

    return Ok(output);
}

#[test]
fn percent_decode_test() {
    // Successes
    assert_eq!(percent_decode("%1a").unwrap(), [26]);
    assert_eq!(percent_decode("%1A").unwrap(), [26]);
    assert_eq!(percent_decode("a").unwrap(), [97]);

    // Failures
    assert!(percent_decode("%").is_err()); //too short
    assert!(percent_decode("%a").is_err()); //too short
    assert!(percent_decode("%ak").is_err()); //not in [0-9a-f]
}

/// Converts bytes to a hexstring
///
/// Lowercase-hex
pub fn hexstring(input: &[u8]) -> String
{
    let mut output = String::new();

    for byte in input {
        output.push_str(&format!("{:02x}", byte));
    }

    return output;
}

#[test]
fn hexstring_test() {
    assert_eq!(hexstring(&[0x5, 0x6]), "0506");
    assert_eq!(hexstring(&[0x0B, 0xf1]), "0bf1");
}

/// Extract query key-values as a HashMap from a RequestUri::AbsolutePath

pub fn query_hashmap(uri: &RequestUri) -> HashMap<&str, &str> {
    let mut retval: HashMap<&str, &str> = HashMap::new();

    let path = match uri {
        &RequestUri::AbsolutePath(ref i) => i,
        _ => return retval,
    };

    let query = match path.find('?') {
        Some(i) => &path[i+1..],
        None => "",
    };

    for component in query.split('&') {
        match component.find('=') {
            Some(position) => {
                retval.insert(&component[..position], &component[position + 1..]);
            }
            None => {},
        }
    }

    return retval;
}

#[test]
pub fn query_hashmap_test() {
    let uri = RequestUri::AbsolutePath("/q?param=hello&n=&m&&something=world".to_owned());
    let hmap = query_hashmap(&uri);

    assert_eq!(*hmap.get("param").unwrap(), "hello");
    assert_eq!(*hmap.get("something").unwrap(), "world");
    assert_eq!(*hmap.get("n").unwrap(), "");
    assert_eq!(hmap.len(), 3);
}

/* The info hash is stored in mongo, which is not binary string safe apparently,
so it needs to be parsed to a hexadecimal string rather than a binary one. */
pub fn parse_info_hash(input: &str) -> Result<String, String> {
    let info_hash_binary = match percent_decode(input) {
        Ok(i) => i,
        Err(j) => return Err(j),
    };

    if info_hash_binary.len() != 20 {
        return Err("Info hash is invalid (too short).".to_owned());
    }

    return Ok(hexstring(&info_hash_binary));
}

#[test]
fn parse_info_hash_test() {
    // Success
    assert_eq!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx%9A").unwrap(), "123456789abcdef123456789abcdef123456789a");

    // Failures
    assert!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx").is_err()); // too short
    assert!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vxab").is_err()); // too long
    assert!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx%ZA").is_err()); // invalid percent encoding
}
<commit_msg>Parse peer id function<commit_after>extern crate url;
use self::url::percent_encoding::from_hex;

extern crate hyper;
use hyper::uri::RequestUri;

use std::collections::HashMap;

/// Percent-decode the input string
///
/// Not using url::percent_decode because it does not throw an error for malformed percent encoding
/// such as %AK
pub fn percent_decode(input: &str) -> Result<Vec<u8>, String> {
    let mut output: Vec<u8> = Vec::new();

    let mut input_iterator = input.as_bytes().into_iter();
    while let Some(i) = input_iterator.next() {
        match i {
            &b'%' => {
                let hexdigit1 = input_iterator.next()
                    .and_then(|h| from_hex(*h));
                let hexdigit2 = input_iterator.next()
                    .and_then(|h| from_hex(*h));
                match (hexdigit1, hexdigit2) {
                    (Some(h1), Some(h2)) => {
                        output.push(h1 * 0x10 + h2);
                    },
                    _ => {
                        return Err("Invalid percent encoding".to_owned());
                    },
                }
            }
            _ => {
                output.push(*i);
            }
        }
    }

    return Ok(output);
}

#[test]
fn percent_decode_test() {
    // Successes
    assert_eq!(percent_decode("%1a").unwrap(), [26]);
    assert_eq!(percent_decode("%1A").unwrap(), [26]);
    assert_eq!(percent_decode("a").unwrap(), [97]);

    // Failures
    assert!(percent_decode("%").is_err()); //too short
    assert!(percent_decode("%a").is_err()); //too short
    assert!(percent_decode("%ak").is_err()); //not in [0-9a-f]
}

/// Converts bytes to a hexstring
///
/// Lowercase-hex
pub fn hexstring(input: &[u8]) -> String
{
    let mut output = String::new();

    for byte in input {
        output.push_str(&format!("{:02x}", byte));
    }

    return output;
}

#[test]
fn hexstring_test() {
    assert_eq!(hexstring(&[0x5, 0x6]), "0506");
    assert_eq!(hexstring(&[0x0B, 0xf1]), "0bf1");
}

/// Extract query key-values as a HashMap from a RequestUri::AbsolutePath

pub fn query_hashmap(uri: &RequestUri) -> HashMap<&str, &str> {
    let mut retval: HashMap<&str, &str> = HashMap::new();

    let path = match uri {
        &RequestUri::AbsolutePath(ref i) => i,
        _ => return retval,
    };

    let query = match path.find('?') {
        Some(i) => &path[i+1..],
        None => "",
    };

    for component in query.split('&') {
        match component.find('=') {
            Some(position) => {
                retval.insert(&component[..position], &component[position + 1..]);
            }
            None => {},
        }
    }

    return retval;
}

#[test]
pub fn query_hashmap_test() {
    let uri = RequestUri::AbsolutePath("/q?param=hello&n=&m&&something=world".to_owned());
    let hmap = query_hashmap(&uri);

    assert_eq!(*hmap.get("param").unwrap(), "hello");
    assert_eq!(*hmap.get("something").unwrap(), "world");
    assert_eq!(*hmap.get("n").unwrap(), "");
    assert_eq!(hmap.len(), 3);
}

/* The info hash is stored in mongo, which is not binary string safe apparently,
so it needs to be parsed to a hexadecimal string rather than a binary one. */
pub fn parse_info_hash(input: &str) -> Result<String, String> {
    let info_hash_binary = match percent_decode(input) {
        Ok(i) => i,
        Err(j) => return Err(j),
    };

    if info_hash_binary.len() != 20 {
        return Err("Info hash is invalid (too short).".to_owned());
    }

    return Ok(hexstring(&info_hash_binary));
}

#[test]
fn parse_info_hash_test() {
    // Success
    assert_eq!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx%9A").unwrap(), "123456789abcdef123456789abcdef123456789a");

    // Failures
    assert!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx").is_err()); // too short
    assert!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vxab").is_err()); // too long
    assert!(parse_info_hash("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx%ZA").is_err()); // invalid percent encoding
}

/* Peer id is only ever used with redis, which is binary string safe. */
pub fn parse_peer_id(input: &str) -> Result<Vec<u8>, String> {
    let peer_id_binary = match percent_decode(input) {
        Ok(i) => i,
        Err(j) => return Err(j),
    };

    if peer_id_binary.len() != 20 {
        return Err("Peer ID is invalid (too short).".to_owned());
    }

    return Ok(peer_id_binary);
}

#[test]
fn parse_peer_id_test() {
    // Success
    assert_eq!(parse_peer_id("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx%9A").unwrap(), [0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf1, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0x12, 0x34, 0x56, 0x78, 0x9a]);

    // Failures
    assert!(parse_peer_id("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx").is_err()); // too short
    assert!(parse_peer_id("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vxab").is_err()); // too long
    assert!(parse_peer_id("%124Vx%9A%BC%DE%F1%23Eg%89%AB%CD%EF%124Vx%ZA").is_err()); // invalid percent encoding
}
<|endoftext|>


<commit_before>use regex_syntax::hir::Hir;
use regex_syntax::ParserBuilder;

/// Determine if a regex pattern contains a literal uppercase character.
pub fn pattern_has_uppercase_char(pattern: &str) -> bool {
    let mut parser = ParserBuilder::new().allow_invalid_utf8(true).build();

    parser
        .parse(pattern)
        .map(|hir| hir_has_uppercase_char(&hir))
        .unwrap_or(false)
}

/// Determine if a regex expression contains a literal uppercase character.
fn hir_has_uppercase_char(hir: &Hir) -> bool {
    use regex_syntax::hir::*;

    match *hir.kind() {
        HirKind::Literal(Literal::Unicode(c)) => c.is_uppercase(),
        HirKind::Literal(Literal::Byte(b)) => char::from(b).is_uppercase(),
        HirKind::Class(Class::Unicode(ref ranges)) => ranges
            .iter()
            .any(|r| r.start().is_uppercase() || r.end().is_uppercase()),
        HirKind::Class(Class::Bytes(ref ranges)) => ranges
            .iter()
            .any(|r| char::from(r.start()).is_uppercase() || char::from(r.end()).is_uppercase()),
        HirKind::Group(Group { ref hir, .. }) | HirKind::Repetition(Repetition { ref hir, .. }) => {
            hir_has_uppercase_char(hir)
        }
        HirKind::Concat(ref hirs) | HirKind::Alternation(ref hirs) => {
            hirs.iter().any(hir_has_uppercase_char)
        }
        _ => false,
    }
}
<commit_msg>Add tests for regex_helper<commit_after>use regex_syntax::hir::Hir;
use regex_syntax::ParserBuilder;

/// Determine if a regex pattern contains a literal uppercase character.
pub fn pattern_has_uppercase_char(pattern: &str) -> bool {
    let mut parser = ParserBuilder::new().allow_invalid_utf8(true).build();

    parser
        .parse(pattern)
        .map(|hir| hir_has_uppercase_char(&hir))
        .unwrap_or(false)
}

/// Determine if a regex expression contains a literal uppercase character.
fn hir_has_uppercase_char(hir: &Hir) -> bool {
    use regex_syntax::hir::*;

    match *hir.kind() {
        HirKind::Literal(Literal::Unicode(c)) => c.is_uppercase(),
        HirKind::Literal(Literal::Byte(b)) => char::from(b).is_uppercase(),
        HirKind::Class(Class::Unicode(ref ranges)) => ranges
            .iter()
            .any(|r| r.start().is_uppercase() || r.end().is_uppercase()),
        HirKind::Class(Class::Bytes(ref ranges)) => ranges
            .iter()
            .any(|r| char::from(r.start()).is_uppercase() || char::from(r.end()).is_uppercase()),
        HirKind::Group(Group { ref hir, .. }) | HirKind::Repetition(Repetition { ref hir, .. }) => {
            hir_has_uppercase_char(hir)
        }
        HirKind::Concat(ref hirs) | HirKind::Alternation(ref hirs) => {
            hirs.iter().any(hir_has_uppercase_char)
        }
        _ => false,
    }
}

#[test]
fn pattern_has_uppercase_char_simple() {
    assert!(pattern_has_uppercase_char("A"));
    assert!(pattern_has_uppercase_char("foo.EXE"));

    assert!(!pattern_has_uppercase_char("a"));
    assert!(!pattern_has_uppercase_char("foo.exe123"));
}

#[test]
fn pattern_has_uppercase_char_advanced() {
    assert!(pattern_has_uppercase_char("foo.[a-zA-Z]"));

    assert!(!pattern_has_uppercase_char(r"\Acargo"));
    assert!(!pattern_has_uppercase_char(r"carg\x6F"));
}
<|endoftext|>


<commit_before>use std::cell::{RefCell, UnsafeCell};
use std::fs;
use std::ops::{Deref, DerefMut};
use std::path::Path;
use std::rc::Rc;
use std::sync::Arc;

use std::io::{Error, Read, Seek, SeekFrom};
use std::fs::File;
use std::io::ErrorKind::{Interrupted, Other, UnexpectedEof};

use super::*;

/// Top-level configuration for the system.
///
/// # Examples
///
/// ```
/// let config = sled::Config::default()
///     .path("/path/to/data".to_owned())
///     .cache_capacity(10_000_000_000)
///     .use_compression(true)
///     .flush_every_ms(Some(1000))
///     .snapshot_after_ops(100_000);
/// ```
#[derive(Debug, Clone)]
pub struct Config {
    inner: Arc<UnsafeCell<ConfigInner>>,
}

unsafe impl Send for Config {}
unsafe impl Sync for Config {}

impl Default for Config {
    fn default() -> Config {
        let now = uptime();
        let nanos = (now.as_secs() * 1_000_000_000) + now.subsec_nanos() as u64;
        let tmp_path = format!("sled.tmp.{}", nanos);
        let inner = Arc::new(UnsafeCell::new(ConfigInner {
            io_bufs: 3,
            io_buf_size: 2 << 22, // 8mb
            blink_fanout: 32,
            page_consolidation_threshold: 10,
            path: tmp_path.to_owned(),
            cache_bits: 6, // 64 shards
            cache_capacity: 1024 * 1024 * 1024, // 1gb
            use_os_cache: true,
            use_compression: true,
            flush_every_ms: None,
            snapshot_after_ops: 1_000_000,
            snapshot_path: None,
            cache_fixup_threshold: 1,
            segment_cleanup_threshold: 0.2,
            min_free_segments: 3,
            zero_copy_storage: false,
            tc: ThreadCache::default(),
            tmp_path: tmp_path.to_owned(),
        }));
        Config {
            inner: inner,
        }
    }
}

impl Deref for Config {
    type Target = ConfigInner;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.inner.get() }
    }
}

impl DerefMut for Config {
    fn deref_mut(&mut self) -> &mut ConfigInner {
        unsafe { &mut *self.inner.get() }
    }
}

impl Config {
    /// create a new `Tree` based on this configuration
    pub fn tree(&self) -> Tree {
        Tree::new(self.clone())
    }

    /// create a new `LockFreeLog` based on this
    /// configuration
    pub fn log(&self) -> LockFreeLog {
        LockFreeLog::start_system(self.clone())
    }
}

#[derive(Debug, Clone)]
pub struct ConfigInner {
    io_bufs: usize,
    io_buf_size: usize,
    blink_fanout: usize,
    page_consolidation_threshold: usize,
    path: String,
    cache_bits: usize,
    cache_capacity: usize,
    use_os_cache: bool,
    use_compression: bool,
    flush_every_ms: Option<u64>,
    snapshot_after_ops: usize,
    snapshot_path: Option<String>,
    cache_fixup_threshold: usize,
    segment_cleanup_threshold: f64,
    min_free_segments: usize,
    zero_copy_storage: bool,
    tc: ThreadCache<fs::File>,
    tmp_path: String,
}

macro_rules! builder {
    ($(($name:ident, $get:ident, $set:ident, $t:ty, $desc:expr)),*) => {
        $(
            #[doc="Get "]
            #[doc=$desc]
            pub fn $get(&self) -> $t {
                self.$name.clone()
            }

            #[doc="Set "]
            #[doc=$desc]
            pub fn $set(&mut self, to: $t) {
                self.$name = to;
            }

            #[doc="Builder, set "]
            #[doc=$desc]
            pub fn $name(&self, to: $t) -> Config {
                let mut ret = self.clone();
                ret.$name = to;
                Config { inner: Arc::new(UnsafeCell::new(ret))}
            }
        )*
    }
}

impl ConfigInner {
    builder!(
        (io_bufs, get_io_bufs, set_io_bufs, usize, "number of io buffers"),
        (io_buf_size, get_io_buf_size, set_io_buf_size, usize, "size of each io flush buffer. MUST be multiple of 512!"),
        (blink_fanout, get_blink_fanout, set_blink_fanout, usize, "b-link node fanout, minimum of 2"),
        (page_consolidation_threshold, get_page_consolidation_threshold, set_page_consolidation_threshold, usize, "page consolidation threshold"),
        (path, get_path, set_path, String, "path for the main storage file"),
        (cache_bits, get_cache_bits, set_cache_bits, usize, "log base 2 of the number of cache shards"),
        (cache_capacity, get_cache_capacity, set_cache_capacity, usize, "maximum size for the system page cache"),
        (use_os_cache, get_use_os_cache, set_use_os_cache, bool, "whether to use the OS page cache"),
        (use_compression, get_use_compression, set_use_compression, bool, "whether to use zstd compression"),
        (flush_every_ms, get_flush_every_ms, set_flush_every_ms, Option<u64>, "number of ms between IO buffer flushes"),
        (snapshot_after_ops, get_snapshot_after_ops, set_snapshot_after_ops, usize, "number of operations between page table snapshots"),
        (snapshot_path, get_snapshot_path, set_snapshot_path, Option<String>, "snapshot file location"),
        (cache_fixup_threshold, get_cache_fixup_threshold, set_cache_fixup_threshold, usize, "the maximum length of a cached page fragment chain"),
        (segment_cleanup_threshold, get_segment_cleanup_threshold, set_segment_cleanup_threshold, f64, "the proportion of remaining valid pages in the segment"),
        (min_free_segments, get_min_free_segments, set_min_free_segments, usize, "the minimum number of free segments to have on-deck before a compaction occurs"),
        (zero_copy_storage, get_zero_copy_storage, set_zero_copy_storage, bool, "disabling of the log segment copy cleaner")
    );

    /// Retrieve a thread-local file handle to the
    /// configured underlying storage,
    /// or create a new one if this is the first time the
    /// thread is accessing it.
    pub fn cached_file(&self) -> Rc<RefCell<fs::File>> {
        self.tc.get_or_else(|| {
            let path = self.get_path();
            let mut options = fs::OpenOptions::new();
            options.create(true);
            options.read(true);
            options.write(true);

            #[cfg(target_os = "linux")]
            {
                if !self.use_os_cache {
                    use std::os::unix::fs::OpenOptionsExt;
                    options.custom_flags(libc::O_DIRECT);
                    panic!("O_DIRECT support not sussed out yet.");
                }
            }

            options.open(path).unwrap()
        })
    }

    pub fn get_tmp_path(&self) -> String {
        self.tmp_path.clone()
    }

    /// returns the current snapshot file prefix
    pub fn snapshot_prefix(&self) -> String {
        let snapshot_path = self.get_snapshot_path();
        let path = self.get_path();
        snapshot_path.unwrap_or(path)
    }

    /// returns the snapshot file paths for this system
    pub fn get_snapshot_files(&self) -> Vec<String> {
        let mut prefix = self.snapshot_prefix();
        prefix.push_str(".");

        let abs_prefix: String = if Path::new(&prefix).is_absolute() {
            prefix
        } else {
            let mut abs_path =
                std::env::current_dir().expect("could not read current dir, maybe deleted?");
            abs_path.push(prefix.clone());
            abs_path.to_str().unwrap().to_owned()
        };

        let filter = |dir_entry: std::io::Result<std::fs::DirEntry>| if let Ok(de) = dir_entry {
            let path_buf = de.path();
            let path = path_buf.as_path();
            let path_str = path.to_str().unwrap();
            if path_str.starts_with(&abs_prefix) && !path_str.ends_with(".in___motion") {
                Some(path_str.to_owned())
            } else {
                None
            }
        } else {
            None
        };

        let snap_dir = Path::new(&abs_prefix).parent().expect(
            "could not read snapshot directory",
        );

        snap_dir
            .read_dir()
            .expect("could not read snapshot directory")
            .filter_map(filter)
            .collect()
    }

    /// read a segment of log messages. Only call after
    /// pausing segment rewriting on the segment accountant!
    pub fn read_segment(&self, offset: LogID) -> std::io::Result<log::Segment> {
        let segment_header_read = self.read(offset)?;
        if segment_header_read.is_corrupt() {
            return Err(Error::new(Other, "corrupt segment"));
        }
        if segment_header_read.is_zeroed() {
            return Err(Error::new(Other, "empty segment"));
        }
        let (lsn, _, _) = segment_header_read.flush().unwrap();
        assert_eq!(lsn % self.get_io_buf_size() as Lsn, 0);

        let mut buf = vec![];
        let cached_f = self.cached_file();
        let mut f = cached_f.borrow_mut();

        f.seek(SeekFrom::Start(offset))?;
        read_up_to(f, &mut buf, self.get_io_buf_size())?;

        Ok(log::Segment {
            buf: buf,
            lsn: lsn,
            read_offset: HEADER_LEN,
            position: offset,
        })
    }

    /// read a buffer from the disk
    pub fn read(&self, id: LogID) -> std::io::Result<LogRead> {
        let start = clock();
        let cached_f = self.cached_file();
        let mut f = cached_f.borrow_mut();
        f.seek(SeekFrom::Start(id))?;

        let mut valid_buf = [0u8; 1];
        f.read_exact(&mut valid_buf)?;
        let valid = valid_buf[0] == 1;

        let mut lsn_buf = [0u8; 8];
        f.read_exact(&mut lsn_buf)?;
        let lsn: Lsn = unsafe { std::mem::transmute(lsn_buf) };

        let mut len_buf = [0u8; 4];
        f.read_exact(&mut len_buf)?;

        let len32: u32 = unsafe { std::mem::transmute(len_buf) };
        let mut len = len32 as usize;
        let max = self.get_io_buf_size() - HEADER_LEN;
        if len > max {
            #[cfg(feature = "log")]
            error!("log read invalid message length, {} should be <= {}", len, max);
            M.read.measure(clock() - start);
            return Ok(LogRead::Corrupted(len));
        } else if len == 0 && !valid {
            // skip to next record, which starts with 1
            loop {
                let mut byte = [0u8; 1];
                if let Err(e) = f.read_exact(&mut byte) {
                    if e.kind() == UnexpectedEof {
                        // we've hit the end of the file
                        break;
                    }
                    panic!("{:?}", e);
                }
                if byte[0] != 1 {
                    debug_assert_eq!(byte[0], 0);
                    len += 1;
                } else {
                    break;
                }
            }
        }

        if !valid {
            M.read.measure(clock() - start);
            // we're 2 short when we started seeking for a non-zero byte (crc16)
            return Ok(LogRead::Zeroed(len + HEADER_LEN - 2));
        } else if (lsn % self.get_io_buf_size() as Lsn) != (id % self.get_io_buf_size() as LogID) {
            return Ok(LogRead::Corrupted(len + HEADER_LEN - 2));
        }

        let mut crc16_buf = [0u8; 2];
        f.read_exact(&mut crc16_buf)?;

        let mut buf = Vec::with_capacity(len);
        unsafe {
            buf.set_len(len);
        }
        f.read_exact(&mut buf)?;

        let checksum = crc16_arr(&buf);
        if checksum != crc16_buf {
            M.read.measure(clock() - start);
            return Ok(LogRead::Corrupted(len));
        }

        #[cfg(feature = "zstd")]
        let res = {
            if self.config().get_use_compression() {
                let start = clock();
                let res = Ok(LogRead::Flush(lsn, decompress(&*buf, max).unwrap(), len));
                M.decompress.measure(clock() - start);
                res
            } else {
                Ok(LogRead::Flush(lsn, buf, len))
            }
        };

        #[cfg(not(feature = "zstd"))]
        let res = Ok(LogRead::Flush(lsn, buf, len));

        M.read.measure(clock() - start);
        res
    }
}

impl Drop for ConfigInner {
    fn drop(&mut self) {
        let ephemeral = self.get_path() == self.get_tmp_path();
        if !ephemeral {
            return;
        }

        // Our files are temporary, so nuke them.

        let _res = fs::remove_file(self.tmp_path.clone());

        let candidates = self.get_snapshot_files();
        for path in candidates {
            if let Err(_e) = std::fs::remove_file(path) {
                    #[cfg(feature = "log")]
                warn!("failed to remove old snapshot file, maybe snapshot race? {}", _e);
            }
        }
    }
}

fn read_up_to(
    mut f: std::cell::RefMut<File>,
    buf: &mut Vec<u8>,
    max: usize,
) -> std::io::Result<usize> {
    buf.reserve(max);
    unsafe {
        buf.set_len(max);
    }
    let mut len = 0;
    loop {
        match f.read(&mut buf[len..max]) {
            Ok(0) => {
                unsafe {
                    buf.set_len(len);
                }
                return Ok(len);
            }
            Ok(n) => {
                len += n;
            }
            Err(ref e) if e.kind() == Interrupted => {}
            Err(e) => return Err(e),
        }
    }
}
<commit_msg>restore periodic flusher by default in config<commit_after>use std::cell::{RefCell, UnsafeCell};
use std::fs;
use std::ops::{Deref, DerefMut};
use std::path::Path;
use std::rc::Rc;
use std::sync::Arc;

use std::io::{Error, Read, Seek, SeekFrom};
use std::fs::File;
use std::io::ErrorKind::{Interrupted, Other, UnexpectedEof};

use super::*;

/// Top-level configuration for the system.
///
/// # Examples
///
/// ```
/// let config = sled::Config::default()
///     .path("/path/to/data".to_owned())
///     .cache_capacity(10_000_000_000)
///     .use_compression(true)
///     .flush_every_ms(Some(1000))
///     .snapshot_after_ops(100_000);
/// ```
#[derive(Debug, Clone)]
pub struct Config {
    inner: Arc<UnsafeCell<ConfigInner>>,
}

unsafe impl Send for Config {}
unsafe impl Sync for Config {}

impl Default for Config {
    fn default() -> Config {
        let now = uptime();
        let nanos = (now.as_secs() * 1_000_000_000) + now.subsec_nanos() as u64;
        let tmp_path = format!("sled.tmp.{}", nanos);
        let inner = Arc::new(UnsafeCell::new(ConfigInner {
            io_bufs: 3,
            io_buf_size: 2 << 22, // 8mb
            blink_fanout: 32,
            page_consolidation_threshold: 10,
            path: tmp_path.to_owned(),
            cache_bits: 6, // 64 shards
            cache_capacity: 1024 * 1024 * 1024, // 1gb
            use_os_cache: true,
            use_compression: true,
            flush_every_ms: Some(500),
            snapshot_after_ops: 1_000_000,
            snapshot_path: None,
            cache_fixup_threshold: 1,
            segment_cleanup_threshold: 0.2,
            min_free_segments: 3,
            zero_copy_storage: false,
            tc: ThreadCache::default(),
            tmp_path: tmp_path.to_owned(),
        }));
        Config {
            inner: inner,
        }
    }
}

impl Deref for Config {
    type Target = ConfigInner;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.inner.get() }
    }
}

impl DerefMut for Config {
    fn deref_mut(&mut self) -> &mut ConfigInner {
        unsafe { &mut *self.inner.get() }
    }
}

impl Config {
    /// create a new `Tree` based on this configuration
    pub fn tree(&self) -> Tree {
        Tree::new(self.clone())
    }

    /// create a new `LockFreeLog` based on this
    /// configuration
    pub fn log(&self) -> LockFreeLog {
        LockFreeLog::start_system(self.clone())
    }
}

#[derive(Debug, Clone)]
pub struct ConfigInner {
    io_bufs: usize,
    io_buf_size: usize,
    blink_fanout: usize,
    page_consolidation_threshold: usize,
    path: String,
    cache_bits: usize,
    cache_capacity: usize,
    use_os_cache: bool,
    use_compression: bool,
    flush_every_ms: Option<u64>,
    snapshot_after_ops: usize,
    snapshot_path: Option<String>,
    cache_fixup_threshold: usize,
    segment_cleanup_threshold: f64,
    min_free_segments: usize,
    zero_copy_storage: bool,
    tc: ThreadCache<fs::File>,
    tmp_path: String,
}

macro_rules! builder {
    ($(($name:ident, $get:ident, $set:ident, $t:ty, $desc:expr)),*) => {
        $(
            #[doc="Get "]
            #[doc=$desc]
            pub fn $get(&self) -> $t {
                self.$name.clone()
            }

            #[doc="Set "]
            #[doc=$desc]
            pub fn $set(&mut self, to: $t) {
                self.$name = to;
            }

            #[doc="Builder, set "]
            #[doc=$desc]
            pub fn $name(&self, to: $t) -> Config {
                let mut ret = self.clone();
                ret.$name = to;
                Config { inner: Arc::new(UnsafeCell::new(ret))}
            }
        )*
    }
}

impl ConfigInner {
    builder!(
        (io_bufs, get_io_bufs, set_io_bufs, usize, "number of io buffers"),
        (io_buf_size, get_io_buf_size, set_io_buf_size, usize, "size of each io flush buffer. MUST be multiple of 512!"),
        (blink_fanout, get_blink_fanout, set_blink_fanout, usize, "b-link node fanout, minimum of 2"),
        (page_consolidation_threshold, get_page_consolidation_threshold, set_page_consolidation_threshold, usize, "page consolidation threshold"),
        (path, get_path, set_path, String, "path for the main storage file"),
        (cache_bits, get_cache_bits, set_cache_bits, usize, "log base 2 of the number of cache shards"),
        (cache_capacity, get_cache_capacity, set_cache_capacity, usize, "maximum size for the system page cache"),
        (use_os_cache, get_use_os_cache, set_use_os_cache, bool, "whether to use the OS page cache"),
        (use_compression, get_use_compression, set_use_compression, bool, "whether to use zstd compression"),
        (flush_every_ms, get_flush_every_ms, set_flush_every_ms, Option<u64>, "number of ms between IO buffer flushes"),
        (snapshot_after_ops, get_snapshot_after_ops, set_snapshot_after_ops, usize, "number of operations between page table snapshots"),
        (snapshot_path, get_snapshot_path, set_snapshot_path, Option<String>, "snapshot file location"),
        (cache_fixup_threshold, get_cache_fixup_threshold, set_cache_fixup_threshold, usize, "the maximum length of a cached page fragment chain"),
        (segment_cleanup_threshold, get_segment_cleanup_threshold, set_segment_cleanup_threshold, f64, "the proportion of remaining valid pages in the segment"),
        (min_free_segments, get_min_free_segments, set_min_free_segments, usize, "the minimum number of free segments to have on-deck before a compaction occurs"),
        (zero_copy_storage, get_zero_copy_storage, set_zero_copy_storage, bool, "disabling of the log segment copy cleaner")
    );

    /// Retrieve a thread-local file handle to the
    /// configured underlying storage,
    /// or create a new one if this is the first time the
    /// thread is accessing it.
    pub fn cached_file(&self) -> Rc<RefCell<fs::File>> {
        self.tc.get_or_else(|| {
            let path = self.get_path();
            let mut options = fs::OpenOptions::new();
            options.create(true);
            options.read(true);
            options.write(true);

            #[cfg(target_os = "linux")]
            {
                if !self.use_os_cache {
                    use std::os::unix::fs::OpenOptionsExt;
                    options.custom_flags(libc::O_DIRECT);
                    panic!("O_DIRECT support not sussed out yet.");
                }
            }

            options.open(path).unwrap()
        })
    }

    pub fn get_tmp_path(&self) -> String {
        self.tmp_path.clone()
    }

    /// returns the current snapshot file prefix
    pub fn snapshot_prefix(&self) -> String {
        let snapshot_path = self.get_snapshot_path();
        let path = self.get_path();
        snapshot_path.unwrap_or(path)
    }

    /// returns the snapshot file paths for this system
    pub fn get_snapshot_files(&self) -> Vec<String> {
        let mut prefix = self.snapshot_prefix();
        prefix.push_str(".");

        let abs_prefix: String = if Path::new(&prefix).is_absolute() {
            prefix
        } else {
            let mut abs_path =
                std::env::current_dir().expect("could not read current dir, maybe deleted?");
            abs_path.push(prefix.clone());
            abs_path.to_str().unwrap().to_owned()
        };

        let filter = |dir_entry: std::io::Result<std::fs::DirEntry>| if let Ok(de) = dir_entry {
            let path_buf = de.path();
            let path = path_buf.as_path();
            let path_str = path.to_str().unwrap();
            if path_str.starts_with(&abs_prefix) && !path_str.ends_with(".in___motion") {
                Some(path_str.to_owned())
            } else {
                None
            }
        } else {
            None
        };

        let snap_dir = Path::new(&abs_prefix).parent().expect(
            "could not read snapshot directory",
        );

        snap_dir
            .read_dir()
            .expect("could not read snapshot directory")
            .filter_map(filter)
            .collect()
    }

    /// read a segment of log messages. Only call after
    /// pausing segment rewriting on the segment accountant!
    pub fn read_segment(&self, offset: LogID) -> std::io::Result<log::Segment> {
        let segment_header_read = self.read(offset)?;
        if segment_header_read.is_corrupt() {
            return Err(Error::new(Other, "corrupt segment"));
        }
        if segment_header_read.is_zeroed() {
            return Err(Error::new(Other, "empty segment"));
        }
        let (lsn, _, _) = segment_header_read.flush().unwrap();
        assert_eq!(lsn % self.get_io_buf_size() as Lsn, 0);

        let mut buf = vec![];
        let cached_f = self.cached_file();
        let mut f = cached_f.borrow_mut();

        f.seek(SeekFrom::Start(offset))?;
        read_up_to(f, &mut buf, self.get_io_buf_size())?;

        Ok(log::Segment {
            buf: buf,
            lsn: lsn,
            read_offset: HEADER_LEN,
            position: offset,
        })
    }

    /// read a buffer from the disk
    pub fn read(&self, id: LogID) -> std::io::Result<LogRead> {
        let start = clock();
        let cached_f = self.cached_file();
        let mut f = cached_f.borrow_mut();
        f.seek(SeekFrom::Start(id))?;

        let mut valid_buf = [0u8; 1];
        f.read_exact(&mut valid_buf)?;
        let valid = valid_buf[0] == 1;

        let mut lsn_buf = [0u8; 8];
        f.read_exact(&mut lsn_buf)?;
        let lsn: Lsn = unsafe { std::mem::transmute(lsn_buf) };

        let mut len_buf = [0u8; 4];
        f.read_exact(&mut len_buf)?;

        let len32: u32 = unsafe { std::mem::transmute(len_buf) };
        let mut len = len32 as usize;
        let max = self.get_io_buf_size() - HEADER_LEN;
        if len > max {
            #[cfg(feature = "log")]
            error!("log read invalid message length, {} should be <= {}", len, max);
            M.read.measure(clock() - start);
            return Ok(LogRead::Corrupted(len));
        } else if len == 0 && !valid {
            // skip to next record, which starts with 1
            loop {
                let mut byte = [0u8; 1];
                if let Err(e) = f.read_exact(&mut byte) {
                    if e.kind() == UnexpectedEof {
                        // we've hit the end of the file
                        break;
                    }
                    panic!("{:?}", e);
                }
                if byte[0] != 1 {
                    debug_assert_eq!(byte[0], 0);
                    len += 1;
                } else {
                    break;
                }
            }
        }

        if !valid {
            M.read.measure(clock() - start);
            // we're 2 short when we started seeking for a non-zero byte (crc16)
            return Ok(LogRead::Zeroed(len + HEADER_LEN - 2));
        } else if (lsn % self.get_io_buf_size() as Lsn) != (id % self.get_io_buf_size() as LogID) {
            return Ok(LogRead::Corrupted(len + HEADER_LEN - 2));
        }

        let mut crc16_buf = [0u8; 2];
        f.read_exact(&mut crc16_buf)?;

        let mut buf = Vec::with_capacity(len);
        unsafe {
            buf.set_len(len);
        }
        f.read_exact(&mut buf)?;

        let checksum = crc16_arr(&buf);
        if checksum != crc16_buf {
            M.read.measure(clock() - start);
            return Ok(LogRead::Corrupted(len));
        }

        #[cfg(feature = "zstd")]
        let res = {
            if self.config().get_use_compression() {
                let start = clock();
                let res = Ok(LogRead::Flush(lsn, decompress(&*buf, max).unwrap(), len));
                M.decompress.measure(clock() - start);
                res
            } else {
                Ok(LogRead::Flush(lsn, buf, len))
            }
        };

        #[cfg(not(feature = "zstd"))]
        let res = Ok(LogRead::Flush(lsn, buf, len));

        M.read.measure(clock() - start);
        res
    }
}

impl Drop for ConfigInner {
    fn drop(&mut self) {
        let ephemeral = self.get_path() == self.get_tmp_path();
        if !ephemeral {
            return;
        }

        // Our files are temporary, so nuke them.

        let _res = fs::remove_file(self.tmp_path.clone());

        let candidates = self.get_snapshot_files();
        for path in candidates {
            if let Err(_e) = std::fs::remove_file(path) {
                    #[cfg(feature = "log")]
                warn!("failed to remove old snapshot file, maybe snapshot race? {}", _e);
            }
        }
    }
}

fn read_up_to(
    mut f: std::cell::RefMut<File>,
    buf: &mut Vec<u8>,
    max: usize,
) -> std::io::Result<usize> {
    buf.reserve(max);
    unsafe {
        buf.set_len(max);
    }
    let mut len = 0;
    loop {
        match f.read(&mut buf[len..max]) {
            Ok(0) => {
                unsafe {
                    buf.set_len(len);
                }
                return Ok(len);
            }
            Ok(n) => {
                len += n;
            }
            Err(ref e) if e.kind() == Interrupted => {}
            Err(e) => return Err(e),
        }
    }
}
<|endoftext|>


<commit_before>use std::ascii::AsciiExt;
use std::borrow::ToOwned;
use std::env;
use std::fs;
use std::io::{self, Read};
use std::ops::Deref;
use std::path::PathBuf;

use csv;
use index::Indexed;
use rustc_serialize::{Decodable, Decoder};

use CliResult;
use select::{SelectColumns, Selection};
use util;

#[derive(Clone, Copy, Debug)]
pub struct Delimiter(pub u8);

/// Delimiter represents values that can be passed from the command line that
/// can be used as a field delimiter in CSV data.
///
/// Its purpose is to ensure that the Unicode character given decodes to a
/// valid ASCII character as required by the CSV parser.
impl Delimiter {
    pub fn as_byte(self) -> u8 {
        self.0
    }
}

impl Decodable for Delimiter {
    fn decode<D: Decoder>(d: &mut D) -> Result<Delimiter, D::Error> {
        let c = d.read_str()?;
        match &*c {
            r"\t" => Ok(Delimiter(b'\t')),
            s => {
                if s.len() != 1 {
                    let msg = format!("Could not convert '{}' to a single \
                                       ASCII character.", s);
                    return Err(d.error(&*msg));
                }
                let c = s.chars().next().unwrap();
                if c.is_ascii() {
                    Ok(Delimiter(c as u8))
                } else {
                    let msg = format!("Could not convert '{}' \
                                       to ASCII delimiter.", c);
                    Err(d.error(&*msg))
                }
            }
        }
    }
}

#[derive(Debug)]
pub struct Config {
    path: Option<PathBuf>, // None implies <stdin>
    idx_path: Option<PathBuf>,
    select_columns: Option<SelectColumns>,
    delimiter: u8,
    pub no_headers: bool,
    flexible: bool,
    terminator: csv::Terminator,
    quote: u8,
    quote_style: csv::QuoteStyle,
    double_quote: bool,
    escape: u8,
    quoting: bool,
}

impl Config {
    pub fn new(path: &Option<String>) -> Config {
        let (path, delim) = match *path {
            None => (None, b','),
            Some(ref s) if s.deref() == "-" => (None, b','),
            Some(ref s) => {
                let path = PathBuf::from(s);
                let delim =
                    if path.extension().map_or(false, |v| v == "tsv") {
                        b'\t'
                    } else {
                        b','
                    };
                (Some(path), delim)
            }
        };
        Config {
            path: path,
            idx_path: None,
            select_columns: None,
            delimiter: delim,
            no_headers: false,
            flexible: false,
            terminator: csv::Terminator::Any(b'\n'),
            quote: b'"',
            quote_style: csv::QuoteStyle::Necessary,
            double_quote: true,
            escape: b'\\',
            quoting: true,
        }
    }

    pub fn delimiter(mut self, d: Option<Delimiter>) -> Config {
        if let Some(d) = d {
            self.delimiter = d.as_byte();
        }
        self
    }

    pub fn no_headers(mut self, mut yes: bool) -> Config {
        if env::var("XSV_TOGGLE_HEADERS").unwrap_or("0".to_owned()) == "1" {
            yes = !yes;
        }
        self.no_headers = yes;
        self
    }

    pub fn flexible(mut self, yes: bool) -> Config {
        self.flexible = yes;
        self
    }

    pub fn crlf(mut self, yes: bool) -> Config {
        if yes {
            self.terminator = csv::Terminator::CRLF;
        } else {
            self.terminator = csv::Terminator::Any(b'\n');
        }
        self
    }

    pub fn terminator(mut self, term: csv::Terminator) -> Config {
        self.terminator = term;
        self
    }

    pub fn quote(mut self, quote: u8) -> Config {
        self.quote = quote;
        self
    }

    pub fn quote_style(mut self, style: csv::QuoteStyle) -> Config {
        self.quote_style = style;
        self
    }

    pub fn double_quote(mut self, yes: bool) -> Config {
        self.double_quote = yes;
        self
    }

    pub fn escape(mut self, escape: u8) -> Config {
        self.escape = escape;
        self
    }

    pub fn quoting(mut self, yes: bool) -> Config {
        self.quoting = yes;
        self
    }

    pub fn select(mut self, sel_cols: SelectColumns) -> Config {
        self.select_columns = Some(sel_cols);
        self
    }

    pub fn is_std(&self) -> bool {
        self.path.is_none()
    }

    pub fn selection(
        &self,
        first_record: &csv::ByteRecord,
    ) -> Result<Selection, String> {
        match self.select_columns {
            None => Err("Config has no 'SelectColums'. Did you call \
                         Config::select?".to_owned()),
            Some(ref sel) => sel.selection(first_record, !self.no_headers),
        }
    }

    pub fn write_headers<R: io::Read, W: io::Write>
                        (&self, r: &mut csv::Reader<R>, w: &mut csv::Writer<W>)
                        -> csv::Result<()> {
        if !self.no_headers {
            let r = r.byte_headers()?;
            if !r.is_empty() {
                w.write_record(r)?;
            }
        }
        Ok(())
    }

    pub fn writer(&self)
                 -> io::Result<csv::Writer<Box<io::Write+'static>>> {
        Ok(self.from_writer(self.io_writer()?))
    }

    pub fn reader(&self)
                 -> io::Result<csv::Reader<Box<io::Read+'static>>> {
        Ok(self.from_reader(self.io_reader()?))
    }

    pub fn reader_file(&self) -> io::Result<csv::Reader<fs::File>> {
        match self.path {
            None => Err(io::Error::new(
                io::ErrorKind::Other, "Cannot use <stdin> here",
            )),
            Some(ref p) => fs::File::open(p).map(|f| self.from_reader(f)),
        }
    }

    pub fn index_files(&self)
           -> io::Result<Option<(csv::Reader<fs::File>, fs::File)>> {
        let (csv_file, idx_file) = match (&self.path, &self.idx_path) {
            (&None, &None) => return Ok(None),
            (&None, &Some(_)) => return Err(io::Error::new(
                io::ErrorKind::Other,
                "Cannot use <stdin> with indexes",
                // Some(format!("index file: {}", p.display()))
            )),
            (&Some(ref p), &None) => {
                // We generally don't want to report an error here, since we're
                // passively trying to find an index.
                let idx_file = match fs::File::open(&util::idx_path(p)) {
                    // TODO: Maybe we should report an error if the file exists
                    // but is not readable.
                    Err(_) => return Ok(None),
                    Ok(f) => f,
                };
                (fs::File::open(p)?, idx_file)
            }
            (&Some(ref p), &Some(ref ip)) => {
                (fs::File::open(p)?, fs::File::open(ip)?)
            }
        };
        // If the CSV data was last modified after the index file was last
        // modified, then return an error and demand the user regenerate the
        // index.
        let data_modified = util::last_modified(&csv_file.metadata()?);
        let idx_modified = util::last_modified(&idx_file.metadata()?);
        if data_modified > idx_modified {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                "The CSV file was modified after the index file. \
                 Please re-create the index.",
            ));
        }
        let csv_rdr = self.from_reader(csv_file);
        Ok(Some((csv_rdr, idx_file)))
    }

    pub fn indexed(&self)
                  -> CliResult<Option<Indexed<fs::File, fs::File>>> {
        match self.index_files()? {
            None => Ok(None),
            Some((r, i)) => Ok(Some(Indexed::open(r, i)?)),
        }
    }

    pub fn io_reader(&self) -> io::Result<Box<io::Read+'static>> {
        Ok(match self.path {
                None => Box::new(io::stdin()),
                Some(ref p) => {
                    match fs::File::open(p){
                        Ok(x) => Box::new(x),
                        Err(err) => {
                            let msg = format!(
                                "failed to open {}: {}", p.display(), err);
                            return Err(io::Error::new(
                                io::ErrorKind::NotFound,
                                msg,
                            ));
                        }
                    }
                },
            })
    }

    pub fn from_reader<R: Read>(&self, rdr: R) -> csv::Reader<R> {
        csv::ReaderBuilder::new()
            .flexible(self.flexible)
            .delimiter(self.delimiter)
            .has_headers(!self.no_headers)
            .quote(self.quote)
            .quoting(self.quoting)
            .from_reader(rdr)
    }

    pub fn io_writer(&self) -> io::Result<Box<io::Write+'static>> {
        Ok(match self.path {
            None => Box::new(io::stdout()),
            Some(ref p) => Box::new(fs::File::create(p)?),
        })
    }

    pub fn from_writer<W: io::Write>(&self, wtr: W) -> csv::Writer<W> {
        csv::WriterBuilder::new()
            .flexible(self.flexible)
            .delimiter(self.delimiter)
            .terminator(self.terminator)
            .quote(self.quote)
            .quote_style(self.quote_style)
            .double_quote(self.double_quote)
            .escape(self.escape)
            .buffer_capacity(32 * (1<<10))
            .from_writer(wtr)
    }
}
<commit_msg>config: propagate escape config<commit_after>use std::ascii::AsciiExt;
use std::borrow::ToOwned;
use std::env;
use std::fs;
use std::io::{self, Read};
use std::ops::Deref;
use std::path::PathBuf;

use csv;
use index::Indexed;
use rustc_serialize::{Decodable, Decoder};

use CliResult;
use select::{SelectColumns, Selection};
use util;

#[derive(Clone, Copy, Debug)]
pub struct Delimiter(pub u8);

/// Delimiter represents values that can be passed from the command line that
/// can be used as a field delimiter in CSV data.
///
/// Its purpose is to ensure that the Unicode character given decodes to a
/// valid ASCII character as required by the CSV parser.
impl Delimiter {
    pub fn as_byte(self) -> u8 {
        self.0
    }
}

impl Decodable for Delimiter {
    fn decode<D: Decoder>(d: &mut D) -> Result<Delimiter, D::Error> {
        let c = d.read_str()?;
        match &*c {
            r"\t" => Ok(Delimiter(b'\t')),
            s => {
                if s.len() != 1 {
                    let msg = format!("Could not convert '{}' to a single \
                                       ASCII character.", s);
                    return Err(d.error(&*msg));
                }
                let c = s.chars().next().unwrap();
                if c.is_ascii() {
                    Ok(Delimiter(c as u8))
                } else {
                    let msg = format!("Could not convert '{}' \
                                       to ASCII delimiter.", c);
                    Err(d.error(&*msg))
                }
            }
        }
    }
}

#[derive(Debug)]
pub struct Config {
    path: Option<PathBuf>, // None implies <stdin>
    idx_path: Option<PathBuf>,
    select_columns: Option<SelectColumns>,
    delimiter: u8,
    pub no_headers: bool,
    flexible: bool,
    terminator: csv::Terminator,
    quote: u8,
    quote_style: csv::QuoteStyle,
    double_quote: bool,
    escape: u8,
    quoting: bool,
}

impl Config {
    pub fn new(path: &Option<String>) -> Config {
        let (path, delim) = match *path {
            None => (None, b','),
            Some(ref s) if s.deref() == "-" => (None, b','),
            Some(ref s) => {
                let path = PathBuf::from(s);
                let delim =
                    if path.extension().map_or(false, |v| v == "tsv") {
                        b'\t'
                    } else {
                        b','
                    };
                (Some(path), delim)
            }
        };
        Config {
            path: path,
            idx_path: None,
            select_columns: None,
            delimiter: delim,
            no_headers: false,
            flexible: false,
            terminator: csv::Terminator::Any(b'\n'),
            quote: b'"',
            quote_style: csv::QuoteStyle::Necessary,
            double_quote: true,
            escape: b'\\',
            quoting: true,
        }
    }

    pub fn delimiter(mut self, d: Option<Delimiter>) -> Config {
        if let Some(d) = d {
            self.delimiter = d.as_byte();
        }
        self
    }

    pub fn no_headers(mut self, mut yes: bool) -> Config {
        if env::var("XSV_TOGGLE_HEADERS").unwrap_or("0".to_owned()) == "1" {
            yes = !yes;
        }
        self.no_headers = yes;
        self
    }

    pub fn flexible(mut self, yes: bool) -> Config {
        self.flexible = yes;
        self
    }

    pub fn crlf(mut self, yes: bool) -> Config {
        if yes {
            self.terminator = csv::Terminator::CRLF;
        } else {
            self.terminator = csv::Terminator::Any(b'\n');
        }
        self
    }

    pub fn terminator(mut self, term: csv::Terminator) -> Config {
        self.terminator = term;
        self
    }

    pub fn quote(mut self, quote: u8) -> Config {
        self.quote = quote;
        self
    }

    pub fn quote_style(mut self, style: csv::QuoteStyle) -> Config {
        self.quote_style = style;
        self
    }

    pub fn double_quote(mut self, yes: bool) -> Config {
        self.double_quote = yes;
        self
    }

    pub fn escape(mut self, escape: u8) -> Config {
        self.escape = escape;
        self
    }

    pub fn quoting(mut self, yes: bool) -> Config {
        self.quoting = yes;
        self
    }

    pub fn select(mut self, sel_cols: SelectColumns) -> Config {
        self.select_columns = Some(sel_cols);
        self
    }

    pub fn is_std(&self) -> bool {
        self.path.is_none()
    }

    pub fn selection(
        &self,
        first_record: &csv::ByteRecord,
    ) -> Result<Selection, String> {
        match self.select_columns {
            None => Err("Config has no 'SelectColums'. Did you call \
                         Config::select?".to_owned()),
            Some(ref sel) => sel.selection(first_record, !self.no_headers),
        }
    }

    pub fn write_headers<R: io::Read, W: io::Write>
                        (&self, r: &mut csv::Reader<R>, w: &mut csv::Writer<W>)
                        -> csv::Result<()> {
        if !self.no_headers {
            let r = r.byte_headers()?;
            if !r.is_empty() {
                w.write_record(r)?;
            }
        }
        Ok(())
    }

    pub fn writer(&self)
                 -> io::Result<csv::Writer<Box<io::Write+'static>>> {
        Ok(self.from_writer(self.io_writer()?))
    }

    pub fn reader(&self)
                 -> io::Result<csv::Reader<Box<io::Read+'static>>> {
        Ok(self.from_reader(self.io_reader()?))
    }

    pub fn reader_file(&self) -> io::Result<csv::Reader<fs::File>> {
        match self.path {
            None => Err(io::Error::new(
                io::ErrorKind::Other, "Cannot use <stdin> here",
            )),
            Some(ref p) => fs::File::open(p).map(|f| self.from_reader(f)),
        }
    }

    pub fn index_files(&self)
           -> io::Result<Option<(csv::Reader<fs::File>, fs::File)>> {
        let (csv_file, idx_file) = match (&self.path, &self.idx_path) {
            (&None, &None) => return Ok(None),
            (&None, &Some(_)) => return Err(io::Error::new(
                io::ErrorKind::Other,
                "Cannot use <stdin> with indexes",
                // Some(format!("index file: {}", p.display()))
            )),
            (&Some(ref p), &None) => {
                // We generally don't want to report an error here, since we're
                // passively trying to find an index.
                let idx_file = match fs::File::open(&util::idx_path(p)) {
                    // TODO: Maybe we should report an error if the file exists
                    // but is not readable.
                    Err(_) => return Ok(None),
                    Ok(f) => f,
                };
                (fs::File::open(p)?, idx_file)
            }
            (&Some(ref p), &Some(ref ip)) => {
                (fs::File::open(p)?, fs::File::open(ip)?)
            }
        };
        // If the CSV data was last modified after the index file was last
        // modified, then return an error and demand the user regenerate the
        // index.
        let data_modified = util::last_modified(&csv_file.metadata()?);
        let idx_modified = util::last_modified(&idx_file.metadata()?);
        if data_modified > idx_modified {
            return Err(io::Error::new(
                io::ErrorKind::Other,
                "The CSV file was modified after the index file. \
                 Please re-create the index.",
            ));
        }
        let csv_rdr = self.from_reader(csv_file);
        Ok(Some((csv_rdr, idx_file)))
    }

    pub fn indexed(&self)
                  -> CliResult<Option<Indexed<fs::File, fs::File>>> {
        match self.index_files()? {
            None => Ok(None),
            Some((r, i)) => Ok(Some(Indexed::open(r, i)?)),
        }
    }

    pub fn io_reader(&self) -> io::Result<Box<io::Read+'static>> {
        Ok(match self.path {
                None => Box::new(io::stdin()),
                Some(ref p) => {
                    match fs::File::open(p){
                        Ok(x) => Box::new(x),
                        Err(err) => {
                            let msg = format!(
                                "failed to open {}: {}", p.display(), err);
                            return Err(io::Error::new(
                                io::ErrorKind::NotFound,
                                msg,
                            ));
                        }
                    }
                },
            })
    }

    pub fn from_reader<R: Read>(&self, rdr: R) -> csv::Reader<R> {
        csv::ReaderBuilder::new()
            .flexible(self.flexible)
            .delimiter(self.delimiter)
            .has_headers(!self.no_headers)
            .quote(self.quote)
            .quoting(self.quoting)
            .escape(Some(self.escape))
            .from_reader(rdr)
    }

    pub fn io_writer(&self) -> io::Result<Box<io::Write+'static>> {
        Ok(match self.path {
            None => Box::new(io::stdout()),
            Some(ref p) => Box::new(fs::File::create(p)?),
        })
    }

    pub fn from_writer<W: io::Write>(&self, wtr: W) -> csv::Writer<W> {
        csv::WriterBuilder::new()
            .flexible(self.flexible)
            .delimiter(self.delimiter)
            .terminator(self.terminator)
            .quote(self.quote)
            .quote_style(self.quote_style)
            .double_quote(self.double_quote)
            .escape(self.escape)
            .buffer_capacity(32 * (1<<10))
            .from_writer(wtr)
    }
}
<|endoftext|>


<commit_before>use etcd;
use std::thread;
use std::fmt;
use std::time::Duration;
use std::sync::Arc;
use std::collections::BTreeMap;
use rustc_serialize::json;
use rustc_serialize::{Decodable,Encodable};

pub struct ConfigClient<T> {
    client: Arc<InnerClient<T>>,
    lease_mgr: thread::JoinHandle<()>,
    watcher: thread::JoinHandle<()>,
}

#[derive(Clone,Debug, Default)]
pub struct ConfigurationView<T> {
    this_node: String,
    members: BTreeMap<String, T>,
    pub epoch: u64,
}

struct InnerClient<T> {
    etcd: etcd::Client,
    data: T,
    lease_time: Duration,
    callback: Box<Fn(ConfigurationView<T>) + Send + Sync + 'static>,
    lease_key: Option<String>,
}

impl<T: 'static + Decodable + Encodable + fmt::Debug + Eq + Clone + Send + Sync> ConfigClient<T> {
    pub fn new<F: Fn(ConfigurationView<T>) + Send + Sync + 'static>(addr: &str, data: T, lease_time: Duration, callback: F)
        -> Result<ConfigClient<T>, ()> {
        let etcd = etcd::Client::new(addr).expect("etcd client");
        let mut client = InnerClient {
            etcd: etcd,
            data: data,
            lease_time: lease_time,
            callback: Box::new(callback),
            lease_key: None,
        };
        let lease = client.setup_lease();

        let client = Arc::new(client);
        let lease_mgr = { let client = client.clone();
            thread::Builder::new().name("etcd config".to_string()).spawn(move || {
                client.run_lease(lease)
            }).expect("etcd thread")
        };
        let watcher = {
            let client = client.clone();
            thread::Builder::new().name("etcd watcher".to_string()).spawn(move || {
                client.run_watch()
            }).expect("etcd watcher")
        };
        Ok(ConfigClient { client: client, lease_mgr: lease_mgr, watcher: watcher })
    }
}

const MEMBERS : &'static str = "/chain/members";
const SEQUENCER : &'static str = "/chain/config_seq";
const KEY_EXISTS : u64 = 105;
const KEY_NOT_EXISTS : u64 = 100;
const COMPARE_FAILED : u64 = 101;

#[derive(Debug,PartialEq,Eq,Default,Clone, RustcEncodable, RustcDecodable)]
struct ConfigSequencer {
    keys: Vec<String>,
    epoch: u64,
}

impl<T: Decodable + Encodable + fmt::Debug + Eq + Clone> InnerClient<T> {

    fn setup_lease(&mut self) -> u64 {
        match self.etcd.create_dir(MEMBERS, None) {
            Ok(res) => info!("Created dir: {}: {:?}: ", MEMBERS, res),
            Err(etcd::Error::Etcd (ref e)) if e.error_code == KEY_EXISTS => info!("Dir exists: {:?}: ", MEMBERS),
            Err(e) => panic!("Unexpected error creating {}: {:?}", MEMBERS, e),
        }

        let seq : ConfigSequencer = Default::default();
        match self.etcd.create(SEQUENCER, &json::encode(&seq).expect("encode epoch"), None) {
            Ok(res) => info!("Created seq: {}: {:?}: ", SEQUENCER, res),
            Err(etcd::Error::Etcd (ref e)) if e.error_code == KEY_EXISTS => info!("Sequencer exists: {:?}: ", SEQUENCER),
            Err(e) => panic!("Unexpected error creating {}: {:?}", SEQUENCER, e),
        }

        let me = self.etcd.create_in_order(MEMBERS, &self.data_json(), Some(self.lease_time.as_secs())).expect("Create unique node");
        info!("My node! {:?}", me);
        self.lease_key = Some(me.node.key.expect("Key name"));
        me.node.modified_index.expect("Lease node version")
    }

    fn data_json(&self) -> String {
        json::encode(&self.data).expect("json encode")
    }

    fn run_lease(&self, mut lease_index: u64) {
        let pausetime = self.lease_time / 2;
        let lease_key = self.lease_key.as_ref().expect("Should have created lease node");

        loop {

            thread::sleep(pausetime);
            trace!("touch node: {:?}={:?}", lease_key, self.data);
            let res = self.etcd.compare_and_swap(&lease_key, &self.data_json(),
                    Some(self.lease_time.as_secs()), None, Some(lease_index))
                .expect("Update lease");
            trace!("Update: {:?}", res);
            lease_index = res.node.modified_index.expect("lease version");
        }
    }

    fn list_members(&self) -> BTreeMap<String, T> {
        let current_listing = self.etcd.get(MEMBERS, true, true, true).expect("List members");
        trace!("Listing: {:?}", current_listing);
        current_listing.node.nodes.expect("Node listing").into_iter()
                .filter_map(|n|
                    if let (Some(k), Some(v)) = (n.key, n.value) {
                        Some((k, json::decode(&v).expect("decode json"))) } else { None })
                .collect::<BTreeMap<String, T>>()
    }

    fn run_watch(&self) {
        info!("Starting etcd watcher");
        let lease_key = self.lease_key.as_ref().expect("Should have created lease node");

        let current_listing = self.etcd.get(MEMBERS, true, true, true).expect("List members");
        debug!("Listing: {:?}", current_listing);
        let mut last_observed_index = current_listing.node.nodes.unwrap_or_else(|| Vec::new()).into_iter()
            .filter_map(|x| x.modified_index).max();

        let mut curr_members = BTreeMap::new();
        loop {
            trace!("Awaiting for {} from {:?}", MEMBERS, last_observed_index);
            let res = self.etcd.watch(MEMBERS, last_observed_index, true).expect("watch");
            trace!("Watch: {:?}", res);
            last_observed_index = res.node.modified_index.map(|x| x+1);

            let members = self.list_members();
            assert!(members.contains_key(lease_key), "I am ostensibly alive");

            let seq = self.verify_epoch(&members);

            trace!("Members: {:?}; seq: {:?}", members, seq);
            if curr_members != members {
                curr_members = members;
                info!("Membership change! {:?}", curr_members);
                (self.callback)(ConfigurationView {
                    this_node: lease_key.clone(),
                    members: curr_members.clone(),
                    epoch: seq.epoch,
                })
            }
        }

    }

    fn verify_epoch(&self, members: &BTreeMap<String, T>) -> ConfigSequencer {
        loop {
            let (seq_index, mut seq) : (u64, ConfigSequencer) = match self.etcd.get(SEQUENCER, false, false, false) {
                Ok(etcd::KeySpaceInfo {
                    node: etcd::keys::Node {
                        modified_index: Some(modified_index),
                        value: Some(value),
                        ..
                    } ,
                    ..
                }) => (modified_index, json::decode(&value).expect("decode epoch")),
                Ok(e) => panic!("Unexpected response reading {}: {:?}", SEQUENCER, e),
                Err(e) => panic!("Unexpected error reading {}: {:?}", SEQUENCER, e),
            };
            trace!("Sequencer: {}/{:?}", seq_index, seq);
            let current_keys = members.keys().cloned().collect::<Vec<_>>();
            if current_keys != seq.keys {
                debug!("Stale! {:?}", seq);

                seq.keys = current_keys;
                seq.epoch += 1;
                match self.etcd.compare_and_swap(SEQUENCER, &json::encode(&seq).expect("encode epoch"), None, None, Some(seq_index)) {
                        Ok(res) => {
                            debug!("verify_epoch: Updated seq: {}: {:?}: ", SEQUENCER, seq);
                            return seq;
                        },
                        Err(etcd::Error::Etcd (ref e)) if e.error_code == COMPARE_FAILED => debug!("verify_epoch: Raced out; retry"),
                        Err(e) => panic!("Unexpected error creating {}: {:?}", SEQUENCER, e),
                }
            } else {
                debug!("verify_epoch: Fresh! {:?}", seq);
                return seq;
            }
        }
    }
}

impl<T: Clone> ConfigurationView<T> {
    fn head_key(&self) -> Option<&str> {
        self.members.keys().next().map(|s| &**s)
    }

    pub fn should_listen_for_clients(&self) -> bool {
        self.head_key() == Some(&self.this_node)
    }
    pub fn should_listen_for_upstream(&self) -> bool {
        self.head_key() != Some(&self.this_node)
    }
    pub fn should_connect_downstream(&self) -> Option<T> {
        let next = self.members.iter().filter_map(|(k, v)| if *k > self.this_node { Some((k, v)) } else { None }).next();
        next.map(|(_, val)| val.clone())
    }
}
<commit_msg>Death watcher for config.<commit_after>use etcd;
use std::thread;
use std::sync::atomic::{Ordering,AtomicUsize};
use std::fmt;
use std::time::Duration;
use std::sync::Arc;
use std::collections::BTreeMap;
use rustc_serialize::json;
use rustc_serialize::{Decodable,Encodable};

pub struct ConfigClient<T> {
    client: Arc<InnerClient<T>>,
    lease_mgr: thread::JoinHandle<()>,
    watcher: thread::JoinHandle<()>,
}

#[derive(Clone,Debug, Default)]
pub struct ConfigurationView<T> {
    this_node: String,
    members: BTreeMap<String, T>,
    pub epoch: u64,
}

struct InnerClient<T> {
    etcd: etcd::Client,
    data: T,
    lease_time: Duration,
    callback: Box<Fn(ConfigurationView<T>) + Send + Sync + 'static>,
    lease_key: Option<String>,
    watch_count: AtomicUsize,
}

struct DeathWatch<'a>(&'a AtomicUsize);

impl<'a> DeathWatch<'a> {
    fn new(ctr: &'a AtomicUsize) -> DeathWatch<'a> {
        let cnt = ctr.fetch_add(1, Ordering::Relaxed);
        let w = DeathWatch(ctr);
        warn!("Deathwatch::new: New death watcher {:?}", w);
        w
    }
    fn count(&self) -> usize {
        self.0.load(Ordering::Relaxed)
    }
}
impl<'a> Drop for DeathWatch<'a> {
    fn drop(&mut self) {
        self.0.fetch_sub(1, Ordering::Relaxed);
        warn!("Exiting: {:?} left", self);
    }
}

impl<'a> fmt::Debug for DeathWatch<'a> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("DeathWatch").field("count", &self.count()).finish()
    }
}

impl<T: 'static + Decodable + Encodable + fmt::Debug + Eq + Clone + Send + Sync> ConfigClient<T> {
    pub fn new<F: Fn(ConfigurationView<T>) + Send + Sync + 'static>(addr: &str, data: T, lease_time: Duration, callback: F)
        -> Result<ConfigClient<T>, ()> {
        let etcd = etcd::Client::new(addr).expect("etcd client");
        let mut client = InnerClient {
            etcd: etcd,
            data: data,
            lease_time: lease_time,
            callback: Box::new(callback),
            lease_key: None,
            watch_count: AtomicUsize::new(0),
        };
        let lease = client.setup_lease();

        let client = Arc::new(client);
        let lease_mgr = { let client = client.clone();
            thread::Builder::new().name("etcd config".to_string()).spawn(move || {
                client.run_lease(lease)
            }).expect("etcd thread")
        };
        let watcher = {
            let client = client.clone();
            thread::Builder::new().name("etcd watcher".to_string()).spawn(move || {
                client.run_watch()
            }).expect("etcd watcher")
        };
        Ok(ConfigClient { client: client, lease_mgr: lease_mgr, watcher: watcher })
    }

    pub fn is_running(&self) -> bool {
        self.client.is_running()
    }
}

const MEMBERS : &'static str = "/chain/members";
const SEQUENCER : &'static str = "/chain/config_seq";
const KEY_EXISTS : u64 = 105;
const KEY_NOT_EXISTS : u64 = 100;
const COMPARE_FAILED : u64 = 101;

#[derive(Debug,PartialEq,Eq,Default,Clone, RustcEncodable, RustcDecodable)]
struct ConfigSequencer {
    keys: Vec<String>,
    epoch: u64,
}

impl<T: Decodable + Encodable + fmt::Debug + Eq + Clone> InnerClient<T> {
    fn setup_lease(&mut self) -> u64 {
        match self.etcd.create_dir(MEMBERS, None) {
            Ok(res) => info!("Created dir: {}: {:?}: ", MEMBERS, res),
            Err(etcd::Error::Etcd (ref e)) if e.error_code == KEY_EXISTS => info!("Dir exists: {:?}: ", MEMBERS),
            Err(e) => panic!("Unexpected error creating {}: {:?}", MEMBERS, e),
        }

        let seq : ConfigSequencer = Default::default();
        match self.etcd.create(SEQUENCER, &json::encode(&seq).expect("encode epoch"), None) {
            Ok(res) => info!("Created seq: {}: {:?}: ", SEQUENCER, res),
            Err(etcd::Error::Etcd (ref e)) if e.error_code == KEY_EXISTS => info!("Sequencer exists: {:?}: ", SEQUENCER),
            Err(e) => panic!("Unexpected error creating {}: {:?}", SEQUENCER, e),
        }

        let me = self.etcd.create_in_order(MEMBERS, &self.data_json(), Some(self.lease_time.as_secs())).expect("Create unique node");
        info!("My node! {:?}", me);
        self.lease_key = Some(me.node.key.expect("Key name"));
        me.node.modified_index.expect("Lease node version")
    }

    fn data_json(&self) -> String {
        json::encode(&self.data).expect("json encode")
    }

    fn run_lease(&self, mut lease_index: u64) {
        let watch = self.start_watch();
        let pausetime = self.lease_time / 2;
        let lease_key = self.lease_key.as_ref().expect("Should have created lease node");

        loop {

            thread::sleep(pausetime);
            trace!("touch node: {:?}={:?}", lease_key, self.data);
            let res = self.etcd.compare_and_swap(&lease_key, &self.data_json(),
                    Some(self.lease_time.as_secs()), None, Some(lease_index))
                .expect("Update lease");
            trace!("Update: {:?}", res);
            lease_index = res.node.modified_index.expect("lease version");
        }
    }

    fn list_members(&self) -> BTreeMap<String, T> {
        let current_listing = self.etcd.get(MEMBERS, true, true, true).expect("List members");
        trace!("Listing: {:?}", current_listing);
        current_listing.node.nodes.expect("Node listing").into_iter()
                .filter_map(|n|
                    if let (Some(k), Some(v)) = (n.key, n.value) {
                        Some((k, json::decode(&v).expect("decode json"))) } else { None })
                .collect::<BTreeMap<String, T>>()
    }

    fn run_watch(&self) {
        let watch = self.start_watch();
        info!("Starting etcd watcher");
        let lease_key = self.lease_key.as_ref().expect("Should have created lease node");

        let current_listing = self.etcd.get(MEMBERS, true, true, true).expect("List members");
        debug!("Listing: {:?}", current_listing);
        let mut last_observed_index = current_listing.node.nodes.unwrap_or_else(|| Vec::new()).into_iter()
            .filter_map(|x| x.modified_index).max();

        let mut curr_members = BTreeMap::new();
        loop {
            trace!("Awaiting for {} from {:?}", MEMBERS, last_observed_index);
            let res = self.etcd.watch(MEMBERS, last_observed_index, true).expect("watch");
            trace!("Watch: {:?}", res);
            last_observed_index = res.node.modified_index.map(|x| x+1);

            let members = self.list_members();
            assert!(members.contains_key(lease_key), "I am ostensibly alive");

            let seq = self.verify_epoch(&members);

            trace!("Members: {:?}; seq: {:?}", members, seq);
            if curr_members != members {
                curr_members = members;
                info!("Membership change! {:?}", curr_members);
                (self.callback)(ConfigurationView {
                    this_node: lease_key.clone(),
                    members: curr_members.clone(),
                    epoch: seq.epoch,
                })
            }
        }

    }

    fn verify_epoch(&self, members: &BTreeMap<String, T>) -> ConfigSequencer {
        loop {
            let (seq_index, mut seq) : (u64, ConfigSequencer) = match self.etcd.get(SEQUENCER, false, false, false) {
                Ok(etcd::KeySpaceInfo {
                    node: etcd::keys::Node {
                        modified_index: Some(modified_index),
                        value: Some(value),
                        ..
                    } ,
                    ..
                }) => (modified_index, json::decode(&value).expect("decode epoch")),
                Ok(e) => panic!("Unexpected response reading {}: {:?}", SEQUENCER, e),
                Err(e) => panic!("Unexpected error reading {}: {:?}", SEQUENCER, e),
            };
            trace!("Sequencer: {}/{:?}", seq_index, seq);
            let current_keys = members.keys().cloned().collect::<Vec<_>>();
            if current_keys != seq.keys {
                debug!("Stale! {:?}", seq);

                seq.keys = current_keys;
                seq.epoch += 1;
                match self.etcd.compare_and_swap(SEQUENCER, &json::encode(&seq).expect("encode epoch"), None, None, Some(seq_index)) {
                        Ok(res) => {
                            debug!("verify_epoch: Updated seq: {}: {:?}: ", SEQUENCER, seq);
                            return seq;
                        },
                        Err(etcd::Error::Etcd (ref e)) if e.error_code == COMPARE_FAILED => debug!("verify_epoch: Raced out; retry"),
                        Err(e) => panic!("Unexpected error creating {}: {:?}", SEQUENCER, e),
                }
            } else {
                debug!("verify_epoch: Fresh! {:?}", seq);
                return seq;
            }
        }
    }

    fn start_watch<'a>(&'a self) -> DeathWatch<'a> {
        warn!("start_watch: New death watcher {:p}: prev: {:?}", &self.watch_count, self.watch_count.load(Ordering::Relaxed));
        DeathWatch::new(&self.watch_count)
    }

    fn is_running(&self) -> bool {
        self.watch_count.load(Ordering::Relaxed) > 0
    }
}

impl<T: Clone> ConfigurationView<T> {
    fn head_key(&self) -> Option<&str> {
        self.members.keys().next().map(|s| &**s)
    }

    pub fn should_listen_for_clients(&self) -> bool {
        self.head_key() == Some(&self.this_node)
    }
    pub fn should_listen_for_upstream(&self) -> bool {
        self.head_key() != Some(&self.this_node)
    }
    pub fn should_connect_downstream(&self) -> Option<T> {
        let next = self.members.iter().filter_map(|(k, v)| if *k > self.this_node { Some((k, v)) } else { None }).next();
        next.map(|(_, val)| val.clone())
    }
}
<|endoftext|>