<commit_before>use nickel::{
    Params,
};
use std;
use std::collections::HashMap;

macro_rules! parse_field {
    ( $field:expr, $vtype:ty ) => {{
        $field.parse::<$vtype>()
    }}
}

macro_rules! validators {
    ( $([$($field:expr),*]: $vtype:ty => $($func:ident$args:tt),*; )+ ) => {{
        let mut rules: ValidatorParams = Vec::new();
        $(
            let rule = ValidatorParam {
                fields: vec![$($field),*],
                validators: vec![$($func),*],
                type_ref: stringify!($vtype),
                validators_arg: validators! ( @validator_func $($args),* ),
            };
            rules.push(rule);
        )+
        rules
    }};
    ( @validator_func $($args:tt),* ) => {{
        vec!( $( validators! ( @validator_func_args $args ) ),* )
    }};
    ( @validator_func_args () ) => {{
        ValidatorFuncArgs::NoParams
    }};
    ( @validator_func_args ($val:expr) ) => {{
        ValidatorFuncArgs::OneParam($val.to_string())
    }};
    ( @validator_func_args ($val1:expr, $val2:expr) ) => {{
        ValidatorFuncArgs::TwoParam($val1.to_string(), $val2.to_string())
    }};
}

pub struct Validators<'a> {
    validators: ValidatorParams<'a>,
}

#[derive(Debug, Clone)]
pub enum ValidatorFuncArgs {
    NoParams,
    OneParam(String),
    TwoParam(String, String),
}

pub type ValidatorParams<'a> = Vec<ValidatorParam<'a>>;

pub struct ValidatorParam<'a> {
    pub fields: Vec<&'a str>,
    pub validators: Vec<fn(ValidatorFuncArgs)>,
    pub validators_arg: Vec<ValidatorFuncArgs>,
    pub type_ref: &'a str,
}

pub type FormValidationErrors<'a> = HashMap<&'a str, Vec<FormValidationError>>;

#[derive(Debug, Clone, PartialEq)]
pub enum FormValidationError {
    FieldNotExist,
    ParseError
}

pub trait FormValidator {
    fn validate(form_data: &Params)
            -> Result<Self, FormValidationErrors>
            where Self: std::marker::Sized {
        let _validators = Self::validators();
        create_validator(&form_data);
        Ok(Self::fill_form(&form_data)?)
    }
    fn validators<'a>() -> ValidatorParams<'a>;
    fn fill_form(form_data: &Params)
            -> Result<Self, FormValidationErrors>
            where Self: std::marker::Sized;
}

impl<'a> Validators<'a> {
    fn new(validators: ValidatorParams) {
//--
    }
}

fn create_validator(form_data: &Params) -> Result<(), FormValidationErrors> {
    let mut validation_errors: FormValidationErrors = HashMap::new();

    let tst = validators! (
        ["tst1", "tst2"]: i32 => required(), max(1);
        ["tst3", "tst4"]: i32 => max(1, "2+");
    );

    for validator in tst {
        for (key, ref val) in validator.validators.iter().enumerate() {
            let args = validator.validators_arg[key].to_owned();
            for field_name in &validator.fields {
                println!("Fields: {:#?}", field_name);
                println!("Args: {:#?}", args);

                // Fetch fields withour any results
                // We iterate over all fields with `field_name`
                // Try parse with specific type
                // We gathering all possible errors, without skiping any posible error
                form_data.all(field_name)
                    // If field not exist - add error
                    .or_else(|| add_error(&mut validation_errors, field_name, FormValidationError::FieldNotExist))
                    .and_then(|v| {
                        // Iterate over fields
                        Some(v.iter()
                            .map(|x| {
                                // Type inference via Parse
                                parse_field!(x, i32).ok()
                                    // If parse Error - add error
                                    .or_else(|| add_error(&mut validation_errors, field_name, FormValidationError::ParseError));
                                // Return same field
                                x
                            }).cloned().collect::<String>())
                    });
            }
        }
    }

    println!("Errors: {:#?}", validation_errors);

    /*
    println!("Validators: {:#?}", tst.len());
    for vl in tst {
        println!("Fields: {:#?}", vl.fields);
        println!("Validators count: {:#?}", vl.validators.len());
        println!("Args count: {:#?}", vl.validators_arg.len());
        for (key, ref val) in vl.validators.iter().enumerate() {
            let args = vl.validators_arg[key].to_owned();
            val(args);
            //println!("Index: {:#?}", (*val)(vl.validators_arg[key]));
            println!("Field: {:#?}", vl.fields[key]);
            println!("Index: {:#?}", vl.validators_arg[key]);
            //val(vl.validators_arg[key]);
        }
    }*/
    if validation_errors.is_empty() { Ok(()) } else { Err(validation_errors) }
}

fn add_error<'a, T>(validation_errors: &mut FormValidationErrors<'a>, field: &'a str, error: FormValidationError) -> Option<T> {
    if validation_errors.contains_key(&field) {
        let mut err = validation_errors.get_mut(&field).unwrap();
        if !err.contains(&error) {
            err.push(error);
        }
    } else {
        validation_errors.insert(&field, vec!(error));
    }
    None
}

pub fn required(_: ValidatorFuncArgs) {
    println!("required");
    // required template
}

pub fn max(_: ValidatorFuncArgs) {
    println!("max");
    // max template
}
<commit_msg>Changed validators init<commit_after>use nickel::{
    Params,
};
use std;
use std::collections::HashMap;

macro_rules! parse_field {
    ( $field:expr, $vtype:ty ) => {{
        $field.parse::<$vtype>()
    }}
}

macro_rules! validators {
    ( $([$($field:expr),*]: $vtype:ty => $($func:ident$args:tt),*; )+ ) => {{
        let mut rules: ValidatorParams = Vec::new();
        $(
            let rule = ValidatorParam {
                fields: vec![$($field),*],
                validators: vec![$($func),*],
                type_ref: stringify!($vtype),
                validators_arg: validators! ( @validator_func $($args),* ),
            };
            rules.push(rule);
        )+
        rules
    }};
    ( @validator_func $($args:tt),* ) => {{
        vec!( $( validators! ( @validator_func_args $args ) ),* )
    }};
    ( @validator_func_args () ) => {{
        ValidatorFuncArgs::NoParams
    }};
    ( @validator_func_args ($val:expr) ) => {{
        ValidatorFuncArgs::OneParam($val.to_string())
    }};
    ( @validator_func_args ($val1:expr, $val2:expr) ) => {{
        ValidatorFuncArgs::TwoParam($val1.to_string(), $val2.to_string())
    }};
}

pub struct Validators<'a> {
    validators: ValidatorParams<'a>,
}

#[derive(Debug, Clone)]
pub enum ValidatorFuncArgs {
    NoParams,
    OneParam(String),
    TwoParam(String, String),
}

pub type ValidatorParams<'a> = Vec<ValidatorParam<'a>>;

pub struct ValidatorParam<'a> {
    pub fields: Vec<&'a str>,
    pub validators: Vec<fn(ValidatorFuncArgs)>,
    pub validators_arg: Vec<ValidatorFuncArgs>,
    pub type_ref: &'a str,
}

pub type FormValidationErrors<'a> = HashMap<&'a str, Vec<FormValidationError>>;

#[derive(Debug, Clone, PartialEq)]
pub enum FormValidationError {
    FieldNotExist,
    ParseError
}

pub trait FormValidator {
    fn validate(form_data: &Params)
            -> Result<Self, FormValidationErrors>
            where Self: std::marker::Sized {
        Validators::new(Self::validators()).create_validator(&form_data);
        Ok(Self::fill_form(&form_data)?)
    }
    // Validation rules and validators
    fn validators<'a>() -> ValidatorParams<'a>;
    // fill form from form_data to struct
    // when invoked method all data already correct and validated
    fn fill_form(form_data: &Params)
            -> Result<Self, FormValidationErrors>
            where Self: std::marker::Sized;
}

impl<'a> Validators<'a> {
    fn new(validators: ValidatorParams<'a>) -> Validators<'a> {
        Validators { validators: validators }
    }

    fn create_validator(&self, form_data: &Params) -> Result<(), FormValidationErrors> {
        let mut validation_errors: FormValidationErrors = HashMap::new();

        for rules in &self.validators {
            for (key, ref validator) in rules.validators.iter().enumerate() {
                let args = rules.validators_arg[key].to_owned();
                for field_name in &rules.fields {
                    println!("Fields: {:#?}", field_name);
                    println!("Args: {:#?}", args);

                    // Fetch fields withour any results
                    // We iterate over all fields with `field_name`
                    // Try parse with specific type
                    // We gathering all possible errors, without skiping any posible error
                    form_data.all(field_name)
                        // If field not exist - add error
                        .or_else(|| {
                            add_error(&mut validation_errors, field_name, FormValidationError::FieldNotExist);
                            None
                        })
                        .and_then(|v| {
                            // Iterate over fields
                            Some(v.iter()
                                .map(|x| {
                                    // Type inference via Parse
                                    parse_field!(x, i32).ok()
                                        // If parse Error - add error
                                        .or_else(|| {
                                            add_error(&mut validation_errors, field_name, FormValidationError::ParseError);
                                            None
                                        });
                                    // Return same field
                                    x
                                }).cloned().collect::<String>())
                        });
                }
            }
        }

        println!("Errors: {:#?}", validation_errors);

        /*
        println!("Validators: {:#?}", tst.len());
        for vl in tst {
            println!("Fields: {:#?}", vl.fields);
            println!("Validators count: {:#?}", vl.validators.len());
            println!("Args count: {:#?}", vl.validators_arg.len());
            for (key, ref val) in vl.validators.iter().enumerate() {
                let args = vl.validators_arg[key].to_owned();
                val(args);
                //println!("Index: {:#?}", (*val)(vl.validators_arg[key]));
                println!("Field: {:#?}", vl.fields[key]);
                println!("Index: {:#?}", vl.validators_arg[key]);
                //val(vl.validators_arg[key]);
            }
        }*/
        if validation_errors.is_empty() { Ok(()) } else { Err(validation_errors) }
    }
}

fn add_error<'a>(validation_errors: &mut FormValidationErrors<'a>, field: &'a str, error: FormValidationError) {
    if validation_errors.contains_key(&field) {
        let mut err = validation_errors.get_mut(&field).unwrap();
        if !err.contains(&error) {
            err.push(error);
        }
    } else {
        validation_errors.insert(&field, vec!(error));
    }
}

pub fn required(_: ValidatorFuncArgs) {
    println!("required");
    // required template
}

pub fn max(_: ValidatorFuncArgs) {
    println!("max");
    // max template
}
<|endoftext|>


<commit_before>use antidote::Mutex;
use bytes::IntoBuf;
use futures::sync::mpsc;
use futures::{AsyncSink, Poll, Sink, Stream};
use postgres_protocol;
use postgres_protocol::message::backend::Message;
use postgres_protocol::message::frontend;
use std::collections::HashMap;
use std::error::Error as StdError;
use std::sync::{Arc, Weak};

use crate::proto::bind::BindFuture;
use crate::proto::connection::{Request, RequestMessages};
use crate::proto::copy_in::{CopyInFuture, CopyInReceiver, CopyMessage};
use crate::proto::copy_out::CopyOutStream;
use crate::proto::execute::ExecuteFuture;
use crate::proto::idle::{IdleGuard, IdleState};
use crate::proto::portal::Portal;
use crate::proto::prepare::PrepareFuture;
use crate::proto::query::QueryStream;
use crate::proto::simple_query::SimpleQueryFuture;
use crate::proto::statement::Statement;
use crate::types::{IsNull, Oid, ToSql, Type};
use crate::Error;

