<commit_before>#[cfg(target_os = "windows")]
use INVALID_UTF8;
use std::ffi::OsStr;
#[cfg(not(target_os = "windows"))]
use std::os::unix::ffi::OsStrExt;

#[cfg(target_os = "windows")]
trait OsStrExt3 {
    fn from_bytes(b: &[u8]) -> &Self;
    fn as_bytes(&self) -> &[u8];
}

#[doc(hidden)]
pub trait OsStrExt2 {
    fn starts_with(&self, s: &[u8]) -> bool;
    fn split_at_byte(&self, b: u8) -> (&OsStr, &OsStr);
    fn split_at(&self, i: usize) -> (&OsStr, &OsStr);
    fn trim_left_matches(&self, b: u8) -> &OsStr;
    fn len_(&self) -> usize;
    fn contains_byte(&self, b: u8) -> bool;
    fn is_empty_(&self) -> bool;
    fn split(&self, b: u8) -> OsSplit;
}

#[cfg(target_os = "windows")]
impl OsStrExt3 for OsStr {
    fn from_bytes(b: &[u8]) -> &Self {
        use std::mem;
        unsafe { mem::transmute(b) }
    }
    fn as_bytes(&self) -> &[u8] { self.to_str().map(|s| s.as_bytes()).expect(INVALID_UTF8) }
}

impl OsStrExt2 for OsStr {
    fn starts_with(&self, s: &[u8]) -> bool { self.as_bytes().starts_with(s) }

    fn is_empty_(&self) -> bool { self.as_bytes().is_empty() }

    fn contains_byte(&self, byte: u8) -> bool {
        for b in self.as_bytes() {
            if b == &byte {
                return true;
            }
        }
        false
    }

    fn split_at_byte(&self, byte: u8) -> (&OsStr, &OsStr) {
        for (i, b) in self.as_bytes().iter().enumerate() {
            if b == &byte {
                return (OsStr::from_bytes(&self.as_bytes()[..i]),
                        OsStr::from_bytes(&self.as_bytes()[i + 1..]));
            }
        }
        (&*self, OsStr::from_bytes(&self.as_bytes()[self.len_()..self.len_()]))
    }

    fn trim_left_matches(&self, byte: u8) -> &OsStr {
        for (i, b) in self.as_bytes().iter().enumerate() {
            if b != &byte {
                return OsStr::from_bytes(&self.as_bytes()[i..]);
            }
        }
        &*self
    }

    fn split_at(&self, i: usize) -> (&OsStr, &OsStr) {
        (OsStr::from_bytes(&self.as_bytes()[..i]), OsStr::from_bytes(&self.as_bytes()[i..]))
    }

    fn len_(&self) -> usize { self.as_bytes().len() }

    fn split(&self, b: u8) -> OsSplit {
        OsSplit {
            sep: b,
            val: self.as_bytes(),
            pos: 0,
        }
    }
}

#[doc(hidden)]
#[derive(Clone, Debug)]
pub struct OsSplit<'a> {
    sep: u8,
    val: &'a [u8],
    pos: usize,
}

impl<'a> Iterator for OsSplit<'a> {
    type Item = &'a OsStr;

    fn next(&mut self) -> Option<&'a OsStr> {
        debugln!("OsSplit::next: self={:?}", self);
        if self.pos == self.val.len() {
            return None;
        }
        let start = self.pos;
        for b in &self.val[start..] {
            self.pos += 1;
            if *b == self.sep {
                return Some(OsStr::from_bytes(&self.val[start..self.pos - 1]));
            }
        }
        Some(OsStr::from_bytes(&self.val[start..]))
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        let mut count = 0;
        for b in &self.val[self.pos..] {
            if *b == self.sep {
                count += 1;
            }
        }
        if count > 0 {
            return (count, Some(count));
        }
        (0, None)
    }
}

impl<'a> DoubleEndedIterator for OsSplit<'a> {
    fn next_back(&mut self) -> Option<&'a OsStr> {
        if self.pos == 0 {
            return None;
        }
        let start = self.pos;
        for b in self.val[..self.pos].iter().rev() {
            self.pos -= 1;
            if *b == self.sep {
                return Some(OsStr::from_bytes(&self.val[self.pos + 1..start]));
            }
        }
        Some(OsStr::from_bytes(&self.val[..start]))
    }
}
<commit_msg>fix: fixes a trait that's marked private accidentlly, but should be crate internal public<commit_after>#[cfg(target_os = "windows")]
use INVALID_UTF8;
use std::ffi::OsStr;
#[cfg(not(target_os = "windows"))]
use std::os::unix::ffi::OsStrExt;

#[cfg(target_os = "windows")]
pub trait OsStrExt3 {
    fn from_bytes(b: &[u8]) -> &Self;
    fn as_bytes(&self) -> &[u8];
}

#[doc(hidden)]
pub trait OsStrExt2 {
    fn starts_with(&self, s: &[u8]) -> bool;
    fn split_at_byte(&self, b: u8) -> (&OsStr, &OsStr);
    fn split_at(&self, i: usize) -> (&OsStr, &OsStr);
    fn trim_left_matches(&self, b: u8) -> &OsStr;
    fn len_(&self) -> usize;
    fn contains_byte(&self, b: u8) -> bool;
    fn is_empty_(&self) -> bool;
    fn split(&self, b: u8) -> OsSplit;
}

#[cfg(target_os = "windows")]
impl OsStrExt3 for OsStr {
    fn from_bytes(b: &[u8]) -> &Self {
        use std::mem;
        unsafe { mem::transmute(b) }
    }
    fn as_bytes(&self) -> &[u8] { self.to_str().map(|s| s.as_bytes()).expect(INVALID_UTF8) }
}

impl OsStrExt2 for OsStr {
    fn starts_with(&self, s: &[u8]) -> bool { self.as_bytes().starts_with(s) }

    fn is_empty_(&self) -> bool { self.as_bytes().is_empty() }

    fn contains_byte(&self, byte: u8) -> bool {
        for b in self.as_bytes() {
            if b == &byte {
                return true;
            }
        }
        false
    }

    fn split_at_byte(&self, byte: u8) -> (&OsStr, &OsStr) {
        for (i, b) in self.as_bytes().iter().enumerate() {
            if b == &byte {
                return (OsStr::from_bytes(&self.as_bytes()[..i]),
                        OsStr::from_bytes(&self.as_bytes()[i + 1..]));
            }
        }
        (&*self, OsStr::from_bytes(&self.as_bytes()[self.len_()..self.len_()]))
    }

    fn trim_left_matches(&self, byte: u8) -> &OsStr {
        for (i, b) in self.as_bytes().iter().enumerate() {
            if b != &byte {
                return OsStr::from_bytes(&self.as_bytes()[i..]);
            }
        }
        &*self
    }

    fn split_at(&self, i: usize) -> (&OsStr, &OsStr) {
        (OsStr::from_bytes(&self.as_bytes()[..i]), OsStr::from_bytes(&self.as_bytes()[i..]))
    }

    fn len_(&self) -> usize { self.as_bytes().len() }

    fn split(&self, b: u8) -> OsSplit {
        OsSplit {
            sep: b,
            val: self.as_bytes(),
            pos: 0,
        }
    }
}

#[doc(hidden)]
#[derive(Clone, Debug)]
pub struct OsSplit<'a> {
    sep: u8,
    val: &'a [u8],
    pos: usize,
}

impl<'a> Iterator for OsSplit<'a> {
    type Item = &'a OsStr;

    fn next(&mut self) -> Option<&'a OsStr> {
        debugln!("OsSplit::next: self={:?}", self);
        if self.pos == self.val.len() {
            return None;
        }
        let start = self.pos;
        for b in &self.val[start..] {
            self.pos += 1;
            if *b == self.sep {
                return Some(OsStr::from_bytes(&self.val[start..self.pos - 1]));
            }
        }
        Some(OsStr::from_bytes(&self.val[start..]))
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
        let mut count = 0;
        for b in &self.val[self.pos..] {
            if *b == self.sep {
                count += 1;
            }
        }
        if count > 0 {
            return (count, Some(count));
        }
        (0, None)
    }
}

impl<'a> DoubleEndedIterator for OsSplit<'a> {
    fn next_back(&mut self) -> Option<&'a OsStr> {
        if self.pos == 0 {
            return None;
        }
        let start = self.pos;
        for b in self.val[..self.pos].iter().rev() {
            self.pos -= 1;
            if *b == self.sep {
                return Some(OsStr::from_bytes(&self.val[self.pos + 1..start]));
            }
        }
        Some(OsStr::from_bytes(&self.val[..start]))
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2016 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! Payloads are self-contained projects that encapsulate a specific
//! feature or system function. Think of them as reusable chunks of
//! code that can be run across multiple hosts. Any time you have a
//! task that you want to repeat, it should probably go into a
//! payload.
//!
//! For example, a payload could handle installing a specific
//! package, such as Nginx. Or, you could create a payload that
//! configures iptables.
//!
//! # Examples
//!
//! ```no_run
//! # use inapi::{Host, Payload};
#![cfg_attr(feature = "local-run", doc = "# let mut host = Host::local(Some(\"nodes/mynode.json\")).unwrap();")]
#![cfg_attr(feature = "remote-run", doc = "# let mut host = Host::connect(\"nodes/mynode.json\").unwrap();")]
//! let payload = Payload::new("nginx::install").unwrap(); // format is "payload::executable"
//! payload.run(&mut host, None).unwrap();
//! ```

pub mod config;
pub mod ffi;

use czmq::{ZMsg, ZPoller, ZSock, SocketType, ZSys};
use error::{Error, Result};
use host::{Host,HostSendRecv};
use project::Language;
use self::config::Config;
use serde_json;
use std::env::{current_dir, set_current_dir};
use std::process;
use std::path::PathBuf;
use std::thread;
use zdaemon::ConfigFile;

/// Container for running a Payload.
pub struct Payload {
    /// Path to the payload directory.
    path: PathBuf,
    /// Name of the executable/source file to run.
    artifact: Option<String>,
    /// Language the payload is written in.
    language: Language,
}

impl Payload {
    /// Create a new Payload using the payload::artifact notation.
    /// This notation is simply "payload" + separator ("::") +
    /// "executable/source file". For example: "nginx::install".
    ///
    /// By default, payloads live in
    /// <project root>/payloads/<payload_name>. Thus the payload name
    /// "nginx" will resolve to <project root>/payloads/nginx/. You
    /// can also specify an absolute path to your payload, which will
    /// override the resolved path.
    ///
    /// ```no_run
    /// # use inapi::Payload;
    /// // Using standard payload/artifact notation...
    /// let payload = Payload::new("iptables::update").unwrap();
    ///
    /// // Using an absolute path...
    /// let payload = Payload::new("/mnt/intecture/payloads/iptables::update").unwrap();
    /// ```
    pub fn new(payload_artifact: &str) -> Result<Payload> {
        let mut parts: Vec<&str> = payload_artifact.split("::").collect();
        let payload = if parts.len() > 0 {
            parts.remove(0)
        } else {
            return Err(Error::Generic("Invalid payload string".into()));
        };

        let mut buf = PathBuf::from("payloads");
        buf.push(payload);

        buf.push("payload.json");
        let config = try!(Config::load(&buf));
        buf.pop();

        // Check dependencies
        if let Some(deps) = config.dependencies {
            try!(Self::check_deps(&deps));
        }

        Ok(Payload {
            path: buf,
            artifact: if parts.len() > 0 {
                Some(parts.remove(0).into())
            } else {
                None
            },
            language: config.language,
        })
    }

