<commit_before>extern crate rand;

use std::string::String;
use std::collections::HashMap;
use std::vec::Vec;
use std::fs::{File};
use std::env;
use std::path::{Path};
use std::io::{self, BufReader, BufRead};

use rand::{random, Closed01};
use rand::distributions::{IndependentSample, Range};

type Bag = HashMap<usize, usize>;

fn load_bags<P: AsRef<Path>>(path: P) -> io::Result<(Vec<usize>, Vec<Bag>)> {
    let mut bags = Vec::new();
    let mut labels = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        let mut iter = line.split_whitespace();
        let label = iter.next().unwrap().parse::<usize>().unwrap();
        let mut bag = Bag::new();
        for elm in iter {
            let mut iter = elm.split(':');
            let index = iter.next().unwrap().parse::<usize>().unwrap();
            let value = iter.next().unwrap().parse::<usize>().unwrap();
            bag.insert(index - 1, value);
        }
        labels.push(label);
        bags.push(bag);
    }
    Ok((labels, bags))
}

fn load_text_vocabulary<P: AsRef<Path>>(path: P) -> io::Result<Vec<String>> {
    let mut vocab = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        vocab.push(line);
    }
    Ok(vocab)
}

fn lda(num_topics: usize, dataset: Vec<Bag>, alpha: Vec<f64>, beta: Vec<f64>, num_samples: usize) {
    // Initialization
    let mut rng = rand::thread_rng();
    let topics = Range::new(0, num_topics);

    let vocab_size: usize = {
        let mut have_some = false;
        let mut max_index = 0;
        for bag in &dataset {
            for &index in bag.keys() {
                have_some = true;
                if index > max_index {
                    max_index = index;
                }
            }
        }
        if have_some {
            max_index + 1
        }
        else {
            0
        }
    };
    println!("K = {}", num_topics);
    println!("M = {}", &dataset.len());
    println!("V = {}", vocab_size);

    let mut w: Vec<Vec<usize>> = Vec::new();
    let mut z: Vec<Vec<usize>> = Vec::new();
    let mut theta: Vec<Vec<f64>> = Vec::new();
    let mut ndk: Vec<Vec<usize>> = Vec::new();
    for bag in &dataset {
        // w
        let mut w_d = Vec::new();
        for (&index, &count) in bag {
            for _ in 0..count {
                w_d.push(index);
            }
        }
        w.push(w_d);
        // z
        let mut z_d = Vec::new();
        for &count in bag.values() {
            for _ in 0..count {
                z_d.push(0);
            }
        }
        let n_d = z_d.len();
        z.push(z_d);
        // theta
        let mut theta_d = Vec::with_capacity(num_topics);
        for _k in 0..num_topics {
            theta_d.push(1.0 / num_topics as f64);
        }
        theta.push(theta_d);
        // ndk
        let mut ndk_d = Vec::with_capacity(num_topics);
        ndk_d.push(n_d);
        for _k in 1..num_topics {
            ndk_d.push(0);
        }
        ndk.push(ndk_d);
    }
    println!("z = {:?}", z);
    println!("theta = {:?}", theta);
    println!("ndk = {:?}", ndk);

    let mut phi: Vec<Vec<f64>> = Vec::new();
    let mut nkv: Vec<Vec<usize>> = Vec::new();
    // phi
    for _k in 0..num_topics {
        let mut phi_k = Vec::new();
        for _v in 0..vocab_size {
            phi_k.push(1.0 / vocab_size as f64);
        }
        phi.push(phi_k);
    }
    // nkv
    let mut nkv_0 = Vec::new();
    for _v in 0..vocab_size {
        nkv_0.push(0);
    }
    for bag in &dataset {
        for (&index, &count) in bag {
            nkv_0[index] += count;
        }
    }
    nkv.push(nkv_0);
    for _k in 1..num_topics {
        let mut nkv_k = Vec::new();
        for _v in 0..vocab_size {
            nkv_k.push(0);
        }
        nkv.push(nkv_k);
    }
    println!("phi = {:?}", phi);
    println!("nkv = {:?}", nkv);

    // Sampling
    for s in 0..num_samples {
        for (d, w_d) in w.iter().enumerate() {
            for (i, &w_di) in w_d.iter().enumerate() {
                let v = w_di;
                // Compute the proportion
                let mut sum: f64 = 0.0;
                let mut prop: Vec<f64> = Vec::new();
                for k in 0..num_topics {
                    let weight = theta[d][k] * phi[k][v];
                    sum += weight;
                    prop.push(weight);
                }
                for k in 0..num_topics {
                    prop[k] /= sum;
                }
                // Sample z_di
                let Closed01(x) = random::<Closed01<f64>>();
                let mut sum = 0.0;
                let mut sample = num_topics - 1;
                for k in 0..num_topics {
                    sum += prop[k];
                    if x < sum {
                        sample = k;
                        break;
                    }
                }
                z[d][i] = sample;
            }
        }
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let (labels, dataset) = load_bags(&args[1]).unwrap();
    let vocab: Option<Vec<String>> = if args.len() < 3 {
        None
    }
    else {
        Some(load_text_vocabulary(&args[2]).unwrap())
    };
    println!("{:?}", labels);
    println!("{:?}", dataset);
    println!("{:?}", vocab);

    lda(3, dataset, vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0], 100);
}
<commit_msg>Add Categorical distribution<commit_after>extern crate rand;

use std::string::String;
use std::collections::HashMap;
use std::vec::Vec;
use std::fs::{File};
use std::env;
use std::path::{Path};
use std::io::{self, BufReader, BufRead};

use rand::{random, Closed01, Rng};
use rand::distributions::{IndependentSample, Range, Sample};

type Bag = HashMap<usize, usize>;

fn load_bags<P: AsRef<Path>>(path: P) -> io::Result<(Vec<usize>, Vec<Bag>)> {
    let mut bags = Vec::new();
    let mut labels = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        let mut iter = line.split_whitespace();
        let label = iter.next().unwrap().parse::<usize>().unwrap();
        let mut bag = Bag::new();
        for elm in iter {
            let mut iter = elm.split(':');
            let index = iter.next().unwrap().parse::<usize>().unwrap();
            let value = iter.next().unwrap().parse::<usize>().unwrap();
            bag.insert(index - 1, value);
        }
        labels.push(label);
        bags.push(bag);
    }
    Ok((labels, bags))
}

fn load_text_vocabulary<P: AsRef<Path>>(path: P) -> io::Result<Vec<String>> {
    let mut vocab = Vec::new();
    let file = try!(File::open(path));
    let file = BufReader::new(file);
    for line in file.lines() {
        let line = line.unwrap();
        vocab.push(line);
    }
    Ok(vocab)
}

#[derive(Clone)]
struct Categorical {
    prop: Vec<f64>,
}

impl Categorical {
    fn new(weights: Vec<f64>) -> Categorical {
        let sum: f64 = weights.iter().sum();
        let prop: Vec<f64> = weights.iter().map(|w| w / sum).collect();
        Categorical {
            prop: prop,
        }
    }
}

impl Sample<usize> for Categorical {
    fn sample<R: Rng>(&mut self, rng: &mut R) -> usize {
        self.ind_sample(rng)
    }
}

impl IndependentSample<usize> for Categorical {
    fn ind_sample<R: Rng>(&self, rng: &mut R) -> usize {
        let Closed01(x) = random::<Closed01<f64>>();

        let mut sum = 0.0;
        let mut ret = self.prop.len() - 1;
        for (k, &p) in self.prop.iter().enumerate() {
            sum += p;
            if x <= sum {
                ret = k;
            }
        }
        ret
    }
}

fn lda(num_topics: usize, dataset: Vec<Bag>, alpha: Vec<f64>, beta: Vec<f64>, num_samples: usize) {
    // Initialization
    let mut rng = rand::thread_rng();
    let topics = Range::new(0, num_topics);

    let vocab_size: usize = {
        let mut have_some = false;
        let mut max_index = 0;
        for bag in &dataset {
            for &index in bag.keys() {
                have_some = true;
                if index > max_index {
                    max_index = index;
                }
            }
        }
        if have_some {
            max_index + 1
        }
        else {
            0
        }
    };
    println!("K = {}", num_topics);
    println!("M = {}", &dataset.len());
    println!("V = {}", vocab_size);

    let mut w: Vec<Vec<usize>> = Vec::new();
    let mut z: Vec<Vec<usize>> = Vec::new();
    let mut theta: Vec<Vec<f64>> = Vec::new();
    let mut ndk: Vec<Vec<usize>> = Vec::new();
    for bag in &dataset {
        // w
        let mut w_d = Vec::new();
        for (&index, &count) in bag {
            for _ in 0..count {
                w_d.push(index);
            }
        }
        w.push(w_d);
        // z
        let mut z_d = Vec::new();
        for &count in bag.values() {
            for _ in 0..count {
                z_d.push(0);
            }
        }
        let n_d = z_d.len();
        z.push(z_d);
        // theta
        let mut theta_d = Vec::with_capacity(num_topics);
        for _k in 0..num_topics {
            theta_d.push(1.0 / num_topics as f64);
        }
        theta.push(theta_d);
        // ndk
        let mut ndk_d = Vec::with_capacity(num_topics);
        ndk_d.push(n_d);
        for _k in 1..num_topics {
            ndk_d.push(0);
        }
        ndk.push(ndk_d);
    }
    println!("z = {:?}", z);
    println!("theta = {:?}", theta);
    println!("ndk = {:?}", ndk);

    let mut phi: Vec<Vec<f64>> = Vec::new();
    let mut nkv: Vec<Vec<usize>> = Vec::new();
    // phi
    for _k in 0..num_topics {
        let mut phi_k = Vec::new();
        for _v in 0..vocab_size {
            phi_k.push(1.0 / vocab_size as f64);
        }
        phi.push(phi_k);
    }
    // nkv
    let mut nkv_0 = Vec::new();
    for _v in 0..vocab_size {
        nkv_0.push(0);
    }
    for bag in &dataset {
        for (&index, &count) in bag {
            nkv_0[index] += count;
        }
    }
    nkv.push(nkv_0);
    for _k in 1..num_topics {
        let mut nkv_k = Vec::new();
        for _v in 0..vocab_size {
            nkv_k.push(0);
        }
        nkv.push(nkv_k);
    }
    println!("phi = {:?}", phi);
    println!("nkv = {:?}", nkv);

    // Sampling
    for s in 0..num_samples {
        for (d, w_d) in w.iter().enumerate() {
            for (i, &w_di) in w_d.iter().enumerate() {
                let v = w_di;
                // Sample z_di
                let mut weights: Vec<f64> = Vec::new();
                for k in 0..num_topics {
                    weights.push(theta[d][k] * phi[k][v]);
                }
                let cat = Categorical::new(weights);
                z[d][i] = cat.ind_sample(&mut rand::thread_rng());
            }
        }
    }
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let (labels, dataset) = load_bags(&args[1]).unwrap();
    let vocab: Option<Vec<String>> = if args.len() < 3 {
        None
    }
    else {
        Some(load_text_vocabulary(&args[2]).unwrap())
    };
    println!("{:?}", labels);
    println!("{:?}", dataset);
    println!("{:?}", vocab);

    lda(3, dataset, vec![1.0, 1.0, 1.0], vec![1.0, 1.0, 1.0], 100);
}
<|endoftext|>


<commit_before>extern crate clap;
extern crate kafka;

use clap::{Arg, App};
use kafka::client::{KafkaClient, ProduceMessage, RequiredAcks};
use std::time::Duration;
use std::io;
use std::io::prelude::*;