pub struct PendingRequest(Result<(RequestMessages, IdleGuard), Error>);

pub struct WeakClient(Weak<Inner>);

impl WeakClient {
    pub fn upgrade(&self) -> Option<Client> {
        self.0.upgrade().map(Client)
    }
}

struct State {
    types: HashMap<Oid, Type>,
    typeinfo_query: Option<Statement>,
    typeinfo_enum_query: Option<Statement>,
    typeinfo_composite_query: Option<Statement>,
}

struct Inner {
    state: Mutex<State>,
    idle: IdleState,
    sender: mpsc::UnboundedSender<Request>,
}

#[derive(Clone)]
pub struct Client(Arc<Inner>);

impl Client {
    pub fn new(sender: mpsc::UnboundedSender<Request>) -> Client {
        Client(Arc::new(Inner {
            state: Mutex::new(State {
                types: HashMap::new(),
                typeinfo_query: None,
                typeinfo_enum_query: None,
                typeinfo_composite_query: None,
            }),
            idle: IdleState::new(),
            sender,
        }))
    }

    pub fn is_closed(&self) -> bool {
        self.0.sender.is_closed()
    }

    pub fn poll_idle(&self) -> Poll<(), Error> {
        self.0.idle.poll_idle()
    }

    pub fn downgrade(&self) -> WeakClient {
        WeakClient(Arc::downgrade(&self.0))
    }

    pub fn cached_type(&self, oid: Oid) -> Option<Type> {
        self.0.state.lock().types.get(&oid).cloned()
    }

    pub fn cache_type(&self, ty: &Type) {
        self.0.state.lock().types.insert(ty.oid(), ty.clone());
    }

    pub fn typeinfo_query(&self) -> Option<Statement> {
        self.0.state.lock().typeinfo_query.clone()
    }

    pub fn set_typeinfo_query(&self, statement: &Statement) {
        self.0.state.lock().typeinfo_query = Some(statement.clone());
    }

    pub fn typeinfo_enum_query(&self) -> Option<Statement> {
        self.0.state.lock().typeinfo_enum_query.clone()
    }

    pub fn set_typeinfo_enum_query(&self, statement: &Statement) {
        self.0.state.lock().typeinfo_enum_query = Some(statement.clone());
    }

    pub fn typeinfo_composite_query(&self) -> Option<Statement> {
        self.0.state.lock().typeinfo_composite_query.clone()
    }

    pub fn set_typeinfo_composite_query(&self, statement: &Statement) {
        self.0.state.lock().typeinfo_composite_query = Some(statement.clone());
    }

    pub fn send(&self, request: PendingRequest) -> Result<mpsc::Receiver<Message>, Error> {
        let (messages, idle) = request.0?;
        let (sender, receiver) = mpsc::channel(0);
        self.0
            .sender
            .unbounded_send(Request {
                messages,
                sender,
                idle: Some(idle),
            })
            .map(|_| receiver)
            .map_err(|_| Error::closed())
    }

    pub fn batch_execute(&self, query: &str) -> SimpleQueryFuture {
        let pending = self.pending(|buf| {
            frontend::query(query, buf).map_err(Error::parse)?;
            Ok(())
        });

        SimpleQueryFuture::new(self.clone(), pending)
    }

    pub fn prepare(&self, name: String, query: &str, param_types: &[Type]) -> PrepareFuture {
        let pending = self.pending(|buf| {
            frontend::parse(&name, query, param_types.iter().map(|t| t.oid()), buf)
                .map_err(Error::parse)?;
            frontend::describe(b'S', &name, buf).map_err(Error::parse)?;
            frontend::sync(buf);
            Ok(())
        });

        PrepareFuture::new(self.clone(), pending, name)
    }

    pub fn execute(&self, statement: &Statement, params: &[&dyn ToSql]) -> ExecuteFuture {
        let pending = PendingRequest(
            self.excecute_message(statement, params)
                .map(|m| (RequestMessages::Single(m), self.0.idle.guard())),
        );
        ExecuteFuture::new(self.clone(), pending, statement.clone())
    }

    pub fn query(&self, statement: &Statement, params: &[&dyn ToSql]) -> QueryStream<Statement> {
        let pending = PendingRequest(
            self.excecute_message(statement, params)
                .map(|m| (RequestMessages::Single(m), self.0.idle.guard())),
        );
        QueryStream::new(self.clone(), pending, statement.clone())
    }

    pub fn bind(&self, statement: &Statement, name: String, params: &[&dyn ToSql]) -> BindFuture {
        let mut buf = self.bind_message(statement, &name, params);
        if let Ok(ref mut buf) = buf {
            frontend::sync(buf);
        }
        let pending =
            PendingRequest(buf.map(|m| (RequestMessages::Single(m), self.0.idle.guard())));
        BindFuture::new(self.clone(), pending, name, statement.clone())
    }

    pub fn query_portal(&self, portal: &Portal, rows: i32) -> QueryStream<Portal> {
        let pending = self.pending(|buf| {
            frontend::execute(portal.name(), rows, buf).map_err(Error::parse)?;
            frontend::sync(buf);
            Ok(())
        });
        QueryStream::new(self.clone(), pending, portal.clone())
    }

    pub fn copy_in<S>(
        &self,
        statement: &Statement,
        params: &[&dyn ToSql],
        stream: S,
    ) -> CopyInFuture<S>
    where
        S: Stream,
        S::Item: IntoBuf,
        <S::Item as IntoBuf>::Buf: Send,
        S::Error: Into<Box<dyn StdError + Sync + Send>>,
    {
        let (mut sender, receiver) = mpsc::channel(1);
        let pending = PendingRequest(self.excecute_message(statement, params).map(|buf| {
            match sender.start_send(CopyMessage::Data(buf)) {
                Ok(AsyncSink::Ready) => {}
                _ => unreachable!("channel should have capacity"),
            }
            (
                RequestMessages::CopyIn {
                    receiver: CopyInReceiver::new(receiver),
                    pending_message: None,
                },
                self.0.idle.guard(),
            )
        }));
        CopyInFuture::new(self.clone(), pending, statement.clone(), stream, sender)
    }

    pub fn copy_out(&self, statement: &Statement, params: &[&dyn ToSql]) -> CopyOutStream {
        let pending = PendingRequest(
            self.excecute_message(statement, params)
                .map(|m| (RequestMessages::Single(m), self.0.idle.guard())),
        );
        CopyOutStream::new(self.clone(), pending, statement.clone())
    }

    pub fn close_statement(&self, name: &str) {
        self.close(b'S', name)
    }

    pub fn close_portal(&self, name: &str) {
        self.close(b'P', name)
    }

    fn close(&self, ty: u8, name: &str) {
        let mut buf = vec![];
        frontend::close(ty, name, &mut buf).expect("statement name not valid");
        frontend::sync(&mut buf);
        let (sender, _) = mpsc::channel(0);
        let _ = self.0.sender.unbounded_send(Request {
            messages: RequestMessages::Single(buf),
            sender,
            idle: None,
        });
    }

    fn bind_message(
        &self,
        statement: &Statement,
        name: &str,
        params: &[&dyn ToSql],
    ) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let r = frontend::bind(
            name,
            statement.name(),
            Some(1),
            params.iter().zip(statement.params()),
            |(param, ty), buf| match param.to_sql_checked(ty, buf) {
                Ok(IsNull::No) => Ok(postgres_protocol::IsNull::No),
                Ok(IsNull::Yes) => Ok(postgres_protocol::IsNull::Yes),
                Err(e) => Err(e),
            },
            Some(1),
            &mut buf,
        );
        match r {
            Ok(()) => Ok(buf),
            Err(frontend::BindError::Conversion(e)) => Err(Error::to_sql(e)),
            Err(frontend::BindError::Serialization(e)) => Err(Error::encode(e)),
        }
    }

    fn excecute_message(
        &self,
        statement: &Statement,
        params: &[&dyn ToSql],
    ) -> Result<Vec<u8>, Error> {
        let mut buf = self.bind_message(statement, "", params)?;
        frontend::execute("", 0, &mut buf).map_err(Error::parse)?;
        frontend::sync(&mut buf);
        Ok(buf)
    }

    fn pending<F>(&self, messages: F) -> PendingRequest
    where
        F: FnOnce(&mut Vec<u8>) -> Result<(), Error>,
    {
        let mut buf = vec![];
        PendingRequest(
            messages(&mut buf).map(|()| (RequestMessages::Single(buf), self.0.idle.guard())),
        )
    }
}
<commit_msg>Guarantee some buffer for messages<commit_after>use antidote::Mutex;
use bytes::IntoBuf;
use futures::sync::mpsc;
use futures::{AsyncSink, Poll, Sink, Stream};
use postgres_protocol;
use postgres_protocol::message::backend::Message;
use postgres_protocol::message::frontend;
use std::collections::HashMap;
use std::error::Error as StdError;
use std::sync::{Arc, Weak};

use crate::proto::bind::BindFuture;
use crate::proto::connection::{Request, RequestMessages};
use crate::proto::copy_in::{CopyInFuture, CopyInReceiver, CopyMessage};
use crate::proto::copy_out::CopyOutStream;
use crate::proto::execute::ExecuteFuture;
use crate::proto::idle::{IdleGuard, IdleState};
use crate::proto::portal::Portal;
use crate::proto::prepare::PrepareFuture;
use crate::proto::query::QueryStream;
use crate::proto::simple_query::SimpleQueryFuture;
use crate::proto::statement::Statement;
use crate::types::{IsNull, Oid, ToSql, Type};
use crate::Error;

pub struct PendingRequest(Result<(RequestMessages, IdleGuard), Error>);

pub struct WeakClient(Weak<Inner>);

impl WeakClient {
    pub fn upgrade(&self) -> Option<Client> {
        self.0.upgrade().map(Client)
    }
}

struct State {
    types: HashMap<Oid, Type>,
    typeinfo_query: Option<Statement>,
    typeinfo_enum_query: Option<Statement>,
    typeinfo_composite_query: Option<Statement>,
}

struct Inner {
    state: Mutex<State>,
    idle: IdleState,
    sender: mpsc::UnboundedSender<Request>,
}

#[derive(Clone)]
pub struct Client(Arc<Inner>);

impl Client {
    pub fn new(sender: mpsc::UnboundedSender<Request>) -> Client {
        Client(Arc::new(Inner {
            state: Mutex::new(State {
                types: HashMap::new(),
                typeinfo_query: None,
                typeinfo_enum_query: None,
                typeinfo_composite_query: None,
            }),
            idle: IdleState::new(),
            sender,
        }))
    }

    pub fn is_closed(&self) -> bool {
        self.0.sender.is_closed()
    }

    pub fn poll_idle(&self) -> Poll<(), Error> {
        self.0.idle.poll_idle()
    }

    pub fn downgrade(&self) -> WeakClient {
        WeakClient(Arc::downgrade(&self.0))
    }

    pub fn cached_type(&self, oid: Oid) -> Option<Type> {
        self.0.state.lock().types.get(&oid).cloned()
    }

    pub fn cache_type(&self, ty: &Type) {
        self.0.state.lock().types.insert(ty.oid(), ty.clone());
    }

    pub fn typeinfo_query(&self) -> Option<Statement> {
        self.0.state.lock().typeinfo_query.clone()
    }