    /// Compile a payload's source code. This function is also called
    /// by Payload::run(), but is useful for precompiling payloads
    /// ahead of time to catch build errors early.
    ///
    /// Note that this is only useful for compiled languages. If this
    /// function is run on a payload that uses an interpreted
    /// language, it will safely be ignored.
    pub fn build(&self) -> Result<()> {
        let mut make_path = self.path.clone();
        make_path.push("Makefile");

        match self.language {
            Language::C | Language::Rust if make_path.exists() && make_path.is_file() => {
                let current_dir = try!(current_dir());
                try!(set_current_dir(&self.path));

                let output = try!(process::Command::new("make")
                                                   .stdout(process::Stdio::inherit())
                                                   .output());

                try!(set_current_dir(&current_dir));

                if !output.status.success() {
                    return Err(Error::BuildFailed(try!(String::from_utf8(output.stderr))).into());
                }
            },
            Language::Rust => {
                let manifest_path = format!("{}/Cargo.toml", self.path.to_str().unwrap());
                let output = try!(process::Command::new("cargo")
                                                   .args(&["build", "--release", "--manifest-path", &manifest_path])
                                                   .stdout(process::Stdio::inherit())
                                                   .output());
                if !output.status.success() {
                    return Err(Error::BuildFailed(try!(String::from_utf8(output.stderr))).into());
                }
            },
            _ => ()
        }

        Ok(())
    }

    /// Execute the payload's artifact.
    ///
    /// For compiled languages, the artifact will be executed
    /// directly. For interpreted languages, the artifact will be
    /// passed as an argument to the interpreter.
    ///
    /// ```no_run
    /// # use inapi::{Host, Payload};
    #[cfg_attr(feature = "local-run", doc = "# let mut host = Host::local(Some(\"nodes/mynode.json\")).unwrap();")]
    #[cfg_attr(feature = "remote-run", doc = "# let mut host = Host::connect(\"nodes/mynode.json\").unwrap();")]
    /// let payload = Payload::new("iptables::configure").unwrap();
    /// payload.run(&mut host, Some(vec![
    ///     "add_rule",
    ///     "..."
    /// ])).unwrap();
    /// ```
    pub fn run(&self, host: &mut Host, user_args: Option<Vec<&str>>) -> Result<()> {
        // XXX This ugly cloning is a wasteful solution. New ZMQ
        // socket types (SERVER & CLIENT) are threadsafe, which will
        // allow us to thread the ZMQ poller and avoid having to
        // clone half the world. When these are implemented, this
        // code should be changed.

        // Build payload to make sure it's up to date
        try!(self.build());

        let artifact_default = self.artifact.as_ref().map(|a| &**a).unwrap_or("main");
        let api_endpoint = format!("ipc://{}/{}_api.ipc", self.path.to_str().unwrap(), artifact_default);
        let mut api_pipe = ZSock::new(SocketType::DEALER);
        try!(api_pipe.bind(&api_endpoint));

        let file_endpoint = format!("ipc://{}/{}_file.ipc", self.path.to_str().unwrap(), artifact_default);
        let mut file_pipe = ZSock::new(SocketType::DEALER);
        try!(file_pipe.bind(&file_endpoint));

        let (mut parent, child) = try!(ZSys::create_pipe());
        let language = self.language.clone();
        let mut payload_path = self.path.clone();
        let artifact = self.artifact.clone();
        let user_args_c: Option<Vec<String>> = match user_args {
            Some(a) => Some(a.into_iter().map(|arg| String::from(arg)).collect()),
            None => None,
        };

        let handle = thread::spawn(move || {
            match language {
                Language::C => {
                    payload_path.push(artifact.as_ref().map(|a| &**a).unwrap_or("main"));

                    let mut args = vec![api_endpoint, file_endpoint];
                    if let Some(mut a) = user_args_c {
                        args.append(&mut a);
                    }

                    let output = try!(process::Command::new(payload_path.to_str().unwrap())
                                                       .args(&args)
                                                       .stdout(process::Stdio::inherit())
                                                       .output());

                    if !output.status.success() {
                        try!(child.signal(0));
                        return Err(Error::RunFailed(try!(String::from_utf8(output.stderr))).into());
                    }
                },
                Language::Php => {
                    payload_path.push("src");
                    payload_path.push(artifact.as_ref().map(|a| &**a).unwrap_or("main"));
                    if payload_path.extension().is_none() {
                        payload_path.set_extension("php");
                    }

                    let mut args = vec![payload_path.to_str().unwrap().into(), api_endpoint, file_endpoint];
                    if let Some(mut a) = user_args_c {
                        args.append(&mut a);
                    }

                    let output = try!(process::Command::new("php")
                                                       .args(&args)
                                                       .stdout(process::Stdio::inherit())
                                                       .output());

                    if !output.status.success() {
                        try!(child.signal(0));
                        return Err(Error::RunFailed(try!(String::from_utf8(output.stderr))).into());
                    }
                },
                Language::Rust => {
                    if let Some(a) = artifact {
                        payload_path.push("target/release");
                        payload_path.push(a);
                    } else {
                        let dirname = try!(try!(payload_path.file_stem().ok_or(Error::RunFailed("Invalid payload path".into())))
                                                            .to_str().ok_or(Error::RunFailed("Invalid payload path".into())))
                                                            .to_owned();
                        payload_path.push("target/release");
                        payload_path.push(&dirname);
                    }

                    let mut args = vec![
                        api_endpoint,
                        file_endpoint
                    ];

                    if let Some(mut a) = user_args_c {
                        args.append(&mut a);
                    }

                    let output = try!(process::Command::new(&payload_path)
                                                       .args(&args)
                                                       .stdout(process::Stdio::inherit())
                                                       .output());

                    if !output.status.success() {
                        try!(child.signal(0));
                        return Err(Error::RunFailed(try!(String::from_utf8(output.stderr))).into());
                    }
                }
            }

            try!(child.signal(0));
            Ok(())
        });

        // Send data to payload
        let json = try!(serde_json::to_string(host.data()));
        try!(api_pipe.send_str(&json));

        let mut poller = try!(ZPoller::new());
        try!(poller.add(&mut parent));
        try!(poller.add(&mut api_pipe));
        try!(poller.add(&mut file_pipe));

        loop {
            let sock: Option<ZSock> = poller.wait(None);
            if let Some(mut s) = sock {
                if s == api_pipe {
                    let req = try!(ZMsg::recv(&mut s));
                    try!(host.send(req));

                    let reply = try!(host.recv_raw());
                    try!(reply.send(&mut s));
                }
                else if s == file_pipe {
                    let req = try!(ZMsg::recv(&mut s));
                    try!(host.send_file(req));

                    let reply = try!(host.recv_file_raw());
                    try!(reply.send(&mut s));
                }
                else if s == parent {
                    break;
                } else {
                    unreachable!();
                }
            }

            if poller.terminated() {
                break;
            }
        }

        let cmd_result: Result<()> = try!(handle.join());
        try!(cmd_result);

        Ok(())
    }