fn main() {
    let matches = App::new("ktee")
                      .version("0.1.0")
                      .about("tee for kafka")
                      .arg(Arg::with_name("broker")
                               .short("b")
                               .long("broker")
                               .value_name("BROKER")
                               .help("Kafka broker")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("topic")
                               .short("t")
                               .long("topic")
                               .value_name("TOPIC")
                               .help("Kafka topic")
                               .takes_value(true)
                               .required(true))
                      .get_matches();

    let broker = matches.value_of("broker").unwrap();
    let topic = matches.value_of("topic").unwrap();

    let mut client = KafkaClient::new(vec![broker.to_owned()]);
    let meta_res = client.load_metadata_all();
    if let Some(err) = meta_res.err() {
        println!("Error fetching metadata: {}", err);
        return;
    }

    let stdin = io::stdin();
    let mut stdout = io::stdout();
    let mut stderr = io::stderr();
    for line in stdin.lock().lines() {
        if line.is_ok() {
            let l = line.unwrap();
            let req = vec![ProduceMessage::new(topic, 0, None, Some(&l.as_bytes()))];
            let res = client.produce_messages(RequiredAcks::One, Duration::from_millis(0), req);
            if let Some(err) = res.err() {
                writeln!(&mut stderr, "Error sending message: {}", err).unwrap();
            }

            let write_res = write!(&mut stdout, "{}\n", &l);
            if let Some(err) = write_res.err() {
                writeln!(&mut stderr, "Error writing to stdout: {}", err).unwrap();
            }
        }
    }
}
<commit_msg>Switched to using match on stdin result<commit_after>extern crate clap;
extern crate kafka;

use clap::{Arg, App};
use kafka::client::{KafkaClient, ProduceMessage, RequiredAcks};
use std::time::Duration;
use std::io;
use std::io::prelude::*;

fn main() {
    let matches = App::new("ktee")
                      .version("0.1.0")
                      .about("tee for kafka")
                      .arg(Arg::with_name("broker")
                               .short("b")
                               .long("broker")
                               .value_name("BROKER")
                               .help("Kafka broker")
                               .takes_value(true)
                               .required(true))
                      .arg(Arg::with_name("topic")
                               .short("t")
                               .long("topic")
                               .value_name("TOPIC")
                               .help("Kafka topic")
                               .takes_value(true)
                               .required(true))
                      .get_matches();

    let broker = matches.value_of("broker").unwrap();
    let topic = matches.value_of("topic").unwrap();

    let mut client = KafkaClient::new(vec![broker.to_owned()]);
    let meta_res = client.load_metadata_all();
    if let Some(err) = meta_res.err() {
        println!("Error fetching metadata: {}", err);
        return;
    }

    let stdin = io::stdin();
    for line in stdin.lock().lines() {
        match line {
            Ok(line) => {
                let req = vec![ProduceMessage::new(topic, 0, None, Some(&line.as_bytes()))];
                let res = client.produce_messages(RequiredAcks::One, Duration::from_millis(0), req);
                if let Some(err) = res.err() {
                    writeln!(&mut io::stderr(), "Error sending message: {}", err).unwrap();
                }

                let write_res = write!(&mut io::stdout(), "{}\n", &line);
                if let Some(err) = write_res.err() {
                    writeln!(&mut io::stderr(), "Error writing to stdout: {}", err).unwrap();
                }
            }
            Err(err) => writeln!(&mut io::stderr(), "Error reading from stdin: {}", err).unwrap(),
        }
    }
}
<|endoftext|>


<commit_before>#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate reqwest;
#[macro_use]
extern crate serde_derive;
extern crate toml;

#[cfg(test)]
#[macro_use]
extern crate hyper;

use std::env;
use std::io::prelude::*;
use std::fs::File;
use clap::App;

mod github;

static GITHUB_TOKEN: &'static str = "RP_GITHUBTOKEN";

fn main() {
    let yaml = load_yaml!("release-party.yml");
    let matches = App::from_yaml(yaml).get_matches();

    let org_url = make_org_url(&matches);
    let token = env::var(GITHUB_TOKEN).expect(&format!("{} should be set", GITHUB_TOKEN));
    let reqwest_client = get_reqwest_client();

    print_party_links(get_pr_links(
        &get_repos_we_care_about(&token, &org_url, &reqwest_client),
        &token,
        &reqwest_client,
        is_dryrun(&matches),
    ));
}

fn is_dryrun(matches: &clap::ArgMatches) -> bool {
    matches.is_present("DRYRUN")
}

fn org_is_just_org(org: &str) -> bool {
    if org.contains("https://api.github.com") {
        return false;
    }
    true
}

fn suggest_org_arg(org: &str) -> Result<String, String> {
    if org.starts_with("https://api.github.com/orgs/") && org.ends_with("/repos") {
        let suggestion = org.replace("https://api.github.com/orgs/", "").replace("/repos", "");
        return Ok(format!("Try this: {}", suggestion).to_owned());
    }
    Err("Can't make a suggestion".to_owned())
}

fn make_org_url(matches: &clap::ArgMatches) -> String {
    let org = matches
        .value_of("ORG")
        .expect("Please specify a github org");

    if !org_is_just_org(&org) {
        match suggest_org_arg(&org) {
            Ok(suggestion) => panic!("Try this for the org value: {}", suggestion),
            Err(_) => panic!("Please make org just the organization name."),
        }
    }

    format!("https://api.github.com/orgs/{}/repos", org)
}

fn get_pr_links(repos: &Vec<github::GithubRepo>, token: &str, reqwest_client: &reqwest::Client, dryrun: bool) -> Vec<Option<String>> {
    let mut pr_links: Vec<Option<String>> = repos
        .into_iter()
        .map(|repo| match get_release_pr_for(&repo, &token, reqwest_client, dryrun) {
            Some(pr_url) => Some(pr_url),
            None => None,
        })
        .collect();
    // only keep the Some(PR_URL) items:
    pr_links.retain(|maybe_pr_link| maybe_pr_link.is_some());
    pr_links
}

fn get_reqwest_client() -> reqwest::Client {
    match reqwest::Client::new() {
        Ok(new_client) => new_client,
        Err(e) => panic!("Couldn't create new reqwest client: {}", e),
    }
}

fn get_repos_we_care_about(token: &str, github_org_url: &str, reqwest_client: &reqwest::Client) -> Vec<github::GithubRepo> {
    let mut repos = match github::get_repos_at(github_org_url, token, reqwest_client) {
        Ok(repos) => repos,
        Err(e) => panic!(format!("Couldn't get repos from github: {}", e)),
    };
    let repos_to_ignore = ignored_repos();

    // remove repos we don't care about:
    repos.retain(|repo| !repos_to_ignore.contains(&repo.name));

    repos
}

fn get_release_pr_for(repo: &github::GithubRepo, token: &str, client: &reqwest::Client, dryrun: bool) -> Option<String> {
    println!("looking for release PR for {}", repo.name);
    match github::existing_release_pr_location(repo, token, client) {
        Some(url) => Some(url),
        None => if !github::is_release_up_to_date_with_master(&repo.url, token, client) {
            if dryrun {
                Some(format!("Dry run: {} would get a release PR.", repo.url))
            } else {
                match github::create_release_pull_request(repo, token, client) {
                    Ok(pr_url) => Some(pr_url),
                    Err(_) => None,
                }
            }
        } else {
            None
        },
    }
}

fn print_party_links(pr_links: Vec<Option<String>>) {
    if !pr_links.is_empty() {
        println!("\nIt's a release party!  PRs to review and approve:");
        for link in pr_links {
            match link {
                Some(pr_link) => println!("{}", pr_link),
                None => println!("Party link is None: this shouldn't happen."),
            }
        }
    } else {
        println!("\nNo party today, all releases are done.");
    }
}

#[derive(Deserialize, Debug)]
struct IgnoredRepo {
    ignore: Option<Vec<String>>,
}

fn ignored_repos() -> Vec<String> {
    let mut f = match File::open("ignoredrepos.toml") {
        Ok(file) => file,
        Err(e) => {
            println!("Couldn't load ignoredrepos.toml, not ignoring any repos. Reason: {}", e);
            return Vec::new();
        }
    };
    let mut buffer = String::new();
    match f.read_to_string(&mut buffer) {
        Ok(_) => (),
        Err(e) => {
            println!("Couldn't read from ignoredrepos.toml, not ignoring any repos. Reason: {}", e);
            return Vec::new();
        }
    }

    let ignored: IgnoredRepo = match toml::from_str(&buffer) {
        Ok(ignored_repos) => ignored_repos,
        Err(e) => {
            println!(
                "Couldn't parse toml from ignoredrepos.toml, not ignoring any repos. Reason: {}",
                e
            );
            return Vec::new();
        }
    };
    match ignored.ignore {
        Some(repos_to_ignore) => repos_to_ignore,
        None => Vec::new(),
    }
}

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

    #[test]
    fn get_ignored_repos_happy_path() {
        let ignored_repositories = vec!["calagator".to_owned(), "moe".to_owned()];
        assert_eq!(ignored_repositories, ignored_repos());
    }

    #[test]
    fn handle_malformed_org() {
        assert_eq!(false, org_is_just_org("https://api.github.com/orgs/ORG-HERE/repos"));
    }

    #[test]
    fn handle_okay_org() {
        assert_eq!(true, org_is_just_org("ORG-HERE"));
    }

    #[test]
    fn suggestion_for_org_happy() {
        assert_eq!("Try this: ORG-HERE", suggest_org_arg("https://api.github.com/orgs/ORG-HERE/repos").unwrap());
    }

    #[test]
    fn suggestion_for_org_sad() {
        assert_eq!(true, suggest_org_arg("https://api.github.com/orgs/ORG-HERE/").is_err());
        assert_eq!(true, suggest_org_arg("http://api.github.com/orgs/ORG-HERE/").is_err());
        assert_eq!(true, suggest_org_arg("api.github.com/orgs/ORG-HERE/repos").is_err());
    }
}<commit_msg>Exit cleaner when not provided correct data, clean up org suggestion message.<commit_after>#![cfg_attr(feature = "clippy", feature(plugin))]
#![cfg_attr(feature = "clippy", plugin(clippy))]

#[macro_use]
extern crate clap;
extern crate reqwest;
#[macro_use]
extern crate serde_derive;
extern crate toml;

#[cfg(test)]
#[macro_use]
extern crate hyper;

use std::env;
use std::io::prelude::*;
use std::fs::File;
use clap::App;

mod github;

static GITHUB_TOKEN: &'static str = "RP_GITHUBTOKEN";

fn main() {
    let yaml = load_yaml!("release-party.yml");
    let matches = App::from_yaml(yaml).get_matches();

    let org_url = make_org_url(&matches);
    let token = match env::var(GITHUB_TOKEN) {
        Ok(env_var) => env_var,
        Err(_) => {
            print_message_and_exit(&format!("{} environment variable should be set", GITHUB_TOKEN), -1);
            unreachable!();
        },
    };
    let reqwest_client = get_reqwest_client();

    print_party_links(get_pr_links(
        &get_repos_we_care_about(&token, &org_url, &reqwest_client),
        &token,
        &reqwest_client,
        is_dryrun(&matches),
    ));
}

fn is_dryrun(matches: &clap::ArgMatches) -> bool {
    matches.is_present("DRYRUN")
}

fn org_is_just_org(org: &str) -> bool {
    if org.contains("https://api.github.com") {
        return false;
    }
    true
}

fn suggest_org_arg(org: &str) -> Result<String, String> {
    if org.starts_with("https://api.github.com/orgs/") && org.ends_with("/repos") {
        let suggestion = org.replace("https://api.github.com/orgs/", "").replace("/repos", "");
        return Ok(format!("{}", suggestion).to_owned());
    }
    Err("Can't make a suggestion".to_owned())
}