    pub fn set_typeinfo_query(&self, statement: &Statement) {
        self.0.state.lock().typeinfo_query = Some(statement.clone());
    }

    pub fn typeinfo_enum_query(&self) -> Option<Statement> {
        self.0.state.lock().typeinfo_enum_query.clone()
    }

    pub fn set_typeinfo_enum_query(&self, statement: &Statement) {
        self.0.state.lock().typeinfo_enum_query = Some(statement.clone());
    }

    pub fn typeinfo_composite_query(&self) -> Option<Statement> {
        self.0.state.lock().typeinfo_composite_query.clone()
    }

    pub fn set_typeinfo_composite_query(&self, statement: &Statement) {
        self.0.state.lock().typeinfo_composite_query = Some(statement.clone());
    }

    pub fn send(&self, request: PendingRequest) -> Result<mpsc::Receiver<Message>, Error> {
        let (messages, idle) = request.0?;
        let (sender, receiver) = mpsc::channel(1);
        self.0
            .sender
            .unbounded_send(Request {
                messages,
                sender,
                idle: Some(idle),
            })
            .map(|_| receiver)
            .map_err(|_| Error::closed())
    }

    pub fn batch_execute(&self, query: &str) -> SimpleQueryFuture {
        let pending = self.pending(|buf| {
            frontend::query(query, buf).map_err(Error::parse)?;
            Ok(())
        });

        SimpleQueryFuture::new(self.clone(), pending)
    }

    pub fn prepare(&self, name: String, query: &str, param_types: &[Type]) -> PrepareFuture {
        let pending = self.pending(|buf| {
            frontend::parse(&name, query, param_types.iter().map(|t| t.oid()), buf)
                .map_err(Error::parse)?;
            frontend::describe(b'S', &name, buf).map_err(Error::parse)?;
            frontend::sync(buf);
            Ok(())
        });

        PrepareFuture::new(self.clone(), pending, name)
    }

    pub fn execute(&self, statement: &Statement, params: &[&dyn ToSql]) -> ExecuteFuture {
        let pending = PendingRequest(
            self.excecute_message(statement, params)
                .map(|m| (RequestMessages::Single(m), self.0.idle.guard())),
        );
        ExecuteFuture::new(self.clone(), pending, statement.clone())
    }

    pub fn query(&self, statement: &Statement, params: &[&dyn ToSql]) -> QueryStream<Statement> {
        let pending = PendingRequest(
            self.excecute_message(statement, params)
                .map(|m| (RequestMessages::Single(m), self.0.idle.guard())),
        );
        QueryStream::new(self.clone(), pending, statement.clone())
    }

    pub fn bind(&self, statement: &Statement, name: String, params: &[&dyn ToSql]) -> BindFuture {
        let mut buf = self.bind_message(statement, &name, params);
        if let Ok(ref mut buf) = buf {
            frontend::sync(buf);
        }
        let pending =
            PendingRequest(buf.map(|m| (RequestMessages::Single(m), self.0.idle.guard())));
        BindFuture::new(self.clone(), pending, name, statement.clone())
    }

    pub fn query_portal(&self, portal: &Portal, rows: i32) -> QueryStream<Portal> {
        let pending = self.pending(|buf| {
            frontend::execute(portal.name(), rows, buf).map_err(Error::parse)?;
            frontend::sync(buf);
            Ok(())
        });
        QueryStream::new(self.clone(), pending, portal.clone())
    }

    pub fn copy_in<S>(
        &self,
        statement: &Statement,
        params: &[&dyn ToSql],
        stream: S,
    ) -> CopyInFuture<S>
    where
        S: Stream,
        S::Item: IntoBuf,
        <S::Item as IntoBuf>::Buf: Send,
        S::Error: Into<Box<dyn StdError + Sync + Send>>,
    {
        let (mut sender, receiver) = mpsc::channel(1);
        let pending = PendingRequest(self.excecute_message(statement, params).map(|buf| {
            match sender.start_send(CopyMessage::Data(buf)) {
                Ok(AsyncSink::Ready) => {}
                _ => unreachable!("channel should have capacity"),
            }
            (
                RequestMessages::CopyIn {
                    receiver: CopyInReceiver::new(receiver),
                    pending_message: None,
                },
                self.0.idle.guard(),
            )
        }));
        CopyInFuture::new(self.clone(), pending, statement.clone(), stream, sender)
    }

    pub fn copy_out(&self, statement: &Statement, params: &[&dyn ToSql]) -> CopyOutStream {
        let pending = PendingRequest(
            self.excecute_message(statement, params)
                .map(|m| (RequestMessages::Single(m), self.0.idle.guard())),
        );
        CopyOutStream::new(self.clone(), pending, statement.clone())
    }

    pub fn close_statement(&self, name: &str) {
        self.close(b'S', name)
    }

    pub fn close_portal(&self, name: &str) {
        self.close(b'P', name)
    }

    fn close(&self, ty: u8, name: &str) {
        let mut buf = vec![];
        frontend::close(ty, name, &mut buf).expect("statement name not valid");
        frontend::sync(&mut buf);
        let (sender, _) = mpsc::channel(0);
        let _ = self.0.sender.unbounded_send(Request {
            messages: RequestMessages::Single(buf),
            sender,
            idle: None,
        });
    }

    fn bind_message(
        &self,
        statement: &Statement,
        name: &str,
        params: &[&dyn ToSql],
    ) -> Result<Vec<u8>, Error> {
        let mut buf = vec![];
        let r = frontend::bind(
            name,
            statement.name(),
            Some(1),
            params.iter().zip(statement.params()),
            |(param, ty), buf| match param.to_sql_checked(ty, buf) {
                Ok(IsNull::No) => Ok(postgres_protocol::IsNull::No),
                Ok(IsNull::Yes) => Ok(postgres_protocol::IsNull::Yes),
                Err(e) => Err(e),
            },
            Some(1),
            &mut buf,
        );
        match r {
            Ok(()) => Ok(buf),
            Err(frontend::BindError::Conversion(e)) => Err(Error::to_sql(e)),
            Err(frontend::BindError::Serialization(e)) => Err(Error::encode(e)),
        }
    }

    fn excecute_message(
        &self,
        statement: &Statement,
        params: &[&dyn ToSql],
    ) -> Result<Vec<u8>, Error> {
        let mut buf = self.bind_message(statement, "", params)?;
        frontend::execute("", 0, &mut buf).map_err(Error::parse)?;
        frontend::sync(&mut buf);
        Ok(buf)
    }

    fn pending<F>(&self, messages: F) -> PendingRequest
    where
        F: FnOnce(&mut Vec<u8>) -> Result<(), Error>,
    {
        let mut buf = vec![];
        PendingRequest(
            messages(&mut buf).map(|()| (RequestMessages::Single(buf), self.0.idle.guard())),
        )
    }
}
<|endoftext|>


<commit_before>//! Easy file downloading

use errors::*;
use notifications::Notification;
use sha2::{Sha256, Digest};
use std::cell::RefCell;
use std::fs;
use std::io;
use std::env;
use std::path::Path;
use url::Url;

pub fn download_file(url: &Url,
                     path: &Path,
                     hasher: Option<&mut Sha256>,
                     notify_handler: &Fn(Notification))
                     -> Result<()> {

    let file = RefCell::new(try!(fs::File::create(&path).chain_err(
        || "error creating file for download")));
    let hasher = RefCell::new(hasher);

    // This callback will write the download to disk and optionally
    // hash the contents, then forward the notification up the stack
    let handler: &Fn(Notification) -> Result<()> = &|notice| {
        match notice {
            Notification::DownloadDataReceived(data) => {
                try!(io::Write::write_all(&mut *file.borrow_mut(), data)
                     .chain_err(|| "unable to write download to disk"));
                if let Some(ref mut h) = *hasher.borrow_mut() {
                    h.input(data);
                }
            }
            _ => ()
        }

        notify_handler(notice);

        Ok(())
    };

    // Download the file
    if env::var_os("RUSTUP_USE_HYPER").is_some() {
        notify_handler(Notification::UsingHyper);
         try!(self::hyper::download_file(url, handler));
    } else {
        notify_handler(Notification::UsingCurl);
        try!(self::curl::download_file(url, handler));
    }

    try!(file.borrow_mut().sync_data().chain_err(|| "unable to sync download to disk"));

    notify_handler(Notification::DownloadFinished);

    Ok(())
}

/// Download via libcurl; encrypt with the native (or OpenSSl) TLS
/// stack via libcurl
mod curl {
    use curl::easy::Easy;
    use errors::*;
    use notifications::Notification;
    use std::cell::RefCell;
    use std::str;
    use std::time::Duration;
    use url::Url;

    pub fn download_file(url: &Url,
                         notify_handler: &Fn(Notification) -> Result<()> )
                         -> Result<()> {
        // Fetch either a cached libcurl handle (which will preserve open
        // connections) or create a new one if it isn't listed.
        //
        // Once we've acquired it, reset the lifetime from 'static to our local
        // scope.
        thread_local!(static EASY: RefCell<Easy> = RefCell::new(Easy::new()));
        EASY.with(|handle| {
            let mut handle = handle.borrow_mut();

            try!(handle.url(&url.to_string()).chain_err(|| "failed to set url"));
            try!(handle.follow_location(true).chain_err(|| "failed to set follow redirects"));

            // Take at most 30s to connect
            try!(handle.connect_timeout(Duration::new(30, 0)).chain_err(|| "failed to set connect timeout"));

            // Fail if less than 10 bytes are transferred every 30 seconds
            try!(handle.low_speed_limit(10).chain_err(|| "failed to set low speed limit"));
            try!(handle.low_speed_time(Duration::new(30, 0)).chain_err(|| "failed to set low speed time"));

            {
                let fserr = RefCell::new(None);
                let mut transfer = handle.transfer();

                // Data callback for libcurl which is called with data that's
                // downloaded. We just feed it into our hasher and also write it out
                // to disk.
                try!(transfer.write_function(|data| {
                    match notify_handler(Notification::DownloadDataReceived(data)) {
                        Ok(()) => Ok(data.len()),
                        Err(e) => {
                            *fserr.borrow_mut() = Some(e);
                            Ok(0)
                        }
                    }
                }).chain_err(|| "failed to set write"));

                // Listen for headers and parse out a `Content-Length` if it comes
                // so we know how much we're downloading.
                try!(transfer.header_function(|header| {
                    if let Ok(data) = str::from_utf8(header) {
                        let prefix = "Content-Length: ";
                        if data.starts_with(prefix) {
                            if let Ok(s) = data[prefix.len()..].trim().parse() {
                                let msg = Notification::DownloadContentLengthReceived(s);
                                match notify_handler(msg) {
                                    Ok(()) => (),
                                    Err(_e) => {
                                        // FIXME: discarding error _e
                                        return false;
                                    }
                                }
                            }
                        }
                    }
                    true
                }).chain_err(|| "failed to set header"));

                // If an error happens check to see if we had a filesystem error up
                // in `fserr`, but we always want to punt it up.
                try!(transfer.perform().or_else(|e| {
                    match fserr.borrow_mut().take() {
                        Some(fs) => Err(fs).chain_err(|| ErrorKind::HttpError(e)),
                        None => Err(ErrorKind::HttpError(e).into())
                    }
                }));
            }

            // If we didn't get a 200 or 0 ("OK" for files) then return an error
            let code = try!(handle.response_code().chain_err(|| "failed to get response code"));
            if code != 200 && code != 0 {
                return Err(ErrorKind::HttpStatus(code).into());
            }

            Ok(())
        })
    }
}

