<commit_before>use clap::{App, Arg, ArgGroup, ArgMatches};
use cluster_test::{
    aws::Aws,
    cluster::Cluster,
    deployment::{DeploymentManager, SOURCE_TAG, TESTED_TAG},
    effects::{Effect, Reboot},
    experiments::{Experiment, RebootRandomValidators},
    health::{DebugPortLogThread, HealthCheckRunner, LogTail},
    slack::SlackClient,
    suite::ExperimentSuite,
};
use failure::{
    self,
    prelude::{bail, format_err},
};
use std::{
    collections::HashSet,
    env,
    sync::mpsc::{self, TryRecvError},
    thread,
    time::{Duration, Instant},
};
use termion::{color, style};

const HEALTH_POLL_INTERVAL: Duration = Duration::from_secs(5);

pub fn main() {
    let matches = arg_matches();
    let mut runner = ClusterTestRunner::setup(&matches);

    if matches.is_present(ARG_RUN) {
        runner.run_suite_in_loop();
    } else if matches.is_present(ARG_RUN_ONCE) {
        let experiment = RebootRandomValidators::new(3, &runner.cluster);
        runner.run_single_experiment(Box::new(experiment)).unwrap();
    } else if matches.is_present(ARG_TAIL_LOGS) {
        runner.tail_logs();
    } else if matches.is_present(ARG_HEALTH_CHECK) {
        runner.run_health_check();
    } else if matches.is_present(ARG_WIPE_ALL_DB) {
        runner.wipe_all_db();
    } else if matches.is_present(ARG_REBOOT) {
        runner.reboot(matches.values_of_lossy(ARG_REBOOT).unwrap());
    }
}

struct ClusterTestRunner {
    logs: LogTail,
    cluster: Cluster,
    health_check_runner: HealthCheckRunner,
    deployment_manager: DeploymentManager,
    experiment_interval: Duration,
    slack: Option<SlackClient>,
}

impl ClusterTestRunner {
    /// Discovers cluster, setup log, etc
    pub fn setup(matches: &ArgMatches) -> Self {
        let workplace = matches
            .value_of(ARG_WORKPLACE)
            .expect("workplace should be set");
        let aws = Aws::new(workplace.into());
        let peers = matches.values_of_lossy(ARG_PEERS);
        let cluster = Cluster::discover(&aws).expect("Failed to discover cluster");
        let cluster = match peers {
            None => cluster,
            Some(peers) => cluster.sub_cluster(peers),
        };
        println!("Discovered {} peers", cluster.instances().len());
        let log_tail_started = Instant::now();
        let logs = DebugPortLogThread::spawn_new(&cluster);
        let log_tail_startup_time = Instant::now() - log_tail_started;
        println!(
            "Log tail thread started in {} ms",
            log_tail_startup_time.as_millis()
        );
        let health_check_runner = HealthCheckRunner::new_all(cluster.clone());
        let experiment_interval_sec = match env::var("EXPERIMENT_INTERVAL") {
            Ok(s) => s.parse().expect("EXPERIMENT_INTERVAL env is not a number"),
            Err(..) => 15,
        };
        let experiment_interval = Duration::from_secs(experiment_interval_sec);
        let deployment_manager = DeploymentManager::new(aws.clone(), cluster.clone());
        let slack = SlackClient::try_new_from_environment();
        Self {
            logs,
            cluster,
            health_check_runner,
            deployment_manager,
            experiment_interval,
            slack,
        }
    }

    pub fn run_suite_in_loop(&mut self) {
        let mut hash_to_tag = None;
        loop {
            if let Some(hash) = self.deployment_manager.latest_hash_changed() {
                self.slack_message(format!(
                    "Deploying new version of `{}` tag: `{}`",
                    SOURCE_TAG, hash
                ));
                if let Err(e) = self.redeploy(hash.clone()) {
                    self.report_failure(format!("Failed to deploy `{}`: {}", hash, e));
                    return;
                }
                self.slack_message(format!("Deployed new version `{}`", hash));
                hash_to_tag = Some(hash);
            }
            let suite = ExperimentSuite::new_pre_release(&self.cluster);
            if let Err(e) = self.run_suite(suite) {
                self.report_failure(format!("{}", e));
                return;
            }
            if let Some(hash_to_tag) = hash_to_tag.take() {
                self.slack_message(format!(
                    "Test suite succeed first time for `{}`",
                    hash_to_tag
                ));
                if let Err(e) = self
                    .deployment_manager
                    .tag_tested_image(hash_to_tag.clone())
                {
                    self.report_failure(format!("Failed to tag tested image: {}", e));
                    return;
                }
                self.slack_message(format!("Tagged `{}` as `{}`", hash_to_tag, TESTED_TAG));
            }
            thread::sleep(self.experiment_interval);
        }
    }

    fn report_failure(&self, msg: String) {
        self.slack_message(msg);
    }

    fn redeploy(&mut self, hash: String) -> failure::Result<()> {
        self.deployment_manager.redeploy(hash);
        println!("Waiting for 60 seconds to allow ECS to restart tasks...");
        thread::sleep(Duration::from_secs(60));
        println!("Waiting until all validators healthy after deployment");
        self.wait_until_all_healthy()
    }

    fn run_suite(&mut self, suite: ExperimentSuite) -> failure::Result<()> {
        println!("Starting suite");
        let suite_started = Instant::now();
        for experiment in suite.experiments {
            let experiment_name = format!("{}", experiment);
            self.run_single_experiment(experiment).map_err(move |e| {
                format_err!("Experiment `{}` failed: `{}`", experiment_name, e)
            })?;
            thread::sleep(self.experiment_interval);
        }
        println!(
            "Suite completed in {:?}",
            Instant::now().duration_since(suite_started)
        );
        Ok(())
    }

    pub fn run_single_experiment(
        &mut self,
        experiment: Box<dyn Experiment>,
    ) -> failure::Result<()> {
        let events = self.logs.recv_all();
        if !self.health_check_runner.run(&events).is_empty() {
            bail!("Some validators are unhealthy before experiment started");
        }

        println!(
            "{}Starting experiment {}{}{}{}",
            style::Bold,
            color::Fg(color::Blue),
            experiment,
            color::Fg(color::Reset),
            style::Reset
        );
        let affected_validators = experiment.affected_validators();
        let (exp_result_sender, exp_result_recv) = mpsc::channel();
        thread::spawn(move || {
            let result = experiment.run();
            exp_result_sender
                .send(result)
                .expect("Failed to send experiment result");
        });

        // We expect experiments completes and cluster go into healthy state within timeout
        let experiment_deadline = Instant::now() + Duration::from_secs(5 * 60);

        loop {
            if Instant::now() > experiment_deadline {
                bail!("Experiment did not complete in time");
            }
            let deadline = Instant::now() + HEALTH_POLL_INTERVAL;
            // Receive all events that arrived to aws log tail within next 1 second
            // This assumes so far that event propagation time is << 1s, this need to be refined
            // in future to account for actual event propagation delay
            let events = self.logs.recv_all_until_deadline(deadline);
            let failed_validators = self.health_check_runner.run(&events);
            for failed in failed_validators {
                if !affected_validators.contains(&failed) {
                    bail!(
                        "Validator {} failed, not expected for this experiment",
                        failed
                    );
                }
            }
            match exp_result_recv.try_recv() {
                Ok(result) => {
                    result.expect("Failed to run experiment");
                    break;
                }
                Err(TryRecvError::Empty) => {
                    // Experiment in progress, continue monitoring health
                }
                Err(TryRecvError::Disconnected) => {
                    panic!("Experiment thread exited without returning result");
                }
            }
        }

        println!(
            "{}Experiment finished, waiting until all affected validators recover{}",
            style::Bold,
            style::Reset
        );

        for validator in affected_validators.iter() {
            self.health_check_runner.invalidate(validator);
        }

        loop {
            if Instant::now() > experiment_deadline {
                bail!("Cluster did not become healthy in time");
            }
            let deadline = Instant::now() + HEALTH_POLL_INTERVAL;
            // Receive all events that arrived to aws log tail within next 1 second
            // This assumes so far that event propagation time is << 1s, this need to be refined
            // in future to account for actual event propagation delay
            let events = self.logs.recv_all_until_deadline(deadline);
            let failed_validators = self.health_check_runner.run(&events);
            let mut still_affected_validator = HashSet::new();
            for failed in failed_validators {
                if !affected_validators.contains(&failed) {
                    bail!(
                        "Validator {} failed, not expected for this experiment",
                        failed
                    );
                }
                still_affected_validator.insert(failed);
            }
            if still_affected_validator.is_empty() {
                break;
            }
        }

        println!("Experiment completed");
        Ok(())
    }

    fn run_health_check(&mut self) {
        loop {
            let deadline = Instant::now() + Duration::from_secs(1);
            // Receive all events that arrived to aws log tail within next 1 second
            // This assumes so far that event propagation time is << 1s, this need to be refined
            // in future to account for actual event propagation delay
            let events = self.logs.recv_all_until_deadline(deadline);
            self.health_check_runner.run(&events);
        }
    }

    fn wait_until_all_healthy(&mut self) -> failure::Result<()> {
        let wait_deadline = Instant::now() + Duration::from_secs(10 * 60);
        for instance in self.cluster.instances() {
            self.health_check_runner.invalidate(instance.short_hash());
        }
        loop {
            let now = Instant::now();
            if now > wait_deadline {
                bail!("Validators did not become healthy after deployment");
            }
            let deadline = now + HEALTH_POLL_INTERVAL;
            let events = self.logs.recv_all_until_deadline(deadline);
            if self.health_check_runner.run(&events).is_empty() {
                break;
            }
        }
        Ok(())
    }

    fn tail_logs(self) {
        for log in self.logs.event_receiver {
            println!("{:?}", log);
        }
    }

    fn slack_message(&self, msg: String) {
        println!("{}", msg);
        if let Some(ref slack) = self.slack {
            if let Err(e) = slack.send_message(&msg) {
                println!("Failed to send slack message: {}", e);
            }
        }
    }

    fn wipe_all_db(self) {
        println!("Going to wipe db on all validators in cluster!");
        println!("Waiting 10 seconds before proceed");
        thread::sleep(Duration::from_secs(10));
        println!("Starting...");
        for instance in self.cluster.instances() {
            if let Err(e) = instance.run_cmd_tee_err(vec!["sudo", "rm", "-rf", "/data/libra/"]) {
                println!("Failed to wipe {}: {:?}", instance, e);
            }
        }
        println!("Done");
    }