fn make_org_url(matches: &clap::ArgMatches) -> String {
    let org = matches
        .value_of("ORG")
        .expect("Please specify a github org");

    if !org_is_just_org(&org) {
        match suggest_org_arg(&org) {
            Ok(suggestion) => print_message_and_exit(&format!("Try this for the org value: {}", suggestion), -1),
            Err(_) => print_message_and_exit(&format!("Please make org just the organization name."), -1),
        }
    }

    format!("https://api.github.com/orgs/{}/repos", org)
}

fn get_pr_links(repos: &Vec<github::GithubRepo>, token: &str, reqwest_client: &reqwest::Client, dryrun: bool) -> Vec<Option<String>> {
    let mut pr_links: Vec<Option<String>> = repos
        .into_iter()
        .map(|repo| match get_release_pr_for(&repo, &token, reqwest_client, dryrun) {
            Some(pr_url) => Some(pr_url),
            None => None,
        })
        .collect();
    // only keep the Some(PR_URL) items:
    pr_links.retain(|maybe_pr_link| maybe_pr_link.is_some());
    pr_links
}

fn get_reqwest_client() -> reqwest::Client {
    match reqwest::Client::new() {
        Ok(new_client) => new_client,
        Err(e) => panic!("Couldn't create new reqwest client: {}", e),
    }
}

fn get_repos_we_care_about(token: &str, github_org_url: &str, reqwest_client: &reqwest::Client) -> Vec<github::GithubRepo> {
    let mut repos = match github::get_repos_at(github_org_url, token, reqwest_client) {
        Ok(repos) => repos,
        Err(e) => panic!(format!("Couldn't get repos from github: {}", e)),
    };
    let repos_to_ignore = ignored_repos();

    // remove repos we don't care about:
    repos.retain(|repo| !repos_to_ignore.contains(&repo.name));

    repos
}

fn get_release_pr_for(repo: &github::GithubRepo, token: &str, client: &reqwest::Client, dryrun: bool) -> Option<String> {
    println!("looking for release PR for {}", repo.name);
    match github::existing_release_pr_location(repo, token, client) {
        Some(url) => Some(url),
        None => if !github::is_release_up_to_date_with_master(&repo.url, token, client) {
            if dryrun {
                Some(format!("Dry run: {} would get a release PR.", repo.url))
            } else {
                match github::create_release_pull_request(repo, token, client) {
                    Ok(pr_url) => Some(pr_url),
                    Err(_) => None,
                }
            }
        } else {
            None
        },
    }
}

fn print_party_links(pr_links: Vec<Option<String>>) {
    if !pr_links.is_empty() {
        println!("\nIt's a release party!  PRs to review and approve:");
        for link in pr_links {
            match link {
                Some(pr_link) => println!("{}", pr_link),
                None => println!("Party link is None: this shouldn't happen."),
            }
        }
    } else {
        println!("\nNo party today, all releases are done.");
    }
}

#[derive(Deserialize, Debug)]
struct IgnoredRepo {
    ignore: Option<Vec<String>>,
}

fn ignored_repos() -> Vec<String> {
    let mut f = match File::open("ignoredrepos.toml") {
        Ok(file) => file,
        Err(e) => {
            println!("Couldn't load ignoredrepos.toml, not ignoring any repos. Reason: {}", e);
            return Vec::new();
        }
    };
    let mut buffer = String::new();
    match f.read_to_string(&mut buffer) {
        Ok(_) => (),
        Err(e) => {
            println!("Couldn't read from ignoredrepos.toml, not ignoring any repos. Reason: {}", e);
            return Vec::new();
        }
    }

    let ignored: IgnoredRepo = match toml::from_str(&buffer) {
        Ok(ignored_repos) => ignored_repos,
        Err(e) => {
            println!(
                "Couldn't parse toml from ignoredrepos.toml, not ignoring any repos. Reason: {}",
                e
            );
            return Vec::new();
        }
    };
    match ignored.ignore {
        Some(repos_to_ignore) => repos_to_ignore,
        None => Vec::new(),
    }
}

fn print_message_and_exit(message: &str, exit_code: i32) {
    println!("{}", message);
    ::std::process::exit(exit_code);
}

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

    #[test]
    fn get_ignored_repos_happy_path() {
        let ignored_repositories = vec!["calagator".to_owned(), "moe".to_owned()];
        assert_eq!(ignored_repositories, ignored_repos());
    }

    #[test]
    fn handle_malformed_org() {
        assert_eq!(false, org_is_just_org("https://api.github.com/orgs/ORG-HERE/repos"));
    }

    #[test]
    fn handle_okay_org() {
        assert_eq!(true, org_is_just_org("ORG-HERE"));
    }

    #[test]
    fn suggestion_for_org_happy() {
        assert_eq!("ORG-HERE", suggest_org_arg("https://api.github.com/orgs/ORG-HERE/repos").unwrap());
    }

    #[test]
    fn suggestion_for_org_sad() {
        assert_eq!(true, suggest_org_arg("https://api.github.com/orgs/ORG-HERE/").is_err());
        assert_eq!(true, suggest_org_arg("http://api.github.com/orgs/ORG-HERE/").is_err());
        assert_eq!(true, suggest_org_arg("api.github.com/orgs/ORG-HERE/repos").is_err());
    }
}<|endoftext|>


<commit_before>use hyper::Client;
use kuchiki::{Html, NodeRef, NodeDataRef, ElementData};

#[derive(Debug)]
pub struct DictDef {
    pub word: String,
    pub def: String,
    pub example: Option<String>
}

impl DictDef {
    pub fn new (word: String, def: String, example: Option<String>) -> DictDef {
        DictDef {
            word: word,
            def: def,
            example: example
        }
    }
}

fn urban_dict_url (word: &str) -> String {
    let mut s = String::new();
    s.push_str("http://www.urbandictionary.com/define.php?term=");
    s.push_str(word);
    s
}

fn get_first_match(node: &NodeRef, selector: &str) -> Option<NodeDataRef<ElementData>> {
    if let Ok(mut eles) = node.select(selector) {
        eles.next()
    } else {
        None
    }
}

pub fn find_on_urban_dict (word: &str) -> Option<DictDef> {
    let url = urban_dict_url(word);

    let client = Client::new();

    let mut response = client.get(&url).send().unwrap();

    if let Ok(html) = Html::from_stream(&mut response) {
        let doc = html.parse();

        if let Some(panel_ele) = get_first_match(&doc, "div.def-panel") {
            let panel = panel_ele.as_node();
            if let Some(word_ele) = get_first_match(panel, "a.word") {
                let word_node = word_ele.as_node();
                let word_text = word_node.text_contents();

                let def_ele = get_first_match(panel, "div.meaning").unwrap();
                let def_node = def_ele.as_node();
                let def_text = def_node.text_contents();

                let example = if let Some(example_ele) = get_first_match(panel, "div.example") {
                    let example_node = example_ele.as_node();
                    Some(example_node.text_contents())
                } else {
                    None
                };

                return Some(DictDef::new(word_text.trim().to_owned(),
                                         def_text.trim().to_owned(),
                                         example));
            }
        }
    }

    None
}
<commit_msg>(refactor) move data extraction to dedicated fn<commit_after>use hyper::Client;
use kuchiki::{Html, NodeRef, NodeDataRef, ElementData};

#[derive(Debug)]
pub struct DictDef {
    pub word: String,
    pub def: String,
    pub example: Option<String>
}

impl DictDef {
    pub fn new (word: String, def: String, example: Option<String>) -> DictDef {
        DictDef {
            word: word,
            def: def,
            example: example
        }
    }
}

fn urban_dict_url (word: &str) -> String {
    let mut s = String::new();
    s.push_str("http://www.urbandictionary.com/define.php?term=");
    s.push_str(word);
    s
}

fn get_first_match(node: &NodeRef, selector: &str) -> Option<NodeDataRef<ElementData>> {
    if let Ok(mut eles) = node.select(selector) {
        eles.next()
    } else {
        None
    }
}

fn get_def_from_page_ele (panel: &NodeRef) -> Option<DictDef> {
    if let Some(word_ele) = get_first_match(panel, "a.word") {
        let word_node = word_ele.as_node();
        let word_text = word_node.text_contents();

        let def_ele = get_first_match(panel, "div.meaning").unwrap();
        let def_node = def_ele.as_node();
        let def_text = def_node.text_contents();

        let example = if let Some(example_ele) = get_first_match(panel, "div.example") {
            let example_node = example_ele.as_node();
            Some(example_node.text_contents())
        } else {
            None
        };

        Some(DictDef::new(word_text.trim().to_owned(),
                          def_text.trim().to_owned(),
                          example))
    } else {
        None
    }
}

pub fn find_on_urban_dict (word: &str) -> Option<DictDef> {
    let url = urban_dict_url(word);

    let client = Client::new();

    let mut response = client.get(&url).send().unwrap();

    if let Ok(html) = Html::from_stream(&mut response) {
        let doc = html.parse();

        if let Some(panel_ele) = get_first_match(&doc, "div.def-panel") {
            let panel = panel_ele.as_node();
            return get_def_from_page_ele(panel);
        }
    }

    None
}
<|endoftext|>


<commit_before>use std::io;
use std::net::SocketAddr;

use futures::{Future, Poll, Async};
use tokio_core::net::TcpStream;
use tokio_proto::{TryRead, TryWrite};

#[derive(Debug)]
enum ConnectionState {
    ClientReading,
    ClientWriting,
    ServerReading,
    ServerWriting,
}

#[must_use = "Must use Pipe"]
pub struct Pipe {
    client_addr: SocketAddr,
    server_addr: SocketAddr,
    client: TcpStream,
    server: TcpStream,
    state: ConnectionState,

    /// The buffer from the client to send to the server
    send_buf: [u8; 1024],
    send_num_bytes: usize,

    /// The buffer from the server to send to the client
    recv_buf: [u8; 1024],
    recv_num_bytes: usize,
}

impl Pipe {
    pub fn new(client_addr: SocketAddr, client: TcpStream, server_addr: SocketAddr, server: TcpStream) -> Pipe {
        Pipe {
            client_addr: client_addr,
            server_addr: server_addr,
            client: client,
            server: server,
            state: ConnectionState::ClientReading,
            send_buf: [0u8; 1024],
            send_num_bytes: 0,
            recv_buf: [0u8; 1024],
            recv_num_bytes: 0,
        }
    }
}

impl Future for Pipe {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        trace!("Polling...");

        loop {

            match self.state {
                ConnectionState::ClientReading => {
                    trace!("Reading from {}", self.client_addr);

                    let bytes = try_ready!(self.client.try_read(&mut self.send_buf));
                    trace!("Read {} bytes from {}", bytes, self.client_addr);

                    if bytes == 0 {
                        self.state = ConnectionState::ServerWriting;
                        trace!("State switched to {:?}", self.state);
                    } else {
                        self.send_num_bytes = bytes;
                    }
                }

                ConnectionState::ServerWriting => {
                    trace!("Writing to {}", self.server_addr);

                    let bytes = try_ready!(self.server.try_write(&mut self.send_buf[0..self.send_num_bytes]));
                    trace!("Wrote {} bytes to {}", bytes, self.server_addr);

                    if bytes == 0 {
                        self.send_num_bytes = 0;
                        self.state = ConnectionState::ServerReading;
                        trace!("State switched to {:?}", self.state);
                    } else {
                        self.send_num_bytes -= bytes;
                    }
                }

                ConnectionState::ServerReading => {
                    trace!("Reading from {}", self.server_addr);

                    let bytes = try_ready!(self.server.try_read(&mut self.recv_buf));
                    trace!("Read {} bytes from {}", bytes, self.server_addr);

                    if bytes == 0 {
                        self.state = ConnectionState::ClientWriting;
                        trace!("State switched to {:?}", self.state);
                    } else {
                        self.recv_num_bytes = bytes;
                    }
                }

                ConnectionState::ClientWriting => {
                    trace!("Writing to {}", self.client_addr);

                    let bytes = try_ready!(self.client.try_write(&mut self.recv_buf[0..self.recv_num_bytes]));
                    trace!("Wrote {} bytes to {}", bytes, self.client_addr);

                    if bytes == 0 {
                        self.recv_num_bytes = 0;
                        self.state = ConnectionState::ClientReading;
                        trace!("State switched to {:?}", self.state);
                        trace!("Bailing out of loop");
                        break;
                    } else {
                        self.recv_num_bytes -= bytes;
                    }
                }
            }
        }