/// Download via hyper; encrypt with the native (or OpenSSl) TLS
/// stack via native-tls
mod hyper {
    use hyper;
    use notifications::Notification;
    use std::fs;
    use std::io;
    use std::time::Duration;
    use url::Url;
    use errors::*;

    fn proxy_from_env(url: &Url) -> Option<(String, u16)> {
        use std::env::var_os;

        let mut maybe_https_proxy = var_os("https_proxy").map(|ref v| v.to_str().unwrap_or("").to_string());
        if maybe_https_proxy.is_none() {
            maybe_https_proxy = var_os("HTTPS_PROXY").map(|ref v| v.to_str().unwrap_or("").to_string());
        }
        let maybe_http_proxy = var_os("http_proxy").map(|ref v| v.to_str().unwrap_or("").to_string());
        let mut maybe_all_proxy = var_os("all_proxy").map(|ref v| v.to_str().unwrap_or("").to_string());
        if maybe_all_proxy.is_none() {
            maybe_all_proxy = var_os("ALL_PROXY").map(|ref v| v.to_str().unwrap_or("").to_string());
        }
        if let Some(url_value) = match url.scheme() {
                                     "https" => maybe_https_proxy.or(maybe_http_proxy.or(maybe_all_proxy)),
                                     "http" => maybe_http_proxy.or(maybe_all_proxy),
                                     _ => maybe_all_proxy,
                                 } {
            if let Ok(proxy_url) = Url::parse(&url_value) {
                if let Some(host) = proxy_url.host_str() {
                    let port = proxy_url.port().unwrap_or(8080);
                    return Some((host.to_string(), port));
                }
            }
        }
        None
    }

    pub fn download_file(url: &Url,
                         notify_handler: &Fn(Notification) -> Result<()>)
                         -> Result<()> {

        // Short-circuit hyper for the "file:" URL scheme
        if try!(download_from_file_url(url, notify_handler)) {
            return Ok(());
        }

        use hyper::client::{Client, ProxyConfig};
        use hyper::error::Result as HyperResult;
        use hyper::header::ContentLength;
        use hyper::net::{SslClient, NetworkStream, HttpsConnector};
        use native_tls;
        use std::io::Result as IoResult;
        use std::io::{Read, Write};
        use std::net::{SocketAddr, Shutdown};
        use std::sync::{Arc, Mutex};

        // The Hyper HTTP client
        let client;

        let maybe_proxy = proxy_from_env(url);
        if url.scheme() == "https" {

            // All the following is adapter code to use native_tls with hyper.

            struct NativeSslClient;

            impl<T: NetworkStream + Send + Clone> SslClient<T> for NativeSslClient {
                type Stream = NativeSslStream<T>;

                fn wrap_client(&self, stream: T, host: &str) -> HyperResult<Self::Stream> {
                    use native_tls::ClientBuilder as TlsClientBuilder;
                    use hyper::error::Error as HyperError;

                    let mut ssl_builder = try!(TlsClientBuilder::new()
                                               .map_err(|e| HyperError::Ssl(Box::new(e))));
                    let ssl_stream = try!(ssl_builder.handshake(host, stream)
                                          .map_err(|e| HyperError::Ssl(Box::new(e))));

                    Ok(NativeSslStream(Arc::new(Mutex::new(ssl_stream))))
                }
            }

            #[derive(Clone)]
            struct NativeSslStream<T>(Arc<Mutex<native_tls::TlsStream<T>>>);

            #[derive(Debug)]
            struct NativeSslPoisonError;

            impl ::std::error::Error for NativeSslPoisonError {
                fn description(&self) -> &str { "mutex poisoned during TLS operation" }
            }

            impl ::std::fmt::Display for NativeSslPoisonError {
                fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
                    f.write_str(::std::error::Error::description(self))
                }
            }

            impl<T> NetworkStream for NativeSslStream<T>
                where T: NetworkStream
            {
                fn peer_addr(&mut self) -> IoResult<SocketAddr> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.get_mut().peer_addr())
                }
                fn set_read_timeout(&self, dur: Option<Duration>) -> IoResult<()> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|t| t.get_ref().set_read_timeout(dur))
                }
                fn set_write_timeout(&self, dur: Option<Duration>) -> IoResult<()> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|t| t.get_ref().set_write_timeout(dur))
                }
                fn close(&mut self, how: Shutdown) -> IoResult<()> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.get_mut().close(how))
                }
            }

            impl<T> Read for NativeSslStream<T>
                where T: Read + Write
            {
                fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.read(buf))
                }
            }

            impl<T> Write for NativeSslStream<T>
                where T: Read + Write
            {
                fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.write(buf))
                }
                fn flush(&mut self) -> IoResult<()> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.flush())
                }
            }

            maybe_init_certs();

            if maybe_proxy.is_none() {
                // Connect with hyper + native_tls
                client = Client::with_connector(HttpsConnector::new(NativeSslClient));
            } else {
                let proxy_host_port = maybe_proxy.unwrap();
                client = Client::with_proxy_config(ProxyConfig(proxy_host_port.0, proxy_host_port.1, NativeSslClient));
            }
        } else if url.scheme() == "http" {
            if maybe_proxy.is_none() {
                client = Client::new();
            } else {
                let proxy_host_port = maybe_proxy.unwrap();
                client = Client::with_http_proxy(proxy_host_port.0, proxy_host_port.1);
            }
        } else {
            return Err(format!("unsupported URL scheme: '{}'", url.scheme()).into());
        }

        let mut res = try!(client.get(url.clone()).send()
                           .chain_err(|| "failed to make network request"));
        if res.status != hyper::Ok {
            return Err(ErrorKind::HttpStatus(res.status.to_u16() as u32).into());
        }

        let buffer_size = 0x10000;
        let mut buffer = vec![0u8; buffer_size];

        if let Some(len) = res.headers.get::<ContentLength>().cloned() {
            try!(notify_handler(Notification::DownloadContentLengthReceived(len.0)));
        }

        loop {
            let bytes_read = try!(io::Read::read(&mut res, &mut buffer)
                                  .chain_err(|| "error reading from socket"));

            if bytes_read != 0 {
                try!(notify_handler(Notification::DownloadDataReceived(&buffer[0..bytes_read])));
            } else {
                return Ok(());
            }
        }
    }

    // Tell our statically-linked OpenSSL where to find root certs
    // cc https://github.com/alexcrichton/git2-rs/blob/master/libgit2-sys/lib.rs#L1267
    #[cfg(not(any(target_os = "windows", target_os = "macos")))]
    fn maybe_init_certs() {
        use std::sync::{Once, ONCE_INIT};
        static INIT: Once = ONCE_INIT;
        INIT.call_once(|| {
            ::openssl_sys::probe::init_ssl_cert_env_vars();
        });
    }

    #[cfg(any(target_os = "windows", target_os = "macos"))]
    fn maybe_init_certs() { }

    fn download_from_file_url(url: &Url,
                              notify_handler: &Fn(Notification) -> Result<()>)
                              -> Result<bool> {
        use raw::is_file;

        // The file scheme is mostly for use by tests to mock the dist server
        if url.scheme() == "file" {
            let src = try!(url.to_file_path()
                           .map_err(|_| Error::from(format!("bogus file url: '{}'", url))));
            if !is_file(&src) {
                // Because some of multirust's logic depends on checking
                // the error when a downloaded file doesn't exist, make
                // the file case return the same error value as the
                // network case.
                return Err(ErrorKind::HttpStatus(hyper::status::StatusCode::NotFound.to_u16() as u32).into());
            }

            let ref mut f = try!(fs::File::open(src)
                                 .chain_err(|| "unable to open downloaded file"));

            let ref mut buffer = vec![0u8; 0x10000];
            loop {
                let bytes_read = try!(io::Read::read(f, buffer)
                                      .chain_err(|| "unable to read downloaded file"));
                if bytes_read == 0 { break }
                try!(notify_handler(Notification::DownloadDataReceived(&buffer[0..bytes_read])));
            }

            Ok(true)
        } else {
            Ok(false)
        }
    }
}
<commit_msg>Give download_file its own notification types<commit_after>//! Easy file downloading

use errors::*;
use notifications::Notification;
use sha2::{Sha256, Digest};
use std::cell::RefCell;
use std::fs;
use std::io;
use std::env;
use std::path::Path;
use url::Url;

pub fn download_file(url: &Url,
                     path: &Path,
                     hasher: Option<&mut Sha256>,
                     notify_handler: &Fn(Notification))
                     -> Result<()> {

    let file = RefCell::new(try!(fs::File::create(&path).chain_err(
        || "error creating file for download")));
    let hasher = RefCell::new(hasher);

    // This callback will write the download to disk and optionally
    // hash the contents, then forward the notification up the stack
    let callback: &Fn(Event) -> Result<()> = &|msg| {
        match msg {
            Event::DownloadDataReceived(data) => {
                try!(io::Write::write_all(&mut *file.borrow_mut(), data)
                     .chain_err(|| "unable to write download to disk"));
                if let Some(ref mut h) = *hasher.borrow_mut() {
                    h.input(data);
                }
            }
            _ => ()
        }

        match msg {
            Event::DownloadContentLengthReceived(len) => {
                notify_handler(Notification::DownloadContentLengthReceived(len));
            }
            Event::DownloadDataReceived(data) => {
                notify_handler(Notification::DownloadDataReceived(data));
            }
        }

        Ok(())
    };

    // Download the file
    if env::var_os("RUSTUP_USE_HYPER").is_some() {
        notify_handler(Notification::UsingHyper);
         try!(self::hyper::download_file(url, callback));
    } else {
        notify_handler(Notification::UsingCurl);
        try!(self::curl::download_file(url, callback));
    }

    try!(file.borrow_mut().sync_data().chain_err(|| "unable to sync download to disk"));

    notify_handler(Notification::DownloadFinished);

    Ok(())
}

#[derive(Debug)]
pub enum Event<'a> {
    /// Received the Content-Length of the to-be downloaded data.
    DownloadContentLengthReceived(u64),
    /// Received some data.
    DownloadDataReceived(&'a [u8]),
}

/// Download via libcurl; encrypt with the native (or OpenSSl) TLS
/// stack via libcurl
mod curl {
    use curl::easy::Easy;
    use errors::*;
    use std::cell::RefCell;
    use std::str;
    use std::time::Duration;
    use url::Url;
    use super::Event;