    fn reboot(self, validators: Vec<String>) {
        let mut reboots = vec![];
        for validator in validators {
            match self.cluster.get_instance(&validator) {
                None => println!("{} not found", validator),
                Some(instance) => {
                    println!("Rebooting {}", validator);
                    let reboot = Reboot::new(instance.clone());
                    if let Err(err) = reboot.apply() {
                        println!("Failed to reboot {}: {:?}", validator, err);
                    } else {
                        reboots.push(reboot);
                    }
                }
            }
        }
        println!("Waiting to complete");
        while reboots.iter().any(|r| !r.is_complete()) {
            thread::sleep(Duration::from_secs(5));
        }
        println!("Completed");
    }
}

const ARG_WORKPLACE: &str = "workplace";
// Actions:
const ARG_TAIL_LOGS: &str = "tail-logs";
const ARG_HEALTH_CHECK: &str = "health-check";
const ARG_RUN: &str = "run";
const ARG_RUN_ONCE: &str = "run-once";
const ARG_WIPE_ALL_DB: &str = "wipe-all-db";
const ARG_REBOOT: &str = "reboot";
const ARG_PEERS: &str = "peers";

fn arg_matches() -> ArgMatches<'static> {
    // Parameters
    let workplace = Arg::with_name(ARG_WORKPLACE)
        .long("--workplace")
        .short("-w")
        .takes_value(true)
        .required(true);
    let peers = Arg::with_name(ARG_PEERS)
        .long("--peers")
        .short("-p")
        .takes_value(true)
        .use_delimiter(true);
    // Actions
    let wipe_all_db = Arg::with_name(ARG_WIPE_ALL_DB).long("--wipe-all-db");
    let run = Arg::with_name(ARG_RUN).long("--run");
    let run_once = Arg::with_name(ARG_RUN_ONCE).long("--run-once");
    let tail_logs = Arg::with_name(ARG_TAIL_LOGS).long("--tail-logs");
    let health_check = Arg::with_name(ARG_HEALTH_CHECK).long("--health-check");
    let reboot = Arg::with_name(ARG_REBOOT)
        .long("--reboot")
        .takes_value(true)
        .use_delimiter(true);
    // This grouping requires one and only one action (tail logs, run test, etc)
    let action_group = ArgGroup::with_name("action")
        .args(&[
            ARG_TAIL_LOGS,
            ARG_RUN,
            ARG_RUN_ONCE,
            ARG_HEALTH_CHECK,
            ARG_WIPE_ALL_DB,
            ARG_REBOOT,
        ])
        .required(true);
    App::new("cluster_test")
        .author("Libra Association <opensource@libra.org>")
        .group(action_group)
        .args(&[
            // parameters
            workplace,
            peers,
            // actions
            run,
            run_once,
            tail_logs,
            health_check,
            wipe_all_db,
            reboot,
        ])
        .get_matches()
}
<commit_msg>[cluster_test] Increase experiment_deadline 5m->10m<commit_after>use clap::{App, Arg, ArgGroup, ArgMatches};
use cluster_test::{
    aws::Aws,
    cluster::Cluster,
    deployment::{DeploymentManager, SOURCE_TAG, TESTED_TAG},
    effects::{Effect, Reboot},
    experiments::{Experiment, RebootRandomValidators},
    health::{DebugPortLogThread, HealthCheckRunner, LogTail},
    slack::SlackClient,
    suite::ExperimentSuite,
};
use failure::{
    self,
    prelude::{bail, format_err},
};
use std::{
    collections::HashSet,
    env,
    sync::mpsc::{self, TryRecvError},
    thread,
    time::{Duration, Instant},
};
use termion::{color, style};

const HEALTH_POLL_INTERVAL: Duration = Duration::from_secs(5);

pub fn main() {
    let matches = arg_matches();
    let mut runner = ClusterTestRunner::setup(&matches);

    if matches.is_present(ARG_RUN) {
        runner.run_suite_in_loop();
    } else if matches.is_present(ARG_RUN_ONCE) {
        let experiment = RebootRandomValidators::new(3, &runner.cluster);
        runner.run_single_experiment(Box::new(experiment)).unwrap();
    } else if matches.is_present(ARG_TAIL_LOGS) {
        runner.tail_logs();
    } else if matches.is_present(ARG_HEALTH_CHECK) {
        runner.run_health_check();
    } else if matches.is_present(ARG_WIPE_ALL_DB) {
        runner.wipe_all_db();
    } else if matches.is_present(ARG_REBOOT) {
        runner.reboot(matches.values_of_lossy(ARG_REBOOT).unwrap());
    }
}

struct ClusterTestRunner {
    logs: LogTail,
    cluster: Cluster,
    health_check_runner: HealthCheckRunner,
    deployment_manager: DeploymentManager,
    experiment_interval: Duration,
    slack: Option<SlackClient>,
}

impl ClusterTestRunner {
    /// Discovers cluster, setup log, etc
    pub fn setup(matches: &ArgMatches) -> Self {
        let workplace = matches
            .value_of(ARG_WORKPLACE)
            .expect("workplace should be set");
        let aws = Aws::new(workplace.into());
        let peers = matches.values_of_lossy(ARG_PEERS);
        let cluster = Cluster::discover(&aws).expect("Failed to discover cluster");
        let cluster = match peers {
            None => cluster,
            Some(peers) => cluster.sub_cluster(peers),
        };
        println!("Discovered {} peers", cluster.instances().len());
        let log_tail_started = Instant::now();
        let logs = DebugPortLogThread::spawn_new(&cluster);
        let log_tail_startup_time = Instant::now() - log_tail_started;
        println!(
            "Log tail thread started in {} ms",
            log_tail_startup_time.as_millis()
        );
        let health_check_runner = HealthCheckRunner::new_all(cluster.clone());
        let experiment_interval_sec = match env::var("EXPERIMENT_INTERVAL") {
            Ok(s) => s.parse().expect("EXPERIMENT_INTERVAL env is not a number"),
            Err(..) => 15,
        };
        let experiment_interval = Duration::from_secs(experiment_interval_sec);
        let deployment_manager = DeploymentManager::new(aws.clone(), cluster.clone());
        let slack = SlackClient::try_new_from_environment();
        Self {
            logs,
            cluster,
            health_check_runner,
            deployment_manager,
            experiment_interval,
            slack,
        }
    }

    pub fn run_suite_in_loop(&mut self) {
        let mut hash_to_tag = None;
        loop {
            if let Some(hash) = self.deployment_manager.latest_hash_changed() {
                self.slack_message(format!(
                    "Deploying new version of `{}` tag: `{}`",
                    SOURCE_TAG, hash
                ));
                if let Err(e) = self.redeploy(hash.clone()) {
                    self.report_failure(format!("Failed to deploy `{}`: {}", hash, e));
                    return;
                }
                self.slack_message(format!("Deployed new version `{}`", hash));
                hash_to_tag = Some(hash);
            }
            let suite = ExperimentSuite::new_pre_release(&self.cluster);
            if let Err(e) = self.run_suite(suite) {
                self.report_failure(format!("{}", e));
                return;
            }
            if let Some(hash_to_tag) = hash_to_tag.take() {
                self.slack_message(format!(
                    "Test suite succeed first time for `{}`",
                    hash_to_tag
                ));
                if let Err(e) = self
                    .deployment_manager
                    .tag_tested_image(hash_to_tag.clone())
                {
                    self.report_failure(format!("Failed to tag tested image: {}", e));
                    return;
                }
                self.slack_message(format!("Tagged `{}` as `{}`", hash_to_tag, TESTED_TAG));
            }
            thread::sleep(self.experiment_interval);
        }
    }

    fn report_failure(&self, msg: String) {
        self.slack_message(msg);
    }

    fn redeploy(&mut self, hash: String) -> failure::Result<()> {
        self.deployment_manager.redeploy(hash);
        println!("Waiting for 60 seconds to allow ECS to restart tasks...");
        thread::sleep(Duration::from_secs(60));
        println!("Waiting until all validators healthy after deployment");
        self.wait_until_all_healthy()
    }

    fn run_suite(&mut self, suite: ExperimentSuite) -> failure::Result<()> {
        println!("Starting suite");
        let suite_started = Instant::now();
        for experiment in suite.experiments {
            let experiment_name = format!("{}", experiment);
            self.run_single_experiment(experiment).map_err(move |e| {
                format_err!("Experiment `{}` failed: `{}`", experiment_name, e)
            })?;
            thread::sleep(self.experiment_interval);
        }
        println!(
            "Suite completed in {:?}",
            Instant::now().duration_since(suite_started)
        );
        Ok(())
    }

    pub fn run_single_experiment(
        &mut self,
        experiment: Box<dyn Experiment>,
    ) -> failure::Result<()> {
        let events = self.logs.recv_all();
        if !self.health_check_runner.run(&events).is_empty() {
            bail!("Some validators are unhealthy before experiment started");
        }

        println!(
            "{}Starting experiment {}{}{}{}",
            style::Bold,
            color::Fg(color::Blue),
            experiment,
            color::Fg(color::Reset),
            style::Reset
        );
        let affected_validators = experiment.affected_validators();
        let (exp_result_sender, exp_result_recv) = mpsc::channel();
        thread::spawn(move || {
            let result = experiment.run();
            exp_result_sender
                .send(result)
                .expect("Failed to send experiment result");
        });

        // We expect experiments completes and cluster go into healthy state within timeout
        let experiment_deadline = Instant::now() + Duration::from_secs(10 * 60);

        loop {
            if Instant::now() > experiment_deadline {
                bail!("Experiment did not complete in time");
            }
            let deadline = Instant::now() + HEALTH_POLL_INTERVAL;
            // Receive all events that arrived to aws log tail within next 1 second
            // This assumes so far that event propagation time is << 1s, this need to be refined
            // in future to account for actual event propagation delay
            let events = self.logs.recv_all_until_deadline(deadline);
            let failed_validators = self.health_check_runner.run(&events);
            for failed in failed_validators {
                if !affected_validators.contains(&failed) {
                    bail!(
                        "Validator {} failed, not expected for this experiment",
                        failed
                    );
                }
            }
            match exp_result_recv.try_recv() {
                Ok(result) => {
                    result.expect("Failed to run experiment");
                    break;
                }
                Err(TryRecvError::Empty) => {
                    // Experiment in progress, continue monitoring health
                }
                Err(TryRecvError::Disconnected) => {
                    panic!("Experiment thread exited without returning result");
                }
            }
        }

        println!(
            "{}Experiment finished, waiting until all affected validators recover{}",
            style::Bold,
            style::Reset
        );

        for validator in affected_validators.iter() {
            self.health_check_runner.invalidate(validator);
        }

        loop {
            if Instant::now() > experiment_deadline {
                bail!("Cluster did not become healthy in time");
            }
            let deadline = Instant::now() + HEALTH_POLL_INTERVAL;
            // Receive all events that arrived to aws log tail within next 1 second
            // This assumes so far that event propagation time is << 1s, this need to be refined
            // in future to account for actual event propagation delay
            let events = self.logs.recv_all_until_deadline(deadline);
            let failed_validators = self.health_check_runner.run(&events);
            let mut still_affected_validator = HashSet::new();
            for failed in failed_validators {
                if !affected_validators.contains(&failed) {
                    bail!(
                        "Validator {} failed, not expected for this experiment",
                        failed
                    );
                }
                still_affected_validator.insert(failed);
            }
            if still_affected_validator.is_empty() {
                break;
            }
        }

        println!("Experiment completed");
        Ok(())
    }