        Ok(Async::Ready(()))
    }
}
<commit_msg>Use only one buffer per pipe<commit_after>use std::io;
use std::net::SocketAddr;

use futures::{Future, Poll, Async};
use tokio_core::net::TcpStream;
use tokio_proto::{TryRead, TryWrite};

#[derive(Debug)]
enum ConnectionState {
    ClientReading,
    ClientWriting,
    ServerReading,
    ServerWriting,
}

#[must_use = "Must use Pipe"]
pub struct Pipe {
    client_addr: SocketAddr,
    server_addr: SocketAddr,
    client: TcpStream,
    server: TcpStream,
    state: ConnectionState,

    /// A buffer to keep track of the data betewen the client and server
    buf: [u8; 1024],
    num_bytes: usize,
}

impl Pipe {
    pub fn new(client_addr: SocketAddr, client: TcpStream, server_addr: SocketAddr, server: TcpStream) -> Pipe {
        Pipe {
            client_addr: client_addr,
            server_addr: server_addr,
            client: client,
            server: server,
            state: ConnectionState::ClientReading,
            buf: [0u8; 1024],
            num_bytes: 0,
        }
    }
}

impl Future for Pipe {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<(), io::Error> {
        trace!("Polling...");

        loop {

            match self.state {
                ConnectionState::ClientReading => {
                    trace!("Reading from {}", self.client_addr);

                    let bytes = try_ready!(self.client.try_read(&mut self.buf));
                    trace!("Read {} bytes from {}", bytes, self.client_addr);

                    if bytes == 0 {
                        self.state = ConnectionState::ServerWriting;
                        trace!("State switched to {:?}", self.state);
                    } else {
                        self.num_bytes = bytes;
                    }
                }

                ConnectionState::ServerWriting => {
                    trace!("Writing to {}", self.server_addr);

                    let bytes = try_ready!(self.server.try_write(&mut self.buf[0..self.num_bytes]));
                    trace!("Wrote {} bytes to {}", bytes, self.server_addr);

                    if bytes == 0 {
                        self.num_bytes = 0;
                        self.state = ConnectionState::ServerReading;
                        trace!("State switched to {:?}", self.state);
                    } else {
                        self.num_bytes -= bytes;
                    }
                }

                ConnectionState::ServerReading => {
                    trace!("Reading from {}", self.server_addr);

                    let bytes = try_ready!(self.server.try_read(&mut self.buf));
                    trace!("Read {} bytes from {}", bytes, self.server_addr);

                    if bytes == 0 {
                        self.state = ConnectionState::ClientWriting;
                        trace!("State switched to {:?}", self.state);
                    } else {
                        self.num_bytes = bytes;
                    }
                }

                ConnectionState::ClientWriting => {
                    trace!("Writing to {}", self.client_addr);

                    let bytes = try_ready!(self.client.try_write(&mut self.buf[0..self.num_bytes]));
                    trace!("Wrote {} bytes to {}", bytes, self.client_addr);

                    if bytes == 0 {
                        self.num_bytes = 0;
                        self.state = ConnectionState::ClientReading;
                        trace!("State switched to {:?}", self.state);
                        trace!("Bailing out of loop");
                        break;
                    } else {
                        self.num_bytes -= bytes;
                    }
                }
            }
        }

        Ok(Async::Ready(()))
    }
}
<|endoftext|>


<commit_before>/* Copyright 2016 Jordan Miner
 *
 * Licensed under the MIT license <LICENSE or
 * http://opensource.org/licenses/MIT>. This file may not be copied,
 * modified, or distributed except according to those terms.
 */

use super::control_prelude::*;
use super::containers::{
    Container,
    NonDialogContainer,
};

#[derive(Clone)]
pub struct TabInfo<'a, 'b> {
    title: &'a str,
    control: &'b Control,
}

#[derive(Clone)]
pub struct Tabs(HandleRc);

impl Tabs {
    pub fn new() -> Self {
        unsafe {
            ::iup_open();
            let ih = IupTabsv(ptr::null_mut());
            Tabs(HandleRc::new(ih))
        }
    }

    pub fn tab_title(&self, pos: usize) -> String {
        get_str_attribute(self.handle(), &format!("TABTITLE{}\0", pos))
    }

    pub fn set_tab_title(&self, pos: usize, title: &str) -> &Self {
        set_str_attribute(self.handle(), &format!("TABTITLE{}\0", pos), title);
        self
    }

    pub fn add_tabs(&self, info: &[TabInfo]) -> &Self {
        let mut index = self.child_count();
        for ti in info {
            self.append(ti.control).expect("failed to append tab");
            self.set_tab_title(index, ti.title);
            index += 1;
        }
        self
    }
}

impl_control_traits!(Tabs);

impl Container for Tabs {}
impl NonDialogContainer for Tabs {}

<commit_msg>Rename add_tabs to append_tabs to match Container::append<commit_after>/* Copyright 2016 Jordan Miner
 *
 * Licensed under the MIT license <LICENSE or
 * http://opensource.org/licenses/MIT>. This file may not be copied,
 * modified, or distributed except according to those terms.
 */

use super::control_prelude::*;
use super::containers::{
    Container,
    NonDialogContainer,
};

#[derive(Clone)]
pub struct TabInfo<'a, 'b> {
    title: &'a str,
    control: &'b Control,
}

#[derive(Clone)]
pub struct Tabs(HandleRc);

impl Tabs {
    pub fn new() -> Self {
        unsafe {
            ::iup_open();
            let ih = IupTabsv(ptr::null_mut());
            Tabs(HandleRc::new(ih))
        }
    }

    pub fn tab_title(&self, pos: usize) -> String {
        get_str_attribute(self.handle(), &format!("TABTITLE{}\0", pos))
    }

    pub fn set_tab_title(&self, pos: usize, title: &str) -> &Self {
        set_str_attribute(self.handle(), &format!("TABTITLE{}\0", pos), title);
        self
    }

    pub fn append_tabs(&self, info: &[TabInfo]) -> &Self {
        let mut index = self.child_count();
        for ti in info {
            self.append(ti.control).expect("failed to append tab");
            self.set_tab_title(index, ti.title);
            index += 1;
        }
        self
    }
}

impl_control_traits!(Tabs);

impl Container for Tabs {}
impl NonDialogContainer for Tabs {}

<|endoftext|>


<commit_before>use serde::de::{Deserialize, Deserializer};
use serde::ser::{Serialize, Serializer};

use std::fmt::Write;
use std::num::FpCategory;
use std::time::Duration;
use std::{cmp, u64};

pub(crate) fn ser_secs<S: Serializer>(
    dur: &Option<Duration>,
    serializer: S,
) -> std::result::Result<S::Ok, S::Error> {
    dur.map(format_secs).serialize(serializer)
}

fn format_secs(dur: Duration) -> String {
    let (secs, nanos) = (dur.as_secs(), dur.subsec_nanos());
    let mut r = secs.to_string();
    if nanos > 0 {
        write!(r, ".{:09}", nanos).unwrap();
        while r.ends_with('0') {
            r.pop();
        }
    }
    r += "s";
    r
}

pub(crate) fn ser_millis<S: Serializer>(
    dur: &Option<Duration>,
    serializer: S,
) -> std::result::Result<S::Ok, S::Error> {
    dur.map(format_millis).serialize(serializer)
}

fn format_millis(dur: Duration) -> String {
    let (secs, nanos) = (dur.as_secs(), dur.subsec_nanos());
    let millis = 1000 * u128::from(secs) + u128::from(nanos) / 1_000_000;
    let submilli_nanos = nanos % 1_000_000;
    let mut r = millis.to_string();
    if submilli_nanos > 0 {
        write!(r, ".{:06}", submilli_nanos).unwrap();
        while r.ends_with('0') {
            r.pop();
        }
    }
    r += "ms";
    r
}

pub(crate) fn de_secs<'de, D: Deserializer<'de>>(
    deserializer: D,
) -> std::result::Result<Option<Duration>, D::Error> {
    Option::<String>::deserialize(deserializer).and_then(|s| match s {
        None => Ok(None),
        Some(s) => parse_secs(&s).map(Some).map_err(serde::de::Error::custom),
    })
}

pub(crate) fn parse_secs(s: &str) -> std::result::Result<Duration, String> {
    fn extract_unit(s: &str) -> std::result::Result<(&str, f64), String> {
        if s.ends_with("das") {
            Ok((&s[..s.len() - 3], 1e10))
        } else if s.ends_with("hs") {
            Ok((&s[..s.len() - 2], 1e11))
        } else if s.ends_with("ks") {
            Ok((&s[..s.len() - 2], 1e12))
        } else if s.ends_with("Ms") {
            Ok((&s[..s.len() - 2], 1e15))
        } else if s.ends_with("Gs") {
            Ok((&s[..s.len() - 2], 1e18))
        } else if s.ends_with("Ts") {
            Ok((&s[..s.len() - 2], 1e21))
        } else if s.ends_with("Ps") {
            Ok((&s[..s.len() - 2], 1e24))
        } else if s.ends_with("Es") {
            Ok((&s[..s.len() - 2], 1e27))
        } else if s.ends_with("Zs") {
            Ok((&s[..s.len() - 2], 1e30))
        } else if s.ends_with("Ys") {
            Ok((&s[..s.len() - 2], 1e33))
        } else if s.ends_with("ds") {
            Ok((&s[..s.len() - 2], 1e8))
        } else if s.ends_with("cs") {
            Ok((&s[..s.len() - 2], 1e7))
        } else if s.ends_with("ms") {
            Ok((&s[..s.len() - 2], 1e6))
        } else if s.ends_with("μs") {
            Ok((&s[..s.len() - 3], 1e3))
        } else if s.ends_with("us") {
            Ok((&s[..s.len() - 2], 1e3))
        } else if s.ends_with("ns") {
            Ok((&s[..s.len() - 2], 1e0))
        } else if s.ends_with("ps") {
            Ok((&s[..s.len() - 2], 1e-3))
        } else if s.ends_with("fs") && !s.ends_with("infs") {
            Ok((&s[..s.len() - 2], 1e-6))
        } else if s.ends_with("as") {
            Ok((&s[..s.len() - 2], 1e-9))
        } else if s.ends_with("zs") {
            Ok((&s[..s.len() - 2], 1e-12))
        } else if s.ends_with("ys") {
            Ok((&s[..s.len() - 2], 1e-15))
        } else if s.ends_with('s') {
            Ok((&s[..s.len() - 1], 1e9))
        } else {
            Err("invalid or missing unit".to_owned())
        }
    }

    let (s, k) = extract_unit(s.trim())?;
    let nanos = k * s.trim_end().parse::<f64>().map_err(|e| e.to_string())?;
    if nanos.is_nan() {
        Err("NaN".to_owned())
    } else if nanos.is_sign_negative() {
        Err("negative".to_owned())
    } else if nanos.classify() == FpCategory::Zero {
        Ok(Duration::new(0, 0))
    } else if nanos.is_infinite() {
        Ok(Duration::new(u64::MAX, 999_999_999))
    } else {
        let nanos = nanos as u128;
        let secs = cmp::min(nanos / 1_000_000_000, u128::from(u64::MAX)) as u64;
        let subsec_nanos = if nanos > 1_000_000_000 * u128::from(u64::MAX) {
            999_999_999
        } else {
            (nanos % 1_000_000_000) as u32
        };
        Ok(Duration::new(secs, subsec_nanos))
    }
}