    fn check_deps(payloads: &[String]) -> Result<()> {
        for payload in payloads {
            try!(Payload::new(payload));
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use czmq::{ZSock, SocketType};
    use host::Host;
    use project::Language;
    use std::fs;
    use std::io::Write;
    use std::path::PathBuf;
    use std::thread;
    use super::*;
    use super::config::Config;
    use tempdir::TempDir;
    use zdaemon::ConfigFile;

    #[test]
    fn test_new_nodeps() {
        let _ = ::_MOCK_ENV.init();

        let tempdir = TempDir::new("test_payload_new_nodeps").unwrap();
        let mut buf = tempdir.path().to_owned();

        create_cargo_proj(&mut buf);

        let conf = Config {
            author: "Dr. Hibbert".into(),
            repository: "https://github.com/dhibbz/hehehe.git".into(),
            language: Language::Rust,
            dependencies: Some(vec!["missing_payload".into()]),
        };

        buf.push("payload.json");
        conf.save(&buf).unwrap();
        buf.pop();

        assert!(Payload::new(buf.to_str().unwrap()).is_err());
    }

    pub fn test_build_c() {
        let tempdir = TempDir::new("test_payload_build_c").unwrap();
        let mut buf = tempdir.path().to_owned();

        buf.push("Makefile");
        let mut fh = fs::File::create(&buf).unwrap();
        fh.write_all(b"all:
\ttouch test").unwrap();
        buf.pop();

        let conf = Config {
            author: "Dr. Hibbert".into(),
            repository: "https://github.com/dhibbz/hehehe.git".into(),
            language: Language::C,
            dependencies: None,
        };

        buf.push("payload.json");
        conf.save(&buf).unwrap();
        buf.pop();

        let payload = Payload::new(buf.to_str().unwrap()).unwrap();
        payload.build().unwrap();

        buf.push("test");
        assert!(buf.exists());
    }

    pub fn test_run() {
        let tempdir = TempDir::new("test_payload_run").unwrap();
        let mut buf = tempdir.path().to_owned();

        create_cargo_proj(&mut buf);

        let conf = Config {
            author: "Dr. Hibbert".into(),
            repository: "https://github.com/dhibbz/hehehe.git".into(),
            language: Language::Rust,
            dependencies: None,
        };

        buf.push("payload.json");
        conf.save(&buf).unwrap();
        buf.pop();

        let payload_name = buf.into_os_string().into_string().unwrap();
        let payload_name_clone = payload_name.clone();

        let handle = thread::spawn(move || {
            let s = ZSock::new(SocketType::DEALER);
            s.connect(&format!("ipc://{}/main_api.ipc", payload_name_clone)).unwrap();
            s.recv_str().unwrap().unwrap();
        });

        let mut host = Host::test_new(None, None, None, None);
        let payload = Payload::new(&payload_name).unwrap();
        payload.run(&mut host, Some(vec!["abc"])).unwrap();

        handle.join().unwrap();
    }

    pub fn create_cargo_proj(buf: &mut PathBuf) {
        buf.push("payload/src");
        fs::create_dir_all(&buf).unwrap();

        buf.push("main.rs");
        let mut fh = fs::File::create(&buf).unwrap();
        fh.write_all(b"fn main() {}").unwrap();
        buf.pop();
        buf.pop();

        buf.push("Cargo.toml");
        let mut fh = fs::File::create(&buf).unwrap();
        fh.write_all(b"[package]
name = \"payload\"
version = \"0.1.0\"
authors = [\"Don Duck <quack@goosehat.rz>\"]").unwrap();
        buf.pop();
    }
}
<commit_msg>Add timeout and error handler to Payload api pipe to protect from hanging due to early exit of child proc<commit_after>// Copyright 2015-2016 Intecture Developers. See the COPYRIGHT file at the
// top-level directory of this distribution and at
// https://intecture.io/COPYRIGHT.
//
// Licensed under the Mozilla Public License 2.0 <LICENSE or
// https://www.tldrlegal.com/l/mpl-2.0>. This file may not be copied,
// modified, or distributed except according to those terms.

//! Payloads are self-contained projects that encapsulate a specific
//! feature or system function. Think of them as reusable chunks of
//! code that can be run across multiple hosts. Any time you have a
//! task that you want to repeat, it should probably go into a
//! payload.
//!
//! For example, a payload could handle installing a specific
//! package, such as Nginx. Or, you could create a payload that
//! configures iptables.
//!
//! # Examples
//!
//! ```no_run
//! # use inapi::{Host, Payload};
#![cfg_attr(feature = "local-run", doc = "# let mut host = Host::local(Some(\"nodes/mynode.json\")).unwrap();")]
#![cfg_attr(feature = "remote-run", doc = "# let mut host = Host::connect(\"nodes/mynode.json\").unwrap();")]
//! let payload = Payload::new("nginx::install").unwrap(); // format is "payload::executable"
//! payload.run(&mut host, None).unwrap();
//! ```

pub mod config;
pub mod ffi;

use czmq::{ZMsg, ZPoller, ZSock, SocketType, ZSys};
use error::{Error, Result};
use host::{Host,HostSendRecv};
use project::Language;
use self::config::Config;
use serde_json;
use std::env::{current_dir, set_current_dir};
use std::process;
use std::path::PathBuf;
use std::thread;
use zdaemon::ConfigFile;

/// Container for running a Payload.
pub struct Payload {
    /// Path to the payload directory.
    path: PathBuf,
    /// Name of the executable/source file to run.
    artifact: Option<String>,
    /// Language the payload is written in.
    language: Language,
}

impl Payload {
    /// Create a new Payload using the payload::artifact notation.
    /// This notation is simply "payload" + separator ("::") +
    /// "executable/source file". For example: "nginx::install".
    ///
    /// By default, payloads live in
    /// <project root>/payloads/<payload_name>. Thus the payload name
    /// "nginx" will resolve to <project root>/payloads/nginx/. You
    /// can also specify an absolute path to your payload, which will
    /// override the resolved path.
    ///
    /// ```no_run
    /// # use inapi::Payload;
    /// // Using standard payload/artifact notation...
    /// let payload = Payload::new("iptables::update").unwrap();
    ///
    /// // Using an absolute path...
    /// let payload = Payload::new("/mnt/intecture/payloads/iptables::update").unwrap();
    /// ```
    pub fn new(payload_artifact: &str) -> Result<Payload> {
        let mut parts: Vec<&str> = payload_artifact.split("::").collect();
        let payload = if parts.len() > 0 {
            parts.remove(0)
        } else {
            return Err(Error::Generic("Invalid payload string".into()));
        };

        let mut buf = PathBuf::from("payloads");
        buf.push(payload);

        buf.push("payload.json");
        let config = try!(Config::load(&buf));
        buf.pop();

        // Check dependencies
        if let Some(deps) = config.dependencies {
            try!(Self::check_deps(&deps));
        }

        Ok(Payload {
            path: buf,
            artifact: if parts.len() > 0 {
                Some(parts.remove(0).into())
            } else {
                None
            },
            language: config.language,
        })
    }

    /// Compile a payload's source code. This function is also called
    /// by Payload::run(), but is useful for precompiling payloads
    /// ahead of time to catch build errors early.
    ///
    /// Note that this is only useful for compiled languages. If this
    /// function is run on a payload that uses an interpreted
    /// language, it will safely be ignored.
    pub fn build(&self) -> Result<()> {
        let mut make_path = self.path.clone();
        make_path.push("Makefile");

        match self.language {
            Language::C | Language::Rust if make_path.exists() && make_path.is_file() => {
                let current_dir = try!(current_dir());
                try!(set_current_dir(&self.path));

                let output = try!(process::Command::new("make")
                                                   .stdout(process::Stdio::inherit())
                                                   .output());

                try!(set_current_dir(&current_dir));

                if !output.status.success() {
                    return Err(Error::BuildFailed(try!(String::from_utf8(output.stderr))).into());
                }
            },
            Language::Rust => {
                let manifest_path = format!("{}/Cargo.toml", self.path.to_str().unwrap());
                let output = try!(process::Command::new("cargo")
                                                   .args(&["build", "--release", "--manifest-path", &manifest_path])
                                                   .stdout(process::Stdio::inherit())
                                                   .output());
                if !output.status.success() {
                    return Err(Error::BuildFailed(try!(String::from_utf8(output.stderr))).into());
                }
            },
            _ => ()
        }

        Ok(())
    }

    /// Execute the payload's artifact.
    ///
    /// For compiled languages, the artifact will be executed
    /// directly. For interpreted languages, the artifact will be
    /// passed as an argument to the interpreter.
    ///
    /// ```no_run
    /// # use inapi::{Host, Payload};
    #[cfg_attr(feature = "local-run", doc = "# let mut host = Host::local(Some(\"nodes/mynode.json\")).unwrap();")]
    #[cfg_attr(feature = "remote-run", doc = "# let mut host = Host::connect(\"nodes/mynode.json\").unwrap();")]
    /// let payload = Payload::new("iptables::configure").unwrap();
    /// payload.run(&mut host, Some(vec![
    ///     "add_rule",
    ///     "..."
    /// ])).unwrap();
    /// ```
    pub fn run(&self, host: &mut Host, user_args: Option<Vec<&str>>) -> Result<()> {
        // XXX This ugly cloning is a wasteful solution. New ZMQ
        // socket types (SERVER & CLIENT) are threadsafe, which will
        // allow us to thread the ZMQ poller and avoid having to
        // clone half the world. When these are implemented, this
        // code should be changed.

        // Build payload to make sure it's up to date
        try!(self.build());

        let artifact_default = self.artifact.as_ref().map(|a| &**a).unwrap_or("main");
        let api_endpoint = format!("ipc://{}/{}_api.ipc", self.path.to_str().unwrap(), artifact_default);
        let mut api_pipe = ZSock::new(SocketType::DEALER);
        api_pipe.set_sndtimeo(Some(1000));
        try!(api_pipe.bind(&api_endpoint));

        let file_endpoint = format!("ipc://{}/{}_file.ipc", self.path.to_str().unwrap(), artifact_default);
        let mut file_pipe = ZSock::new(SocketType::DEALER);
        try!(file_pipe.bind(&file_endpoint));

        let (mut parent, child) = try!(ZSys::create_pipe());

        let language = self.language.clone();
        let mut payload_path = self.path.clone();
        let artifact = self.artifact.clone();
        let user_args_c: Option<Vec<String>> = match user_args {
            Some(a) => Some(a.into_iter().map(|arg| String::from(arg)).collect()),
            None => None,
        };

        let handle = thread::spawn(move || {
            match language {
                Language::C => {
                    payload_path.push(artifact.as_ref().map(|a| &**a).unwrap_or("main"));

                    let mut args = vec![api_endpoint, file_endpoint];
                    if let Some(mut a) = user_args_c {
                        args.append(&mut a);
                    }

                    let output = try!(process::Command::new(payload_path.to_str().unwrap())
                                                       .args(&args)
                                                       .stdout(process::Stdio::inherit())
                                                       .output());

                    if !output.status.success() {
                        try!(child.signal(0));
                        return Err(Error::RunFailed(try!(String::from_utf8(output.stderr))).into());
                    }
                },
                Language::Php => {
                    payload_path.push("src");
                    payload_path.push(artifact.as_ref().map(|a| &**a).unwrap_or("main"));
                    if payload_path.extension().is_none() {
                        payload_path.set_extension("php");
                    }

                    let mut args = vec![payload_path.to_str().unwrap().into(), api_endpoint, file_endpoint];
                    if let Some(mut a) = user_args_c {
                        args.append(&mut a);
                    }

                    let output = try!(process::Command::new("php")
                                                       .args(&args)
                                                       .stdout(process::Stdio::inherit())
                                                       .output());

                    if !output.status.success() {
                        try!(child.signal(0));
                        return Err(Error::RunFailed(try!(String::from_utf8(output.stderr))).into());
                    }
                },
                Language::Rust => {
                    if let Some(a) = artifact {
                        payload_path.push("target/release");
                        payload_path.push(a);
                    } else {
                        let dirname = try!(try!(payload_path.file_stem().ok_or(Error::RunFailed("Invalid payload path".into())))
                                                            .to_str().ok_or(Error::RunFailed("Invalid payload path".into())))
                                                            .to_owned();
                        payload_path.push("target/release");
                        payload_path.push(&dirname);
                    }

                    let mut args = vec![
                        api_endpoint,
                        file_endpoint
                    ];

                    if let Some(mut a) = user_args_c {
                        args.append(&mut a);
                    }

                    let output = try!(process::Command::new(&payload_path)
                                                       .args(&args)
                                                       .stdout(process::Stdio::inherit())
                                                       .output());

                    if !output.status.success() {
                        try!(child.signal(0));
                        return Err(Error::RunFailed(try!(String::from_utf8(output.stderr))).into());
                    }
                }
            }

            try!(child.signal(0));
            Ok(())
        });

        // Send data to payload
        let json = try!(serde_json::to_string(host.data()));
        match api_pipe.send_str(&json) {
            Ok(_) => (),
            // If we can't send data to the child proc, something has
            // gone very wrong. In order not to lose the child error,
            // we'll attempt to receive it now, or return CZMQ error
            // otherwise.
            Err(e) => {
                parent.set_rcvtimeo(Some(500));
                if parent.wait().is_ok() {
                    let cmd_result: Result<()> = handle.join()?;
                    cmd_result?;
                }

                return Err(e.into());
            }
        }

        let mut poller = try!(ZPoller::new());
        try!(poller.add(&mut parent));
        try!(poller.add(&mut api_pipe));
        try!(poller.add(&mut file_pipe));

        loop {
            let sock: Option<ZSock> = poller.wait(None);
            if let Some(mut s) = sock {
                if s == api_pipe {
                    let req = try!(ZMsg::recv(&mut s));
                    try!(host.send(req));

                    let reply = try!(host.recv_raw());
                    try!(reply.send(&mut s));
                }
                else if s == file_pipe {
                    let req = try!(ZMsg::recv(&mut s));
                    try!(host.send_file(req));

                    let reply = try!(host.recv_file_raw());
                    try!(reply.send(&mut s));
                }
                else if s == parent {
                    break;
                } else {
                    unreachable!();
                }
            }

            if poller.terminated() {
                break;
            }
        }

        let cmd_result: Result<()> = try!(handle.join());
        try!(cmd_result);

        Ok(())
    }

    fn check_deps(payloads: &[String]) -> Result<()> {
        for payload in payloads {
            try!(Payload::new(payload));
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use czmq::{ZSock, SocketType};
    use host::Host;
    use project::Language;
    use std::fs;
    use std::io::Write;
    use std::path::PathBuf;
    use std::thread;
    use super::*;
    use super::config::Config;
    use tempdir::TempDir;
    use zdaemon::ConfigFile;

    #[test]
    fn test_new_nodeps() {
        let _ = ::_MOCK_ENV.init();

        let tempdir = TempDir::new("test_payload_new_nodeps").unwrap();
        let mut buf = tempdir.path().to_owned();

        create_cargo_proj(&mut buf);

        let conf = Config {
            author: "Dr. Hibbert".into(),
            repository: "https://github.com/dhibbz/hehehe.git".into(),
            language: Language::Rust,
            dependencies: Some(vec!["missing_payload".into()]),
        };

        buf.push("payload.json");
        conf.save(&buf).unwrap();
        buf.pop();

        assert!(Payload::new(buf.to_str().unwrap()).is_err());
    }

    pub fn test_build_c() {
        let tempdir = TempDir::new("test_payload_build_c").unwrap();
        let mut buf = tempdir.path().to_owned();

        buf.push("Makefile");
        let mut fh = fs::File::create(&buf).unwrap();
        fh.write_all(b"all:
\ttouch test").unwrap();
        buf.pop();

        let conf = Config {
            author: "Dr. Hibbert".into(),
            repository: "https://github.com/dhibbz/hehehe.git".into(),
            language: Language::C,
            dependencies: None,
        };

        buf.push("payload.json");
        conf.save(&buf).unwrap();
        buf.pop();

        let payload = Payload::new(buf.to_str().unwrap()).unwrap();
        payload.build().unwrap();

        buf.push("test");
        assert!(buf.exists());
    }

    pub fn test_run() {
        let tempdir = TempDir::new("test_payload_run").unwrap();
        let mut buf = tempdir.path().to_owned();

        create_cargo_proj(&mut buf);

        let conf = Config {
            author: "Dr. Hibbert".into(),
            repository: "https://github.com/dhibbz/hehehe.git".into(),
            language: Language::Rust,
            dependencies: None,
        };

        buf.push("payload.json");
        conf.save(&buf).unwrap();
        buf.pop();

        let payload_name = buf.into_os_string().into_string().unwrap();
        let payload_name_clone = payload_name.clone();

        let handle = thread::spawn(move || {
            let s = ZSock::new(SocketType::DEALER);
            s.connect(&format!("ipc://{}/main_api.ipc", payload_name_clone)).unwrap();
            s.recv_str().unwrap().unwrap();
        });

        let mut host = Host::test_new(None, None, None, None);
        let payload = Payload::new(&payload_name).unwrap();
        payload.run(&mut host, Some(vec!["abc"])).unwrap();

        handle.join().unwrap();
    }

    pub fn create_cargo_proj(buf: &mut PathBuf) {
        buf.push("payload/src");
        fs::create_dir_all(&buf).unwrap();

        buf.push("main.rs");
        let mut fh = fs::File::create(&buf).unwrap();
        fh.write_all(b"fn main() {}").unwrap();
        buf.pop();
        buf.pop();

        buf.push("Cargo.toml");
        let mut fh = fs::File::create(&buf).unwrap();
        fh.write_all(b"[package]
name = \"payload\"
version = \"0.1.0\"
authors = [\"Don Duck <quack@goosehat.rz>\"]").unwrap();
        buf.pop();
    }
}
<|endoftext|>


<commit_before>use std::io::Read;
use std::io::BufReader;

use color::{ColorType};
use image::{DecodingResult, ImageDecoder, ImageResult, ImageError};
extern crate byteorder;
use self::byteorder::{BigEndian, ByteOrder};

/// PPM decoder
pub struct PPMDecoder<R> {
    reader: BufReader<R>,
    width: u32,
    height: u32,
    maxwhite: u32,
}

impl<R: Read> PPMDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(r: R) -> ImageResult<PPMDecoder<R>> {
        let mut buf = BufReader::new(r);
        try!(PPMDecoder::read_next_string(&mut buf)); // Skip P6
        let width = try!(PPMDecoder::read_next_u32(&mut buf));
        let height = try!(PPMDecoder::read_next_u32(&mut buf));
        let maxwhite = try!(PPMDecoder::read_next_u32(&mut buf));
        Ok(PPMDecoder {
            reader: buf,
            width: width,
            height: height,
            maxwhite: maxwhite,
        })
    }

  fn read_next_string(reader: &mut BufReader<R>) -> ImageResult<String> {
      let mut bytes = Vec::new();
      loop {
          let mut buffer: [u8;1] = [0];
          try!(reader.read_exact(&mut buffer));
          match buffer[0] {
              b'\n' | b' ' | b'\r' | b'\t' => {
                  if bytes.len() > 0 {
                    break // We're done as we already have some content
                  }
              },
              _ => {
                  bytes.push(buffer[0]);
              },
          }
      }

      match String::from_utf8(bytes) {
          Ok(s) => Ok(s),
          Err(_) => Err(ImageError::FormatError("Couldn't read preamble".to_string())),
      }
  }

  fn read_next_u32(reader: &mut BufReader<R>) -> ImageResult<u32> {
      let s = try!(PPMDecoder::read_next_string(reader));
      match s.parse::<u32>() {
          Ok(v) => Ok(v),
          Err(_) => Err(ImageError::FormatError("Couldn't read preamble".to_string())),
      }
  }
}

impl<R: Read> ImageDecoder for PPMDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        Ok((self.width, self.height))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match self.bytewidth() {
            1 => Ok(ColorType::RGB(8)),
            2 => Ok(ColorType::RGB(16)),
            _ => Err(ImageError::FormatError("Don't know how to decode PPM with more than 16 bits".to_string())),
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        Ok((self.width*3*self.bytewidth()) as usize)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        let mut data = vec![0 as u8; (self.width*self.height*3*self.bytewidth()) as usize];
        match self.reader.read_exact(&mut data) {
            Ok(_) => {},
            Err(e) => return Err(ImageError::IoError(e)),
        };

        if self.bytewidth() == 1 {
            Ok(DecodingResult::U8(data))
        } else {
            let mut out = vec![0 as u16; (self.width*self.height*3) as usize];
            for (o, i) in out.chunks_mut(1).zip(data.chunks(2)) {
                o[0] = BigEndian::read_u16(i);
            }
            Ok(DecodingResult::U16(out))
        }
    }
}

impl<R: Read> PPMDecoder<R> {
    fn bytewidth(&self) -> u32 {
        if self.maxwhite < 256 { 1 } else { 2 }
    }
}
<commit_msg>Simplify the PPM header decode with an iterator<commit_after>use std::io::Read;
use std::io::BufReader;

use color::{ColorType};
use image::{DecodingResult, ImageDecoder, ImageResult, ImageError};
extern crate byteorder;
use self::byteorder::{BigEndian, ByteOrder};

/// PPM decoder
pub struct PPMDecoder<R> {
    reader: BufReader<R>,
    width: u32,
    height: u32,
    maxwhite: u32,
}

impl<R: Read> PPMDecoder<R> {
    /// Create a new decoder that decodes from the stream ```r```
    pub fn new(r: R) -> ImageResult<PPMDecoder<R>> {
        let mut buf = BufReader::new(r);
        try!(PPMDecoder::read_next_string(&mut buf)); // Skip P6
        let width = try!(PPMDecoder::read_next_u32(&mut buf));
        let height = try!(PPMDecoder::read_next_u32(&mut buf));
        let maxwhite = try!(PPMDecoder::read_next_u32(&mut buf));
        Ok(PPMDecoder {
            reader: buf,
            width: width,
            height: height,
            maxwhite: maxwhite,
        })
    }

  fn read_next_string(reader: &mut BufReader<R>) -> ImageResult<String> {
      let mut bytes = Vec::new();
      for byte in reader.bytes() {
          match byte {
              Ok(b'\n') | Ok(b' ') | Ok(b'\r') | Ok(b'\t') => {
                  if bytes.len() > 0 {
                    break // We're done as we already have some content
                  }
              },
              Ok(byte) => {
                  bytes.push(byte);
              },
              Err(_) => break,
          }
      }

      match String::from_utf8(bytes) {
          Ok(s) => Ok(s),
          Err(_) => Err(ImageError::FormatError("Couldn't read preamble".to_string())),
      }
  }

  fn read_next_u32(reader: &mut BufReader<R>) -> ImageResult<u32> {
      let s = try!(PPMDecoder::read_next_string(reader));
      match s.parse::<u32>() {
          Ok(v) => Ok(v),
          Err(_) => Err(ImageError::FormatError("Couldn't read preamble".to_string())),
      }
  }
}

impl<R: Read> ImageDecoder for PPMDecoder<R> {
    fn dimensions(&mut self) -> ImageResult<(u32, u32)> {
        Ok((self.width, self.height))
    }

    fn colortype(&mut self) -> ImageResult<ColorType> {
        match self.bytewidth() {
            1 => Ok(ColorType::RGB(8)),
            2 => Ok(ColorType::RGB(16)),
            _ => Err(ImageError::FormatError("Don't know how to decode PPM with more than 16 bits".to_string())),
        }
    }

    fn row_len(&mut self) -> ImageResult<usize> {
        Ok((self.width*3*self.bytewidth()) as usize)
    }

    fn read_scanline(&mut self, _buf: &mut [u8]) -> ImageResult<u32> {
        unimplemented!();
    }

    fn read_image(&mut self) -> ImageResult<DecodingResult> {
        let mut data = vec![0 as u8; (self.width*self.height*3*self.bytewidth()) as usize];
        match self.reader.read_exact(&mut data) {
            Ok(_) => {},
            Err(e) => return Err(ImageError::IoError(e)),
        };

        if self.bytewidth() == 1 {
            Ok(DecodingResult::U8(data))
        } else {
            let mut out = vec![0 as u16; (self.width*self.height*3) as usize];
            for (o, i) in out.chunks_mut(1).zip(data.chunks(2)) {
                o[0] = BigEndian::read_u16(i);
            }
            Ok(DecodingResult::U16(out))
        }
    }
}

impl<R: Read> PPMDecoder<R> {
    fn bytewidth(&self) -> u32 {
        if self.maxwhite < 256 { 1 } else { 2 }
    }
}
<|endoftext|>


<commit_before>use std::path::Path;
use std::io::prelude::*;
use std::fs::File;
use yaml_rust::{YamlLoader, Yaml};
#[cfg(not(test))] use std::env::home_dir;

pub mod parser;
pub mod processor;

static MUXED_FOLDER: &'static str = "muxed";

pub fn main(project_name: String) -> Vec<Yaml> {
}

pub fn path_string(project_name: String) -> String {
    format!("{}/.{}/{}.yml", homedir_string(), &MUXED_FOLDER.to_string(), project_name)
}

fn read(config_str: String) -> String {
    let path = Path::new(&config_str);
    let mut s = String::new();
    let _ = File::open(path).expect("Config Read error").read_to_string(&mut s);

    return s
}

#[cfg(not(test))] fn homedir_string() -> String {
    match home_dir() {
        Some(dir) => format!("{}", dir.display()),
        None      => panic!("Impossible to get your home dir!")
    }
}

#[cfg(test)] fn homedir_string() -> String {
  String::from("/tmp")
}

#[test]
pub fn path_string_returns_muxed_inside_homedir() {
    let path = format!("{}", path_string("test".to_string()));
    let new  = format!("{}", Path::new("/tmp/.muxed/test.yml").display());
    assert_eq!(path, new)
}
<commit_msg>Rename function for symantics.<commit_after>use std::path::Path;
use std::io::prelude::*;
use std::fs::File;
use yaml_rust::{YamlLoader, Yaml};
#[cfg(not(test))] use std::env::home_dir;

pub mod parser;
pub mod processor;

static MUXED_FOLDER: &'static str = "muxed";

pub fn main(project_name: String) -> Vec<Yaml> {
    YamlLoader::load_from_str(&read(project_path(project_name))).unwrap()
}

pub fn project_path(project_name: String) -> String {
    format!("{}/.{}/{}.yml", homedir_string(), &MUXED_FOLDER.to_string(), project_name)
}

fn read(config_str: String) -> String {
    let path = Path::new(&config_str);
    let mut s = String::new();
    let _ = File::open(path).expect("Config Read error").read_to_string(&mut s);

    return s
}

#[cfg(not(test))] fn homedir_string() -> String {
    match home_dir() {
        Some(dir) => format!("{}", dir.display()),
        None      => panic!("Impossible to get your home dir!")
    }
}

#[cfg(test)] fn homedir_string() -> String {
  String::from("/tmp")
}

#[test]
pub fn project_path_returns_muxed_in_homedir() {
    let path = format!("{}", project_path("test".to_string()));
    let new  = format!("{}", Path::new("/tmp/.muxed/test.yml").display());
    assert_eq!(path, new)
}
<|endoftext|>


<commit_before>extern crate rand;

use self::rand::{thread_rng, Rng};

/**
* Chooses n samples by their weights. The greater their weights the more likely they get chosen.
*
* @invariant sum of weights must not overflow.
* @param samples The to be selected samples
* @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
* @param n Number of randomly chosen indices by weight.
* @return randomly selected samples by their weight
*/
pub fn random_choice<'a, T>(samples: &'a Vec<T>, weights: &Vec<f64>, n: usize) -> Vec<&'a T>{
    // TODO Check, if weight.len() > 0 and n > 0