    fn run_health_check(&mut self) {
        loop {
            let deadline = Instant::now() + Duration::from_secs(1);
            // Receive all events that arrived to aws log tail within next 1 second
            // This assumes so far that event propagation time is << 1s, this need to be refined
            // in future to account for actual event propagation delay
            let events = self.logs.recv_all_until_deadline(deadline);
            self.health_check_runner.run(&events);
        }
    }

    fn wait_until_all_healthy(&mut self) -> failure::Result<()> {
        let wait_deadline = Instant::now() + Duration::from_secs(10 * 60);
        for instance in self.cluster.instances() {
            self.health_check_runner.invalidate(instance.short_hash());
        }
        loop {
            let now = Instant::now();
            if now > wait_deadline {
                bail!("Validators did not become healthy after deployment");
            }
            let deadline = now + HEALTH_POLL_INTERVAL;
            let events = self.logs.recv_all_until_deadline(deadline);
            if self.health_check_runner.run(&events).is_empty() {
                break;
            }
        }
        Ok(())
    }

    fn tail_logs(self) {
        for log in self.logs.event_receiver {
            println!("{:?}", log);
        }
    }

    fn slack_message(&self, msg: String) {
        println!("{}", msg);
        if let Some(ref slack) = self.slack {
            if let Err(e) = slack.send_message(&msg) {
                println!("Failed to send slack message: {}", e);
            }
        }
    }

    fn wipe_all_db(self) {
        println!("Going to wipe db on all validators in cluster!");
        println!("Waiting 10 seconds before proceed");
        thread::sleep(Duration::from_secs(10));
        println!("Starting...");
        for instance in self.cluster.instances() {
            if let Err(e) = instance.run_cmd_tee_err(vec!["sudo", "rm", "-rf", "/data/libra/"]) {
                println!("Failed to wipe {}: {:?}", instance, e);
            }
        }
        println!("Done");
    }

    fn reboot(self, validators: Vec<String>) {
        let mut reboots = vec![];
        for validator in validators {
            match self.cluster.get_instance(&validator) {
                None => println!("{} not found", validator),
                Some(instance) => {
                    println!("Rebooting {}", validator);
                    let reboot = Reboot::new(instance.clone());
                    if let Err(err) = reboot.apply() {
                        println!("Failed to reboot {}: {:?}", validator, err);
                    } else {
                        reboots.push(reboot);
                    }
                }
            }
        }
        println!("Waiting to complete");
        while reboots.iter().any(|r| !r.is_complete()) {
            thread::sleep(Duration::from_secs(5));
        }
        println!("Completed");
    }
}

const ARG_WORKPLACE: &str = "workplace";
// Actions:
const ARG_TAIL_LOGS: &str = "tail-logs";
const ARG_HEALTH_CHECK: &str = "health-check";
const ARG_RUN: &str = "run";
const ARG_RUN_ONCE: &str = "run-once";
const ARG_WIPE_ALL_DB: &str = "wipe-all-db";
const ARG_REBOOT: &str = "reboot";
const ARG_PEERS: &str = "peers";

fn arg_matches() -> ArgMatches<'static> {
    // Parameters
    let workplace = Arg::with_name(ARG_WORKPLACE)
        .long("--workplace")
        .short("-w")
        .takes_value(true)
        .required(true);
    let peers = Arg::with_name(ARG_PEERS)
        .long("--peers")
        .short("-p")
        .takes_value(true)
        .use_delimiter(true);
    // Actions
    let wipe_all_db = Arg::with_name(ARG_WIPE_ALL_DB).long("--wipe-all-db");
    let run = Arg::with_name(ARG_RUN).long("--run");
    let run_once = Arg::with_name(ARG_RUN_ONCE).long("--run-once");
    let tail_logs = Arg::with_name(ARG_TAIL_LOGS).long("--tail-logs");
    let health_check = Arg::with_name(ARG_HEALTH_CHECK).long("--health-check");
    let reboot = Arg::with_name(ARG_REBOOT)
        .long("--reboot")
        .takes_value(true)
        .use_delimiter(true);
    // This grouping requires one and only one action (tail logs, run test, etc)
    let action_group = ArgGroup::with_name("action")
        .args(&[
            ARG_TAIL_LOGS,
            ARG_RUN,
            ARG_RUN_ONCE,
            ARG_HEALTH_CHECK,
            ARG_WIPE_ALL_DB,
            ARG_REBOOT,
        ])
        .required(true);
    App::new("cluster_test")
        .author("Libra Association <opensource@libra.org>")
        .group(action_group)
        .args(&[
            // parameters
            workplace,
            peers,
            // actions
            run,
            run_once,
            tail_logs,
            health_check,
            wipe_all_db,
            reboot,
        ])
        .get_matches()
}
<|endoftext|>


<commit_before>#![allow(dead_code)]

//! This module is meant to directly expose the C nif api. The functions in here
//! should be unsafe, and should not do any validation.
//! Currently we are using erlang_nif_sys, but we do want to be able to easily
//! switch to something else in the future if required.
//!
//! While reexporting from erlang_nif_sys directly would remove a lot of code, it
//! is preferred to keep the function signatures in here for future use/reference.

extern crate erlang_nif_sys;
extern crate std;
//extern crate libc;
//pub use libc::{ c_int, c_uint, c_double, size_t, c_void, c_uchar };
pub use std::os::raw::{ c_int, c_uint, c_double, c_void, c_uchar };
pub type size_t = usize;

pub type NIF_ENV = *mut erlang_nif_sys::ErlNifEnv;
pub type NIF_TERM = size_t;
pub type NIF_BINARY = *mut erlang_nif_sys::ErlNifBinary;
pub type NIF_RESOURCE_TYPE = *const erlang_nif_sys::ErlNifResourceType;

pub type NIF_RESOURCE_HANDLE = *const c_void;
pub type MUTABLE_NIF_RESOURCE_HANDLE = *mut c_void;

pub type NifResourceDtor = extern "C" fn(r_env: NIF_ENV, obj: MUTABLE_NIF_RESOURCE_HANDLE) -> ();
pub type NifResourceFlags = erlang_nif_sys::ErlNifResourceFlags;

pub enum NIF_ERROR {
    BAD_ARG
}

pub type DEF_NIF_FUNC = erlang_nif_sys::ErlNifFunc;
pub type DEF_NIF_ENTRY = erlang_nif_sys::ErlNifEntry;
pub use self::erlang_nif_sys::NIF_MAJOR_VERSION;
pub use self::erlang_nif_sys::NIF_MINOR_VERSION;
pub use self::erlang_nif_sys::ErlNifResourceFlags as NIF_RESOURCE_FLAGS;

#[repr(C)]
pub enum ErlNifTaskFlags {
    ERL_NIF_NORMAL_JOB = 0,
    ERL_NIF_DIRTY_JOB_CPU_BOUND = 1,
    ERL_NIF_DIRTY_JOB_IO_BOUND = 2
}

pub unsafe fn enif_make_badarg(env: NIF_ENV) -> NIF_TERM {
    erlang_nif_sys::enif_make_badarg(env)
}

pub unsafe fn enif_alloc_env() -> NIF_ENV {
    erlang_nif_sys::enif_alloc_env()
}

pub unsafe fn enif_raise_exception(env: NIF_ENV, reason: NIF_TERM) -> NIF_TERM {
    erlang_nif_sys::enif_raise_exception(env, reason)
}

pub unsafe fn enif_make_copy(dest_env: NIF_ENV, source_term: NIF_TERM) -> NIF_TERM {
    erlang_nif_sys::enif_make_copy(dest_env, source_term)
}

// Term type checks
pub unsafe fn enif_is_atom(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_atom(env, term)
}
pub unsafe fn enif_is_list(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_list(env, term)
}
pub unsafe fn enif_is_empty_list(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_empty_list(env, term)
}

// Atoms
pub unsafe fn enif_make_atom_len(env: NIF_ENV, string: *const u8, length: size_t) -> NIF_TERM {
    erlang_nif_sys::enif_make_atom_len(env, string, length)
}

// Binaries
pub unsafe fn enif_release_binary(bin_ref: NIF_BINARY) {
    erlang_nif_sys::enif_release_binary(bin_ref)
}
pub unsafe fn enif_alloc_binary(size: size_t, bin_ref: NIF_BINARY) -> c_int {
    erlang_nif_sys::enif_alloc_binary(size, bin_ref)
}
pub unsafe fn enif_make_binary(env: NIF_ENV, bin_ref: NIF_BINARY) -> NIF_TERM {
    erlang_nif_sys::enif_make_binary(env, bin_ref)
}
pub unsafe fn enif_inspect_binary(env: NIF_ENV, term: NIF_TERM, bin_ref: NIF_BINARY) -> c_int {
    erlang_nif_sys::enif_inspect_binary(env, term, bin_ref)
}
pub unsafe fn enif_make_sub_binary(env: NIF_ENV, bin_term: NIF_TERM, pos: size_t, size: size_t) -> NIF_TERM {
    erlang_nif_sys::enif_make_sub_binary(env, bin_term, pos, size)
}