pub(crate) trait MillisRoundedUp {
    fn millis_rounded_up(self) -> u128;
}

impl MillisRoundedUp for Duration {
    fn millis_rounded_up(self) -> u128 {
        (1_000_000_000 * u128::from(self.as_secs()) + u128::from(self.subsec_nanos()) + 999_999)
            / 1_000_000
    }
}

#[cfg(test)]
mod tests {
    use crate::time::MillisRoundedUp;

    use std::time::Duration;
    use std::u64;

    #[test]
    fn test_parse_secs() {
        fn test(s: &str, expected: std::result::Result<Duration, String>) {
            assert_eq!(super::parse_secs(s), expected.map_err(Into::into));
        }

        test("NaNs", Err("NaN".to_owned()));
        test("+NaNs", Err("NaN".to_owned()));
        test("-NaNs", Err("NaN".to_owned()));
        test("-infs", Err("negative".to_owned()));
        test("-42s", Err("negative".to_owned()));
        test("1S", Err("invalid or missing unit".to_owned()));

        test("infs", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("+infs", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("1Ys", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("1Zs", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("100Es", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("1Ps", Ok(Duration::new(999_999_999_999_999, 983_222_784)));
        test("100Ts", Ok(Duration::new(99_999_999_999_999, 991_611_392)));
        test("10Ts", Ok(Duration::new(10_000_000_000_000, 0)));
        test("1Gs", Ok(Duration::new(1_000_000_000, 0)));
        test("1Ms", Ok(Duration::new(1_000_000, 0)));
        test("1ks", Ok(Duration::new(1000, 0)));
        test("1das", Ok(Duration::new(10, 0)));
        test("0s", Ok(Duration::new(0, 0)));
        test(" 10 s ", Ok(Duration::new(10, 0)));
        test("42s", Ok(Duration::new(42, 0)));
        test("123.456s", Ok(Duration::new(123, 456_000_000)));
        test("1E3s", Ok(Duration::new(1000, 0)));
        test("1ds", Ok(Duration::new(0, 100_000_000)));
        test("1cs", Ok(Duration::new(0, 10_000_000)));
        test("42ms", Ok(Duration::new(0, 42_000_000)));
        test("123.456ms", Ok(Duration::new(0, 123_456_000)));
        test("1E3ms", Ok(Duration::new(1, 0)));
        test("42μs", Ok(Duration::new(0, 42_000)));
        test("123.456μs", Ok(Duration::new(0, 123_456)));
        test("1E3μs", Ok(Duration::new(0, 1_000_000)));
        test("1us", Ok(Duration::new(0, 1000)));
        test("42ns", Ok(Duration::new(0, 42)));
        test("123.456ns", Ok(Duration::new(0, 123)));
        test("1E3ns", Ok(Duration::new(0, 1000)));
        test("1000ps", Ok(Duration::new(0, 1)));
        test("100ps", Ok(Duration::new(0, 0)));
        test("1E6fs", Ok(Duration::new(0, 1)));
        test("1E9as", Ok(Duration::new(0, 1)));
        test("1E12zs", Ok(Duration::new(0, 1)));
        test("1E15ys", Ok(Duration::new(0, 1)));
    }

    #[test]
    fn test_format_secs() {
        fn test(dur: Duration, expected: &str) {
            assert_eq!(super::format_secs(dur), expected);
        }

        test(Duration::from_secs(42), "42s");
        test(Duration::from_millis(42), "0.042s");
        test(Duration::from_nanos(42), "0.000000042s");
        test(Duration::new(1234, 567_800_000), "1234.5678s");
    }

    #[test]
    fn test_format_millis() {
        fn test(dur: Duration, expected: &str) {
            assert_eq!(super::format_millis(dur), expected);
        }

        test(Duration::from_secs(42), "42000ms");
        test(Duration::from_millis(42), "42ms");
        test(Duration::from_nanos(42), "0.000042ms");
        test(Duration::new(1234, 567_800_000), "1234567.8ms");
    }

    #[test]
    fn test_millis_rounded_up() {
        fn test(dur: Duration, expected: u128) {
            assert_eq!(dur.millis_rounded_up(), expected);
        }

        test(Duration::from_nanos(0), 0);
        test(Duration::from_nanos(1), 1);
        test(Duration::from_nanos(999_999), 1);
        test(Duration::from_nanos(1_000_000), 1);
        test(Duration::from_nanos(1_000_001), 2);
        test(Duration::from_secs(1), 1000);
    }
}
<commit_msg>Cover ::time<commit_after>use serde::de::{Deserialize, Deserializer};
use serde::ser::{Serialize, Serializer};

use std::fmt::Write;
use std::num::FpCategory;
use std::time::Duration;
use std::{cmp, u64};

pub(crate) fn ser_secs<S: Serializer>(
    dur: &Option<Duration>,
    serializer: S,
) -> std::result::Result<S::Ok, S::Error> {
    fn format_secs(dur: Duration) -> String {
        let (secs, nanos) = (dur.as_secs(), dur.subsec_nanos());
        let mut r = secs.to_string();
        if nanos > 0 {
            write!(r, ".{:09}", nanos).unwrap();
            while r.ends_with('0') {
                r.pop();
            }
        }
        r += "s";
        r
    }

    dur.map(format_secs).serialize(serializer)
}

pub(crate) fn ser_millis<S: Serializer>(
    dur: &Option<Duration>,
    serializer: S,
) -> std::result::Result<S::Ok, S::Error> {
    fn format_millis(dur: Duration) -> String {
        let (secs, nanos) = (dur.as_secs(), dur.subsec_nanos());
        let millis = 1000 * u128::from(secs) + u128::from(nanos) / 1_000_000;
        let submilli_nanos = nanos % 1_000_000;
        let mut r = millis.to_string();
        if submilli_nanos > 0 {
            write!(r, ".{:06}", submilli_nanos).unwrap();
            while r.ends_with('0') {
                r.pop();
            }
        }
        r += "ms";
        r
    }

    dur.map(format_millis).serialize(serializer)
}

pub(crate) fn de_secs<'de, D: Deserializer<'de>>(
    deserializer: D,
) -> std::result::Result<Option<Duration>, D::Error> {
    Option::<String>::deserialize(deserializer).and_then(|s| match s {
        None => Ok(None),
        Some(s) => parse_secs(&s).map(Some).map_err(serde::de::Error::custom),
    })
}

pub(crate) fn parse_secs(s: &str) -> std::result::Result<Duration, String> {
    fn extract_unit(s: &str) -> std::result::Result<(&str, f64), String> {
        if s.ends_with("das") {
            Ok((&s[..s.len() - 3], 1e10))
        } else if s.ends_with("hs") {
            Ok((&s[..s.len() - 2], 1e11))
        } else if s.ends_with("ks") {
            Ok((&s[..s.len() - 2], 1e12))
        } else if s.ends_with("Ms") {
            Ok((&s[..s.len() - 2], 1e15))
        } else if s.ends_with("Gs") {
            Ok((&s[..s.len() - 2], 1e18))
        } else if s.ends_with("Ts") {
            Ok((&s[..s.len() - 2], 1e21))
        } else if s.ends_with("Ps") {
            Ok((&s[..s.len() - 2], 1e24))
        } else if s.ends_with("Es") {
            Ok((&s[..s.len() - 2], 1e27))
        } else if s.ends_with("Zs") {
            Ok((&s[..s.len() - 2], 1e30))
        } else if s.ends_with("Ys") {
            Ok((&s[..s.len() - 2], 1e33))
        } else if s.ends_with("ds") {
            Ok((&s[..s.len() - 2], 1e8))
        } else if s.ends_with("cs") {
            Ok((&s[..s.len() - 2], 1e7))
        } else if s.ends_with("ms") {
            Ok((&s[..s.len() - 2], 1e6))
        } else if s.ends_with("μs") {
            Ok((&s[..s.len() - 3], 1e3))
        } else if s.ends_with("us") {
            Ok((&s[..s.len() - 2], 1e3))
        } else if s.ends_with("ns") {
            Ok((&s[..s.len() - 2], 1e0))
        } else if s.ends_with("ps") {
            Ok((&s[..s.len() - 2], 1e-3))
        } else if s.ends_with("fs") && !s.ends_with("infs") {
            Ok((&s[..s.len() - 2], 1e-6))
        } else if s.ends_with("as") {
            Ok((&s[..s.len() - 2], 1e-9))
        } else if s.ends_with("zs") {
            Ok((&s[..s.len() - 2], 1e-12))
        } else if s.ends_with("ys") {
            Ok((&s[..s.len() - 2], 1e-15))
        } else if s.ends_with('s') {
            Ok((&s[..s.len() - 1], 1e9))
        } else {
            Err("invalid or missing unit".to_owned())
        }
    }

    let (s, k) = extract_unit(s.trim())?;
    let nanos = k * s.trim_end().parse::<f64>().map_err(|e| e.to_string())?;
    if nanos.is_nan() {
        Err("NaN".to_owned())
    } else if nanos.is_sign_negative() {
        Err("negative".to_owned())
    } else if nanos.classify() == FpCategory::Zero {
        Ok(Duration::new(0, 0))
    } else if nanos.is_infinite() {
        Ok(Duration::new(u64::MAX, 999_999_999))
    } else {
        let nanos = nanos as u128;
        let secs = cmp::min(nanos / 1_000_000_000, u128::from(u64::MAX)) as u64;
        let subsec_nanos = if nanos > 1_000_000_000 * u128::from(u64::MAX) {
            999_999_999
        } else {
            (nanos % 1_000_000_000) as u32
        };
        Ok(Duration::new(secs, subsec_nanos))
    }
}

pub(crate) trait MillisRoundedUp {
    fn millis_rounded_up(self) -> u128;
}

impl MillisRoundedUp for Duration {
    fn millis_rounded_up(self) -> u128 {
        (1_000_000_000 * u128::from(self.as_secs()) + u128::from(self.subsec_nanos()) + 999_999)
            / 1_000_000
    }
}

#[cfg(test)]
mod tests {
    use crate::time::MillisRoundedUp;

    use serde_derive::{Deserialize, Serialize};

    use std::time::Duration;
    use std::u64;

    #[test]
    fn test_parse_secs() {
        fn test(s: &str, expected: std::result::Result<Duration, String>) {
            assert_eq!(super::parse_secs(s), expected.map_err(Into::into));
        }

        test("NaNs", Err("NaN".to_owned()));
        test("+NaNs", Err("NaN".to_owned()));
        test("-NaNs", Err("NaN".to_owned()));
        test("-infs", Err("negative".to_owned()));
        test("-42s", Err("negative".to_owned()));
        test("1S", Err("invalid or missing unit".to_owned()));

        test("infs", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("+infs", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("1Ys", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("1Zs", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("100Es", Ok(Duration::new(u64::MAX, 999_999_999)));
        test("1Ps", Ok(Duration::new(999_999_999_999_999, 983_222_784)));
        test("100Ts", Ok(Duration::new(99_999_999_999_999, 991_611_392)));
        test("10Ts", Ok(Duration::new(10_000_000_000_000, 0)));
        test("1Gs", Ok(Duration::new(1_000_000_000, 0)));
        test("1Ms", Ok(Duration::new(1_000_000, 0)));
        test("1ks", Ok(Duration::new(1000, 0)));
        test("1hs", Ok(Duration::new(100, 0)));
        test("1das", Ok(Duration::new(10, 0)));
        test("0s", Ok(Duration::new(0, 0)));
        test(" 10 s ", Ok(Duration::new(10, 0)));
        test("42s", Ok(Duration::new(42, 0)));
        test("123.456s", Ok(Duration::new(123, 456_000_000)));
        test("1E3s", Ok(Duration::new(1000, 0)));
        test("1ds", Ok(Duration::new(0, 100_000_000)));
        test("1cs", Ok(Duration::new(0, 10_000_000)));
        test("42ms", Ok(Duration::new(0, 42_000_000)));
        test("123.456ms", Ok(Duration::new(0, 123_456_000)));
        test("1E3ms", Ok(Duration::new(1, 0)));
        test("42μs", Ok(Duration::new(0, 42_000)));
        test("123.456μs", Ok(Duration::new(0, 123_456)));
        test("1E3μs", Ok(Duration::new(0, 1_000_000)));
        test("1us", Ok(Duration::new(0, 1000)));
        test("42ns", Ok(Duration::new(0, 42)));
        test("123.456ns", Ok(Duration::new(0, 123)));
        test("1E3ns", Ok(Duration::new(0, 1000)));
        test("1000ps", Ok(Duration::new(0, 1)));
        test("100ps", Ok(Duration::new(0, 0)));
        test("1E6fs", Ok(Duration::new(0, 1)));
        test("1E9as", Ok(Duration::new(0, 1)));
        test("1E12zs", Ok(Duration::new(0, 1)));
        test("1E15ys", Ok(Duration::new(0, 1)));
    }

    #[test]
    fn test_ser_secs() -> serde_json::Result<()> {
        fn test(dur: Duration, expected: &str) -> serde_json::Result<()> {
            #[derive(Serialize)]
            struct S {
                #[serde(serialize_with = "super::ser_secs")]
                value: Option<Duration>,
            }

            #[derive(Deserialize)]
            struct D {
                value: Option<String>,
            }

            let actual = serde_json::to_string(&S { value: Some(dur) })?;
            let actual = serde_json::from_str::<D>(&actual)?.value;
            assert_eq!(actual.as_ref().map(AsRef::as_ref), Some(expected));
            Ok(())
        }

        test(Duration::from_secs(42), "42s")?;
        test(Duration::from_millis(42), "0.042s")?;
        test(Duration::from_nanos(42), "0.000000042s")?;
        test(Duration::new(1234, 567_800_000), "1234.5678s")
    }

    #[test]
    fn test_ser_millis() -> serde_json::Result<()> {
        fn test(dur: Duration, expected: &str) -> serde_json::Result<()> {
            #[derive(Serialize)]
            struct S {
                #[serde(serialize_with = "super::ser_millis")]
                value: Option<Duration>,
            }

            #[derive(Deserialize)]
            struct D {
                value: Option<String>,
            }

            let actual = serde_json::to_string(&S { value: Some(dur) })?;
            let actual = serde_json::from_str::<D>(&actual)?.value;
            assert_eq!(actual.as_ref().map(AsRef::as_ref), Some(expected));
            Ok(())
        }

        test(Duration::from_secs(42), "42000ms")?;
        test(Duration::from_millis(42), "42ms")?;
        test(Duration::from_nanos(42), "0.000042ms")?;
        test(Duration::new(1234, 567_800_000), "1234567.8ms")
    }

    #[test]
    fn test_millis_rounded_up() {
        fn test(dur: Duration, expected: u128) {
            assert_eq!(dur.millis_rounded_up(), expected);
        }

        test(Duration::from_nanos(0), 0);
        test(Duration::from_nanos(1), 1);
        test(Duration::from_nanos(999_999), 1);
        test(Duration::from_nanos(1_000_000), 1);
        test(Duration::from_nanos(1_000_001), 2);
        test(Duration::from_secs(1), 1000);
    }
}
<|endoftext|>


<commit_before>//! Unix handling of child processes
//!
//! Right now the only "fancy" thing about this is how we implement the
//! `Future` implementation on `Child` to get the exit status. Unix offers
//! no way to register a child with epoll, and the only real way to get a
//! notification when a process exits is the SIGCHLD signal.
//!
//! Signal handling in general is *super* hairy and complicated, and it's even
//! more complicated here with the fact that signals are coalesced, so we may
//! not get a SIGCHLD-per-child.
//!
//! Our best approximation here is to check *all spawned processes* for all
//! SIGCHLD signals received. To do that we create a `Signal`, implemented in
//! the `tokio-signal` crate, which is a stream over signals being received.
//!
//! Later when we poll the process's exit status we simply check to see if a
//! SIGCHLD has happened since we last checked, and while that returns "yes" we
//! keep trying.
//!
//! Note that this means that this isn't really scalable, but then again
//! processes in general aren't scalable (e.g. millions) so it shouldn't be that
//! bad in theory...

extern crate libc;
extern crate tokio_signal;

use std::io;
use std::os::unix::prelude::*;
use std::process::{self, ExitStatus};

use futures::future::FlattenStream;
use futures::{Future, Poll, Async, Stream};
use mio::unix::{EventedFd, UnixReady};
use mio::{PollOpt, Ready, Token};
use mio::event::Evented;
use mio;
use self::tokio_signal::unix::Signal;
use std::fmt;
use tokio_io::IoFuture;
use tokio_reactor::{Handle, PollEvented};

#[must_use = "futures do nothing unless polled"]
pub struct Child {
    inner: process::Child,
    reaped: bool,
    sigchld: FlattenStream<IoFuture<Signal>>,
}

impl fmt::Debug for Child {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Child")
            .field("pid", &self.inner.id())
            .field("inner", &self.inner)
            .field("reaped", &self.reaped)
            .field("sigchld", &"..")
            .finish()
    }
}

impl Child {
    pub fn new(inner: process::Child, handle: &Handle) -> Child {
        Child {
            inner: inner,
            reaped: false,
            sigchld: Signal::with_handle(libc::SIGCHLD, handle).flatten_stream(),
        }
    }

    pub fn register_stdin(&mut self, handle: &Handle)
                          -> io::Result<Option<ChildStdin>> {
        stdio(self.inner.stdin.take(), handle)
    }

    pub fn register_stdout(&mut self, handle: &Handle)
                           -> io::Result<Option<ChildStdout>> {
        stdio(self.inner.stdout.take(), handle)
    }

    pub fn register_stderr(&mut self, handle: &Handle)
                           -> io::Result<Option<ChildStderr>> {
        stdio(self.inner.stderr.take(), handle)
    }

    pub fn id(&self) -> u32 {
        self.inner.id()
    }

    pub fn kill(&mut self) -> io::Result<()> {
        if !self.reaped {
            // NB: SIGKILL cannnot be caught, so the process will definitely exit immediately.
            // We're not waiting for the process itself but for the kernel to execute the kill.
            self.inner.kill()?;
            let _ = self.try_wait(true);
        }

        Ok(())
    }

    pub fn poll_exit(&mut self) -> Poll<ExitStatus, io::Error> {
        loop {
            // Ensure that once we've successfully waited we won't try to
            // `kill` above.
            if let Some(e) = try!(self.try_wait(false)) {
                return Ok(e.into())
            }

            // If the child hasn't exited yet, then it's our responsibility to
            // ensure the current task gets notified when it might be able to
            // make progress.
            //
            // As described in `spawn` above, we just indicate that we can
            // next make progress once a SIGCHLD is received.
            if try!(self.sigchld.poll()).is_not_ready() {
                return Ok(Async::NotReady)
            }
        }
    }

    fn try_wait(&mut self, block_on_wait: bool) -> io::Result<Option<ExitStatus>> {
        assert!(!self.reaped);
        let exit = try!(try_wait_process(self.id() as libc::pid_t, block_on_wait));

        if let Some(_) = exit {
            self.reaped = true;
        }

        Ok(exit)
    }
}

fn try_wait_process(id: libc::pid_t, block_on_wait: bool) -> io::Result<Option<ExitStatus>> {
    let wait_flags = if block_on_wait { 0 } else { libc::WNOHANG };
    let mut status = 0;

    loop {
        match unsafe { libc::waitpid(id, &mut status, wait_flags) } {
            0 => return Ok(None),
            n if n < 0 => {
                let err = io::Error::last_os_error();
                if err.kind() == io::ErrorKind::Interrupted {
                    continue
                }
                return Err(err)
            }
            n => {
                assert_eq!(n, id);
                return Ok(Some(ExitStatus::from_raw(status)))
            }
        }
    }
}

#[derive(Debug)]
pub struct Fd<T>(T);

impl<T: io::Read> io::Read for Fd<T> {
    fn read(&mut self, bytes: &mut [u8]) -> io::Result<usize> {
        self.0.read(bytes)
    }
}

impl<T: io::Write> io::Write for Fd<T> {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        self.0.write(bytes)
    }

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

pub type ChildStdin = PollEvented<Fd<process::ChildStdin>>;
pub type ChildStdout = PollEvented<Fd<process::ChildStdout>>;
pub type ChildStderr = PollEvented<Fd<process::ChildStderr>>;

impl<T> Evented for Fd<T> where T: AsRawFd {
    fn register(&self,
                poll: &mio::Poll,
                token: Token,
                interest: Ready,
                opts: PollOpt)
                -> io::Result<()> {
        EventedFd(&self.0.as_raw_fd()).register(poll,
                                                token,
                                                interest | UnixReady::hup(),
                                                opts)
    }

    fn reregister(&self,
                  poll: &mio::Poll,
                  token: Token,
                  interest: Ready,
                  opts: PollOpt)
                  -> io::Result<()> {
        EventedFd(&self.0.as_raw_fd()).reregister(poll,
                                                  token,
                                                  interest | UnixReady::hup(),
                                                  opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        EventedFd(&self.0.as_raw_fd()).deregister(poll)
    }
}

fn stdio<T>(option: Option<T>, handle: &Handle)
            -> io::Result<Option<PollEvented<Fd<T>>>>
    where T: AsRawFd
{
    let io = match option {
        Some(io) => io,
        None => return Ok(None),
    };

    // Set the fd to nonblocking before we pass it to the event loop
    unsafe {
        let fd = io.as_raw_fd();
        let r = libc::fcntl(fd, libc::F_GETFL);
        if r == -1 {
            return Err(io::Error::last_os_error())
        }
        let r = libc::fcntl(fd, libc::F_SETFL, r | libc::O_NONBLOCK);
        if r == -1 {
            return Err(io::Error::last_os_error())
        }
    }
    let io = try!(PollEvented::new_with_handle(Fd(io), handle));
    Ok(Some(io))
}
<commit_msg>process: Unix: preregister for signal notifications before polling child<commit_after>//! Unix handling of child processes
//!
//! Right now the only "fancy" thing about this is how we implement the
//! `Future` implementation on `Child` to get the exit status. Unix offers
//! no way to register a child with epoll, and the only real way to get a
//! notification when a process exits is the SIGCHLD signal.
//!
//! Signal handling in general is *super* hairy and complicated, and it's even
//! more complicated here with the fact that signals are coalesced, so we may
//! not get a SIGCHLD-per-child.
//!
//! Our best approximation here is to check *all spawned processes* for all
//! SIGCHLD signals received. To do that we create a `Signal`, implemented in
//! the `tokio-signal` crate, which is a stream over signals being received.
//!
//! Later when we poll the process's exit status we simply check to see if a
//! SIGCHLD has happened since we last checked, and while that returns "yes" we
//! keep trying.
//!
//! Note that this means that this isn't really scalable, but then again
//! processes in general aren't scalable (e.g. millions) so it shouldn't be that
//! bad in theory...

extern crate libc;
extern crate tokio_signal;

use std::io;
use std::os::unix::prelude::*;
use std::process::{self, ExitStatus};

use futures::future::FlattenStream;
use futures::{Future, Poll, Async, Stream};
use mio::unix::{EventedFd, UnixReady};
use mio::{PollOpt, Ready, Token};
use mio::event::Evented;
use mio;
use self::tokio_signal::unix::Signal;
use std::fmt;
use tokio_io::IoFuture;
use tokio_reactor::{Handle, PollEvented};

#[must_use = "futures do nothing unless polled"]
pub struct Child {
    inner: process::Child,
    reaped: bool,
    sigchld: FlattenStream<IoFuture<Signal>>,
}

impl fmt::Debug for Child {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("Child")
            .field("pid", &self.inner.id())
            .field("inner", &self.inner)
            .field("reaped", &self.reaped)
            .field("sigchld", &"..")
            .finish()
    }
}

impl Child {
    pub fn new(inner: process::Child, handle: &Handle) -> Child {
        Child {
            inner: inner,
            reaped: false,
            sigchld: Signal::with_handle(libc::SIGCHLD, handle).flatten_stream(),
        }
    }

    pub fn register_stdin(&mut self, handle: &Handle)
                          -> io::Result<Option<ChildStdin>> {
        stdio(self.inner.stdin.take(), handle)
    }

    pub fn register_stdout(&mut self, handle: &Handle)
                           -> io::Result<Option<ChildStdout>> {
        stdio(self.inner.stdout.take(), handle)
    }

    pub fn register_stderr(&mut self, handle: &Handle)
                           -> io::Result<Option<ChildStderr>> {
        stdio(self.inner.stderr.take(), handle)
    }

    pub fn id(&self) -> u32 {
        self.inner.id()
    }

    pub fn kill(&mut self) -> io::Result<()> {
        if !self.reaped {
            // NB: SIGKILL cannnot be caught, so the process will definitely exit immediately.
            // We're not waiting for the process itself but for the kernel to execute the kill.
            self.inner.kill()?;
            let _ = self.try_wait(true);
        }

        Ok(())
    }

    pub fn poll_exit(&mut self) -> Poll<ExitStatus, io::Error> {
        loop {
            // Ensure we don't register for additional notifications
            // if the child has already finished.
            if self.reaped {
                return Ok(Async::NotReady);
            }

            // If the child hasn't exited yet, then it's our responsibility to
            // ensure the current task gets notified when it might be able to
            // make progress.
            //
            // As described in `spawn` above, we just indicate that we can
            // next make progress once a SIGCHLD is received.
            //
            // However, we will register for a notification on the next signal
            // BEFORE we poll the child. Otherwise it is possible that the child
            // can exit and the signal can arrive after we last polled the child,
            // but before we've registered for a notification on the next signal
            // (this can cause a deadlock if there are no more spawned children
            // which can generate a different signal for us). A side effect of
            // pre-registering for signal notifications is that when the child
            // exits, we will have already registered for an additional
            // notification we don't need to consume. If another signal arrives,
            // this future's task will be notified/woken up again. Since the
            // futures model allows for spurious wake ups this extra wakeup
            // should not cause significant issues with parent futures.
            let registered_interest = try!(self.sigchld.poll()).is_not_ready();

            if let Some(e) = try!(self.try_wait(false)) {
                return Ok(e.into());
            }

            // If our attempt to poll for the next signal was not ready, then
            // we've arranged for our task to get notified and we can bail out.
            if registered_interest {
                return Ok(Async::NotReady);
            } else {
                // Otherwise, if the signal stream delivered a signal to us, we
                // won't get notified at the next signal, so we'll loop and try
                // again.
                continue;
            }
        }
    }

    fn try_wait(&mut self, block_on_wait: bool) -> io::Result<Option<ExitStatus>> {
        assert!(!self.reaped);
        let exit = try!(try_wait_process(self.id() as libc::pid_t, block_on_wait));

        if let Some(_) = exit {
            self.reaped = true;
        }

        Ok(exit)
    }
}

fn try_wait_process(id: libc::pid_t, block_on_wait: bool) -> io::Result<Option<ExitStatus>> {
    let wait_flags = if block_on_wait { 0 } else { libc::WNOHANG };
    let mut status = 0;

    loop {
        match unsafe { libc::waitpid(id, &mut status, wait_flags) } {
            0 => return Ok(None),
            n if n < 0 => {
                let err = io::Error::last_os_error();
                if err.kind() == io::ErrorKind::Interrupted {
                    continue
                }
                return Err(err)
            }
            n => {
                assert_eq!(n, id);
                return Ok(Some(ExitStatus::from_raw(status)))
            }
        }
    }
}

#[derive(Debug)]
pub struct Fd<T>(T);

impl<T: io::Read> io::Read for Fd<T> {
    fn read(&mut self, bytes: &mut [u8]) -> io::Result<usize> {
        self.0.read(bytes)
    }
}

impl<T: io::Write> io::Write for Fd<T> {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        self.0.write(bytes)
    }

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

pub type ChildStdin = PollEvented<Fd<process::ChildStdin>>;
pub type ChildStdout = PollEvented<Fd<process::ChildStdout>>;
pub type ChildStderr = PollEvented<Fd<process::ChildStderr>>;

impl<T> Evented for Fd<T> where T: AsRawFd {
    fn register(&self,
                poll: &mio::Poll,
                token: Token,
                interest: Ready,
                opts: PollOpt)
                -> io::Result<()> {
        EventedFd(&self.0.as_raw_fd()).register(poll,
                                                token,
                                                interest | UnixReady::hup(),
                                                opts)
    }

    fn reregister(&self,
                  poll: &mio::Poll,
                  token: Token,
                  interest: Ready,
                  opts: PollOpt)
                  -> io::Result<()> {
        EventedFd(&self.0.as_raw_fd()).reregister(poll,
                                                  token,
                                                  interest | UnixReady::hup(),
                                                  opts)
    }

    fn deregister(&self, poll: &mio::Poll) -> io::Result<()> {
        EventedFd(&self.0.as_raw_fd()).deregister(poll)
    }
}

fn stdio<T>(option: Option<T>, handle: &Handle)
            -> io::Result<Option<PollEvented<Fd<T>>>>
    where T: AsRawFd
{
    let io = match option {
        Some(io) => io,
        None => return Ok(None),
    };

    // Set the fd to nonblocking before we pass it to the event loop
    unsafe {
        let fd = io.as_raw_fd();
        let r = libc::fcntl(fd, libc::F_GETFL);
        if r == -1 {
            return Err(io::Error::last_os_error())
        }
        let r = libc::fcntl(fd, libc::F_SETFL, r | libc::O_NONBLOCK);
        if r == -1 {
            return Err(io::Error::last_os_error())
        }
    }
    let io = try!(PollEvented::new_with_handle(Fd(io), handle));
    Ok(Some(io))
}
<|endoftext|>


<commit_before>extern crate libc;

use std::ffi::CString;
use std::fs::File;
use std::io::{Error, ErrorKind, Result};
use std::mem;
use std::os::unix::ffi::OsStrExt;
use std::os::unix::fs::MetadataExt;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::os::unix::raw::off_t;
use std::path::Path;

pub fn duplicate(file: &File) -> Result<File> {
    unsafe {
        let fd = libc::dup(file.as_raw_fd());

        if fd < 0 {
            Err(Error::last_os_error())
        } else {
            Ok(File::from_raw_fd(fd))
        }
    }
}

pub fn lock_shared(file: &File) -> Result<()> {
    flock(file, libc::LOCK_SH)
}

pub fn lock_exclusive(file: &File) -> Result<()> {
    flock(file, libc::LOCK_EX)
}

pub fn try_lock_shared(file: &File) -> Result<()> {
    flock(file, libc::LOCK_SH | libc::LOCK_NB)
}

pub fn try_lock_exclusive(file: &File) -> Result<()> {
    flock(file, libc::LOCK_EX | libc::LOCK_NB)
}

pub fn unlock(file: &File) -> Result<()> {
    flock(file, libc::LOCK_UN)
}

pub fn lock_error() -> Error {
    Error::from_raw_os_error(libc::EWOULDBLOCK)
}

fn flock(file: &File, flag: libc::c_int) -> Result<()> {
    let ret = unsafe { libc::flock(file.as_raw_fd(), flag) };
    if ret < 0 { Err(Error::last_os_error()) } else { Ok(()) }
}

pub fn allocated_size(file: &File) -> Result<u64> {
    file.metadata().map(|m| m.blocks() as u64 * 512)
}

#[cfg(any(target_os = "linux",
          target_os = "freebsd",
          target_os = "android",
          target_os = "nacl"))]
pub fn allocate(file: &File, len: u64) -> Result<()> {
    let ret = unsafe { libc::posix_fallocate(file.as_raw_fd(), 0, len as off_t) };
    if ret == 0 { Ok(()) } else { Err(Error::last_os_error()) }
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
pub fn allocate(file: &File, len: u64) -> Result<()> {
    let stat = try!(file.metadata());

    if len > stat.blocks() as u64 * 512 {
        let mut fstore = libc::fstore_t {
            fst_flags: libc::F_ALLOCATECONTIG,
            fst_posmode: libc::F_PEOFPOSMODE,
            fst_offset: 0,
            fst_length: len as off_t,
            fst_bytesalloc: 0,
        };

        let ret = unsafe { libc::fcntl(file.as_raw_fd(), libc::F_PREALLOCATE, &fstore) };
        if ret == -1 {
            // Unable to allocate contiguous disk space; attempt to allocate non-contiguously.
            fstore.fst_flags = libc::F_ALLOCATEALL;
            let ret = unsafe { libc::fcntl(file.as_raw_fd(), libc::F_PREALLOCATE, &fstore) };
            if ret == -1 {
                return Err(Error::last_os_error());
            }
        }
    }

    if len > stat.size() as u64 {
        file.set_len(len)
    } else {
        Ok(())
    }
}

#[cfg(any(target_os = "openbsd",
          target_os = "netbsd"))]
pub fn allocate(file: &File, len: u64) -> Result<()> {
    // No file allocation API available, just set the length if necessary.
    if len > try!(file.metadata()).len() as u64 {
        file.set_len(len)
    } else {
        Ok(())
    }
}

fn statvfs<P>(path: P) -> Result<libc::statvfs> where P: AsRef<Path> {
    let cstr = match CString::new(path.as_ref().as_os_str().as_bytes()) {
        Ok(cstr) => cstr,
        Err(..) => return Err(Error::new(ErrorKind::InvalidInput, "path contained a null")),
    };

    unsafe {
        let mut stat: libc::statvfs = mem::zeroed();
        // danburkert/fs2-rs#1: cast is necessary for platforms where c_char != u8.
        if libc::statvfs(cstr.as_ptr() as *const _, &mut stat) == -1 {
            Err(Error::last_os_error())
        } else {
            Ok(stat)
        }
    }
}

pub fn free_space<P>(path: P) -> Result<u64> where P: AsRef<Path> {
    statvfs(path).map(|statvfs| statvfs.f_frsize as u64 * statvfs.f_bfree as u64)
}

pub fn available_space<P>(path: P) -> Result<u64> where P: AsRef<Path> {
    statvfs(path).map(|statvfs| statvfs.f_frsize as u64 * statvfs.f_bavail as u64)
}

pub fn total_space<P>(path: P) -> Result<u64> where P: AsRef<Path> {
    statvfs(path).map(|statvfs| statvfs.f_frsize as u64 * statvfs.f_blocks as u64)
}

pub fn allocation_granularity<P>(path: P) -> Result<u64> where P: AsRef<Path> {
    statvfs(path).map(|statvfs| statvfs.f_frsize as u64)
}

#[cfg(test)]
mod test {
    extern crate tempdir;
    extern crate libc;

    use std::fs::{self, File};
    use std::os::unix::io::AsRawFd;

    use {FileExt, lock_contended_error};

    /// The duplicate method returns a file with a new file descriptor.
    #[test]
    fn duplicate_new_fd() {
        let tempdir = tempdir::TempDir::new("fs2").unwrap();
        let path = tempdir.path().join("fs2");
        let file1 = fs::OpenOptions::new().write(true).create(true).open(&path).unwrap();
        let file2 = file1.duplicate().unwrap();
        assert!(file1.as_raw_fd() != file2.as_raw_fd());
    }

    /// The duplicate method should preservesthe close on exec flag.
    #[test]
    fn duplicate_cloexec() {

        fn flags(file: &File) -> libc::c_int {
            unsafe { libc::fcntl(file.as_raw_fd(), libc::F_GETFL, 0) }
        }

        let tempdir = tempdir::TempDir::new("fs2").unwrap();
        let path = tempdir.path().join("fs2");
        let file1 = fs::OpenOptions::new().write(true).create(true).open(&path).unwrap();
        let file2 = file1.duplicate().unwrap();

        assert_eq!(flags(&file1), flags(&file2));
    }

    /// Tests that locking a file descriptor will replace any existing locks
    /// held on the file descriptor.
    #[test]
    fn lock_replace() {
        let tempdir = tempdir::TempDir::new("fs2").unwrap();
        let path = tempdir.path().join("fs2");
        let file1 = fs::OpenOptions::new().create(true).open(&path).unwrap();
        let file2 = fs::OpenOptions::new().create(true).open(&path).unwrap();

        // Creating a shared lock will drop an exclusive lock.
        file1.lock_exclusive().unwrap();
        file1.lock_shared().unwrap();
        file2.lock_shared().unwrap();

        // Attempting to replace a shared lock with an exclusive lock will fail
        // with multiple lock holders, and remove the original shared lock.
        assert_eq!(file2.try_lock_exclusive().unwrap_err().raw_os_error(),
                   lock_contended_error().raw_os_error());
        file1.lock_shared().unwrap();
    }

    /// Tests that locks are shared among duplicated file descriptors.
    #[test]
    fn lock_duplicate() {
        let tempdir = tempdir::TempDir::new("fs2").unwrap();
        let path = tempdir.path().join("fs2");
        let file1 = fs::OpenOptions::new().create(true).open(&path).unwrap();
        let file2 = file1.duplicate().unwrap();
        let file3 = fs::OpenOptions::new().create(true).open(&path).unwrap();

        // Create a lock through fd1, then replace it through fd2.
        file1.lock_shared().unwrap();
        file2.lock_exclusive().unwrap();
        assert_eq!(file3.try_lock_shared().unwrap_err().raw_os_error(),
                   lock_contended_error().raw_os_error());

        // Either of the file descriptors should be able to unlock.
        file1.unlock().unwrap();
        file3.lock_shared().unwrap();
    }
}
<commit_msg>Use libc::off_t instead of std::*::off_t<commit_after>extern crate libc;

use std::ffi::CString;
use std::fs::File;
use std::io::{Error, ErrorKind, Result};
use std::mem;
use std::os::unix::ffi::OsStrExt;
use std::os::unix::fs::MetadataExt;
use std::os::unix::io::{AsRawFd, FromRawFd};
use std::path::Path;

pub fn duplicate(file: &File) -> Result<File> {
    unsafe {
        let fd = libc::dup(file.as_raw_fd());

        if fd < 0 {
            Err(Error::last_os_error())
        } else {
            Ok(File::from_raw_fd(fd))
        }
    }
}

pub fn lock_shared(file: &File) -> Result<()> {
    flock(file, libc::LOCK_SH)
}

pub fn lock_exclusive(file: &File) -> Result<()> {
    flock(file, libc::LOCK_EX)
}

pub fn try_lock_shared(file: &File) -> Result<()> {
    flock(file, libc::LOCK_SH | libc::LOCK_NB)
}

pub fn try_lock_exclusive(file: &File) -> Result<()> {
    flock(file, libc::LOCK_EX | libc::LOCK_NB)
}

pub fn unlock(file: &File) -> Result<()> {
    flock(file, libc::LOCK_UN)
}

pub fn lock_error() -> Error {
    Error::from_raw_os_error(libc::EWOULDBLOCK)
}

fn flock(file: &File, flag: libc::c_int) -> Result<()> {
    let ret = unsafe { libc::flock(file.as_raw_fd(), flag) };
    if ret < 0 { Err(Error::last_os_error()) } else { Ok(()) }
}

pub fn allocated_size(file: &File) -> Result<u64> {
    file.metadata().map(|m| m.blocks() as u64 * 512)
}

#[cfg(any(target_os = "linux",
          target_os = "freebsd",
          target_os = "android",
          target_os = "nacl"))]