    let sum:f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
    let spoke_gap: f64 = sum / n as f64;

    // next_f64() ∈ [0.0, 1.0)
    let spin = thread_rng().next_f64() * spoke_gap;

    let mut i: usize = 0;
    let mut accumulated_weights = weights[0];
    let mut choices: Vec<&T> = Vec::with_capacity(n);
    let mut current_spoke: f64 = spin;

    for _ in 0 .. n {
        while accumulated_weights < current_spoke {
            i += 1;
            accumulated_weights += weights[i];
        }
        choices.push(&samples[i]);
        current_spoke += spoke_gap;
    }

    choices
}

#[cfg(test)]
mod tests {
    #[test]
    fn basics() {
        let test_vec: Vec<f64> = vec![3.1, 1.2, 1.3, 2.0];
        //let sum:f64 = test_vec.iter().fold(0.0, |acc, &i| acc + i);

        //assert_eq!(6.0, sum);

        let choices = super::random_choice(&test_vec, &test_vec, 4 as usize);

        for choice in choices {
            print!("{}, ", choice);
        }
    }
}<commit_msg>in place variant added<commit_after>/// Implementation of Stochastic universal sampling
/// https://en.wikipedia.org/wiki/Stochastic_universal_sampling
/// Runtime: O(n)
/// Memory Usage: O(n) or O(1) (in place variant)

extern crate rand;

use self::rand::{thread_rng, Rng};

/// Chooses n samples by their weights. The greater their weights the more likely they get chosen.
///
/// @invariant sum of weights must not overflow.
/// @param samples The to be selected samples
/// @param weights Weights that get chosen by their weight/probability. One weight can be greater 1.
/// @param n Number of randomly chosen indices by weight.
/// @return randomly selected samples by their weights
pub fn random_choice<'a, T>(samples: &'a Vec<T>, weights: &Vec<f64>, n: usize) -> Vec<&'a T>{
    // TODO Check, if weight.len() > 0 and n > 0

    let sum:f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
    let spoke_gap: f64 = sum / n as f64;

    // next_f64() ∈ [0.0, 1.0)
    let spin = thread_rng().next_f64() * spoke_gap;

    let mut i: usize = 0;
    let mut accumulated_weights = weights[0];
    let mut choices: Vec<&T> = Vec::with_capacity(n);
    let mut current_spoke: f64 = spin;

    for _ in 0 .. n {
        while accumulated_weights < current_spoke {
            i += 1;
            accumulated_weights += weights[i];
        }
        choices.push(&samples[i]);
        current_spoke += spoke_gap;
    }

    choices
}