// Maps
pub unsafe fn enif_get_map_value(env: NIF_ENV, map: NIF_TERM, key: NIF_TERM, value: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_get_map_value(env, map, key, value)
}
pub unsafe fn enif_get_map_size(env: NIF_ENV, map: NIF_TERM, size: *mut size_t) -> c_int {
    erlang_nif_sys::enif_get_map_size(env, map, size)
}
pub unsafe fn enif_make_new_map(env: NIF_ENV) -> NIF_TERM {
    erlang_nif_sys::enif_make_new_map(env)
}
pub unsafe fn enif_make_map_put(env: NIF_ENV, map_in: NIF_TERM, key: NIF_TERM, value: NIF_TERM, map_out: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_make_map_put(env, map_in, key, value, map_out)
}
pub unsafe fn enif_make_map_update(env: NIF_ENV, map_in: NIF_TERM, key: NIF_TERM, value: NIF_TERM, map_out: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_make_map_update(env, map_in, key, value, map_out)
}
pub unsafe fn enif_make_map_remove(env: NIF_ENV, map_in: NIF_TERM, key: NIF_TERM, map_out: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_make_map_remove(env, map_in, key, map_out)
}

// Tuple
pub unsafe fn enif_get_tuple(env: NIF_ENV, term: NIF_TERM, arity: *mut c_int, array: *mut *const NIF_TERM) -> c_int {
    erlang_nif_sys::enif_get_tuple(env, term, arity, array)
}
pub unsafe fn enif_make_tuple_from_array(env: NIF_ENV, terms: *const NIF_TERM, terms_len: c_uint) -> NIF_TERM {
    erlang_nif_sys::enif_make_tuple_from_array(env, terms, terms_len)
}

// List
pub unsafe fn enif_get_list_cell(env: NIF_ENV, list: NIF_TERM, head: *mut NIF_TERM, tail: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_get_list_cell(env, list, head, tail)
}
pub unsafe fn enif_get_list_length(env: NIF_ENV, list: NIF_TERM, length: *mut c_uint) -> c_int {
    erlang_nif_sys::enif_get_list_length(env, list, length)
}
pub unsafe fn enif_make_list_from_array(env: NIF_ENV, arr: *const NIF_TERM, cnt: c_uint) -> NIF_TERM {
    erlang_nif_sys::enif_make_list_from_array(env, arr, cnt)
}

// Resources
pub unsafe fn enif_open_resource_type(env: NIF_ENV, module_str: *const c_uchar, name: *const c_uchar,
                                      dtor: Option<NifResourceDtor>, flags: NifResourceFlags, tried: *mut NifResourceFlags
                                      ) -> NIF_RESOURCE_TYPE {
    erlang_nif_sys::enif_open_resource_type(env, module_str, name, dtor, flags, tried)
}
pub unsafe fn enif_alloc_resource(typ: NIF_RESOURCE_TYPE, size: usize) -> NIF_RESOURCE_HANDLE {
    erlang_nif_sys::enif_alloc_resource(typ, size)
}
pub unsafe fn enif_make_resource(env: NIF_ENV, obj: NIF_RESOURCE_HANDLE) -> NIF_TERM {
    erlang_nif_sys::enif_make_resource(env, obj)
}
pub unsafe fn enif_get_resource(env: NIF_ENV, term: NIF_TERM, typ: NIF_RESOURCE_TYPE, objp: *mut NIF_RESOURCE_HANDLE) -> c_int {
    erlang_nif_sys::enif_get_resource(env, term, typ, objp)
}
pub unsafe fn enif_release_resource(obj: NIF_RESOURCE_HANDLE) {
    erlang_nif_sys::enif_release_resource(obj)
}
pub unsafe fn enif_keep_resource(obj: NIF_RESOURCE_HANDLE) {
    erlang_nif_sys::enif_keep_resource(obj)
}

// Numbers
macro_rules! wrap_number {
    ($typ: ty, $encode: ident, $decode: ident) => {
        pub unsafe fn $encode(env: NIF_ENV, val: $typ) -> NIF_TERM {
            erlang_nif_sys::$encode(env, val)
        }
        pub unsafe fn $decode(env: NIF_ENV, term: NIF_TERM, dest: *mut $typ) -> c_int {
            erlang_nif_sys::$decode(env, term, dest)
        }
    }
}
wrap_number!(c_int, enif_make_int, enif_get_int);
wrap_number!(c_uint, enif_make_uint, enif_get_uint);
wrap_number!(i64, enif_make_int64, enif_get_int64);
wrap_number!(u64, enif_make_uint64, enif_get_uint64);
wrap_number!(c_double, enif_make_double, enif_get_double);
<commit_msg>Add missing `enif_is_*` functions<commit_after>#![allow(dead_code)]

//! This module is meant to directly expose the C nif api. The functions in here
//! should be unsafe, and should not do any validation.
//! Currently we are using erlang_nif_sys, but we do want to be able to easily
//! switch to something else in the future if required.
//!
//! While reexporting from erlang_nif_sys directly would remove a lot of code, it
//! is preferred to keep the function signatures in here for future use/reference.

extern crate erlang_nif_sys;
extern crate std;
//extern crate libc;
//pub use libc::{ c_int, c_uint, c_double, size_t, c_void, c_uchar };
pub use std::os::raw::{ c_int, c_uint, c_double, c_void, c_uchar };
pub type size_t = usize;

pub type NIF_ENV = *mut erlang_nif_sys::ErlNifEnv;
pub type NIF_TERM = size_t;
pub type NIF_BINARY = *mut erlang_nif_sys::ErlNifBinary;
pub type NIF_RESOURCE_TYPE = *const erlang_nif_sys::ErlNifResourceType;

pub type NIF_RESOURCE_HANDLE = *const c_void;
pub type MUTABLE_NIF_RESOURCE_HANDLE = *mut c_void;

pub type NifResourceDtor = extern "C" fn(r_env: NIF_ENV, obj: MUTABLE_NIF_RESOURCE_HANDLE) -> ();
pub type NifResourceFlags = erlang_nif_sys::ErlNifResourceFlags;

pub enum NIF_ERROR {
    BAD_ARG
}

pub type DEF_NIF_FUNC = erlang_nif_sys::ErlNifFunc;
pub type DEF_NIF_ENTRY = erlang_nif_sys::ErlNifEntry;
pub use self::erlang_nif_sys::NIF_MAJOR_VERSION;
pub use self::erlang_nif_sys::NIF_MINOR_VERSION;
pub use self::erlang_nif_sys::ErlNifResourceFlags as NIF_RESOURCE_FLAGS;

#[repr(C)]
pub enum ErlNifTaskFlags {
    ERL_NIF_NORMAL_JOB = 0,
    ERL_NIF_DIRTY_JOB_CPU_BOUND = 1,
    ERL_NIF_DIRTY_JOB_IO_BOUND = 2
}

pub unsafe fn enif_make_badarg(env: NIF_ENV) -> NIF_TERM {
    erlang_nif_sys::enif_make_badarg(env)
}

pub unsafe fn enif_alloc_env() -> NIF_ENV {
    erlang_nif_sys::enif_alloc_env()
}

pub unsafe fn enif_raise_exception(env: NIF_ENV, reason: NIF_TERM) -> NIF_TERM {
    erlang_nif_sys::enif_raise_exception(env, reason)
}

pub unsafe fn enif_make_copy(dest_env: NIF_ENV, source_term: NIF_TERM) -> NIF_TERM {
    erlang_nif_sys::enif_make_copy(dest_env, source_term)
}

// Term type checks
pub unsafe fn enif_is_atom(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_atom(env, term)
}
pub unsafe fn enif_is_binary(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_binary(env, term)
}
pub unsafe fn enif_is_empty_list(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_empty_list(env, term)
}
pub unsafe fn enif_is_exception(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_exception(env, term)
}
pub unsafe fn enif_is_fun(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_fun(env, term)
}
pub unsafe fn enif_is_list(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_list(env, term)
}
pub unsafe fn enif_is_map(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_map(env, term)
}
pub unsafe fn enif_is_pid(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_pid(env, term)
}
pub unsafe fn enif_is_port(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_port(env, term)
}
pub unsafe fn enif_is_ref(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_ref(env, term)
}
pub unsafe fn enif_is_tuple(env: NIF_ENV, term: NIF_TERM) -> c_int {
    erlang_nif_sys::enif_is_tuple(env, term)
}

// Atoms
pub unsafe fn enif_make_atom_len(env: NIF_ENV, string: *const u8, length: size_t) -> NIF_TERM {
    erlang_nif_sys::enif_make_atom_len(env, string, length)
}

// Binaries
pub unsafe fn enif_release_binary(bin_ref: NIF_BINARY) {
    erlang_nif_sys::enif_release_binary(bin_ref)
}
pub unsafe fn enif_alloc_binary(size: size_t, bin_ref: NIF_BINARY) -> c_int {
    erlang_nif_sys::enif_alloc_binary(size, bin_ref)
}
pub unsafe fn enif_make_binary(env: NIF_ENV, bin_ref: NIF_BINARY) -> NIF_TERM {
    erlang_nif_sys::enif_make_binary(env, bin_ref)
}
pub unsafe fn enif_inspect_binary(env: NIF_ENV, term: NIF_TERM, bin_ref: NIF_BINARY) -> c_int {
    erlang_nif_sys::enif_inspect_binary(env, term, bin_ref)
}
pub unsafe fn enif_make_sub_binary(env: NIF_ENV, bin_term: NIF_TERM, pos: size_t, size: size_t) -> NIF_TERM {
    erlang_nif_sys::enif_make_sub_binary(env, bin_term, pos, size)
}

// Maps
pub unsafe fn enif_get_map_value(env: NIF_ENV, map: NIF_TERM, key: NIF_TERM, value: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_get_map_value(env, map, key, value)
}
pub unsafe fn enif_get_map_size(env: NIF_ENV, map: NIF_TERM, size: *mut size_t) -> c_int {
    erlang_nif_sys::enif_get_map_size(env, map, size)
}
pub unsafe fn enif_make_new_map(env: NIF_ENV) -> NIF_TERM {
    erlang_nif_sys::enif_make_new_map(env)
}
pub unsafe fn enif_make_map_put(env: NIF_ENV, map_in: NIF_TERM, key: NIF_TERM, value: NIF_TERM, map_out: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_make_map_put(env, map_in, key, value, map_out)
}
pub unsafe fn enif_make_map_update(env: NIF_ENV, map_in: NIF_TERM, key: NIF_TERM, value: NIF_TERM, map_out: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_make_map_update(env, map_in, key, value, map_out)
}
pub unsafe fn enif_make_map_remove(env: NIF_ENV, map_in: NIF_TERM, key: NIF_TERM, map_out: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_make_map_remove(env, map_in, key, map_out)
}