pub fn allocate(file: &File, len: u64) -> Result<()> {
    let ret = unsafe { libc::posix_fallocate(file.as_raw_fd(), 0, len as libc::off_t) };
    if ret == 0 { Ok(()) } else { Err(Error::last_os_error()) }
}

#[cfg(any(target_os = "macos", target_os = "ios"))]
pub fn allocate(file: &File, len: u64) -> Result<()> {
    let stat = try!(file.metadata());

    if len > stat.blocks() as u64 * 512 {
        let mut fstore = libc::fstore_t {
            fst_flags: libc::F_ALLOCATECONTIG,
            fst_posmode: libc::F_PEOFPOSMODE,
            fst_offset: 0,
            fst_length: len as libc::off_t,
            fst_bytesalloc: 0,
        };

        let ret = unsafe { libc::fcntl(file.as_raw_fd(), libc::F_PREALLOCATE, &fstore) };
        if ret == -1 {
            // Unable to allocate contiguous disk space; attempt to allocate non-contiguously.
            fstore.fst_flags = libc::F_ALLOCATEALL;
            let ret = unsafe { libc::fcntl(file.as_raw_fd(), libc::F_PREALLOCATE, &fstore) };
            if ret == -1 {
                return Err(Error::last_os_error());
            }
        }
    }

    if len > stat.size() as u64 {
        file.set_len(len)
    } else {
        Ok(())
    }
}