pub fn random_choice_in_place<T: Copy>(samples: &mut Vec<T>, weights: &Vec<f64>) {
    // TODO Check, if weight.len() > 0 and n > 0
    let sum:f64 = weights.iter().fold(0.0, |acc, &i| acc + i);
    let n: usize = weights.len();
    let spoke_gap: f64 = sum / n as f64;

    // next_f64() ∈ [0.0, 1.0)
    let spin = thread_rng().next_f64() * spoke_gap;

    let mut j: usize = 0;
    let mut accumulated_weights = weights[0];
    let mut current_spoke: f64 = spin;

    for i in 0 .. n {
        while accumulated_weights < current_spoke {
            j += 1;
            accumulated_weights += weights[j];
        }
        samples[i] = samples[j];
        current_spoke += spoke_gap;
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_random_choice() {
        let test_vec: Vec<f64> = vec![3.1, 1.2, 1.3, 2.0];
        //let sum:f64 = test_vec.iter().fold(0.0, |acc, &i| acc + i);

        //assert_eq!(6.0, sum);

        let choices = super::random_choice(&test_vec, &test_vec, 4 as usize);

        for choice in choices {
            print!("{}, ", choice);
        }
    }
}<|endoftext|>


<commit_before>//! Validate checks some basic safety conditions.

use super::{NormResult, NormError};
use super::norm_util::{self, Symbols};

use grammar::consts::*;
use grammar::parse_tree::*;
use grammar::repr as r;
use intern::{intern, InternedString};
use collections::{Multimap, set};
use util::Sep;

#[cfg(test)]
mod test;

pub fn validate(grammar: &Grammar) -> NormResult<()> {
    let match_token: Option<&MatchToken> =
        grammar.items
               .iter()
               .filter_map(|item| item.as_match_token())
               .next();

    let extern_token: Option<&ExternToken> =
        grammar.items
               .iter()
               .filter_map(|item| item.as_extern_token())
               .next();

    let validator = Validator {
        grammar: grammar,
        match_token: match_token,
        extern_token: extern_token,
    };

    validator.validate()
}

struct Validator<'grammar> {
    grammar: &'grammar Grammar,
    match_token: Option<&'grammar MatchToken>,
    extern_token: Option<&'grammar ExternToken>,
}

impl<'grammar> Validator<'grammar> {
    fn validate(&self) -> NormResult<()> {
        let allowed_names = vec![intern(LALR),
                                 intern(TABLE_DRIVEN),
                                 intern(RECURSIVE_ASCENT),
                                 intern(TEST_ALL)];
        for annotation in &self.grammar.annotations {
            if !allowed_names.contains(&annotation.id) {
                return_err!(annotation.id_span,
                            "unrecognized annotation `{}`",
                            annotation.id);
            }
        }

        for item in &self.grammar.items {
            match *item {
                GrammarItem::Use(..) => { }

                GrammarItem::MatchToken(ref data) => {
                    if data.span != self.match_token.unwrap().span {
                        return_err!(
                            data.span,
                            "multiple match definitions are not permitted");
                    }

                    // We may want to allow a limited extern to coexist with match in the future
                    if let Some(d) = self.extern_token {
                        return_err!(
                            d.span,
                            "extern and match definitions are mutually exclusive");
                    }

                    // Ensure that the catch all is final item of final block
                    for (contents_idx, match_contents) in data.contents.iter().enumerate() {
                        for (item_idx, item) in match_contents.items.iter().enumerate() {
                            if item.is_catch_all() && (contents_idx != &data.contents.len()-1 || item_idx != &match_contents.items.len()-1) {
                                return_err!(
                                    item.span(),
                                    "Catch all must be final item");
                            } else {
                                println!("ok");
                            }
                        }
                    }
                }

                GrammarItem::ExternToken(ref data) => {
                    if data.span != self.extern_token.unwrap().span {
                        return_err!(
                            data.span,
                            "multiple extern definitions are not permitted");
                    }

                    // We may want to allow a limited extern to coexist with match in the future
                    if let Some(d) = self.match_token {
                        return_err!(
                            d.span,
                            "match and extern definitions are mutually exclusive");
                    }

                    let allowed_names = vec![intern(LOCATION), intern(ERROR)];
                    let mut new_names = set();
                    for associated_type in &data.associated_types {
                        if !allowed_names.contains(&associated_type.type_name) {
                            return_err!(
                                associated_type.type_span,
                                "associated type `{}` not recognized, \
                                 try one of the following: {}",
                                associated_type.type_name,
                                Sep(", ", &allowed_names));
                        } else if !new_names.insert(associated_type.type_name) {
                            return_err!(
                                associated_type.type_span,
                                "associated type `{}` already specified",
                                associated_type.type_name);
                        }
                    }
                }
                GrammarItem::Nonterminal(ref data) => {
                    let inline_annotation = intern(INLINE);
                    let known_annotations = vec![inline_annotation];
                    let mut found_annotations = set();
                    for annotation in &data.annotations {
                        if !known_annotations.contains(&annotation.id) {
                            return_err!(annotation.id_span,
                                        "unrecognized annotation `{}`",
                                        annotation.id);
                        } else if !found_annotations.insert(annotation.id) {
                            return_err!(annotation.id_span,
                                        "duplicate annotation `{}`",
                                        annotation.id);
                        } else if annotation.id == inline_annotation && data.public {
                            return_err!(annotation.id_span,
                                        "public items cannot be marked #[inline]");
                        }
                    }

                    for alternative in &data.alternatives {
                        try!(self.validate_alternative(alternative));
                    }
                }
                GrammarItem::InternToken(..) => { }
            }
        }
        Ok(())
    }