// Tuple
pub unsafe fn enif_get_tuple(env: NIF_ENV, term: NIF_TERM, arity: *mut c_int, array: *mut *const NIF_TERM) -> c_int {
    erlang_nif_sys::enif_get_tuple(env, term, arity, array)
}
pub unsafe fn enif_make_tuple_from_array(env: NIF_ENV, terms: *const NIF_TERM, terms_len: c_uint) -> NIF_TERM {
    erlang_nif_sys::enif_make_tuple_from_array(env, terms, terms_len)
}

// List
pub unsafe fn enif_get_list_cell(env: NIF_ENV, list: NIF_TERM, head: *mut NIF_TERM, tail: *mut NIF_TERM) -> c_int {
    erlang_nif_sys::enif_get_list_cell(env, list, head, tail)
}
pub unsafe fn enif_get_list_length(env: NIF_ENV, list: NIF_TERM, length: *mut c_uint) -> c_int {
    erlang_nif_sys::enif_get_list_length(env, list, length)
}
pub unsafe fn enif_make_list_from_array(env: NIF_ENV, arr: *const NIF_TERM, cnt: c_uint) -> NIF_TERM {
    erlang_nif_sys::enif_make_list_from_array(env, arr, cnt)
}

// Resources
pub unsafe fn enif_open_resource_type(env: NIF_ENV, module_str: *const c_uchar, name: *const c_uchar,
                                      dtor: Option<NifResourceDtor>, flags: NifResourceFlags, tried: *mut NifResourceFlags
                                      ) -> NIF_RESOURCE_TYPE {
    erlang_nif_sys::enif_open_resource_type(env, module_str, name, dtor, flags, tried)
}
pub unsafe fn enif_alloc_resource(typ: NIF_RESOURCE_TYPE, size: usize) -> NIF_RESOURCE_HANDLE {
    erlang_nif_sys::enif_alloc_resource(typ, size)
}
pub unsafe fn enif_make_resource(env: NIF_ENV, obj: NIF_RESOURCE_HANDLE) -> NIF_TERM {
    erlang_nif_sys::enif_make_resource(env, obj)
}
pub unsafe fn enif_get_resource(env: NIF_ENV, term: NIF_TERM, typ: NIF_RESOURCE_TYPE, objp: *mut NIF_RESOURCE_HANDLE) -> c_int {
    erlang_nif_sys::enif_get_resource(env, term, typ, objp)
}
pub unsafe fn enif_release_resource(obj: NIF_RESOURCE_HANDLE) {
    erlang_nif_sys::enif_release_resource(obj)
}
pub unsafe fn enif_keep_resource(obj: NIF_RESOURCE_HANDLE) {
    erlang_nif_sys::enif_keep_resource(obj)
}

// Numbers
macro_rules! wrap_number {
    ($typ: ty, $encode: ident, $decode: ident) => {
        pub unsafe fn $encode(env: NIF_ENV, val: $typ) -> NIF_TERM {
            erlang_nif_sys::$encode(env, val)
        }
        pub unsafe fn $decode(env: NIF_ENV, term: NIF_TERM, dest: *mut $typ) -> c_int {
            erlang_nif_sys::$decode(env, term, dest)
        }
    }
}
wrap_number!(c_int, enif_make_int, enif_get_int);
wrap_number!(c_uint, enif_make_uint, enif_get_uint);
wrap_number!(i64, enif_make_int64, enif_get_int64);
wrap_number!(u64, enif_make_uint64, enif_get_uint64);
wrap_number!(c_double, enif_make_double, enif_get_double);
<|endoftext|>


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

use std::ffi::OsString;
use std::path::PathBuf;
use std::process::{self, Command};
use std::ptr;
use std::io;
use time::{Duration, SteadyTime};

use kernel32;
use winapi;

use error::{Error, Result};

use super::{HabExitStatus, ExitStatusExt, ShutdownMethod};

const STILL_ACTIVE: u32 = 259;

pub fn become_command(command: PathBuf, args: Vec<OsString>) -> Result<()> {
    become_child_command(command, args)
}

/// Get process identifier of calling process.
pub fn current_pid() -> u32 {
    unsafe { kernel32::GetCurrentProcessId() as u32 }
}

/// Determines if a process is running with the given process identifier.
pub fn is_alive(pid: u32) -> bool {
    handle_from_pid(pid).is_some()
}

/// Executes a command as a child process and exits with the child's exit code.
///
/// Note that if successful, this function will not return.
///
/// # Failures
///
/// * If the child process cannot be created
fn become_child_command(command: PathBuf, args: Vec<OsString>) -> Result<()> {
    debug!("Calling child process: ({:?}) {:?}",
           command.display(),
           &args);
    let status = try!(Command::new(command).args(&args).status());
    // Let's honor the exit codes from the child process we finished running
    process::exit(status.code().unwrap())
}

fn handle_from_pid(pid: u32) -> Option<winapi::HANDLE> {
    unsafe {
        let proc_handle = kernel32::OpenProcess(winapi::PROCESS_QUERY_LIMITED_INFORMATION |
                                                winapi::PROCESS_TERMINATE,
                                                winapi::FALSE,
                                                pid as winapi::DWORD);

        // we expect this to happen if the process died
        // before OpenProcess completes
        if proc_handle == ptr::null_mut() {
            return None;
        } else {
            return Some(proc_handle);
        }
    }
}

pub struct Child {
    handle: Option<winapi::HANDLE>,
    last_status: Option<u32>,
    pid: u32,
}

impl Child {
    // On windows we need the process handle to capture status
    // Here we will attempt to get the handle from the pid but if the
    // process dies before we can get it, we will just wait() on the
    // std::process::Child and cache the exit_status which we will return
    // when status is called.
    pub fn new(child: &mut process::Child) -> Result<Child> {
        let (win_handle, status) = match handle_from_pid(child.id()) {
            Some(handle) => (Some(handle), Ok(None)),
            _ => {
                (None,
                 {
                     match child.wait() {
                         Ok(exit) => Ok(Some(exit.code().unwrap() as u32)),
                         Err(e) => {
                             Err(format!("Failed to retrieve exit code for pid {} : {}",
                                         child.id(),
                                         e))
                         }
                     }
                 })
            }
        };

        match status {
            Ok(status) => {
                Ok(Child {
                       handle: win_handle,
                       last_status: status,
                       pid: child.id(),
                   })
            }
            Err(e) => Err(Error::GetHabChildFailed(e)),
        }
    }

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

    pub fn status(&mut self) -> Result<HabExitStatus> {
        if self.last_status.is_some() {
            return Ok(HabExitStatus { status: Some(self.last_status.unwrap()) });
        }

        let mut exit_status: u32 = 0;

        unsafe {
            let ret = kernel32::GetExitCodeProcess(self.handle.unwrap(),
                                                   &mut exit_status as winapi::LPDWORD);
            if ret == 0 {
                return Err(Error::GetExitCodeProcessFailed(format!("Failed to retrieve Exit \
                                                                    Code for pid {}: {}",
                                                                   self.pid,
                                                                   io::Error::last_os_error())));
            }
        }

        if exit_status == STILL_ACTIVE {
            return Ok(HabExitStatus { status: None });
        };

        Ok(HabExitStatus { status: Some(exit_status) })
    }

    pub fn kill(&mut self) -> Result<ShutdownMethod> {
        if self.last_status.is_some() {
            return Ok(ShutdownMethod::AlreadyExited);
        }

        let mut ret;
        unsafe {
            // Turn off ctrl-C handling for current process
            ret = kernel32::SetConsoleCtrlHandler(None, winapi::TRUE);
            if ret == 0 {
                debug!("Failed to call SetConsoleCtrlHandler on pid {}: {}",
                       self.pid,
                       io::Error::last_os_error());
            }

            if ret != 0 {
                // Send a ctrl-C
                ret = kernel32::GenerateConsoleCtrlEvent(0, 0);
                if ret == 0 {
                    debug!("Failed to send ctrl-c to pid {}: {}",
                           self.pid,
                           io::Error::last_os_error());
                }
            }
        }

        let stop_time = SteadyTime::now() + Duration::seconds(8);

        let result;
        loop {
            if ret == 0 || SteadyTime::now() > stop_time {
                unsafe {
                    ret = kernel32::TerminateProcess(self.handle.unwrap(), 1);
                    if ret == 0 {
                        result = Err(Error::TerminateProcessFailed(format!("Failed to call \
                                                                       terminate pid {}: {}",
                                                                      self.pid,
                                                                      io::Error::last_os_error())));
                    } else {
                        result = Ok(ShutdownMethod::Killed);
                    }
                    break;
                }
            }

            match self.status() {
                Ok(status) => {
                    if !status.no_status() {
                        result = Ok(ShutdownMethod::GracefulTermination);
                        break;
                    }
                }
                _ => {}
            }
        }

        // turn Ctrl-C handling back on for current process
        ret = unsafe { kernel32::SetConsoleCtrlHandler(None, winapi::FALSE) };
        if ret == 0 {
            debug!("Failed to call SetConsoleCtrlHandler on pid {}: {}",
                   self.pid,
                   io::Error::last_os_error());
        }

        result
    }
}

// Have to implement these due to our HANDLE field
unsafe impl Send for Child {}
unsafe impl Sync for Child {}

impl Drop for Child {
    fn drop(&mut self) {
        match self.handle {
            None => {}
            Some(handle) => unsafe {
                let _ = kernel32::CloseHandle(handle);
            },
        }
    }
}

impl ExitStatusExt for HabExitStatus {
    fn code(&self) -> Option<u32> {
        self.status
    }