    pub fn download_file(url: &Url,
                         callback: &Fn(Event) -> Result<()> )
                         -> Result<()> {
        // Fetch either a cached libcurl handle (which will preserve open
        // connections) or create a new one if it isn't listed.
        //
        // Once we've acquired it, reset the lifetime from 'static to our local
        // scope.
        thread_local!(static EASY: RefCell<Easy> = RefCell::new(Easy::new()));
        EASY.with(|handle| {
            let mut handle = handle.borrow_mut();

            try!(handle.url(&url.to_string()).chain_err(|| "failed to set url"));
            try!(handle.follow_location(true).chain_err(|| "failed to set follow redirects"));

            // Take at most 30s to connect
            try!(handle.connect_timeout(Duration::new(30, 0)).chain_err(|| "failed to set connect timeout"));

            // Fail if less than 10 bytes are transferred every 30 seconds
            try!(handle.low_speed_limit(10).chain_err(|| "failed to set low speed limit"));
            try!(handle.low_speed_time(Duration::new(30, 0)).chain_err(|| "failed to set low speed time"));

            {
                let fserr = RefCell::new(None);
                let mut transfer = handle.transfer();

                // Data callback for libcurl which is called with data that's
                // downloaded. We just feed it into our hasher and also write it out
                // to disk.
                try!(transfer.write_function(|data| {
                    match callback(Event::DownloadDataReceived(data)) {
                        Ok(()) => Ok(data.len()),
                        Err(e) => {
                            *fserr.borrow_mut() = Some(e);
                            Ok(0)
                        }
                    }
                }).chain_err(|| "failed to set write"));

                // Listen for headers and parse out a `Content-Length` if it comes
                // so we know how much we're downloading.
                try!(transfer.header_function(|header| {
                    if let Ok(data) = str::from_utf8(header) {
                        let prefix = "Content-Length: ";
                        if data.starts_with(prefix) {
                            if let Ok(s) = data[prefix.len()..].trim().parse() {
                                let msg = Event::DownloadContentLengthReceived(s);
                                match callback(msg) {
                                    Ok(()) => (),
                                    Err(_e) => {
                                        // FIXME: discarding error _e
                                        return false;
                                    }
                                }
                            }
                        }
                    }
                    true
                }).chain_err(|| "failed to set header"));

                // If an error happens check to see if we had a filesystem error up
                // in `fserr`, but we always want to punt it up.
                try!(transfer.perform().or_else(|e| {
                    match fserr.borrow_mut().take() {
                        Some(fs) => Err(fs).chain_err(|| ErrorKind::HttpError(e)),
                        None => Err(ErrorKind::HttpError(e).into())
                    }
                }));
            }

            // If we didn't get a 200 or 0 ("OK" for files) then return an error
            let code = try!(handle.response_code().chain_err(|| "failed to get response code"));
            if code != 200 && code != 0 {
                return Err(ErrorKind::HttpStatus(code).into());
            }

            Ok(())
        })
    }
}

/// Download via hyper; encrypt with the native (or OpenSSl) TLS
/// stack via native-tls
mod hyper {
    use hyper;
    use super::Event;
    use std::fs;
    use std::io;
    use std::time::Duration;
    use url::Url;
    use errors::*;

    fn proxy_from_env(url: &Url) -> Option<(String, u16)> {
        use std::env::var_os;

        let mut maybe_https_proxy = var_os("https_proxy").map(|ref v| v.to_str().unwrap_or("").to_string());
        if maybe_https_proxy.is_none() {
            maybe_https_proxy = var_os("HTTPS_PROXY").map(|ref v| v.to_str().unwrap_or("").to_string());
        }
        let maybe_http_proxy = var_os("http_proxy").map(|ref v| v.to_str().unwrap_or("").to_string());
        let mut maybe_all_proxy = var_os("all_proxy").map(|ref v| v.to_str().unwrap_or("").to_string());
        if maybe_all_proxy.is_none() {
            maybe_all_proxy = var_os("ALL_PROXY").map(|ref v| v.to_str().unwrap_or("").to_string());
        }
        if let Some(url_value) = match url.scheme() {
                                     "https" => maybe_https_proxy.or(maybe_http_proxy.or(maybe_all_proxy)),
                                     "http" => maybe_http_proxy.or(maybe_all_proxy),
                                     _ => maybe_all_proxy,
                                 } {
            if let Ok(proxy_url) = Url::parse(&url_value) {
                if let Some(host) = proxy_url.host_str() {
                    let port = proxy_url.port().unwrap_or(8080);
                    return Some((host.to_string(), port));
                }
            }
        }
        None
    }

    pub fn download_file(url: &Url,
                         callback: &Fn(Event) -> Result<()>)
                         -> Result<()> {

        // Short-circuit hyper for the "file:" URL scheme
        if try!(download_from_file_url(url, callback)) {
            return Ok(());
        }

        use hyper::client::{Client, ProxyConfig};
        use hyper::error::Result as HyperResult;
        use hyper::header::ContentLength;
        use hyper::net::{SslClient, NetworkStream, HttpsConnector};
        use native_tls;
        use std::io::Result as IoResult;
        use std::io::{Read, Write};
        use std::net::{SocketAddr, Shutdown};
        use std::sync::{Arc, Mutex};

        // The Hyper HTTP client
        let client;

        let maybe_proxy = proxy_from_env(url);
        if url.scheme() == "https" {

            // All the following is adapter code to use native_tls with hyper.

            struct NativeSslClient;

            impl<T: NetworkStream + Send + Clone> SslClient<T> for NativeSslClient {
                type Stream = NativeSslStream<T>;

                fn wrap_client(&self, stream: T, host: &str) -> HyperResult<Self::Stream> {
                    use native_tls::ClientBuilder as TlsClientBuilder;
                    use hyper::error::Error as HyperError;

                    let mut ssl_builder = try!(TlsClientBuilder::new()
                                               .map_err(|e| HyperError::Ssl(Box::new(e))));
                    let ssl_stream = try!(ssl_builder.handshake(host, stream)
                                          .map_err(|e| HyperError::Ssl(Box::new(e))));

                    Ok(NativeSslStream(Arc::new(Mutex::new(ssl_stream))))
                }
            }

            #[derive(Clone)]
            struct NativeSslStream<T>(Arc<Mutex<native_tls::TlsStream<T>>>);

            #[derive(Debug)]
            struct NativeSslPoisonError;

            impl ::std::error::Error for NativeSslPoisonError {
                fn description(&self) -> &str { "mutex poisoned during TLS operation" }
            }

            impl ::std::fmt::Display for NativeSslPoisonError {
                fn fmt(&self, f: &mut ::std::fmt::Formatter) -> ::std::result::Result<(), ::std::fmt::Error> {
                    f.write_str(::std::error::Error::description(self))
                }
            }

            impl<T> NetworkStream for NativeSslStream<T>
                where T: NetworkStream
            {
                fn peer_addr(&mut self) -> IoResult<SocketAddr> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.get_mut().peer_addr())
                }
                fn set_read_timeout(&self, dur: Option<Duration>) -> IoResult<()> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|t| t.get_ref().set_read_timeout(dur))
                }
                fn set_write_timeout(&self, dur: Option<Duration>) -> IoResult<()> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|t| t.get_ref().set_write_timeout(dur))
                }
                fn close(&mut self, how: Shutdown) -> IoResult<()> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.get_mut().close(how))
                }
            }

            impl<T> Read for NativeSslStream<T>
                where T: Read + Write
            {
                fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.read(buf))
                }
            }

            impl<T> Write for NativeSslStream<T>
                where T: Read + Write
            {
                fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.write(buf))
                }
                fn flush(&mut self) -> IoResult<()> {
                    self.0.lock()
                        .map_err(|_| io::Error::new(io::ErrorKind::Other, NativeSslPoisonError))
                        .and_then(|mut t| t.flush())
                }
            }

            maybe_init_certs();

            if maybe_proxy.is_none() {
                // Connect with hyper + native_tls
                client = Client::with_connector(HttpsConnector::new(NativeSslClient));
            } else {
                let proxy_host_port = maybe_proxy.unwrap();
                client = Client::with_proxy_config(ProxyConfig(proxy_host_port.0, proxy_host_port.1, NativeSslClient));
            }
        } else if url.scheme() == "http" {
            if maybe_proxy.is_none() {
                client = Client::new();
            } else {
                let proxy_host_port = maybe_proxy.unwrap();
                client = Client::with_http_proxy(proxy_host_port.0, proxy_host_port.1);
            }
        } else {
            return Err(format!("unsupported URL scheme: '{}'", url.scheme()).into());
        }

        let mut res = try!(client.get(url.clone()).send()
                           .chain_err(|| "failed to make network request"));
        if res.status != hyper::Ok {
            return Err(ErrorKind::HttpStatus(res.status.to_u16() as u32).into());
        }

        let buffer_size = 0x10000;
        let mut buffer = vec![0u8; buffer_size];

        if let Some(len) = res.headers.get::<ContentLength>().cloned() {
            try!(callback(Event::DownloadContentLengthReceived(len.0)));
        }

        loop {
            let bytes_read = try!(io::Read::read(&mut res, &mut buffer)
                                  .chain_err(|| "error reading from socket"));

            if bytes_read != 0 {
                try!(callback(Event::DownloadDataReceived(&buffer[0..bytes_read])));
            } else {
                return Ok(());
            }
        }
    }

    // Tell our statically-linked OpenSSL where to find root certs
    // cc https://github.com/alexcrichton/git2-rs/blob/master/libgit2-sys/lib.rs#L1267
    #[cfg(not(any(target_os = "windows", target_os = "macos")))]
    fn maybe_init_certs() {
        use std::sync::{Once, ONCE_INIT};
        static INIT: Once = ONCE_INIT;
        INIT.call_once(|| {
            ::openssl_sys::probe::init_ssl_cert_env_vars();
        });
    }

    #[cfg(any(target_os = "windows", target_os = "macos"))]
    fn maybe_init_certs() { }

    fn download_from_file_url(url: &Url,
                              callback: &Fn(Event) -> Result<()>)
                              -> Result<bool> {
        use raw::is_file;

        // The file scheme is mostly for use by tests to mock the dist server
        if url.scheme() == "file" {
            let src = try!(url.to_file_path()
                           .map_err(|_| Error::from(format!("bogus file url: '{}'", url))));
            if !is_file(&src) {
                // Because some of multirust's logic depends on checking
                // the error when a downloaded file doesn't exist, make
                // the file case return the same error value as the
                // network case.
                return Err(ErrorKind::HttpStatus(hyper::status::StatusCode::NotFound.to_u16() as u32).into());
            }

            let ref mut f = try!(fs::File::open(src)
                                 .chain_err(|| "unable to open downloaded file"));

            let ref mut buffer = vec![0u8; 0x10000];
            loop {
                let bytes_read = try!(io::Read::read(f, buffer)
                                      .chain_err(|| "unable to read downloaded file"));
                if bytes_read == 0 { break }
                try!(callback(Event::DownloadDataReceived(&buffer[0..bytes_read])));
            }

            Ok(true)
        } else {
            Ok(false)
        }
    }
}
<|endoftext|>


<commit_before>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

use std::fs;
use std::io;
use std::sync::atomic::Ordering;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::*;

use fail;
use raft::eraftpb::MessageType;

use raftstore::store::*;
use test_raftstore::*;
use tikv_util::config::*;
use tikv_util::HandyRwLock;

#[test]
fn test_overlap_cleanup() {
    let mut cluster = new_node_cluster(0, 3);
    // Disable raft log gc in this test case.
    cluster.cfg.raft_store.raft_log_gc_tick_interval = ReadableDuration::secs(60);

    let gen_snapshot_fp = "region_gen_snap";

    let pd_client = Arc::clone(&cluster.pd_client);
    // Disable default max peer count check.
    pd_client.disable_default_operator();

    let region_id = cluster.run_conf_change();
    pd_client.must_add_peer(region_id, new_peer(2, 2));

    cluster.must_put(b"k1", b"v1");
    must_get_equal(&cluster.get_engine(2), b"k1", b"v1");

    cluster.must_transfer_leader(region_id, new_peer(2, 2));
    // This will only pause the bootstrapped region, so the split region
    // can still work as expected.
    fail::cfg(gen_snapshot_fp, "pause").unwrap();
    pd_client.must_add_peer(region_id, new_peer(3, 3));
    cluster.must_put(b"k3", b"v3");
    assert_snapshot(&cluster.get_snap_dir(2), region_id, true);
    let region1 = cluster.get_region(b"k1");
    cluster.must_split(&region1, b"k2");
    // Wait till the snapshot of split region is applied, whose range is ["", "k2").
    must_get_equal(&cluster.get_engine(3), b"k1", b"v1");
    // Resume the fail point and pause it again. So only the paused snapshot is generated.
    // And the paused snapshot's range is ["", ""), hence overlap.
    fail::cfg(gen_snapshot_fp, "pause").unwrap();
    // Overlap snapshot should be deleted.
    assert_snapshot(&cluster.get_snap_dir(3), region_id, false);
    fail::remove(gen_snapshot_fp);
}