    fn validate_alternative(&self,
                            alternative: &Alternative)
                            -> NormResult<()> {
        try!(self.validate_expr(&alternative.expr));

        match norm_util::analyze_expr(&alternative.expr) {
            Symbols::Named(syms) => {
                if alternative.action.is_none() {
                    let sym =
                        syms.iter()
                            .map(|&(_, _, sym)| sym)
                            .next()
                            .unwrap();
                    return_err!(
                        sym.span,
                        "named symbols (like `{}`) require a custom action",
                        sym);
                }
            }
            Symbols::Anon(syms) => {
                let empty_string = "".to_string();
                let action = {
                    match alternative.action {
                        Some(ActionKind::User(ref action)) => action,
                        Some(ActionKind::Fallible(ref action)) => action,
                        _ => &empty_string
                    }
                };
                if norm_util::check_funky_expression(action).is_in_curly_brackets() {
                    let sym =
                        syms.iter()
                            .map(|&(_, sym)| sym)
                            .next()
                            .unwrap();
                    return_err!(
                        sym.span,
                        "the `<>` expression requires to explicitly assign fields' names to values");
                }
            }
        }

        Ok(())
    }

    fn validate_expr(&self,
                     expr: &ExprSymbol)
                     -> NormResult<()> {
        for symbol in &expr.symbols {
            try!(self.validate_symbol(symbol));
        }

        let chosen: Vec<&Symbol> =
            expr.symbols.iter()
                        .filter(|sym| match sym.kind {
                            SymbolKind::Choose(_) => true,
                            _ => false,
                        })
                        .collect();

        let named: Multimap<InternedString, Vec<&Symbol>> =
            expr.symbols.iter()
                        .filter_map(|sym| match sym.kind {
                            SymbolKind::Name(nt, _) => Some((nt, sym)),
                            _ => None,
                        })
                        .collect();

        if !chosen.is_empty() && !named.is_empty() {
            return_err!(chosen[0].span,
                        "anonymous symbols like this one cannot be combined with \
                         named symbols like `{}`",
                        named.into_iter().next().unwrap().1[0]);
        }

        for (name, syms) in named.into_iter() {
            if syms.len() > 1{
                return_err!(syms[1].span,
                            "multiple symbols named `{}` are not permitted",
                            name);
            }
        }

        Ok(())
    }

    fn validate_symbol(&self,
                       symbol: &Symbol)
                       -> NormResult<()> {
        match symbol.kind {
            SymbolKind::Expr(ref expr) => {
                try!(self.validate_expr(expr));
            }
            SymbolKind::AmbiguousId(_) => {
                /* see resolve */
            }
            SymbolKind::Terminal(_) => {
                /* see postvalidate! */
            }
            SymbolKind::Nonterminal(_) => {
                /* see resolve */
            }
            SymbolKind::Error => {
                let mut algorithm = r::Algorithm::default();
                read_algorithm(&self.grammar.annotations, &mut algorithm);
                if algorithm.codegen == r::LrCodeGeneration::RecursiveAscent {
                    return_err!(symbol.span, "error recovery is not yet supported by recursive ascent parsers");
                }
            }
            SymbolKind::Macro(ref msym) => {
                debug_assert!(msym.args.len() > 0);
                for arg in &msym.args {
                    try!(self.validate_symbol(arg));
                }
            }
            SymbolKind::Repeat(ref repeat) => {
                try!(self.validate_symbol(&repeat.symbol));
            }
            SymbolKind::Choose(ref sym) | SymbolKind::Name(_, ref sym) => {
                try!(self.validate_symbol(sym));
            }
            SymbolKind::Lookahead | SymbolKind::Lookbehind => {
                // if using an internal tokenizer, lookahead/lookbehind are ok.
                if let Some(extern_token) = self.extern_token {
                    if extern_token.enum_token.is_some() {
                        // otherwise, the Location type must be specified.
                        let loc = intern(LOCATION);
                        if self.extern_token.unwrap().associated_type(loc).is_none() {
                            return_err!(
                                symbol.span,
                                "lookahead/lookbehind require you to declare the type of \
                                 a location; add a `type {} = ..` statement to the extern token \
                                 block",
                                LOCATION);
                        }
                    }
                }
            }
        }

        Ok(())
    }
}
<commit_msg>Better span for error message<commit_after>//! Validate checks some basic safety conditions.

use super::{NormResult, NormError};
use super::norm_util::{self, Symbols};

use grammar::consts::*;
use grammar::parse_tree::*;
use grammar::repr as r;
use intern::{intern, InternedString};
use collections::{Multimap, set};
use util::Sep;

#[cfg(test)]
mod test;

pub fn validate(grammar: &Grammar) -> NormResult<()> {
    let match_token: Option<&MatchToken> =
        grammar.items
               .iter()
               .filter_map(|item| item.as_match_token())
               .next();

    let extern_token: Option<&ExternToken> =
        grammar.items
               .iter()
               .filter_map(|item| item.as_extern_token())
               .next();

    let validator = Validator {
        grammar: grammar,
        match_token: match_token,
        extern_token: extern_token,
    };

    validator.validate()
}

struct Validator<'grammar> {
    grammar: &'grammar Grammar,
    match_token: Option<&'grammar MatchToken>,
    extern_token: Option<&'grammar ExternToken>,
}

impl<'grammar> Validator<'grammar> {
    fn validate(&self) -> NormResult<()> {
        let allowed_names = vec![intern(LALR),
                                 intern(TABLE_DRIVEN),
                                 intern(RECURSIVE_ASCENT),
                                 intern(TEST_ALL)];
        for annotation in &self.grammar.annotations {
            if !allowed_names.contains(&annotation.id) {
                return_err!(annotation.id_span,
                            "unrecognized annotation `{}`",
                            annotation.id);
            }
        }

        for item in &self.grammar.items {
            match *item {
                GrammarItem::Use(..) => { }

                GrammarItem::MatchToken(ref data) => {
                    if data.span != self.match_token.unwrap().span {
                        return_err!(
                            data.span,
                            "multiple match definitions are not permitted");
                    }

                    // We may want to allow a limited extern to coexist with match in the future
                    if let Some(d) = self.extern_token {
                        return_err!(
                            d.span,
                            "extern and match definitions are mutually exclusive");
                    }

                    // Ensure that the catch all is final item of final block
                    for (contents_idx, match_contents) in data.contents.iter().enumerate() {
                        for (item_idx, item) in match_contents.items.iter().enumerate() {
                            if item.is_catch_all() && (contents_idx != &data.contents.len()-1 || item_idx != &match_contents.items.len()-1) {
                                return_err!(
                                    item.span(),
                                    "Catch all must be final item");
                            } else {
                                println!("ok");
                            }
                        }
                    }
                }

                GrammarItem::ExternToken(ref data) => {
                    if data.span != self.extern_token.unwrap().span {
                        return_err!(
                            data.span,
                            "multiple extern definitions are not permitted");
                    }

                    // We may want to allow a limited extern to coexist with match in the future
                    if let Some(d) = self.match_token {
                        return_err!(
                            d.span,
                            "match and extern definitions are mutually exclusive");
                    }

                    let allowed_names = vec![intern(LOCATION), intern(ERROR)];
                    let mut new_names = set();
                    for associated_type in &data.associated_types {
                        if !allowed_names.contains(&associated_type.type_name) {
                            return_err!(
                                associated_type.type_span,
                                "associated type `{}` not recognized, \
                                 try one of the following: {}",
                                associated_type.type_name,
                                Sep(", ", &allowed_names));
                        } else if !new_names.insert(associated_type.type_name) {
                            return_err!(
                                associated_type.type_span,
                                "associated type `{}` already specified",
                                associated_type.type_name);
                        }
                    }
                }
                GrammarItem::Nonterminal(ref data) => {
                    let inline_annotation = intern(INLINE);
                    let known_annotations = vec![inline_annotation];
                    let mut found_annotations = set();
                    for annotation in &data.annotations {
                        if !known_annotations.contains(&annotation.id) {
                            return_err!(annotation.id_span,
                                        "unrecognized annotation `{}`",
                                        annotation.id);
                        } else if !found_annotations.insert(annotation.id) {
                            return_err!(annotation.id_span,
                                        "duplicate annotation `{}`",
                                        annotation.id);
                        } else if annotation.id == inline_annotation && data.public {
                            return_err!(annotation.id_span,
                                        "public items cannot be marked #[inline]");
                        }
                    }

                    for alternative in &data.alternatives {
                        try!(self.validate_alternative(alternative));
                    }
                }
                GrammarItem::InternToken(..) => { }
            }
        }
        Ok(())
    }

    fn validate_alternative(&self,
                            alternative: &Alternative)
                            -> NormResult<()> {
        try!(self.validate_expr(&alternative.expr));

        match norm_util::analyze_expr(&alternative.expr) {
            Symbols::Named(syms) => {
                if alternative.action.is_none() {
                    let sym =
                        syms.iter()
                            .map(|&(_, _, sym)| sym)
                            .next()
                            .unwrap();
                    return_err!(
                        sym.span,
                        "named symbols (like `{}`) require a custom action",
                        sym);
                }
            }
            Symbols::Anon(syms) => {
                let empty_string = "".to_string();
                let action = {
                    match alternative.action {
                        Some(ActionKind::User(ref action)) => action,
                        Some(ActionKind::Fallible(ref action)) => action,
                        _ => &empty_string
                    }
                };
                if norm_util::check_funky_expression(action).is_in_curly_brackets() {
                    return_err!(
                        alternative.span,
                        "the `<>` expression requires to explicitly assign fields' names to values");
                }
            }
        }

        Ok(())
    }