    fn signal(&self) -> Option<u32> {
        None
    }
}

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

    #[test]
    fn running_process_returns_no_exit_status() {
        let mut cmd = Command::new("C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.\
                                    exe");
        cmd.arg("-noprofile").arg("-command").arg("while($true) { Start-Sleep 1 }");
        let mut child = cmd.spawn().unwrap();

        let mut hab_child = HabChild::from(&mut child).unwrap();

        assert!(hab_child.status().unwrap().no_status())
    }

    #[test]
    fn successfully_run_process_exits_zero() {
        let mut cmd = Command::new("C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.\
                                    exe");
        cmd.arg("-noprofile").arg("-command").arg("$a='b'");
        let mut child = cmd.spawn().unwrap();

        let mut hab_child = HabChild::from(&mut child).unwrap();
        let _ = child.wait();

        assert_eq!(hab_child.status().unwrap().code(), Some(0))
    }

    #[test]
    fn terminated_process_returns_non_zero_exit() {
        let mut cmd = Command::new("C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.\
                                    exe");
        cmd.arg("-noprofile").arg("-command").arg("while($true) { Start-Sleep 1 }");
        let mut child = cmd.spawn().unwrap();

        let mut hab_child = HabChild::from(&mut child).unwrap();
        let _ = child.kill();

        assert!(hab_child.status().unwrap().code() != Some(0))
    }

    #[test]
    fn process_that_exits_with_specific_code_has_same_exit_code() {
        let mut cmd = Command::new("C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.\
                                    exe");
        cmd.arg("-noprofile").arg("-command").arg("exit 5000");
        let mut child = cmd.spawn().unwrap();

        let mut hab_child = HabChild::from(&mut child).unwrap();
        let _ = child.wait();

        assert_eq!(hab_child.status().unwrap().code(), Some(5000))
    }
}
<commit_msg>Fix file handle leak in `handle_from_pid()<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::ffi::OsString;
use std::path::PathBuf;
use std::process::{self, Command};
use std::ptr;
use std::io;
use time::{Duration, SteadyTime};

use kernel32;
use winapi;

use error::{Error, Result};

use super::{HabExitStatus, ExitStatusExt, ShutdownMethod};

const STILL_ACTIVE: u32 = 259;

pub fn become_command(command: PathBuf, args: Vec<OsString>) -> Result<()> {
    become_child_command(command, args)
}

/// Get process identifier of calling process.
pub fn current_pid() -> u32 {
    unsafe { kernel32::GetCurrentProcessId() as u32 }
}

/// Determines if a process is running with the given process identifier.
pub fn is_alive(pid: u32) -> bool {
    match handle_from_pid(pid) {
        Some(handle) => {
            unsafe {
                let _ = kernel32::CloseHandle(handle);
            }
            true
        }
        None => false,
    }
}

/// Executes a command as a child process and exits with the child's exit code.
///
/// Note that if successful, this function will not return.
///
/// # Failures
///
/// * If the child process cannot be created
fn become_child_command(command: PathBuf, args: Vec<OsString>) -> Result<()> {
    debug!("Calling child process: ({:?}) {:?}",
           command.display(),
           &args);
    let status = try!(Command::new(command).args(&args).status());
    // Let's honor the exit codes from the child process we finished running
    process::exit(status.code().unwrap())
}

fn handle_from_pid(pid: u32) -> Option<winapi::HANDLE> {
    unsafe {
        let proc_handle = kernel32::OpenProcess(winapi::PROCESS_QUERY_LIMITED_INFORMATION |
                                                winapi::PROCESS_TERMINATE,
                                                winapi::FALSE,
                                                pid as winapi::DWORD);

        // we expect this to happen if the process died
        // before OpenProcess completes
        if proc_handle == ptr::null_mut() {
            return None;
        } else {
            return Some(proc_handle);
        }
    }
}

pub struct Child {
    handle: Option<winapi::HANDLE>,
    last_status: Option<u32>,
    pid: u32,
}

impl Child {
    // On windows we need the process handle to capture status
    // Here we will attempt to get the handle from the pid but if the
    // process dies before we can get it, we will just wait() on the
    // std::process::Child and cache the exit_status which we will return
    // when status is called.
    pub fn new(child: &mut process::Child) -> Result<Child> {
        let (win_handle, status) = match handle_from_pid(child.id()) {
            Some(handle) => (Some(handle), Ok(None)),
            _ => {
                (None,
                 {
                     match child.wait() {
                         Ok(exit) => Ok(Some(exit.code().unwrap() as u32)),
                         Err(e) => {
                             Err(format!("Failed to retrieve exit code for pid {} : {}",
                                         child.id(),
                                         e))
                         }
                     }
                 })
            }
        };

        match status {
            Ok(status) => {
                Ok(Child {
                       handle: win_handle,
                       last_status: status,
                       pid: child.id(),
                   })
            }
            Err(e) => Err(Error::GetHabChildFailed(e)),
        }
    }

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

    pub fn status(&mut self) -> Result<HabExitStatus> {
        if self.last_status.is_some() {
            return Ok(HabExitStatus { status: Some(self.last_status.unwrap()) });
        }

        let mut exit_status: u32 = 0;

        unsafe {
            let ret = kernel32::GetExitCodeProcess(self.handle.unwrap(),
                                                   &mut exit_status as winapi::LPDWORD);
            if ret == 0 {
                return Err(Error::GetExitCodeProcessFailed(format!("Failed to retrieve Exit \
                                                                    Code for pid {}: {}",
                                                                   self.pid,
                                                                   io::Error::last_os_error())));
            }
        }

        if exit_status == STILL_ACTIVE {
            return Ok(HabExitStatus { status: None });
        };

        Ok(HabExitStatus { status: Some(exit_status) })
    }

    pub fn kill(&mut self) -> Result<ShutdownMethod> {
        if self.last_status.is_some() {
            return Ok(ShutdownMethod::AlreadyExited);
        }

        let mut ret;
        unsafe {
            // Turn off ctrl-C handling for current process
            ret = kernel32::SetConsoleCtrlHandler(None, winapi::TRUE);
            if ret == 0 {
                debug!("Failed to call SetConsoleCtrlHandler on pid {}: {}",
                       self.pid,
                       io::Error::last_os_error());
            }

            if ret != 0 {
                // Send a ctrl-C
                ret = kernel32::GenerateConsoleCtrlEvent(0, 0);
                if ret == 0 {
                    debug!("Failed to send ctrl-c to pid {}: {}",
                           self.pid,
                           io::Error::last_os_error());
                }
            }
        }

        let stop_time = SteadyTime::now() + Duration::seconds(8);

        let result;
        loop {
            if ret == 0 || SteadyTime::now() > stop_time {
                unsafe {
                    ret = kernel32::TerminateProcess(self.handle.unwrap(), 1);
                    if ret == 0 {
                        result = Err(Error::TerminateProcessFailed(format!("Failed to call \
                                                                       terminate pid {}: {}",
                                                                      self.pid,
                                                                      io::Error::last_os_error())));
                    } else {
                        result = Ok(ShutdownMethod::Killed);
                    }
                    break;
                }
            }

            match self.status() {
                Ok(status) => {
                    if !status.no_status() {
                        result = Ok(ShutdownMethod::GracefulTermination);
                        break;
                    }
                }
                _ => {}
            }
        }

        // turn Ctrl-C handling back on for current process
        ret = unsafe { kernel32::SetConsoleCtrlHandler(None, winapi::FALSE) };
        if ret == 0 {
            debug!("Failed to call SetConsoleCtrlHandler on pid {}: {}",
                   self.pid,
                   io::Error::last_os_error());
        }

        result
    }
}

// Have to implement these due to our HANDLE field
unsafe impl Send for Child {}
unsafe impl Sync for Child {}

impl Drop for Child {
    fn drop(&mut self) {
        match self.handle {
            None => {}
            Some(handle) => unsafe {
                let _ = kernel32::CloseHandle(handle);
            },
        }
    }
}

impl ExitStatusExt for HabExitStatus {
    fn code(&self) -> Option<u32> {
        self.status
    }

    fn signal(&self) -> Option<u32> {
        None
    }
}

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

    #[test]
    fn running_process_returns_no_exit_status() {
        let mut cmd = Command::new("C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.\
                                    exe");
        cmd.arg("-noprofile").arg("-command").arg("while($true) { Start-Sleep 1 }");
        let mut child = cmd.spawn().unwrap();

        let mut hab_child = HabChild::from(&mut child).unwrap();

        assert!(hab_child.status().unwrap().no_status())
    }

    #[test]
    fn successfully_run_process_exits_zero() {
        let mut cmd = Command::new("C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.\
                                    exe");
        cmd.arg("-noprofile").arg("-command").arg("$a='b'");
        let mut child = cmd.spawn().unwrap();

        let mut hab_child = HabChild::from(&mut child).unwrap();
        let _ = child.wait();

        assert_eq!(hab_child.status().unwrap().code(), Some(0))
    }

    #[test]
    fn terminated_process_returns_non_zero_exit() {
        let mut cmd = Command::new("C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.\
                                    exe");
        cmd.arg("-noprofile").arg("-command").arg("while($true) { Start-Sleep 1 }");
        let mut child = cmd.spawn().unwrap();

        let mut hab_child = HabChild::from(&mut child).unwrap();
        let _ = child.kill();

        assert!(hab_child.status().unwrap().code() != Some(0))
    }

    #[test]
    fn process_that_exits_with_specific_code_has_same_exit_code() {
        let mut cmd = Command::new("C:\\Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.\
                                    exe");
        cmd.arg("-noprofile").arg("-command").arg("exit 5000");
        let mut child = cmd.spawn().unwrap();

        let mut hab_child = HabChild::from(&mut child).unwrap();
        let _ = child.wait();

        assert_eq!(hab_child.status().unwrap().code(), Some(5000))
    }
}
<|endoftext|>


<commit_before>use crate::to_statement::private::{Sealed, ToStatementType};
use crate::Statement;

mod private {
    use crate::{Client, Error, Statement};

    pub trait Sealed {}

    pub enum ToStatementType<'a> {
        Statement(&'a Statement),
        Query(&'a str),
    }

    impl<'a> ToStatementType<'a> {
        pub async fn into_statement(self, client: &Client) -> Result<Statement, Error> {
            match self {
                ToStatementType::Statement(s) => Ok(s.clone()),
                ToStatementType::Query(s) => client.prepare(s).await,
            }
        }
    }
}

/// A trait abstracting over prepared and unprepared statements.
///
/// Many methods are generic over this bound, so that they support both a raw query string as well as a statement which
/// was prepared previously.
///
/// This trait is "sealed" and cannot be implemented by anything outside this crate.
pub trait ToStatement: Sealed {
    #[doc(hidden)]
    fn __convert(&self) -> ToStatementType<'_>;
}

impl ToStatement for Statement {
    fn __convert(&self) -> ToStatementType<'_> {
        ToStatementType::Statement(self)
    }
}

impl Sealed for Statement {}

impl ToStatement for str {
    fn __convert(&self) -> ToStatementType<'_> {
        ToStatementType::Query(self)
    }
}

impl Sealed for str {}
<commit_msg>Implement ToStatement for String<commit_after>use crate::to_statement::private::{Sealed, ToStatementType};
use crate::Statement;

mod private {
    use crate::{Client, Error, Statement};

    pub trait Sealed {}

    pub enum ToStatementType<'a> {
        Statement(&'a Statement),
        Query(&'a str),
    }

    impl<'a> ToStatementType<'a> {
        pub async fn into_statement(self, client: &Client) -> Result<Statement, Error> {
            match self {
                ToStatementType::Statement(s) => Ok(s.clone()),
                ToStatementType::Query(s) => client.prepare(s).await,
            }
        }
    }
}

/// A trait abstracting over prepared and unprepared statements.
///
/// Many methods are generic over this bound, so that they support both a raw query string as well as a statement which
/// was prepared previously.
///
/// This trait is "sealed" and cannot be implemented by anything outside this crate.
pub trait ToStatement: Sealed {
    #[doc(hidden)]
    fn __convert(&self) -> ToStatementType<'_>;
}

impl ToStatement for Statement {
    fn __convert(&self) -> ToStatementType<'_> {
        ToStatementType::Statement(self)
    }
}

impl Sealed for Statement {}

impl ToStatement for str {
    fn __convert(&self) -> ToStatementType<'_> {
        ToStatementType::Query(self)
    }
}

impl Sealed for str {}

impl ToStatement for String {
    fn __convert(&self) -> ToStatementType<'_> {
        ToStatementType::Query(self)
    }
}

impl Sealed for String {}
<|endoftext|>


<commit_before>//
// tui.rs
// Copyright (C) 2017 Szymon Urbaś <szymon.urbas@aol.com>
// Distributed under terms of the BSD (2-clause) license.
//
// Created on: 16 Mar 2017 22:14:17 +0100 (CET)
//

use cursive::Cursive;
use cursive::view::*;
use cursive::views::*;

use process::*;

pub struct Tui {
  renderer: Cursive,
}

pub fn new() -> Tui {
  let mut renderer = Cursive::new();

  let process_list = ListView::new().with_id("process_list");
  let results = TextView::new("").with_id("results").full_height();
  let header = TextView::new("").with_id("header");
  let debug = TextView::new("").scroll_strategy(ScrollStrategy::StickToBottom).with_id("debug").full_height();

  let mut layout = LinearLayout::horizontal();

  layout.add_child(Dialog::around(process_list).title("Process list"));
  layout.add_child(
    LinearLayout::vertical()
      .child(Dialog::around(header).title("Header"))
      .child(Dialog::around(results).title("Results"))
      .child(Dialog::around(debug).title("Debug info"))
  );

  renderer.set_fps(60);
  renderer.add_layer(layout);
  renderer.add_global_callback('q', |tui| tui.quit());

  Tui {
    renderer: renderer,
  }
}

impl Tui {
  pub fn update(&mut self) {
    self.renderer.step();
  }

  pub fn draw_process_list(&mut self, process_list: &[Process]) {
    let mut process_list_view = self.renderer.find_id::<ListView>("process_list").unwrap();

    process_list_view.clear();

    for p in process_list.iter() {
      let progress_bar_value = ((p.execution_time as f64 / p.burst_time as f64) as f64 * 100f64) as usize;

      process_list_view.add_child(p.name.clone().as_str(),
        ProgressBar::new().with_value(Counter::new(progress_bar_value)));
    }
  }

  pub fn set_header(&mut self, text: String) {
    let mut header = self.renderer.find_id::<TextView>("header").unwrap();

    header.set_content(text.clone());
  }

  pub fn debug(&mut self, text: String) {
    let mut debug = self.renderer.find_id::<TextView>("debug").unwrap();

    debug.append_content(&text);
    debug.append_content("\n");
  }
}

/*
 * vi: ts=2 sw=2 expandtab
 */

<commit_msg>Make the process list a bit wider<commit_after>//
// tui.rs
// Copyright (C) 2017 Szymon Urbaś <szymon.urbas@aol.com>
// Distributed under terms of the BSD (2-clause) license.
//
// Created on: 16 Mar 2017 22:14:17 +0100 (CET)
//

use cursive::Cursive;
use cursive::view::*;
use cursive::views::*;

use process::*;

pub struct Tui {
  renderer: Cursive,
}

pub fn new() -> Tui {
  let mut renderer = Cursive::new();

  let process_list = ListView::new().with_id("process_list").min_width(32);
  let results = TextView::new("").with_id("results").full_height();
  let header = TextView::new("").with_id("header");
  let debug = TextView::new("").scroll_strategy(ScrollStrategy::StickToBottom).with_id("debug").full_height();

  let mut layout = LinearLayout::horizontal();

  layout.add_child(Dialog::around(process_list).title("Process list"));
  layout.add_child(
    LinearLayout::vertical()
      .child(Dialog::around(header).title("Header"))
      .child(Dialog::around(results).title("Results"))
      .child(Dialog::around(debug).title("Debug info"))
  );

  renderer.set_fps(60);
  renderer.add_layer(layout);
  renderer.add_global_callback('q', |tui| tui.quit());

  Tui {
    renderer: renderer,
  }
}

impl Tui {
  pub fn update(&mut self) {
    self.renderer.step();
  }

  pub fn draw_process_list(&mut self, process_list: &[Process]) {
    let mut process_list_view = self.renderer.find_id::<ListView>("process_list").unwrap();

    process_list_view.clear();

    for p in process_list.iter() {
      let progress_bar_value = ((p.execution_time as f64 / p.burst_time as f64) as f64 * 100f64) as usize;

      process_list_view.add_child(p.name.clone().as_str(),
        ProgressBar::new().with_value(Counter::new(progress_bar_value)));
    }
  }

  pub fn set_header(&mut self, text: String) {
    let mut header = self.renderer.find_id::<TextView>("header").unwrap();

    header.set_content(text.clone());
  }

  pub fn debug(&mut self, text: String) {
    let mut debug = self.renderer.find_id::<TextView>("debug").unwrap();

    debug.append_content(&text);
    debug.append_content("\n");
  }
}

/*
 * vi: ts=2 sw=2 expandtab
 */

<|endoftext|>


<commit_before>#![cfg(not(debug_assertions))]

#[macro_use]
extern crate lazy_static;

use beacon_chain::{
    test_utils::{AttestationStrategy, BeaconChainHarness, BlockStrategy},
    StateSkipConfig,
};
use tree_hash::TreeHash;
use types::{AggregateSignature, EthSpec, Hash256, Keypair, MainnetEthSpec, RelativeEpoch, Slot};

pub const VALIDATOR_COUNT: usize = 16;

lazy_static! {
    /// A cached set of keys.
    static ref KEYPAIRS: Vec<Keypair> = types::test_utils::generate_deterministic_keypairs(VALIDATOR_COUNT);
}

/// This test builds a chain that is just long enough to finalize an epoch then it produces an
/// attestation at each slot from genesis through to three epochs past the head.
///
/// It checks the produced attestation against some locally computed values.
#[test]
fn produces_attestations() {
    let num_blocks_produced = MainnetEthSpec::slots_per_epoch() * 4;

    let harness = BeaconChainHarness::new(MainnetEthSpec, KEYPAIRS[..].to_vec());

    // Skip past the genesis slot.
    harness.advance_slot();

    harness.extend_chain(
        num_blocks_produced as usize,
        BlockStrategy::OnCanonicalHead,
        AttestationStrategy::AllValidators,
    );

    let chain = &harness.chain;

    let state = &harness.chain.head().expect("should get head").beacon_state;
    assert_eq!(state.slot, num_blocks_produced, "head should have updated");
    assert!(
        state.finalized_checkpoint.epoch > 0,
        "head should have updated"
    );

    let current_slot = chain.slot().expect("should get slot");

    // Test all valid committee indices for all slots in the chain.
    for slot in 0..=current_slot.as_u64() + MainnetEthSpec::slots_per_epoch() * 3 {
        let slot = Slot::from(slot);
        let state = chain
            .state_at_slot(slot, StateSkipConfig::WithStateRoots)
            .expect("should get state");

        let block_slot = if slot > current_slot {
            current_slot
        } else {
            slot
        };
        let block = chain
            .block_at_slot(block_slot)
            .expect("should get block")
            .expect("block should not be skipped");
        let block_root = Hash256::from_slice(&block.message.tree_hash_root());

        let epoch_boundary_slot = state
            .current_epoch()
            .start_slot(MainnetEthSpec::slots_per_epoch());
        let target_root = if state.slot == epoch_boundary_slot {
            block_root
        } else {
            *state
                .get_block_root(epoch_boundary_slot)
                .expect("should get target block root")
        };

        let committee_cache = state
            .committee_cache(RelativeEpoch::Current)
            .expect("should get committee_cache");

        let committee_count = committee_cache.committees_per_slot();

        for index in 0..committee_count {
            let committee_len = committee_cache
                .get_beacon_committee(slot, index)
                .expect("should get committee for slot")
                .committee
                .len();

            let attestation = chain
                .produce_attestation(slot, index)
                .expect("should produce attestation");

            let data = &attestation.data;

            assert_eq!(
                attestation.aggregation_bits.len(),
                committee_len,
                "bad committee len"
            );
            assert!(
                attestation.aggregation_bits.is_zero(),
                "some committee bits are set"
            );
            assert_eq!(
                attestation.signature,
                AggregateSignature::new(),
                "bad signature"
            );
            assert_eq!(data.index, index, "bad index");
            assert_eq!(data.slot, slot, "bad slot");
            assert_eq!(data.beacon_block_root, block_root, "bad block root");
            assert_eq!(
                data.source, state.current_justified_checkpoint,
                "bad source"
            );
            assert_eq!(
                data.source, state.current_justified_checkpoint,
                "bad source"
            );
            assert_eq!(data.target.epoch, state.current_epoch(), "bad target epoch");
            assert_eq!(data.target.root, target_root, "bad target root");
        }
    }
}
<commit_msg>Fix compile issue (#893)<commit_after>#![cfg(not(debug_assertions))]

#[macro_use]
extern crate lazy_static;

use beacon_chain::{
    test_utils::{AttestationStrategy, BeaconChainHarness, BlockStrategy},
    StateSkipConfig,
};
use tree_hash::TreeHash;
use types::{AggregateSignature, EthSpec, Keypair, MainnetEthSpec, RelativeEpoch, Slot};

pub const VALIDATOR_COUNT: usize = 16;

lazy_static! {
    /// A cached set of keys.
    static ref KEYPAIRS: Vec<Keypair> = types::test_utils::generate_deterministic_keypairs(VALIDATOR_COUNT);
}

/// This test builds a chain that is just long enough to finalize an epoch then it produces an
/// attestation at each slot from genesis through to three epochs past the head.
///
/// It checks the produced attestation against some locally computed values.
#[test]
fn produces_attestations() {
    let num_blocks_produced = MainnetEthSpec::slots_per_epoch() * 4;

    let harness = BeaconChainHarness::new(MainnetEthSpec, KEYPAIRS[..].to_vec());

    // Skip past the genesis slot.
    harness.advance_slot();

    harness.extend_chain(
        num_blocks_produced as usize,
        BlockStrategy::OnCanonicalHead,
        AttestationStrategy::AllValidators,
    );

    let chain = &harness.chain;

    let state = &harness.chain.head().expect("should get head").beacon_state;
    assert_eq!(state.slot, num_blocks_produced, "head should have updated");
    assert!(
        state.finalized_checkpoint.epoch > 0,
        "head should have updated"
    );

    let current_slot = chain.slot().expect("should get slot");

    // Test all valid committee indices for all slots in the chain.
    for slot in 0..=current_slot.as_u64() + MainnetEthSpec::slots_per_epoch() * 3 {
        let slot = Slot::from(slot);
        let state = chain
            .state_at_slot(slot, StateSkipConfig::WithStateRoots)
            .expect("should get state");

        let block_slot = if slot > current_slot {
            current_slot
        } else {
            slot
        };
        let block = chain
            .block_at_slot(block_slot)
            .expect("should get block")
            .expect("block should not be skipped");
        let block_root = block.message.tree_hash_root();

        let epoch_boundary_slot = state
            .current_epoch()
            .start_slot(MainnetEthSpec::slots_per_epoch());
        let target_root = if state.slot == epoch_boundary_slot {
            block_root
        } else {
            *state
                .get_block_root(epoch_boundary_slot)
                .expect("should get target block root")
        };

        let committee_cache = state
            .committee_cache(RelativeEpoch::Current)
            .expect("should get committee_cache");

        let committee_count = committee_cache.committees_per_slot();

        for index in 0..committee_count {
            let committee_len = committee_cache
                .get_beacon_committee(slot, index)
                .expect("should get committee for slot")
                .committee
                .len();

            let attestation = chain
                .produce_attestation(slot, index)
                .expect("should produce attestation");

            let data = &attestation.data;

            assert_eq!(
                attestation.aggregation_bits.len(),
                committee_len,
                "bad committee len"
            );
            assert!(
                attestation.aggregation_bits.is_zero(),
                "some committee bits are set"
            );
            assert_eq!(
                attestation.signature,
                AggregateSignature::new(),
                "bad signature"
            );
            assert_eq!(data.index, index, "bad index");
            assert_eq!(data.slot, slot, "bad slot");
            assert_eq!(data.beacon_block_root, block_root, "bad block root");
            assert_eq!(
                data.source, state.current_justified_checkpoint,
                "bad source"
            );
            assert_eq!(
                data.source, state.current_justified_checkpoint,
                "bad source"
            );
            assert_eq!(data.target.epoch, state.current_epoch(), "bad target epoch");
            assert_eq!(data.target.root, target_root, "bad target root");
        }
    }
}
<|endoftext|>


<commit_before>//! Lower-level types and re-exports.
//!
//! Most users will not have to interact with the types in this module, but it is useful for those
//! writing extractors, middleware, libraries, or interacting with the service API directly.

pub use crate::config::{AppConfig, AppService};
#[doc(hidden)]
pub use crate::handler::Handler;
pub use crate::info::{ConnectionInfo, PeerAddr};
pub use crate::rmap::ResourceMap;
pub use crate::service::{HttpServiceFactory, ServiceRequest, ServiceResponse, WebService};

pub use crate::types::form::UrlEncoded;
pub use crate::types::json::JsonBody;
pub use crate::types::readlines::Readlines;

pub use actix_http::{Extensions, Payload, PayloadStream, RequestHead, Response, ResponseHead};
pub use actix_router::{Path, ResourceDef, ResourcePath, Url};
pub use actix_server::{Server, ServerHandle};
pub use actix_service::{
    always_ready, fn_factory, fn_service, forward_ready, Service, ServiceFactory, Transform,
};

#[cfg(feature = "__compress")]
pub use actix_http::encoding::Decoder as Decompress;

use crate::http::header::ContentEncoding;

use actix_router::Patterns;

pub(crate) fn ensure_leading_slash(mut patterns: Patterns) -> Patterns {
    match &mut patterns {
        Patterns::Single(pat) => {
            if !pat.is_empty() && !pat.starts_with('/') {
                pat.insert(0, '/');
            };
        }
        Patterns::List(pats) => {
            for pat in pats {
                if !pat.is_empty() && !pat.starts_with('/') {
                    pat.insert(0, '/');
                };
            }
        }
    }

    patterns
}
struct Enc(ContentEncoding);

/// Helper trait that allows to set specific encoding for response.
pub trait BodyEncoding {
    /// Get content encoding
    fn get_encoding(&self) -> Option<ContentEncoding>;

    /// Set content encoding
    ///
    /// Must be used with [`crate::middleware::Compress`] to take effect.
    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self;
}

impl BodyEncoding for actix_http::ResponseBuilder {
    fn get_encoding(&self) -> Option<ContentEncoding> {
        self.extensions().get::<Enc>().map(|enc| enc.0)
    }

    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self {
        self.extensions_mut().insert(Enc(encoding));
        self
    }
}

impl<B> BodyEncoding for actix_http::Response<B> {
    fn get_encoding(&self) -> Option<ContentEncoding> {
        self.extensions().get::<Enc>().map(|enc| enc.0)
    }

    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self {
        self.extensions_mut().insert(Enc(encoding));
        self
    }
}

impl BodyEncoding for crate::HttpResponseBuilder {
    fn get_encoding(&self) -> Option<ContentEncoding> {
        self.extensions().get::<Enc>().map(|enc| enc.0)
    }

    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self {
        self.extensions_mut().insert(Enc(encoding));
        self
    }
}

impl<B> BodyEncoding for crate::HttpResponse<B> {
    fn get_encoding(&self) -> Option<ContentEncoding> {
        self.extensions().get::<Enc>().map(|enc| enc.0)
    }

    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self {
        self.extensions_mut().insert(Enc(encoding));
        self
    }
}

// TODO: remove this if it doesn't appear to be needed

#[allow(dead_code)]
#[derive(Debug)]
pub(crate) enum AnyBody {
    None,
    Full { body: crate::web::Bytes },
    Boxed { body: actix_http::body::BoxBody },
}

impl crate::body::MessageBody for AnyBody {
    type Error = crate::BoxError;

    /// Body size hint.
    fn size(&self) -> crate::body::BodySize {
        match self {
            AnyBody::None => crate::body::BodySize::None,
            AnyBody::Full { body } => body.size(),
            AnyBody::Boxed { body } => body.size(),
        }
    }

    /// Attempt to pull out the next chunk of body bytes.
    fn poll_next(
        self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Result<crate::web::Bytes, Self::Error>>> {
        match self.get_mut() {
            AnyBody::None => std::task::Poll::Ready(None),
            AnyBody::Full { body } => {
                let bytes = std::mem::take(body);
                std::task::Poll::Ready(Some(Ok(bytes)))
            }
            AnyBody::Boxed { body } => body.as_pin_mut().poll_next(cx),
        }
    }

    fn try_into_bytes(self) -> Result<crate::web::Bytes, Self> {
        match self {
            AnyBody::None => Ok(crate::web::Bytes::new()),
            AnyBody::Full { body } => Ok(body),
            _ => Err(self),
        }
    }
}
<commit_msg>remove hidden anybody<commit_after>//! Lower-level types and re-exports.
//!
//! Most users will not have to interact with the types in this module, but it is useful for those
//! writing extractors, middleware, libraries, or interacting with the service API directly.

pub use crate::config::{AppConfig, AppService};
#[doc(hidden)]
pub use crate::handler::Handler;
pub use crate::info::{ConnectionInfo, PeerAddr};
pub use crate::rmap::ResourceMap;
pub use crate::service::{HttpServiceFactory, ServiceRequest, ServiceResponse, WebService};

pub use crate::types::form::UrlEncoded;
pub use crate::types::json::JsonBody;
pub use crate::types::readlines::Readlines;

pub use actix_http::{Extensions, Payload, PayloadStream, RequestHead, Response, ResponseHead};
pub use actix_router::{Path, ResourceDef, ResourcePath, Url};
pub use actix_server::{Server, ServerHandle};
pub use actix_service::{
    always_ready, fn_factory, fn_service, forward_ready, Service, ServiceFactory, Transform,
};

#[cfg(feature = "__compress")]
pub use actix_http::encoding::Decoder as Decompress;

use crate::http::header::ContentEncoding;

use actix_router::Patterns;

pub(crate) fn ensure_leading_slash(mut patterns: Patterns) -> Patterns {
    match &mut patterns {
        Patterns::Single(pat) => {
            if !pat.is_empty() && !pat.starts_with('/') {
                pat.insert(0, '/');
            };
        }
        Patterns::List(pats) => {
            for pat in pats {
                if !pat.is_empty() && !pat.starts_with('/') {
                    pat.insert(0, '/');
                };
            }
        }
    }

    patterns
}
struct Enc(ContentEncoding);

/// Helper trait that allows to set specific encoding for response.
pub trait BodyEncoding {
    /// Get content encoding
    fn get_encoding(&self) -> Option<ContentEncoding>;

    /// Set content encoding
    ///
    /// Must be used with [`crate::middleware::Compress`] to take effect.
    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self;
}

impl BodyEncoding for actix_http::ResponseBuilder {
    fn get_encoding(&self) -> Option<ContentEncoding> {
        self.extensions().get::<Enc>().map(|enc| enc.0)
    }

    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self {
        self.extensions_mut().insert(Enc(encoding));
        self
    }
}

impl<B> BodyEncoding for actix_http::Response<B> {
    fn get_encoding(&self) -> Option<ContentEncoding> {
        self.extensions().get::<Enc>().map(|enc| enc.0)
    }

    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self {
        self.extensions_mut().insert(Enc(encoding));
        self
    }
}

impl BodyEncoding for crate::HttpResponseBuilder {
    fn get_encoding(&self) -> Option<ContentEncoding> {
        self.extensions().get::<Enc>().map(|enc| enc.0)
    }

    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self {
        self.extensions_mut().insert(Enc(encoding));
        self
    }
}

impl<B> BodyEncoding for crate::HttpResponse<B> {
    fn get_encoding(&self) -> Option<ContentEncoding> {
        self.extensions().get::<Enc>().map(|enc| enc.0)
    }

    fn encoding(&mut self, encoding: ContentEncoding) -> &mut Self {
        self.extensions_mut().insert(Enc(encoding));
        self
    }
}
<|endoftext|>