// When resolving remote address, all messages will be dropped and
// report unreachable. However unreachable won't reset follower's
// progress if it's in Snapshot state. So trying to send a snapshot
// when the address is being resolved will leave follower's progress
// stay in Snapshot forever.
#[test]
fn test_server_snapshot_on_resolve_failure() {
    let mut cluster = new_server_cluster(1, 4);
    configure_for_snapshot(&mut cluster);

    let on_resolve_fp = "transport_snapshot_on_resolve";
    let on_send_store_fp = "transport_on_send_store";

    let pd_client = Arc::clone(&cluster.pd_client);
    // Disable default max peer count check.
    pd_client.disable_default_operator();
    cluster.run();

    cluster.must_transfer_leader(1, new_peer(1, 1));
    pd_client.must_remove_peer(1, new_peer(4, 4));
    cluster.must_put(b"k1", b"v1");

    let ready_notify = Arc::default();
    let (notify_tx, notify_rx) = mpsc::channel();
    cluster.sim.write().unwrap().add_send_filter(
        1,
        Box::new(MessageTypeNotifier::new(
            MessageType::MsgSnapshot,
            notify_tx,
            Arc::clone(&ready_notify),
        )),
    );

    let (drop_snapshot_tx, drop_snapshot_rx) = mpsc::channel();
    cluster
        .sim
        .write()
        .unwrap()
        .add_recv_filter(4, Box::new(DropSnapshotFilter::new(drop_snapshot_tx)));

    pd_client.add_peer(1, new_peer(4, 5));

    // The leader is trying to send snapshots, but the filter drops snapshots.
    drop_snapshot_rx
        .recv_timeout(Duration::from_secs(3))
        .unwrap();

    // "return(4)" those failure occurs if TiKV resolves or sends to store 4.
    fail::cfg(on_resolve_fp, "return(4)").unwrap();
    fail::cfg(on_send_store_fp, "return(4)").unwrap();

    // We are ready to recv notify.
    ready_notify.store(true, Ordering::SeqCst);
    notify_rx.recv_timeout(Duration::from_secs(3)).unwrap();

    let engine4 = cluster.get_engine(4);
    must_get_none(&engine4, b"k1");
    cluster.sim.write().unwrap().clear_recv_filters(4);

    // Remove the on_send_store_fp.
    // Now it will resolve the store 4's address via heartbeat messages,
    // so snapshots works fine.
    //
    // But keep the on_resolve_fp.
    // Any snapshot messages that has been sent before will meet the
    // injected resolve failure eventually.
    // It perverts a race condition, remove the on_resolve_fp before snapshot
    // messages meet the failpoint, that fails the test.
    fail::remove(on_send_store_fp);

    notify_rx.recv_timeout(Duration::from_secs(3)).unwrap();
    cluster.must_put(b"k2", b"v2");
    must_get_equal(&engine4, b"k1", b"v1");
    must_get_equal(&engine4, b"k2", b"v2");

    // Clean up.
    fail::remove(on_resolve_fp);
}

#[test]
fn test_generate_snapshot() {
    let mut cluster = new_server_cluster(1, 5);
    configure_for_snapshot(&mut cluster);
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    cluster.run();
    cluster.must_transfer_leader(1, new_peer(1, 1));
    cluster.stop_node(4);
    cluster.stop_node(5);
    (0..10).for_each(|_| cluster.must_put(b"k2", b"v2"));
    // Sleep for a while to ensure all logs are compacted.
    thread::sleep(Duration::from_millis(100));

    fail::cfg("snapshot_delete_after_send", "pause").unwrap();

    // Let store 4 inform leader to generate a snapshot.
    cluster.run_node(4).unwrap();
    must_get_equal(&cluster.get_engine(4), b"k2", b"v2");

    fail::cfg("snapshot_enter_do_build", "pause").unwrap();
    cluster.run_node(5).unwrap();
    thread::sleep(Duration::from_millis(100));

    fail::cfg("snapshot_delete_after_send", "off").unwrap();
    must_empty_dir(cluster.get_snap_dir(1));

    // The task is droped so that we can't get the snapshot on store 5.
    fail::cfg("snapshot_enter_do_build", "pause").unwrap();
    must_get_none(&cluster.get_engine(5), b"k2");

    fail::cfg("snapshot_enter_do_build", "off").unwrap();
    must_get_equal(&cluster.get_engine(5), b"k2", b"v2");

    fail::remove("snapshot_enter_do_build");
    fail::remove("snapshot_delete_after_send");
}

fn must_empty_dir(path: String) {
    for _ in 0..500 {
        thread::sleep(Duration::from_millis(10));
        if fs::read_dir(&path).unwrap().count() == 0 {
            return;
        }
    }

    let entries = fs::read_dir(&path)
        .and_then(|dir| dir.collect::<io::Result<Vec<_>>>())
        .unwrap();
    if !entries.is_empty() {
        panic!(
            "the directory {:?} should be empty, but has entries: {:?}",
            path, entries
        );
    }
}

fn assert_snapshot(snap_dir: &str, region_id: u64, exist: bool) {
    let region_id = format!("{}", region_id);
    let timer = Instant::now();
    loop {
        for p in fs::read_dir(&snap_dir).unwrap() {
            let name = p.unwrap().file_name().into_string().unwrap();
            let mut parts = name.split('_');
            parts.next();
            if parts.next().unwrap() == region_id && exist {
                return;
            }
        }
        if !exist {
            return;
        }

        if timer.elapsed() < Duration::from_secs(6) {
            thread::sleep(Duration::from_millis(20));
        } else {
            panic!(
                "assert snapshot [exist: {}, region: {}] fail",
                exist, region_id
            );
        }
    }
}

#[test]
fn test_node_request_snapshot_on_split() {
    let mut cluster = new_node_cluster(0, 3);
    configure_for_request_snapshot(&mut cluster);
    cluster.run();

    let region = cluster.get_region(b"");
    // Make sure peer 2 does not in the pending state.
    cluster.must_transfer_leader(1, new_peer(2, 2));
    for _ in 0..100 {
        cluster.must_put(&[7; 100], &[7; 100]);
    }
    cluster.must_transfer_leader(1, new_peer(3, 3));

    let split_fp = "apply_before_split_1_3";
    fail::cfg(split_fp, "pause").unwrap();
    let (split_tx, split_rx) = mpsc::channel();
    cluster.split_region(
        &region,
        b"k1",
        Callback::Write(Box::new(move |_| {
            split_tx.send(()).unwrap();
        })),
    );
    // Split is stopped on peer3.
    split_rx
        .recv_timeout(Duration::from_millis(100))
        .unwrap_err();

    // Request snapshot.
    let committed_index = cluster.must_request_snapshot(2, region.get_id());

    // Install snapshot filter after requesting snapshot.
    let (tx, rx) = mpsc::channel();
    let notifier = Mutex::new(Some(tx));
    cluster.sim.wl().add_recv_filter(
        2,
        Box::new(RecvSnapshotFilter {
            notifier,
            region_id: region.get_id(),
        }),
    );
    // There is no snapshot as long as we pause the split.
    rx.recv_timeout(Duration::from_millis(500)).unwrap_err();

    // Continue split.
    fail::remove(split_fp);
    split_rx.recv().unwrap();
    let mut m = rx.recv().unwrap();
    let snapshot = m.take_message().take_snapshot();

    // Requested snapshot_index >= committed_index.
    assert!(
        snapshot.get_metadata().get_index() >= committed_index,
        "{:?} | {}",
        m,
        committed_index
    );
}