    fn validate_expr(&self,
                     expr: &ExprSymbol)
                     -> NormResult<()> {
        for symbol in &expr.symbols {
            try!(self.validate_symbol(symbol));
        }

        let chosen: Vec<&Symbol> =
            expr.symbols.iter()
                        .filter(|sym| match sym.kind {
                            SymbolKind::Choose(_) => true,
                            _ => false,
                        })
                        .collect();

        let named: Multimap<InternedString, Vec<&Symbol>> =
            expr.symbols.iter()
                        .filter_map(|sym| match sym.kind {
                            SymbolKind::Name(nt, _) => Some((nt, sym)),
                            _ => None,
                        })
                        .collect();

        if !chosen.is_empty() && !named.is_empty() {
            return_err!(chosen[0].span,
                        "anonymous symbols like this one cannot be combined with \
                         named symbols like `{}`",
                        named.into_iter().next().unwrap().1[0]);
        }

        for (name, syms) in named.into_iter() {
            if syms.len() > 1{
                return_err!(syms[1].span,
                            "multiple symbols named `{}` are not permitted",
                            name);
            }
        }

        Ok(())
    }

    fn validate_symbol(&self,
                       symbol: &Symbol)
                       -> NormResult<()> {
        match symbol.kind {
            SymbolKind::Expr(ref expr) => {
                try!(self.validate_expr(expr));
            }
            SymbolKind::AmbiguousId(_) => {
                /* see resolve */
            }
            SymbolKind::Terminal(_) => {
                /* see postvalidate! */
            }
            SymbolKind::Nonterminal(_) => {
                /* see resolve */
            }
            SymbolKind::Error => {
                let mut algorithm = r::Algorithm::default();
                read_algorithm(&self.grammar.annotations, &mut algorithm);
                if algorithm.codegen == r::LrCodeGeneration::RecursiveAscent {
                    return_err!(symbol.span, "error recovery is not yet supported by recursive ascent parsers");
                }
            }
            SymbolKind::Macro(ref msym) => {
                debug_assert!(msym.args.len() > 0);
                for arg in &msym.args {
                    try!(self.validate_symbol(arg));
                }
            }
            SymbolKind::Repeat(ref repeat) => {
                try!(self.validate_symbol(&repeat.symbol));
            }
            SymbolKind::Choose(ref sym) | SymbolKind::Name(_, ref sym) => {
                try!(self.validate_symbol(sym));
            }
            SymbolKind::Lookahead | SymbolKind::Lookbehind => {
                // if using an internal tokenizer, lookahead/lookbehind are ok.
                if let Some(extern_token) = self.extern_token {
                    if extern_token.enum_token.is_some() {
                        // otherwise, the Location type must be specified.
                        let loc = intern(LOCATION);
                        if self.extern_token.unwrap().associated_type(loc).is_none() {
                            return_err!(
                                symbol.span,
                                "lookahead/lookbehind require you to declare the type of \
                                 a location; add a `type {} = ..` statement to the extern token \
                                 block",
                                LOCATION);
                        }
                    }
                }
            }
        }

        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::{self, Display, Formatter};
use std::net::SocketAddr;
use std::boxed::FnBox;
use std::time::Instant;
use std::sync::{Arc, RwLock};

use mio::Token;
use futures::{Async, Future, Poll, Stream};
use futures::stream::{self, Once};
use grpc::{ChannelBuilder, Environment, WriteFlags};
use kvproto::raft_serverpb::SnapshotChunk;
use kvproto::raft_serverpb::RaftMessage;
use kvproto::tikvpb_grpc::TikvClient;

use raftstore::store::{SnapEntry, SnapKey, SnapManager, Snapshot};
use util::threadpool::{DefaultContext, ThreadPool, ThreadPoolBuilder};
use util::worker::Runnable;
use util::buf::PipeBuffer;
use util::collections::{HashMap, HashMapEntry as Entry};
use util::HandyRwLock;

use super::metrics::*;
use super::{Error, Result};
use super::transport::RaftStoreRouter;

pub type Callback = Box<FnBox(Result<()>) + Send>;

const DEFAULT_SENDER_POOL_SIZE: usize = 3;

/// `Task` that `Runner` can handle.
///
/// `Register` register a pending snapshot file with token;
/// `Write` write data to snapshot file;
/// `Close` save the snapshot file;
/// `Discard` discard all the unsaved changes made to snapshot file;
/// `SendTo` send the snapshot file to specified address.
pub enum Task {
    Register(Token, RaftMessage),
    Write(Token, PipeBuffer),
    Close(Token),
    Discard(Token),
    SendTo {
        addr: SocketAddr,
        msg: RaftMessage,
        cb: Callback,
    },
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            Task::Register(token, ref meta) => write!(f, "Register {:?} token: {:?}", meta, token),
            Task::Write(token, _) => write!(f, "Write snap for {:?}", token),
            Task::Close(token) => write!(f, "Close file {:?}", token),
            Task::Discard(token) => write!(f, "Discard file {:?}", token),
            Task::SendTo {
                ref addr, ref msg, ..
            } => write!(f, "SendTo Snap[to: {}, snap: {:?}]", addr, msg),
        }
    }
}

struct SnapChunk {
    snap: Arc<RwLock<Box<Snapshot>>>,
    remain_bytes: usize,
}

const SNAP_CHUNK_LEN: usize = 1024 * 1024;

impl Stream for SnapChunk {
    type Item = (SnapshotChunk, WriteFlags);
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Error> {
        let mut buf = match self.remain_bytes {
            0 => return Ok(Async::Ready(None)),
            n if n > SNAP_CHUNK_LEN => vec![0; SNAP_CHUNK_LEN],
            n => vec![0; n],
        };
        match self.snap.wl().read_exact(buf.as_mut_slice()) {
            Ok(_) => {
                self.remain_bytes -= buf.len();
                let mut chunk = SnapshotChunk::new();
                chunk.set_data(buf);
                Ok(Async::Ready(
                    Some((chunk, WriteFlags::default().buffer_hint(true))),
                ))
            }
            Err(e) => Err(box_err!("failed to read snapshot chunk: {}", e)),
        }
    }
}

/// Send the snapshot to specified address.
///
/// It will first send the normal raft snapshot message and then send the snapshot file.
fn send_snap(
    env: Arc<Environment>,
    mgr: SnapManager,
    addr: SocketAddr,
    msg: RaftMessage,
) -> Result<()> {
    assert!(msg.get_message().has_snapshot());
    let timer = Instant::now();

    let send_timer = SEND_SNAP_HISTOGRAM.start_coarse_timer();

    let key = {
        let snap = msg.get_message().get_snapshot();
        SnapKey::from_snap(snap)?
    };
    mgr.register(key.clone(), SnapEntry::Sending);
    defer!({
        mgr.deregister(&key, &SnapEntry::Sending);
    });
    let s = box_try!(mgr.get_snapshot_for_sending(&key));
    if !s.exists() {
        return Err(box_err!("missing snap file: {:?}", s.path()));
    }
    let total_size = s.total_size()?;

    // snapshot file has been validated when created, so no need to validate again.
    let s = Arc::new(RwLock::new(s));

    let chunks = {
        let snap_chunk = SnapChunk {
            snap: s.clone(),
            remain_bytes: total_size as usize,
        };
        let first: Once<(SnapshotChunk, _), Error> = stream::once({
            let mut chunk = SnapshotChunk::new();
            chunk.set_message(msg);
            Ok((chunk, WriteFlags::default().buffer_hint(true)))
        });
        first.chain(snap_chunk)
    };

    let channel = ChannelBuilder::new(env).connect(&format!("{}", addr));
    let client = TikvClient::new(channel);
    let (sink, receiver) = client.snapshot();
    let send = chunks.forward(sink);
    let res = send.and_then(|_| receiver.map_err(Error::from))
        .and_then(|_| {
            info!(
                "[region {}] sent snapshot {} [size: {}, dur: {:?}]",
                key.region_id,
                key,
                total_size,
                timer.elapsed()
            );
            s.wl().delete();
            Ok(())
        })
        .wait()
        .map_err(Error::from);

    send_timer.observe_duration();
    res
}

pub struct Runner<R: RaftStoreRouter + 'static> {
    env: Arc<Environment>,
    snap_mgr: SnapManager,
    files: HashMap<Token, (Box<Snapshot>, RaftMessage)>,
    pool: ThreadPool<DefaultContext>,
    raft_router: R,
}

impl<R: RaftStoreRouter + 'static> Runner<R> {
    pub fn new(env: Arc<Environment>, snap_mgr: SnapManager, r: R) -> Runner<R> {
        Runner {
            env: env,
            snap_mgr: snap_mgr,
            files: map![],
            pool: ThreadPoolBuilder::with_default_factory(thd_name!("snap sender"))
                .thread_count(DEFAULT_SENDER_POOL_SIZE)
                .build(),
            raft_router: r,
        }
    }
}