#[cfg(any(target_os = "openbsd",
          target_os = "netbsd"))]
pub fn allocate(file: &File, len: u64) -> Result<()> {
    // No file allocation API available, just set the length if necessary.
    if len > try!(file.metadata()).len() as u64 {
        file.set_len(len)
    } else {
        Ok(())
    }
}

fn statvfs<P>(path: P) -> Result<libc::statvfs> where P: AsRef<Path> {
    let cstr = match CString::new(path.as_ref().as_os_str().as_bytes()) {
        Ok(cstr) => cstr,
        Err(..) => return Err(Error::new(ErrorKind::InvalidInput, "path contained a null")),
    };

    unsafe {
        let mut stat: libc::statvfs = mem::zeroed();
        // danburkert/fs2-rs#1: cast is necessary for platforms where c_char != u8.
        if libc::statvfs(cstr.as_ptr() as *const _, &mut stat) == -1 {
            Err(Error::last_os_error())
        } else {
            Ok(stat)
        }
    }
}

pub fn free_space<P>(path: P) -> Result<u64> where P: AsRef<Path> {
    statvfs(path).map(|statvfs| statvfs.f_frsize as u64 * statvfs.f_bfree as u64)
}

pub fn available_space<P>(path: P) -> Result<u64> where P: AsRef<Path> {
    statvfs(path).map(|statvfs| statvfs.f_frsize as u64 * statvfs.f_bavail as u64)
}