// A peer on store 3 is isolated and is applying snapshot. (add failpoint so it's always pending)
// Then two conf change happens, this peer is removed and a new peer is added on store 3.
// Then isolation clear, this peer will be destroyed because of a bigger peer id in msg.
// Peerfsm can be destroyed synchronously because snapshot state is pending and can be canceled.
// I.e. async_remove is false.
#[test]
fn test_destroy_peer_on_pending_snapshot() {
    let mut cluster = new_server_cluster(0, 4);
    configure_for_snapshot(&mut cluster);
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    let r1 = cluster.run_conf_change();
    pd_client.must_add_peer(r1, new_peer(2, 2));
    pd_client.must_add_peer(r1, new_peer(3, 3));

    cluster.must_put(b"k1", b"v1");
    // Ensure peer 3 is initialized.
    must_get_equal(&cluster.get_engine(3), b"k1", b"v1");

    cluster.must_transfer_leader(1, new_peer(1, 1));

    cluster.add_send_filter(IsolationFilterFactory::new(3));

    for i in 0..20 {
        cluster.must_put(format!("k1{}", i).as_bytes(), b"v1");
    }

    let apply_snapshot_fp = "apply_pending_snapshot";
    fail::cfg(apply_snapshot_fp, "return()").unwrap();

    cluster.clear_send_filters();
    // Wait for leader send snapshot.
    sleep_ms(100);

    cluster.add_send_filter(IsolationFilterFactory::new(3));

    pd_client.must_remove_peer(r1, new_peer(3, 3));
    pd_client.must_add_peer(r1, new_peer(4, 4));

    pd_client.must_remove_peer(r1, new_peer(4, 4));
    pd_client.must_add_peer(r1, new_peer(3, 5));

    let destroy_peer_fp = "destroy_peer";
    fail::cfg(destroy_peer_fp, "pause").unwrap();
    cluster.clear_send_filters();
    // Wait for leader send msg to peer 3.
    // Then destroy peer 3 and create peer 5.
    sleep_ms(100);
    fail::remove(destroy_peer_fp);

    fail::remove(apply_snapshot_fp);
    // After peer 5 has applied snapshot, data should be got.
    must_get_equal(&cluster.get_engine(3), b"k119", b"v1");
}
<commit_msg>tests: make snapshot generate stable (#7304)<commit_after>// Copyright 2017 TiKV Project Authors. Licensed under Apache-2.0.

use std::fs;
use std::io;
use std::sync::atomic::Ordering;
use std::sync::mpsc;
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::*;

use fail;
use raft::eraftpb::MessageType;

use raftstore::store::*;
use test_raftstore::*;
use tikv_util::config::*;
use tikv_util::HandyRwLock;

#[test]
fn test_overlap_cleanup() {
    let mut cluster = new_node_cluster(0, 3);
    // Disable raft log gc in this test case.
    cluster.cfg.raft_store.raft_log_gc_tick_interval = ReadableDuration::secs(60);

    let gen_snapshot_fp = "region_gen_snap";

    let pd_client = Arc::clone(&cluster.pd_client);
    // Disable default max peer count check.
    pd_client.disable_default_operator();

    let region_id = cluster.run_conf_change();
    pd_client.must_add_peer(region_id, new_peer(2, 2));

    cluster.must_put(b"k1", b"v1");
    must_get_equal(&cluster.get_engine(2), b"k1", b"v1");

    cluster.must_transfer_leader(region_id, new_peer(2, 2));
    // This will only pause the bootstrapped region, so the split region
    // can still work as expected.
    fail::cfg(gen_snapshot_fp, "pause").unwrap();
    pd_client.must_add_peer(region_id, new_peer(3, 3));
    cluster.must_put(b"k3", b"v3");
    assert_snapshot(&cluster.get_snap_dir(2), region_id, true);
    let region1 = cluster.get_region(b"k1");
    cluster.must_split(&region1, b"k2");
    // Wait till the snapshot of split region is applied, whose range is ["", "k2").
    must_get_equal(&cluster.get_engine(3), b"k1", b"v1");
    // Resume the fail point and pause it again. So only the paused snapshot is generated.
    // And the paused snapshot's range is ["", ""), hence overlap.
    fail::cfg(gen_snapshot_fp, "pause").unwrap();
    // Overlap snapshot should be deleted.
    assert_snapshot(&cluster.get_snap_dir(3), region_id, false);
    fail::remove(gen_snapshot_fp);
}

// When resolving remote address, all messages will be dropped and
// report unreachable. However unreachable won't reset follower's
// progress if it's in Snapshot state. So trying to send a snapshot
// when the address is being resolved will leave follower's progress
// stay in Snapshot forever.
#[test]
fn test_server_snapshot_on_resolve_failure() {
    let mut cluster = new_server_cluster(1, 4);
    configure_for_snapshot(&mut cluster);

    let on_resolve_fp = "transport_snapshot_on_resolve";
    let on_send_store_fp = "transport_on_send_store";

    let pd_client = Arc::clone(&cluster.pd_client);
    // Disable default max peer count check.
    pd_client.disable_default_operator();
    cluster.run();

    cluster.must_transfer_leader(1, new_peer(1, 1));
    pd_client.must_remove_peer(1, new_peer(4, 4));
    cluster.must_put(b"k1", b"v1");

    let ready_notify = Arc::default();
    let (notify_tx, notify_rx) = mpsc::channel();
    cluster.sim.write().unwrap().add_send_filter(
        1,
        Box::new(MessageTypeNotifier::new(
            MessageType::MsgSnapshot,
            notify_tx,
            Arc::clone(&ready_notify),
        )),
    );

    let (drop_snapshot_tx, drop_snapshot_rx) = mpsc::channel();
    cluster
        .sim
        .write()
        .unwrap()
        .add_recv_filter(4, Box::new(DropSnapshotFilter::new(drop_snapshot_tx)));

    pd_client.add_peer(1, new_peer(4, 5));

    // The leader is trying to send snapshots, but the filter drops snapshots.
    drop_snapshot_rx
        .recv_timeout(Duration::from_secs(3))
        .unwrap();

    // "return(4)" those failure occurs if TiKV resolves or sends to store 4.
    fail::cfg(on_resolve_fp, "return(4)").unwrap();
    fail::cfg(on_send_store_fp, "return(4)").unwrap();

    // We are ready to recv notify.
    ready_notify.store(true, Ordering::SeqCst);
    notify_rx.recv_timeout(Duration::from_secs(3)).unwrap();

    let engine4 = cluster.get_engine(4);
    must_get_none(&engine4, b"k1");
    cluster.sim.write().unwrap().clear_recv_filters(4);

    // Remove the on_send_store_fp.
    // Now it will resolve the store 4's address via heartbeat messages,
    // so snapshots works fine.
    //
    // But keep the on_resolve_fp.
    // Any snapshot messages that has been sent before will meet the
    // injected resolve failure eventually.
    // It perverts a race condition, remove the on_resolve_fp before snapshot
    // messages meet the failpoint, that fails the test.
    fail::remove(on_send_store_fp);

    notify_rx.recv_timeout(Duration::from_secs(3)).unwrap();
    cluster.must_put(b"k2", b"v2");
    must_get_equal(&engine4, b"k1", b"v1");
    must_get_equal(&engine4, b"k2", b"v2");

    // Clean up.
    fail::remove(on_resolve_fp);
}

#[test]
fn test_generate_snapshot() {
    let mut cluster = new_server_cluster(1, 5);
    cluster.cfg.raft_store.raft_log_gc_tick_interval = ReadableDuration::millis(20);
    cluster.cfg.raft_store.raft_log_gc_count_limit = 8;
    cluster.cfg.raft_store.merge_max_log_gap = 3;
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    cluster.run();
    cluster.must_transfer_leader(1, new_peer(1, 1));
    cluster.stop_node(4);
    cluster.stop_node(5);
    (0..10).for_each(|_| cluster.must_put(b"k2", b"v2"));
    // Sleep for a while to ensure all logs are compacted.
    thread::sleep(Duration::from_millis(100));

    fail::cfg("snapshot_delete_after_send", "pause").unwrap();

    // Let store 4 inform leader to generate a snapshot.
    cluster.run_node(4).unwrap();
    must_get_equal(&cluster.get_engine(4), b"k2", b"v2");

    fail::cfg("snapshot_enter_do_build", "pause").unwrap();
    cluster.run_node(5).unwrap();
    thread::sleep(Duration::from_millis(100));

    fail::cfg("snapshot_delete_after_send", "off").unwrap();
    must_empty_dir(cluster.get_snap_dir(1));

    // The task is droped so that we can't get the snapshot on store 5.
    fail::cfg("snapshot_enter_do_build", "pause").unwrap();
    must_get_none(&cluster.get_engine(5), b"k2");

    fail::cfg("snapshot_enter_do_build", "off").unwrap();
    must_get_equal(&cluster.get_engine(5), b"k2", b"v2");

    fail::remove("snapshot_enter_do_build");
    fail::remove("snapshot_delete_after_send");
}

fn must_empty_dir(path: String) {
    for _ in 0..500 {
        thread::sleep(Duration::from_millis(10));
        if fs::read_dir(&path).unwrap().count() == 0 {
            return;
        }
    }

    let entries = fs::read_dir(&path)
        .and_then(|dir| dir.collect::<io::Result<Vec<_>>>())
        .unwrap();
    if !entries.is_empty() {
        panic!(
            "the directory {:?} should be empty, but has entries: {:?}",
            path, entries
        );
    }
}

fn assert_snapshot(snap_dir: &str, region_id: u64, exist: bool) {
    let region_id = format!("{}", region_id);
    let timer = Instant::now();
    loop {
        for p in fs::read_dir(&snap_dir).unwrap() {
            let name = p.unwrap().file_name().into_string().unwrap();
            let mut parts = name.split('_');
            parts.next();
            if parts.next().unwrap() == region_id && exist {
                return;
            }
        }
        if !exist {
            return;
        }

        if timer.elapsed() < Duration::from_secs(6) {
            thread::sleep(Duration::from_millis(20));
        } else {
            panic!(
                "assert snapshot [exist: {}, region: {}] fail",
                exist, region_id
            );
        }
    }
}

#[test]
fn test_node_request_snapshot_on_split() {
    let mut cluster = new_node_cluster(0, 3);
    configure_for_request_snapshot(&mut cluster);
    cluster.run();

    let region = cluster.get_region(b"");
    // Make sure peer 2 does not in the pending state.
    cluster.must_transfer_leader(1, new_peer(2, 2));
    for _ in 0..100 {
        cluster.must_put(&[7; 100], &[7; 100]);
    }
    cluster.must_transfer_leader(1, new_peer(3, 3));

    let split_fp = "apply_before_split_1_3";
    fail::cfg(split_fp, "pause").unwrap();
    let (split_tx, split_rx) = mpsc::channel();
    cluster.split_region(
        &region,
        b"k1",
        Callback::Write(Box::new(move |_| {
            split_tx.send(()).unwrap();
        })),
    );
    // Split is stopped on peer3.
    split_rx
        .recv_timeout(Duration::from_millis(100))
        .unwrap_err();

    // Request snapshot.
    let committed_index = cluster.must_request_snapshot(2, region.get_id());

    // Install snapshot filter after requesting snapshot.
    let (tx, rx) = mpsc::channel();
    let notifier = Mutex::new(Some(tx));
    cluster.sim.wl().add_recv_filter(
        2,
        Box::new(RecvSnapshotFilter {
            notifier,
            region_id: region.get_id(),
        }),
    );
    // There is no snapshot as long as we pause the split.
    rx.recv_timeout(Duration::from_millis(500)).unwrap_err();

    // Continue split.
    fail::remove(split_fp);
    split_rx.recv().unwrap();
    let mut m = rx.recv().unwrap();
    let snapshot = m.take_message().take_snapshot();

    // Requested snapshot_index >= committed_index.
    assert!(
        snapshot.get_metadata().get_index() >= committed_index,
        "{:?} | {}",
        m,
        committed_index
    );
}

// A peer on store 3 is isolated and is applying snapshot. (add failpoint so it's always pending)
// Then two conf change happens, this peer is removed and a new peer is added on store 3.
// Then isolation clear, this peer will be destroyed because of a bigger peer id in msg.
// Peerfsm can be destroyed synchronously because snapshot state is pending and can be canceled.
// I.e. async_remove is false.
#[test]
fn test_destroy_peer_on_pending_snapshot() {
    let mut cluster = new_server_cluster(0, 4);
    configure_for_snapshot(&mut cluster);
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    let r1 = cluster.run_conf_change();
    pd_client.must_add_peer(r1, new_peer(2, 2));
    pd_client.must_add_peer(r1, new_peer(3, 3));

    cluster.must_put(b"k1", b"v1");
    // Ensure peer 3 is initialized.
    must_get_equal(&cluster.get_engine(3), b"k1", b"v1");

    cluster.must_transfer_leader(1, new_peer(1, 1));

    cluster.add_send_filter(IsolationFilterFactory::new(3));

    for i in 0..20 {
        cluster.must_put(format!("k1{}", i).as_bytes(), b"v1");
    }

    let apply_snapshot_fp = "apply_pending_snapshot";
    fail::cfg(apply_snapshot_fp, "return()").unwrap();

    cluster.clear_send_filters();
    // Wait for leader send snapshot.
    sleep_ms(100);

    cluster.add_send_filter(IsolationFilterFactory::new(3));

    pd_client.must_remove_peer(r1, new_peer(3, 3));
    pd_client.must_add_peer(r1, new_peer(4, 4));

    pd_client.must_remove_peer(r1, new_peer(4, 4));
    pd_client.must_add_peer(r1, new_peer(3, 5));

    let destroy_peer_fp = "destroy_peer";
    fail::cfg(destroy_peer_fp, "pause").unwrap();
    cluster.clear_send_filters();
    // Wait for leader send msg to peer 3.
    // Then destroy peer 3 and create peer 5.
    sleep_ms(100);
    fail::remove(destroy_peer_fp);

    fail::remove(apply_snapshot_fp);
    // After peer 5 has applied snapshot, data should be got.
    must_get_equal(&cluster.get_engine(3), b"k119", b"v1");
}
<|endoftext|>


<commit_before>use std::num::NonZeroU32;
use std::time::{Duration, Instant};
use {algorithms::Algorithm, NegativeMultiDecision};

/// A representation of a bare in-memory algorithm, without any bucket
/// attached.
#[derive(Debug)]
pub struct AlgorithmForTest(A);

impl<'a, A> AlgorithmForTest
where
    A: Algorithm,
{
    pub fn new<U: Into<Option<NonZeroU32>>, D: Into<Option<Duration>>>(
        cap: NonZeroU32,
        weight: U,
        duration: D,
    ) -> Self {
        AlgorithmForTest(
            A::construct(
                cap,
                weight.into().unwrap_or(nonzero!(1u32)),
                duration
                    .into()
                    .unwrap_or(::std::time::Duration::from_secs(1)),
            ).unwrap(),
        )
    }

    pub fn algorithm(&'a self) -> &'a A {
        &self.0
    }

    pub fn state(&self) -> A::BucketState {
        A::BucketState::default()
    }

    pub fn check(&self, state: &A::BucketState, t0: Instant) -> Result<(), A::NegativeDecision> {
        self.0.test_and_update(state, t0)
    }

    pub fn check_n(
        &self,
        state: &A::BucketState,
        n: u32,
        t0: Instant,
    ) -> Result<(), NegativeMultiDecision> {
        self.0.test_n_and_update(state, n, t0)
    }
}

impl Default for AlgorithmForTest
where
    A: Algorithm,
{
    fn default() -> Self {
        Self::new(nonzero!(1u32), None, None)
    }
}

#[macro_export]
macro_rules! bench_with_algorithm_variants {
    ($variant:expr, $var:ident, $code:block) => {
        match $variant {
            $crate::test_utilities::variants::Variant::GCRA => {
                let mut $var = $crate::test_utilities::algorithms::AlgorithmForTest::<
                    $crate::GCRA,
                >::default();
                $code
            }
            $crate::test_utilities::variants::Variant::LeakyBucket => {
                let mut $var = $crate::test_utilities::algorithms::AlgorithmForTest::<
                    $crate::LeakyBucket,
                >::default();
                $code
            }
        }
    };
}
<commit_msg>Rustfmt to make the build happy.<commit_after>use std::num::NonZeroU32;
use std::time::{Duration, Instant};
use {algorithms::Algorithm, NegativeMultiDecision};

/// A representation of a bare in-memory algorithm, without any bucket
/// attached.
#[derive(Debug)]
pub struct AlgorithmForTest(A);

impl<'a, A> AlgorithmForTest
where
    A: Algorithm,
{
    pub fn new<U: Into<Option<NonZeroU32>>, D: Into<Option<Duration>>>(
        cap: NonZeroU32,
        weight: U,
        duration: D,
    ) -> Self {
        AlgorithmForTest(
            A::construct(
                cap,
                weight.into().unwrap_or(nonzero!(1u32)),
                duration
                    .into()
                    .unwrap_or(::std::time::Duration::from_secs(1)),
            ).unwrap(),
        )
    }

    pub fn algorithm(&'a self) -> &'a A {
        &self.0
    }

    pub fn state(&self) -> A::BucketState {
        A::BucketState::default()
    }

    pub fn check(&self, state: &A::BucketState, t0: Instant) -> Result<(), A::NegativeDecision> {
        self.0.test_and_update(state, t0)
    }

    pub fn check_n(
        &self,
        state: &A::BucketState,
        n: u32,
        t0: Instant,
    ) -> Result<(), NegativeMultiDecision> {
        self.0.test_n_and_update(state, n, t0)
    }
}

impl Default for AlgorithmForTest
where
    A: Algorithm,
{
    fn default() -> Self {
        Self::new(nonzero!(1u32), None, None)
    }
}

#[macro_export]
macro_rules! bench_with_algorithm_variants {
    ($variant:expr, $var:ident, $code:block) => {
        match $variant {
            $crate::test_utilities::variants::Variant::GCRA => {
                let mut $var =
                    $crate::test_utilities::algorithms::AlgorithmForTest::<$crate::GCRA>::default();
                $code
            }
            $crate::test_utilities::variants::Variant::LeakyBucket => {
                let mut $var = $crate::test_utilities::algorithms::AlgorithmForTest::<
                    $crate::LeakyBucket,
                >::default();
                $code
            }
        }
    };
}
<|endoftext|>


<commit_before>// FIXME: Needs additional cocoa setup on OS X. rust-cocoa should probably just
// be a dependency

#[test]
#[ignore(cfg(target_os = "macos"))]
pub fn test_everything() {

    do task::spawn {
        assert sdl::init(~[sdl::InitVideo, sdl::InitTimer]) == true;
        run_tests(~[
            general::test_was_init,
            general::test_set_error,
            general::test_error,
            general::test_clear_error,
            video::test_set_video_mode,
            // FIXME: Doesn't work when called from a directory that
            // doesn't contain the test image file
            video::test_blit, //Comment out before merge
            test_event::test_poll_event_none,
            // FIXME: This test is interactive
            test_event::test_keyboard, //Comment out before merge

        ]);
        sdl::quit();
    }
}

fn run_tests(tests: &[fn()]) {
    for tests.each |test| {
        (*test)();
    }
}

mod general {
    pub fn test_was_init() {
        assert vec::contains(~[sdl::InitTimer], &sdl::InitTimer);
    }

    pub fn test_set_error() {
        sdl::set_error(~"test");
        assert sdl::get_error() == ~"test";
    }

    pub fn test_error() {
        sdl::clear_error();
        assert str::is_empty(sdl::get_error());
        sdl::error(sdl::ENoMem);
        assert str::is_not_empty(sdl::get_error());
    }

    pub fn test_clear_error() {
        sdl::set_error(~"test");
        sdl::clear_error();
        assert str::is_empty(sdl::get_error());
    }
}

mod test_event {
    pub fn test_poll_event_none() {
        ::event::poll_event(|event| assert event == ::event::NoEvent);
    }

    pub fn test_keyboard() {
        io::println(~"press a key in the window");
        let maybe_surface = ::video::set_video_mode(320, 200, 32,
            ~[::video::SWSurface], ~[::video::DoubleBuf, ::video::Resizable]);

        match maybe_surface {
            result::Ok(_) => {
                let mut keydown = false;
                let mut keyup = false;
                while !keydown || !keyup {
                    ::event::poll_event(|event| {
                        match event {
                          event::KeyUpEvent(keyboard) => {
                              keyup = true;
                              assert (keyboard.keycode as u64 != 0);
                          },
                          event::KeyDownEvent(keyboard) => {
                              keydown = true;
                              //assert (keyboard.keycode != keyboard::SDLKUnknown); //FIXME: This is causing a stack error. Rust bug or code error?
                              assert (keyboard.keycode as u64 != 0);
                          },
                          event::QuitEvent => fail,
                          _ => { }
                        }
                    })
                }
            }
            result::Err(_) => {
                fail;
            }
        }
    }
}

mod video {

    pub fn test_set_video_mode() {
        let maybe_surface = ::video::set_video_mode(320, 200, 32,
            ~[::video::HWSurface], ~[::video::DoubleBuf]);

        assert result::is_ok(&maybe_surface);
    }

    pub fn test_blit() {
        let maybe_screen = ::video::set_video_mode(320, 200, 32,
            ~[::video::SWSurface], ~[::video::DoubleBuf]);

        assert result::is_ok(&maybe_screen);

        match maybe_screen {
            result::Ok(screen) => {
                // FIXME: We need to load this from the crate instead of
                // off the filesystem
                let maybe_image = match ::video::load_bmp(~"rust-logo-128x128-blk.bmp") {
                    result::Ok(raw_image) => {
                        raw_image.display_format()
                    },
                    result::Err(err) => result::Err(err)
                };

                assert result::is_ok(&maybe_image);

                match maybe_image {
                    result::Ok(image) => {
                        for iter::repeat(1u) || {
                            screen.blit_surface(image);
                            screen.flip();
                            ::event::poll_event(|_event| {})
                        };
                    },
                    result::Err(_) => ()
                };
            },
            result::Err(_) => ()
        };
    }
}
<commit_msg>More bug notes<commit_after>// FIXME: Needs additional cocoa setup on OS X. rust-cocoa should probably just
// be a dependency

#[test]
#[ignore(cfg(target_os = "macos"))]
pub fn test_everything() {

    do task::spawn {
        assert sdl::init(~[sdl::InitVideo, sdl::InitTimer]) == true;
        run_tests(~[
            general::test_was_init,
            general::test_set_error,
            general::test_error,
            general::test_clear_error,
            video::test_set_video_mode,
            // FIXME: Doesn't work when called from a directory that
            // doesn't contain the test image file
            video::test_blit, //Comment out before merge
            test_event::test_poll_event_none,
            // FIXME: This test is interactive
            test_event::test_keyboard, //Comment out before merge

        ]);
        sdl::quit();
    }
}

fn run_tests(tests: &[fn()]) {
    for tests.each |test| {
        (*test)();
    }
}

mod general {
    pub fn test_was_init() {
        assert vec::contains(~[sdl::InitTimer], &sdl::InitTimer);
    }

    pub fn test_set_error() {
        sdl::set_error(~"test");
        assert sdl::get_error() == ~"test";
    }

    pub fn test_error() {
        sdl::clear_error();
        assert str::is_empty(sdl::get_error());
        sdl::error(sdl::ENoMem);
        assert str::is_not_empty(sdl::get_error());
    }

    pub fn test_clear_error() {
        sdl::set_error(~"test");
        sdl::clear_error();
        assert str::is_empty(sdl::get_error());
    }
}

mod test_event {
    pub fn test_poll_event_none() {
        ::event::poll_event(|event| assert event == ::event::NoEvent);
    }

    pub fn test_keyboard() {
        io::println(~"press a key in the window");
        let maybe_surface = ::video::set_video_mode(320, 200, 32,
            ~[::video::SWSurface], ~[::video::DoubleBuf, ::video::Resizable]);

        match maybe_surface {
            result::Ok(_) => {
                let mut keydown = false;
                let mut keyup = false;
                while !keydown || !keyup {
                    ::event::poll_event(|event| {
                        match event {
                          event::KeyUpEvent(keyboard) => {
                              keyup = true;
                              assert (keyboard.keycode as u64 != 0);
                          },
                          event::KeyDownEvent(keyboard) => {
                              keydown = true;
                              //assert (keyboard.keycode != keyboard::SDLKUnknown); //FIXME: This is causing a stack error. Rust bug or code error? Seems fine when compared via Match
                              assert (keyboard.keycode as u64 != 0);
                          },
                          event::QuitEvent => fail,
                          _ => { }
                        }
                    })
                }
            }
            result::Err(_) => {
                fail;
            }
        }
    }
}

mod video {

    pub fn test_set_video_mode() {
        let maybe_surface = ::video::set_video_mode(320, 200, 32,
            ~[::video::HWSurface], ~[::video::DoubleBuf]);

        assert result::is_ok(&maybe_surface);
    }

    pub fn test_blit() {
        let maybe_screen = ::video::set_video_mode(320, 200, 32,
            ~[::video::SWSurface], ~[::video::DoubleBuf]);

        assert result::is_ok(&maybe_screen);

        match maybe_screen {
            result::Ok(screen) => {
                // FIXME: We need to load this from the crate instead of
                // off the filesystem
                let maybe_image = match ::video::load_bmp(~"rust-logo-128x128-blk.bmp") {
                    result::Ok(raw_image) => {
                        raw_image.display_format()
                    },
                    result::Err(err) => result::Err(err)
                };

                assert result::is_ok(&maybe_image);

                match maybe_image {
                    result::Ok(image) => {
                        for iter::repeat(1u) || {
                            screen.blit_surface(image);
                            screen.flip();
                            ::event::poll_event(|_event| {})
                        };
                    },
                    result::Err(_) => ()
                };
            },
            result::Err(_) => ()
        };
    }
}
<|endoftext|>