impl<R: RaftStoreRouter + 'static> Runnable<Task> for Runner<R> {
    fn run(&mut self, task: Task) {
        match task {
            Task::Register(token, meta) => {
                SNAP_TASK_COUNTER.with_label_values(&["register"]).inc();
                let mgr = self.snap_mgr.clone();
                let key = match SnapKey::from_snap(meta.get_message().get_snapshot()) {
                    Ok(k) => k,
                    Err(e) => {
                        error!("failed to create snap key for token {:?}: {:?}", token, e);
                        return;
                    }
                };
                match mgr.get_snapshot_for_receiving(
                    &key,
                    meta.get_message().get_snapshot().get_data(),
                ) {
                    Ok(snap) => {
                        if snap.exists() {
                            info!(
                                "snapshot file {} already exists, skip receiving.",
                                snap.path()
                            );
                            if let Err(e) = self.raft_router.send_raft_msg(meta) {
                                error!("send snapshot for key {} token {:?}: {:?}", key, token, e);
                            }
                            return;
                        }
                        debug!("begin to receive snap {:?}", meta);
                        mgr.register(key, SnapEntry::Receiving);
                        self.files.insert(token, (snap, meta));
                    }
                    Err(e) => {
                        error!(
                            "failed to create snapshot file for token {:?}: {:?}",
                            token,
                            e
                        );
                        return;
                    }
                }
            }
            Task::Write(token, mut data) => {
                SNAP_TASK_COUNTER.with_label_values(&["write"]).inc();
                match self.files.entry(token) {
                    Entry::Occupied(mut e) => {
                        if let Err(err) = data.write_all_to(&mut e.get_mut().0) {
                            error!(
                                "failed to write data to snapshot file {} for token {:?}: {:?}",
                                e.get_mut().0.path(),
                                token,
                                err
                            );
                            let (_, msg) = e.remove();
                            let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                            self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                        }
                    }
                    Entry::Vacant(_) => error!("invalid snap token {:?}", token),
                }
            }
            Task::Close(token) => {
                SNAP_TASK_COUNTER.with_label_values(&["close"]).inc();
                match self.files.remove(&token) {
                    Some((mut snap, msg)) => {
                        let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                        info!("saving snapshot to {}", snap.path());
                        defer!({
                            self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                        });
                        if let Err(e) = snap.save() {
                            error!(
                                "failed to save snapshot file {} for token {:?}: {:?}",
                                snap.path(),
                                token,
                                e
                            );
                            return;
                        }
                        if let Err(e) = self.raft_router.send_raft_msg(msg) {
                            error!("send snapshot for token {:?} err {:?}", token, e);
                        }
                    }
                    None => error!("invalid snap token {:?}", token),
                }
            }
            Task::Discard(token) => {
                SNAP_TASK_COUNTER.with_label_values(&["discard"]).inc();
                if let Some((_, msg)) = self.files.remove(&token) {
                    debug!("discard snapshot: {:?}", msg);
                    // because token is inserted, following can't panic.
                    let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                    self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                }
            }
            Task::SendTo { addr, msg, cb } => {
                SNAP_TASK_COUNTER.with_label_values(&["send"]).inc();
                let env = self.env.clone();
                let mgr = self.snap_mgr.clone();
                self.pool.execute(move |_| {
                    let res = send_snap(env, mgr, addr, msg);
                    if res.is_err() {
                        error!("failed to send snap to {}: {:?}", addr, res);
                    }
                    cb(res)
                });
            }
        }
    }
}
<commit_msg>snap: shorten write guard lifetime (#2490)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::fmt::{self, Display, Formatter};
use std::net::SocketAddr;
use std::boxed::FnBox;
use std::time::Instant;
use std::sync::{Arc, RwLock};

use mio::Token;
use futures::{Async, Future, Poll, Stream};
use futures::stream::{self, Once};
use grpc::{ChannelBuilder, Environment, WriteFlags};
use kvproto::raft_serverpb::SnapshotChunk;
use kvproto::raft_serverpb::RaftMessage;
use kvproto::tikvpb_grpc::TikvClient;

use raftstore::store::{SnapEntry, SnapKey, SnapManager, Snapshot};
use util::threadpool::{DefaultContext, ThreadPool, ThreadPoolBuilder};
use util::worker::Runnable;
use util::buf::PipeBuffer;
use util::collections::{HashMap, HashMapEntry as Entry};
use util::HandyRwLock;

use super::metrics::*;
use super::{Error, Result};
use super::transport::RaftStoreRouter;

pub type Callback = Box<FnBox(Result<()>) + Send>;

const DEFAULT_SENDER_POOL_SIZE: usize = 3;

/// `Task` that `Runner` can handle.
///
/// `Register` register a pending snapshot file with token;
/// `Write` write data to snapshot file;
/// `Close` save the snapshot file;
/// `Discard` discard all the unsaved changes made to snapshot file;
/// `SendTo` send the snapshot file to specified address.
pub enum Task {
    Register(Token, RaftMessage),
    Write(Token, PipeBuffer),
    Close(Token),
    Discard(Token),
    SendTo {
        addr: SocketAddr,
        msg: RaftMessage,
        cb: Callback,
    },
}

impl Display for Task {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        match *self {
            Task::Register(token, ref meta) => write!(f, "Register {:?} token: {:?}", meta, token),
            Task::Write(token, _) => write!(f, "Write snap for {:?}", token),
            Task::Close(token) => write!(f, "Close file {:?}", token),
            Task::Discard(token) => write!(f, "Discard file {:?}", token),
            Task::SendTo {
                ref addr, ref msg, ..
            } => write!(f, "SendTo Snap[to: {}, snap: {:?}]", addr, msg),
        }
    }
}

struct SnapChunk {
    snap: Arc<RwLock<Box<Snapshot>>>,
    remain_bytes: usize,
}

const SNAP_CHUNK_LEN: usize = 1024 * 1024;

impl Stream for SnapChunk {
    type Item = (SnapshotChunk, WriteFlags);
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Error> {
        let mut buf = match self.remain_bytes {
            0 => return Ok(Async::Ready(None)),
            n if n > SNAP_CHUNK_LEN => vec![0; SNAP_CHUNK_LEN],
            n => vec![0; n],
        };
        let result = self.snap.wl().read_exact(buf.as_mut_slice());
        match result {
            Ok(_) => {
                self.remain_bytes -= buf.len();
                let mut chunk = SnapshotChunk::new();
                chunk.set_data(buf);
                Ok(Async::Ready(
                    Some((chunk, WriteFlags::default().buffer_hint(true))),
                ))
            }
            Err(e) => Err(box_err!("failed to read snapshot chunk: {}", e)),
        }
    }
}

/// Send the snapshot to specified address.
///
/// It will first send the normal raft snapshot message and then send the snapshot file.
fn send_snap(
    env: Arc<Environment>,
    mgr: SnapManager,
    addr: SocketAddr,
    msg: RaftMessage,
) -> Result<()> {
    assert!(msg.get_message().has_snapshot());
    let timer = Instant::now();

    let send_timer = SEND_SNAP_HISTOGRAM.start_coarse_timer();

    let key = {
        let snap = msg.get_message().get_snapshot();
        SnapKey::from_snap(snap)?
    };
    mgr.register(key.clone(), SnapEntry::Sending);
    defer!({
        mgr.deregister(&key, &SnapEntry::Sending);
    });
    let s = box_try!(mgr.get_snapshot_for_sending(&key));
    if !s.exists() {
        return Err(box_err!("missing snap file: {:?}", s.path()));
    }
    let total_size = s.total_size()?;

    // snapshot file has been validated when created, so no need to validate again.
    let s = Arc::new(RwLock::new(s));

    let chunks = {
        let snap_chunk = SnapChunk {
            snap: s.clone(),
            remain_bytes: total_size as usize,
        };
        let first: Once<(SnapshotChunk, _), Error> = stream::once({
            let mut chunk = SnapshotChunk::new();
            chunk.set_message(msg);
            Ok((chunk, WriteFlags::default().buffer_hint(true)))
        });
        first.chain(snap_chunk)
    };

    let channel = ChannelBuilder::new(env).connect(&format!("{}", addr));
    let client = TikvClient::new(channel);
    let (sink, receiver) = client.snapshot();
    let send = chunks.forward(sink);
    let res = send.and_then(|_| receiver.map_err(Error::from))
        .and_then(|_| {
            info!(
                "[region {}] sent snapshot {} [size: {}, dur: {:?}]",
                key.region_id,
                key,
                total_size,
                timer.elapsed()
            );
            s.wl().delete();
            Ok(())
        })
        .wait()
        .map_err(Error::from);

    send_timer.observe_duration();
    res
}

pub struct Runner<R: RaftStoreRouter + 'static> {
    env: Arc<Environment>,
    snap_mgr: SnapManager,
    files: HashMap<Token, (Box<Snapshot>, RaftMessage)>,
    pool: ThreadPool<DefaultContext>,
    raft_router: R,
}

impl<R: RaftStoreRouter + 'static> Runner<R> {
    pub fn new(env: Arc<Environment>, snap_mgr: SnapManager, r: R) -> Runner<R> {
        Runner {
            env: env,
            snap_mgr: snap_mgr,
            files: map![],
            pool: ThreadPoolBuilder::with_default_factory(thd_name!("snap sender"))
                .thread_count(DEFAULT_SENDER_POOL_SIZE)
                .build(),
            raft_router: r,
        }
    }
}

impl<R: RaftStoreRouter + 'static> Runnable<Task> for Runner<R> {
    fn run(&mut self, task: Task) {
        match task {
            Task::Register(token, meta) => {
                SNAP_TASK_COUNTER.with_label_values(&["register"]).inc();
                let mgr = self.snap_mgr.clone();
                let key = match SnapKey::from_snap(meta.get_message().get_snapshot()) {
                    Ok(k) => k,
                    Err(e) => {
                        error!("failed to create snap key for token {:?}: {:?}", token, e);
                        return;
                    }
                };
                match mgr.get_snapshot_for_receiving(
                    &key,
                    meta.get_message().get_snapshot().get_data(),
                ) {
                    Ok(snap) => {
                        if snap.exists() {
                            info!(
                                "snapshot file {} already exists, skip receiving.",
                                snap.path()
                            );
                            if let Err(e) = self.raft_router.send_raft_msg(meta) {
                                error!("send snapshot for key {} token {:?}: {:?}", key, token, e);
                            }
                            return;
                        }
                        debug!("begin to receive snap {:?}", meta);
                        mgr.register(key, SnapEntry::Receiving);
                        self.files.insert(token, (snap, meta));
                    }
                    Err(e) => {
                        error!(
                            "failed to create snapshot file for token {:?}: {:?}",
                            token,
                            e
                        );
                        return;
                    }
                }
            }
            Task::Write(token, mut data) => {
                SNAP_TASK_COUNTER.with_label_values(&["write"]).inc();
                match self.files.entry(token) {
                    Entry::Occupied(mut e) => {
                        if let Err(err) = data.write_all_to(&mut e.get_mut().0) {
                            error!(
                                "failed to write data to snapshot file {} for token {:?}: {:?}",
                                e.get_mut().0.path(),
                                token,
                                err
                            );
                            let (_, msg) = e.remove();
                            let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                            self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                        }
                    }
                    Entry::Vacant(_) => error!("invalid snap token {:?}", token),
                }
            }
            Task::Close(token) => {
                SNAP_TASK_COUNTER.with_label_values(&["close"]).inc();
                match self.files.remove(&token) {
                    Some((mut snap, msg)) => {
                        let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                        info!("saving snapshot to {}", snap.path());
                        defer!({
                            self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                        });
                        if let Err(e) = snap.save() {
                            error!(
                                "failed to save snapshot file {} for token {:?}: {:?}",
                                snap.path(),
                                token,
                                e
                            );
                            return;
                        }
                        if let Err(e) = self.raft_router.send_raft_msg(msg) {
                            error!("send snapshot for token {:?} err {:?}", token, e);
                        }
                    }
                    None => error!("invalid snap token {:?}", token),
                }
            }
            Task::Discard(token) => {
                SNAP_TASK_COUNTER.with_label_values(&["discard"]).inc();
                if let Some((_, msg)) = self.files.remove(&token) {
                    debug!("discard snapshot: {:?}", msg);
                    // because token is inserted, following can't panic.
                    let key = SnapKey::from_snap(msg.get_message().get_snapshot()).unwrap();
                    self.snap_mgr.deregister(&key, &SnapEntry::Receiving);
                }
            }
            Task::SendTo { addr, msg, cb } => {
                SNAP_TASK_COUNTER.with_label_values(&["send"]).inc();
                let env = self.env.clone();
                let mgr = self.snap_mgr.clone();
                self.pool.execute(move |_| {
                    let res = send_snap(env, mgr, addr, msg);
                    if res.is_err() {
                        error!("failed to send snap to {}: {:?}", addr, res);
                    }
                    cb(res)
                });
            }
        }
    }
}
<|endoftext|>