pub fn total_space<P>(path: P) -> Result<u64> where P: AsRef<Path> {
    statvfs(path).map(|statvfs| statvfs.f_frsize as u64 * statvfs.f_blocks as u64)
}

pub fn allocation_granularity<P>(path: P) -> Result<u64> where P: AsRef<Path> {
    statvfs(path).map(|statvfs| statvfs.f_frsize as u64)
}

#[cfg(test)]
mod test {
    extern crate tempdir;
    extern crate libc;

    use std::fs::{self, File};
    use std::os::unix::io::AsRawFd;

    use {FileExt, lock_contended_error};

    /// The duplicate method returns a file with a new file descriptor.
    #[test]
    fn duplicate_new_fd() {
        let tempdir = tempdir::TempDir::new("fs2").unwrap();
        let path = tempdir.path().join("fs2");
        let file1 = fs::OpenOptions::new().write(true).create(true).open(&path).unwrap();
        let file2 = file1.duplicate().unwrap();
        assert!(file1.as_raw_fd() != file2.as_raw_fd());
    }

    /// The duplicate method should preservesthe close on exec flag.
    #[test]
    fn duplicate_cloexec() {

        fn flags(file: &File) -> libc::c_int {
            unsafe { libc::fcntl(file.as_raw_fd(), libc::F_GETFL, 0) }
        }

        let tempdir = tempdir::TempDir::new("fs2").unwrap();
        let path = tempdir.path().join("fs2");
        let file1 = fs::OpenOptions::new().write(true).create(true).open(&path).unwrap();
        let file2 = file1.duplicate().unwrap();

        assert_eq!(flags(&file1), flags(&file2));
    }

    /// Tests that locking a file descriptor will replace any existing locks
    /// held on the file descriptor.
    #[test]
    fn lock_replace() {
        let tempdir = tempdir::TempDir::new("fs2").unwrap();
        let path = tempdir.path().join("fs2");
        let file1 = fs::OpenOptions::new().create(true).open(&path).unwrap();
        let file2 = fs::OpenOptions::new().create(true).open(&path).unwrap();

        // Creating a shared lock will drop an exclusive lock.
        file1.lock_exclusive().unwrap();
        file1.lock_shared().unwrap();
        file2.lock_shared().unwrap();

        // Attempting to replace a shared lock with an exclusive lock will fail
        // with multiple lock holders, and remove the original shared lock.
        assert_eq!(file2.try_lock_exclusive().unwrap_err().raw_os_error(),
                   lock_contended_error().raw_os_error());
        file1.lock_shared().unwrap();
    }

    /// Tests that locks are shared among duplicated file descriptors.
    #[test]
    fn lock_duplicate() {
        let tempdir = tempdir::TempDir::new("fs2").unwrap();
        let path = tempdir.path().join("fs2");
        let file1 = fs::OpenOptions::new().create(true).open(&path).unwrap();
        let file2 = file1.duplicate().unwrap();
        let file3 = fs::OpenOptions::new().create(true).open(&path).unwrap();

        // Create a lock through fd1, then replace it through fd2.
        file1.lock_shared().unwrap();
        file2.lock_exclusive().unwrap();
        assert_eq!(file3.try_lock_shared().unwrap_err().raw_os_error(),
                   lock_contended_error().raw_os_error());

        // Either of the file descriptors should be able to unlock.
        file1.unlock().unwrap();
        file3.lock_shared().unwrap();
    }
}
<|endoftext|>