<commit_before>use fallible_iterator::FallibleIterator;
use futures::sink;
use futures::sync::mpsc;
use futures::{try_ready, Future, Poll, Sink, Stream};
use postgres_protocol::authentication;
use postgres_protocol::authentication::sasl::{self, ScramSha256};
use postgres_protocol::message::backend::Message;
use postgres_protocol::message::frontend;
use state_machine_future::{transition, RentToOwn, StateMachineFuture};
use std::collections::HashMap;
use std::io;
use tokio_codec::Framed;
use tokio_io::{AsyncRead, AsyncWrite};

use crate::proto::{Client, Connection, PostgresCodec, TlsFuture};
use crate::{CancelData, ChannelBinding, Config, Error, TlsMode};

#[derive(StateMachineFuture)]
pub enum Handshake<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsMode<S>,
{
    #[state_machine_future(start, transitions(SendingStartup))]
    Start {
        future: TlsFuture<S, T>,
        config: Config,
    },
    #[state_machine_future(transitions(ReadingAuth))]
    SendingStartup {
        future: sink::Send<Framed<T::Stream, PostgresCodec>>,
        config: Config,
        channel_binding: ChannelBinding,
    },
    #[state_machine_future(transitions(ReadingInfo, SendingPassword, SendingSasl))]
    ReadingAuth {
        stream: Framed<T::Stream, PostgresCodec>,
        config: Config,
        channel_binding: ChannelBinding,
    },
    #[state_machine_future(transitions(ReadingAuthCompletion))]
    SendingPassword {
        future: sink::Send<Framed<T::Stream, PostgresCodec>>,
    },
    #[state_machine_future(transitions(ReadingSasl))]
    SendingSasl {
        future: sink::Send<Framed<T::Stream, PostgresCodec>>,
        scram: ScramSha256,
    },
    #[state_machine_future(transitions(SendingSasl, ReadingAuthCompletion))]
    ReadingSasl {
        stream: Framed<T::Stream, PostgresCodec>,
        scram: ScramSha256,
    },
    #[state_machine_future(transitions(ReadingInfo))]
    ReadingAuthCompletion {
        stream: Framed<T::Stream, PostgresCodec>,
    },
    #[state_machine_future(transitions(Finished))]
    ReadingInfo {
        stream: Framed<T::Stream, PostgresCodec>,
        cancel_data: Option<CancelData>,
        parameters: HashMap<String, String>,
    },
    #[state_machine_future(ready)]
    Finished((Client, Connection<T::Stream>)),
    #[state_machine_future(error)]
    Failed(Error),
}

impl<S, T> PollHandshake<S, T> for Handshake<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsMode<S>,
{
    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<S, T>>) -> Poll<AfterStart<S, T>, Error> {
        let (stream, channel_binding) = try_ready!(state.future.poll());
        let state = state.take();

        let mut params = vec![("client_encoding", "UTF8"), ("timezone", "GMT")];
        if let Some(user) = &state.config.0.user {
            params.push(("user", &**user));
        }
        if let Some(dbname) = &state.config.0.dbname {
            params.push(("database", &**dbname));
        }
        if let Some(options) = &state.config.0.options {
            params.push(("options", &**options));
        }
        if let Some(application_name) = &state.config.0.application_name {
            params.push(("application_name", &**application_name));
        }

        let mut buf = vec![];
        frontend::startup_message(params.iter().map(|e| (e.0, e.1)), &mut buf)
            .map_err(Error::encode)?;

        let stream = Framed::new(stream, PostgresCodec);

        transition!(SendingStartup {
            future: stream.send(buf),
            config: state.config,
            channel_binding,
        })
    }

    fn poll_sending_startup<'a>(
        state: &'a mut RentToOwn<'a, SendingStartup<S, T>>,
    ) -> Poll<AfterSendingStartup<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingAuth {
            stream,
            config: state.config,
            channel_binding: state.channel_binding,
        })
    }

    fn poll_reading_auth<'a>(
        state: &'a mut RentToOwn<'a, ReadingAuth<S, T>>,
    ) -> Poll<AfterReadingAuth<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let state = state.take();

        match message {
            Some(Message::AuthenticationOk) => transition!(ReadingInfo {
                stream: state.stream,
                cancel_data: None,
                parameters: HashMap::new(),
            }),
            Some(Message::AuthenticationCleartextPassword) => {
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;
                let mut buf = vec![];
                frontend::password_message(pass, &mut buf).map_err(Error::encode)?;
                transition!(SendingPassword {
                    future: state.stream.send(buf)
                })
            }
            Some(Message::AuthenticationMd5Password(body)) => {
                let user = state
                    .config
                    .0
                    .user
                    .as_ref()
                    .ok_or_else(|| Error::config("user missing".into()))?;
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;
                let output = authentication::md5_hash(user.as_bytes(), pass, body.salt());
                let mut buf = vec![];
                frontend::password_message(output.as_bytes(), &mut buf).map_err(Error::encode)?;
                transition!(SendingPassword {
                    future: state.stream.send(buf)
                })
            }
            Some(Message::AuthenticationSasl(body)) => {
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;

                let mut has_scram = false;
                let mut has_scram_plus = false;
                let mut mechanisms = body.mechanisms();
                while let Some(mechanism) = mechanisms.next().map_err(Error::parse)? {
                    match mechanism {
                        sasl::SCRAM_SHA_256 => has_scram = true,
                        sasl::SCRAM_SHA_256_PLUS => has_scram_plus = true,
                        _ => {}
                    }
                }

                let channel_binding = if let Some(tls_server_end_point) =
                    state.channel_binding.tls_server_end_point
                {
                    Some(sasl::ChannelBinding::tls_server_end_point(
                        tls_server_end_point,
                    ))
                } else {
                    None
                };

                let (channel_binding, mechanism) = if has_scram_plus {
                    match channel_binding {
                        Some(channel_binding) => (channel_binding, sasl::SCRAM_SHA_256_PLUS),
                        None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
                    }
                } else if has_scram {
                    match channel_binding {
                        Some(_) => (sasl::ChannelBinding::unrequested(), sasl::SCRAM_SHA_256),
                        None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
                    }
                } else {
                    return Err(Error::authentication(
                        "unsupported authentication method".into(),
                    ));
                };

                let scram = ScramSha256::new(pass, channel_binding);

                let mut buf = vec![];
                frontend::sasl_initial_response(mechanism, scram.message(), &mut buf)
                    .map_err(Error::encode)?;

                transition!(SendingSasl {
                    future: state.stream.send(buf),
                    scram,
                })
            }
            Some(Message::AuthenticationKerberosV5)
            | Some(Message::AuthenticationScmCredential)
            | Some(Message::AuthenticationGss)
            | Some(Message::AuthenticationSspi) => Err(Error::authentication(
                "unsupported authentication method".into(),
            )),
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_sending_password<'a>(
        state: &'a mut RentToOwn<'a, SendingPassword<S, T>>,
    ) -> Poll<AfterSendingPassword<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        transition!(ReadingAuthCompletion { stream })
    }

    fn poll_sending_sasl<'a>(
        state: &'a mut RentToOwn<'a, SendingSasl<S, T>>,
    ) -> Poll<AfterSendingSasl<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingSasl {
            stream,
            scram: state.scram,
        })
    }

    fn poll_reading_sasl<'a>(
        state: &'a mut RentToOwn<'a, ReadingSasl<S, T>>,
    ) -> Poll<AfterReadingSasl<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let mut state = state.take();

        match message {
            Some(Message::AuthenticationSaslContinue(body)) => {
                state
                    .scram
                    .update(body.data())
                    .map_err(|e| Error::authentication(Box::new(e)))?;
                let mut buf = vec![];
                frontend::sasl_response(state.scram.message(), &mut buf).map_err(Error::encode)?;
                transition!(SendingSasl {
                    future: state.stream.send(buf),
                    scram: state.scram,
                })
            }
            Some(Message::AuthenticationSaslFinal(body)) => {
                state
                    .scram
                    .finish(body.data())
                    .map_err(|e| Error::authentication(Box::new(e)))?;
                transition!(ReadingAuthCompletion {
                    stream: state.stream
                })
            }
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_reading_auth_completion<'a>(
        state: &'a mut RentToOwn<'a, ReadingAuthCompletion<S, T>>,
    ) -> Poll<AfterReadingAuthCompletion<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let state = state.take();

        match message {
            Some(Message::AuthenticationOk) => transition!(ReadingInfo {
                stream: state.stream,
                cancel_data: None,
                parameters: HashMap::new()
            }),
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_reading_info<'a>(
        state: &'a mut RentToOwn<'a, ReadingInfo<S, T>>,
    ) -> Poll<AfterReadingInfo<S, T>, Error> {
        loop {
            let message = try_ready!(state.stream.poll().map_err(Error::io));
            match message {
                Some(Message::BackendKeyData(body)) => {
                    state.cancel_data = Some(CancelData {
                        process_id: body.process_id(),
                        secret_key: body.secret_key(),
                    });
                }
                Some(Message::ParameterStatus(body)) => {
                    state.parameters.insert(
                        body.name().map_err(Error::parse)?.to_string(),
                        body.value().map_err(Error::parse)?.to_string(),
                    );
                }
                Some(Message::ReadyForQuery(_)) => {
                    let state = state.take();
                    let cancel_data = state.cancel_data.ok_or_else(|| {
                        Error::parse(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "BackendKeyData message missing",
                        ))
                    })?;
                    let (sender, receiver) = mpsc::unbounded();
                    let client = Client::new(sender);
                    let connection =
                        Connection::new(state.stream, cancel_data, state.parameters, receiver);
                    transition!(Finished((client, connection)))
                }
                Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
                Some(Message::NoticeResponse(_)) => {}
                Some(_) => return Err(Error::unexpected_message()),
                None => return Err(Error::closed()),
            }
        }
    }
}

impl<S, T> HandshakeFuture<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsMode<S>,
{
    pub fn new(stream: S, tls_mode: T, config: Config) -> HandshakeFuture<S, T> {
        Handshake::start(TlsFuture::new(stream, tls_mode), config)
    }
}
<commit_msg>cleanup<commit_after>use fallible_iterator::FallibleIterator;
use futures::sink;
use futures::sync::mpsc;
use futures::{try_ready, Future, Poll, Sink, Stream};
use postgres_protocol::authentication;
use postgres_protocol::authentication::sasl::{self, ScramSha256};
use postgres_protocol::message::backend::Message;
use postgres_protocol::message::frontend;
use state_machine_future::{transition, RentToOwn, StateMachineFuture};
use std::collections::HashMap;
use std::io;
use tokio_codec::Framed;
use tokio_io::{AsyncRead, AsyncWrite};

use crate::proto::{Client, Connection, PostgresCodec, TlsFuture};
use crate::{CancelData, ChannelBinding, Config, Error, TlsMode};

#[derive(StateMachineFuture)]
pub enum Handshake<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsMode<S>,
{
    #[state_machine_future(start, transitions(SendingStartup))]
    Start {
        future: TlsFuture<S, T>,
        config: Config,
    },
    #[state_machine_future(transitions(ReadingAuth))]
    SendingStartup {
        future: sink::Send<Framed<T::Stream, PostgresCodec>>,
        config: Config,
        channel_binding: ChannelBinding,
    },
    #[state_machine_future(transitions(ReadingInfo, SendingPassword, SendingSasl))]
    ReadingAuth {
        stream: Framed<T::Stream, PostgresCodec>,
        config: Config,
        channel_binding: ChannelBinding,
    },
    #[state_machine_future(transitions(ReadingAuthCompletion))]
    SendingPassword {
        future: sink::Send<Framed<T::Stream, PostgresCodec>>,
    },
    #[state_machine_future(transitions(ReadingSasl))]
    SendingSasl {
        future: sink::Send<Framed<T::Stream, PostgresCodec>>,
        scram: ScramSha256,
    },
    #[state_machine_future(transitions(SendingSasl, ReadingAuthCompletion))]
    ReadingSasl {
        stream: Framed<T::Stream, PostgresCodec>,
        scram: ScramSha256,
    },
    #[state_machine_future(transitions(ReadingInfo))]
    ReadingAuthCompletion {
        stream: Framed<T::Stream, PostgresCodec>,
    },
    #[state_machine_future(transitions(Finished))]
    ReadingInfo {
        stream: Framed<T::Stream, PostgresCodec>,
        cancel_data: Option<CancelData>,
        parameters: HashMap<String, String>,
    },
    #[state_machine_future(ready)]
    Finished((Client, Connection<T::Stream>)),
    #[state_machine_future(error)]
    Failed(Error),
}

impl<S, T> PollHandshake<S, T> for Handshake<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsMode<S>,
{
    fn poll_start<'a>(state: &'a mut RentToOwn<'a, Start<S, T>>) -> Poll<AfterStart<S, T>, Error> {
        let (stream, channel_binding) = try_ready!(state.future.poll());
        let state = state.take();

        let mut params = vec![("client_encoding", "UTF8"), ("timezone", "GMT")];
        if let Some(user) = &state.config.0.user {
            params.push(("user", &**user));
        }
        if let Some(dbname) = &state.config.0.dbname {
            params.push(("database", &**dbname));
        }
        if let Some(options) = &state.config.0.options {
            params.push(("options", &**options));
        }
        if let Some(application_name) = &state.config.0.application_name {
            params.push(("application_name", &**application_name));
        }

        let mut buf = vec![];
        frontend::startup_message(params, &mut buf).map_err(Error::encode)?;

        let stream = Framed::new(stream, PostgresCodec);

        transition!(SendingStartup {
            future: stream.send(buf),
            config: state.config,
            channel_binding,
        })
    }

    fn poll_sending_startup<'a>(
        state: &'a mut RentToOwn<'a, SendingStartup<S, T>>,
    ) -> Poll<AfterSendingStartup<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingAuth {
            stream,
            config: state.config,
            channel_binding: state.channel_binding,
        })
    }

    fn poll_reading_auth<'a>(
        state: &'a mut RentToOwn<'a, ReadingAuth<S, T>>,
    ) -> Poll<AfterReadingAuth<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let state = state.take();

        match message {
            Some(Message::AuthenticationOk) => transition!(ReadingInfo {
                stream: state.stream,
                cancel_data: None,
                parameters: HashMap::new(),
            }),
            Some(Message::AuthenticationCleartextPassword) => {
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;
                let mut buf = vec![];
                frontend::password_message(pass, &mut buf).map_err(Error::encode)?;
                transition!(SendingPassword {
                    future: state.stream.send(buf)
                })
            }
            Some(Message::AuthenticationMd5Password(body)) => {
                let user = state
                    .config
                    .0
                    .user
                    .as_ref()
                    .ok_or_else(|| Error::config("user missing".into()))?;
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;
                let output = authentication::md5_hash(user.as_bytes(), pass, body.salt());
                let mut buf = vec![];
                frontend::password_message(output.as_bytes(), &mut buf).map_err(Error::encode)?;
                transition!(SendingPassword {
                    future: state.stream.send(buf)
                })
            }
            Some(Message::AuthenticationSasl(body)) => {
                let pass = state
                    .config
                    .0
                    .password
                    .as_ref()
                    .ok_or_else(|| Error::config("password missing".into()))?;

                let mut has_scram = false;
                let mut has_scram_plus = false;
                let mut mechanisms = body.mechanisms();
                while let Some(mechanism) = mechanisms.next().map_err(Error::parse)? {
                    match mechanism {
                        sasl::SCRAM_SHA_256 => has_scram = true,
                        sasl::SCRAM_SHA_256_PLUS => has_scram_plus = true,
                        _ => {}
                    }
                }

                let channel_binding = if let Some(tls_server_end_point) =
                    state.channel_binding.tls_server_end_point
                {
                    Some(sasl::ChannelBinding::tls_server_end_point(
                        tls_server_end_point,
                    ))
                } else {
                    None
                };

                let (channel_binding, mechanism) = if has_scram_plus {
                    match channel_binding {
                        Some(channel_binding) => (channel_binding, sasl::SCRAM_SHA_256_PLUS),
                        None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
                    }
                } else if has_scram {
                    match channel_binding {
                        Some(_) => (sasl::ChannelBinding::unrequested(), sasl::SCRAM_SHA_256),
                        None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
                    }
                } else {
                    return Err(Error::authentication(
                        "unsupported authentication method".into(),
                    ));
                };

                let scram = ScramSha256::new(pass, channel_binding);

                let mut buf = vec![];
                frontend::sasl_initial_response(mechanism, scram.message(), &mut buf)
                    .map_err(Error::encode)?;

                transition!(SendingSasl {
                    future: state.stream.send(buf),
                    scram,
                })
            }
            Some(Message::AuthenticationKerberosV5)
            | Some(Message::AuthenticationScmCredential)
            | Some(Message::AuthenticationGss)
            | Some(Message::AuthenticationSspi) => Err(Error::authentication(
                "unsupported authentication method".into(),
            )),
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_sending_password<'a>(
        state: &'a mut RentToOwn<'a, SendingPassword<S, T>>,
    ) -> Poll<AfterSendingPassword<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        transition!(ReadingAuthCompletion { stream })
    }

    fn poll_sending_sasl<'a>(
        state: &'a mut RentToOwn<'a, SendingSasl<S, T>>,
    ) -> Poll<AfterSendingSasl<S, T>, Error> {
        let stream = try_ready!(state.future.poll().map_err(Error::io));
        let state = state.take();
        transition!(ReadingSasl {
            stream,
            scram: state.scram,
        })
    }

    fn poll_reading_sasl<'a>(
        state: &'a mut RentToOwn<'a, ReadingSasl<S, T>>,
    ) -> Poll<AfterReadingSasl<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let mut state = state.take();

        match message {
            Some(Message::AuthenticationSaslContinue(body)) => {
                state
                    .scram
                    .update(body.data())
                    .map_err(|e| Error::authentication(Box::new(e)))?;
                let mut buf = vec![];
                frontend::sasl_response(state.scram.message(), &mut buf).map_err(Error::encode)?;
                transition!(SendingSasl {
                    future: state.stream.send(buf),
                    scram: state.scram,
                })
            }
            Some(Message::AuthenticationSaslFinal(body)) => {
                state
                    .scram
                    .finish(body.data())
                    .map_err(|e| Error::authentication(Box::new(e)))?;
                transition!(ReadingAuthCompletion {
                    stream: state.stream
                })
            }
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_reading_auth_completion<'a>(
        state: &'a mut RentToOwn<'a, ReadingAuthCompletion<S, T>>,
    ) -> Poll<AfterReadingAuthCompletion<S, T>, Error> {
        let message = try_ready!(state.stream.poll().map_err(Error::io));
        let state = state.take();

        match message {
            Some(Message::AuthenticationOk) => transition!(ReadingInfo {
                stream: state.stream,
                cancel_data: None,
                parameters: HashMap::new()
            }),
            Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
            Some(_) => Err(Error::unexpected_message()),
            None => Err(Error::closed()),
        }
    }

    fn poll_reading_info<'a>(
        state: &'a mut RentToOwn<'a, ReadingInfo<S, T>>,
    ) -> Poll<AfterReadingInfo<S, T>, Error> {
        loop {
            let message = try_ready!(state.stream.poll().map_err(Error::io));
            match message {
                Some(Message::BackendKeyData(body)) => {
                    state.cancel_data = Some(CancelData {
                        process_id: body.process_id(),
                        secret_key: body.secret_key(),
                    });
                }
                Some(Message::ParameterStatus(body)) => {
                    state.parameters.insert(
                        body.name().map_err(Error::parse)?.to_string(),
                        body.value().map_err(Error::parse)?.to_string(),
                    );
                }
                Some(Message::ReadyForQuery(_)) => {
                    let state = state.take();
                    let cancel_data = state.cancel_data.ok_or_else(|| {
                        Error::parse(io::Error::new(
                            io::ErrorKind::InvalidData,
                            "BackendKeyData message missing",
                        ))
                    })?;
                    let (sender, receiver) = mpsc::unbounded();
                    let client = Client::new(sender);
                    let connection =
                        Connection::new(state.stream, cancel_data, state.parameters, receiver);
                    transition!(Finished((client, connection)))
                }
                Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
                Some(Message::NoticeResponse(_)) => {}
                Some(_) => return Err(Error::unexpected_message()),
                None => return Err(Error::closed()),
            }
        }
    }
}

impl<S, T> HandshakeFuture<S, T>
where
    S: AsyncRead + AsyncWrite,
    T: TlsMode<S>,
{
    pub fn new(stream: S, tls_mode: T, config: Config) -> HandshakeFuture<S, T> {
        Handshake::start(TlsFuture::new(stream, tls_mode), config)
    }
}
<|endoftext|>


<commit_before>
use super::core::*;
use super::queue_state::*;
use super::job_queue::*;

use std::sync::*;
use futures::task::{ArcWake};

///
/// Waker that will wake the specified queue in the specified scheduler core
///
pub (super) struct WakeQueue(pub (super) Arc<JobQueue>, pub (super) Arc<SchedulerCore>);

impl ArcWake for WakeQueue {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        // Decompose this structure
        let WakeQueue(ref queue, ref core) = **arc_self;

        // Move the queue to the idle state if we can
        {
            let mut queue_core = queue.core.lock().unwrap();

            // Queue can be woken if it's in the WaitingForWake state
            match queue_core.state {
                QueueState::WaitingForUnpark    => { panic!("WakeQueue cannot unpark a parked queue") },

                QueueState::WaitingForWake      => queue_core.state = QueueState::Idle,
                QueueState::Running             => queue_core.state = QueueState::AwokenWhileRunning,
                other_state                     => queue_core.state = other_state
            }
        }

        // Cause the core to reschedule its events
        core.reschedule_queue(queue, Arc::clone(core));
    }
}
<commit_msg>The double-waker makes for a situation where it's possible to reach a 'WakeQueue' when the future is running synchronously and is in the 'WakeThread' state<commit_after>
use super::core::*;
use super::queue_state::*;
use super::job_queue::*;

use std::sync::*;
use futures::task::{ArcWake};

///
/// Waker that will wake the specified queue in the specified scheduler core
///
pub (super) struct WakeQueue(pub (super) Arc<JobQueue>, pub (super) Arc<SchedulerCore>);

impl ArcWake for WakeQueue {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        // Decompose this structure
        let WakeQueue(ref queue, ref core) = **arc_self;

        // Move the queue to the idle state if we can
        {
            let mut queue_core = queue.core.lock().unwrap();

            // Queue can be woken if it's in the WaitingForWake state
            match queue_core.state {
                QueueState::WaitingForUnpark    => {
                    // Assume that this was part of a DoubleWake that has become stale (woke up, reached a park, created a WakeThread for notifications)
                    queue_core.state = QueueState::WaitingForUnpark;
                    return;
                },

                QueueState::WaitingForWake      => queue_core.state = QueueState::Idle,
                QueueState::Running             => queue_core.state = QueueState::AwokenWhileRunning,
                other_state                     => queue_core.state = other_state
            }
        }

        // Cause the core to reschedule its events
        core.reschedule_queue(queue, Arc::clone(core));
    }
}
<|endoftext|>


<commit_before>use std::collections::HashSet;
use std::env::args;

#[derive(PartialEq, Eq, Hash)]
struct Point {
    x: i32,
    y: i32,
}

impl Point {
    fn update(&self, instr: char) -> Point {
        match instr {
            '^' => Point{x: self.x, y: self.y + 1},
            'v' => Point{x: self.x, y: self.y - 1},
            '<' => Point{x: self.x - 1, y: self.y},
            '>' => Point{x: self.x + 1, y: self.y},
            _   => unreachable!()
        }
    }
}

fn main() {
    let instructions = match args().nth(1) {
        Some(v) => v,
        None    => "".to_string(),
    };
    let starting_point = Point{x: 0, y: 0};
    let mut points = vec![starting_point];
    for tuple in instructions.chars().zip(0..(instructions.len())) {
        let (instr, i) = match tuple { (a@_, b@_) => (a, b) };
        let next_point = points[i].update(instr);
        points.push(next_point);
    }

    let mut point_set = HashSet::new();
    for point in points {
        point_set.insert(point);
    }
    println!("{}", point_set.len());
}
<commit_msg>day 3, part 2<commit_after>use std::collections::HashSet;
use std::env::args;

#[derive(PartialEq, Eq, Hash)]
struct Point {
    x: i32,
    y: i32,
}

impl Point {
    fn update(&self, instr: char) -> Point {
        match instr {
            '^' => Point{x: self.x, y: self.y + 1},
            'v' => Point{x: self.x, y: self.y - 1},
            '<' => Point{x: self.x - 1, y: self.y},
            '>' => Point{x: self.x + 1, y: self.y},
            _   => unreachable!()
        }
    }
}

fn main() {
    let instructions = match args().nth(1) {
        Some(v) => v,
        None    => "".to_string(),
    };
    let starting_point = Point{x: 0, y: 0};
    let mut points = vec![starting_point, Point{x: 0, y:0}];
    for tuple in instructions.chars().zip(1..(instructions.len() + 1)) {
        let (instr, i) = match tuple { (a@_, b@_) => (a, b) };
        let next_point = points[i - 1].update(instr);
        points.push(next_point);
    }

    let mut point_set = HashSet::new();
    for point in points {
        point_set.insert(point);
    }
    println!("{}", point_set.len());
}
<|endoftext|>


<commit_before>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::io;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::fs::File;
use std::io::ErrorKind;
use std::io::{Seek, Write, SeekFrom};
use std::fs::{OpenOptions, read_dir};
use std::os::unix::prelude::AsRawFd;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::thread;
use std::time::Duration;

use time::Timespec;
use devicemapper::Device;
use uuid::Uuid;

use types::{Bytes, Sectors};
use engine::{DevUuid, EngineResult, EngineError, ErrorEnum, PoolUuid};

use consts::*;
use super::metadata::{StaticHeader, BDA, validate_mda_size};
use super::engine::DevOwnership;
pub use super::BlockDevSave;

const MIN_DEV_SIZE: Bytes = Bytes(IEC::Gi as u64);

ioctl!(read blkgetsize64 with 0x12, 114; u64);

pub fn blkdev_size(file: &File) -> EngineResult<Bytes> {
    let mut val: u64 = 0;

    match unsafe { blkgetsize64(file.as_raw_fd(), &mut val) } {
        Err(x) => Err(EngineError::Nix(x)),
        Ok(_) => Ok(Bytes(val)),
    }
}

/// Resolve a list of Paths of some sort to a set of unique Devices.
/// Return an IOError if there was a problem resolving any particular device.
// FIXME: BTreeSet -> HashSet once Device is hashable
pub fn resolve_devices(paths: &[&Path]) -> io::Result<BTreeSet<Device>> {
    let mut devices = BTreeSet::new();
    for path in paths {
        let dev = try!(Device::from_str(&path.to_string_lossy()));
        devices.insert(dev);
    }
    Ok(devices)
}

/// Find all Stratis Blockdevs.
///
/// Returns a map of pool uuids to maps of blockdev uuids to blockdevs.
pub fn find_all() -> EngineResult<HashMap<PoolUuid, HashMap<DevUuid, BlockDev>>> {

    /// If a Path refers to a valid Stratis blockdev, return a BlockDev
    /// struct. Otherwise, return None. Return an error if there was
    /// a problem inspecting the device.
    fn setup(devnode: &Path) -> EngineResult<Option<BlockDev>> {
        let dev = try!(Device::from_str(&devnode.to_string_lossy()));

        let mut f = try!(OpenOptions::new()
            .read(true)
            .open(devnode));

        let bda = try!(BDA::load(&mut f));

        Ok(Some(BlockDev {
            dev: dev,
            devnode: devnode.to_owned(),
            bda: bda,
        }))
    }

    let mut pool_map = HashMap::new();
    for dir_e in try!(read_dir("/dev")) {
        let devnode = match dir_e {
            Ok(d) => d.path(),
            Err(_) => continue,
        };

        match setup(&devnode) {
            Ok(Some(blockdev)) => {
                pool_map.entry(blockdev.pool_uuid().clone())
                    .or_insert_with(HashMap::new)
                    .insert(blockdev.uuid().clone(), blockdev);
            }
            _ => continue,
        };
    }

    Ok(pool_map)
}

/// Zero sectors at the given offset
pub fn wipe_sectors(path: &Path, offset: Sectors, sector_count: Sectors) -> EngineResult<()> {
    let mut f = try!(OpenOptions::new()
        .write(true)
        .open(path));

    let zeroed = [0u8; SECTOR_SIZE];

    // set the start point to the offset
    try!(f.seek(SeekFrom::Start(*offset)));
    for _ in 0..*sector_count {
        try!(f.write_all(&zeroed));
    }
    try!(f.flush());
    Ok(())
}

/// Initialize multiple blockdevs at once. This allows all of them
/// to be checked for usability before writing to any of them.
// FIXME: BTreeSet -> HashSet once Device is hashable
pub fn initialize(pool_uuid: &PoolUuid,
                  devices: BTreeSet<Device>,
                  mda_size: Sectors,
                  force: bool)
                  -> EngineResult<Vec<BlockDev>> {
    /// Gets device information, returns an error if problem with obtaining
    /// that information.
    /// Returns a tuple with the blockdev's path, its size in bytes,
    /// its ownership as determined by calling determine_ownership(),
    /// and an open File handle, all of which are needed later.
    fn dev_info(dev: &Device) -> EngineResult<(PathBuf, Bytes, DevOwnership, File)> {
        let devnode = try!(dev.path().ok_or_else(|| {
            io::Error::new(ErrorKind::InvalidInput,
                           format!("could not get device node from dev {}", dev.dstr()))
        }));
        let mut f = try!(OpenOptions::new()
            .read(true)
            .write(true)
            .open(&devnode)
            .map_err(|_| {
                io::Error::new(ErrorKind::PermissionDenied,
                               format!("Could not open {}", devnode.display()))
            }));

        let dev_size = try!(blkdev_size(&f));

        let ownership = match StaticHeader::determine_ownership(&mut f) {
            Ok(ownership) => ownership,
            Err(err) => {
                let error_message = format!("{} for device {}", err, devnode.display());
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            }
        };

        Ok((devnode, dev_size, ownership, f))
    }

    /// Filter devices for admission to pool based on dev_infos.
    /// If there is an error finding out the info, return that error.
    /// Also, return an error if a device is not appropriate for this pool.
    fn filter_devs<I>(dev_infos: I,
                      pool_uuid: &PoolUuid,
                      force: bool)
                      -> EngineResult<Vec<(Device, (PathBuf, Bytes, File))>>
        where I: Iterator<Item = (Device, EngineResult<(PathBuf, Bytes, DevOwnership, File)>)>
    {
        let mut add_devs = Vec::new();
        for (dev, dev_result) in dev_infos {
            let (devnode, dev_size, ownership, f) = try!(dev_result);
            if dev_size < MIN_DEV_SIZE {
                let error_message = format!("{} too small, minimum {} bytes",
                                            devnode.display(),
                                            MIN_DEV_SIZE);
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            };
            match ownership {
                DevOwnership::Unowned => add_devs.push((dev, (devnode, dev_size, f))),
                DevOwnership::Theirs => {
                    if !force {
                        let error_str = format!("First 8K of {} not zeroed", devnode.display());
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    } else {
                        add_devs.push((dev, (devnode, dev_size, f)))
                    }
                }
                DevOwnership::Ours(uuid) => {
                    if *pool_uuid != uuid {
                        let error_str = format!("Device {} already belongs to Stratis pool {}",
                                                devnode.display(),
                                                uuid);
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    }
                }
            }
        }
        Ok(add_devs)
    }

    try!(validate_mda_size(mda_size));

    let dev_infos = devices.into_iter().map(|d: Device| (d, dev_info(&d)));

    let add_devs = try!(filter_devs(dev_infos, pool_uuid, force));

    // TODO: Fix this code.  We should deal with any number of blockdevs
    //
    if add_devs.len() < 2 {
        return Err(EngineError::Engine(ErrorEnum::Error,
                                       "Need at least 2 blockdevs to create a pool".into()));
    }

    let mut bds = Vec::new();
    for (dev, (devnode, dev_size, mut f)) in add_devs {

        let bda = try!(BDA::initialize(&mut f,
                                       pool_uuid,
                                       &Uuid::new_v4(),
                                       mda_size,
                                       dev_size.sectors()));

        let bd = BlockDev {
            dev: dev,
            devnode: devnode.clone(),
            bda: bda,
        };
        bds.push(bd);
    }
    Ok(bds)
}

#[derive(Debug)]
pub struct BlockDev {
    dev: Device,
    pub devnode: PathBuf,
    bda: BDA,
}

impl BlockDev {
    pub fn to_save(&self) -> BlockDevSave {
        BlockDevSave {
            devnode: self.devnode.clone(),
            total_size: self.size(),
        }
    }

    pub fn wipe_metadata(&self) -> EngineResult<()> {
        let mut f = try!(OpenOptions::new().write(true).open(&self.devnode));
        BDA::wipe(&mut f)
    }

    /// Get the "x:y" device string for this blockdev
    pub fn dstr(&self) -> String {
        self.dev.dstr()
    }

    pub fn save_state(&mut self, time: &Timespec, metadata: &[u8]) -> EngineResult<()> {
        let mut f = try!(OpenOptions::new().write(true).open(&self.devnode));
        try!(self.bda.save_state(time, metadata, &mut f));

        Ok(())
    }

    pub fn load_state(&self) -> EngineResult<Option<Vec<u8>>> {
        let mut f = try!(OpenOptions::new().read(true).open(&self.devnode));
        self.bda.load_state(&mut f)
    }

    /// List the available-for-upper-layer-use range in this blockdev.
    pub fn avail_range(&self) -> (Sectors, Sectors) {
        let start = self.bda.size();
        let size = self.size();
        // Blockdev size is at least MIN_DEV_SIZE, so this can fail only if
        // size of metadata area exceeds 1 GiB. Initial metadata area size
        // is 4 MiB.
        assert!(start <= size);
        (start, size - start)
    }

    /// The /dev/mapper/<name> device is not immediately available for use.
    /// TODO: Implement wait for event or poll.
    pub fn wait_for_dm() {
        thread::sleep(Duration::from_millis(500))
    }

    /// The device's UUID.
    pub fn uuid(&self) -> &DevUuid {
        self.bda.dev_uuid()
    }

    /// The device's pool's UUID.
    pub fn pool_uuid(&self) -> &PoolUuid {
        self.bda.pool_uuid()
    }

    /// The device's size.
    pub fn size(&self) -> Sectors {
        self.bda.dev_size()
    }

    /// Last time metadata was written to this device.
    pub fn last_update_time(&self) -> &Option<Timespec> {
        self.bda.last_update_time()
    }
}
<commit_msg>Remove an extra try! and Ok(())<commit_after>// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

use std::io;
use std::collections::BTreeSet;
use std::collections::HashMap;
use std::fs::File;
use std::io::ErrorKind;
use std::io::{Seek, Write, SeekFrom};
use std::fs::{OpenOptions, read_dir};
use std::os::unix::prelude::AsRawFd;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use std::thread;
use std::time::Duration;

use time::Timespec;
use devicemapper::Device;
use uuid::Uuid;

use types::{Bytes, Sectors};
use engine::{DevUuid, EngineResult, EngineError, ErrorEnum, PoolUuid};

use consts::*;
use super::metadata::{StaticHeader, BDA, validate_mda_size};
use super::engine::DevOwnership;
pub use super::BlockDevSave;

const MIN_DEV_SIZE: Bytes = Bytes(IEC::Gi as u64);

ioctl!(read blkgetsize64 with 0x12, 114; u64);

pub fn blkdev_size(file: &File) -> EngineResult<Bytes> {
    let mut val: u64 = 0;

    match unsafe { blkgetsize64(file.as_raw_fd(), &mut val) } {
        Err(x) => Err(EngineError::Nix(x)),
        Ok(_) => Ok(Bytes(val)),
    }
}

/// Resolve a list of Paths of some sort to a set of unique Devices.
/// Return an IOError if there was a problem resolving any particular device.
// FIXME: BTreeSet -> HashSet once Device is hashable
pub fn resolve_devices(paths: &[&Path]) -> io::Result<BTreeSet<Device>> {
    let mut devices = BTreeSet::new();
    for path in paths {
        let dev = try!(Device::from_str(&path.to_string_lossy()));
        devices.insert(dev);
    }
    Ok(devices)
}

/// Find all Stratis Blockdevs.
///
/// Returns a map of pool uuids to maps of blockdev uuids to blockdevs.
pub fn find_all() -> EngineResult<HashMap<PoolUuid, HashMap<DevUuid, BlockDev>>> {

    /// If a Path refers to a valid Stratis blockdev, return a BlockDev
    /// struct. Otherwise, return None. Return an error if there was
    /// a problem inspecting the device.
    fn setup(devnode: &Path) -> EngineResult<Option<BlockDev>> {
        let dev = try!(Device::from_str(&devnode.to_string_lossy()));

        let mut f = try!(OpenOptions::new()
            .read(true)
            .open(devnode));

        let bda = try!(BDA::load(&mut f));

        Ok(Some(BlockDev {
            dev: dev,
            devnode: devnode.to_owned(),
            bda: bda,
        }))
    }

    let mut pool_map = HashMap::new();
    for dir_e in try!(read_dir("/dev")) {
        let devnode = match dir_e {
            Ok(d) => d.path(),
            Err(_) => continue,
        };

        match setup(&devnode) {
            Ok(Some(blockdev)) => {
                pool_map.entry(blockdev.pool_uuid().clone())
                    .or_insert_with(HashMap::new)
                    .insert(blockdev.uuid().clone(), blockdev);
            }
            _ => continue,
        };
    }

    Ok(pool_map)
}

/// Zero sectors at the given offset
pub fn wipe_sectors(path: &Path, offset: Sectors, sector_count: Sectors) -> EngineResult<()> {
    let mut f = try!(OpenOptions::new()
        .write(true)
        .open(path));

    let zeroed = [0u8; SECTOR_SIZE];

    // set the start point to the offset
    try!(f.seek(SeekFrom::Start(*offset)));
    for _ in 0..*sector_count {
        try!(f.write_all(&zeroed));
    }
    try!(f.flush());
    Ok(())
}

/// Initialize multiple blockdevs at once. This allows all of them
/// to be checked for usability before writing to any of them.
// FIXME: BTreeSet -> HashSet once Device is hashable
pub fn initialize(pool_uuid: &PoolUuid,
                  devices: BTreeSet<Device>,
                  mda_size: Sectors,
                  force: bool)
                  -> EngineResult<Vec<BlockDev>> {
    /// Gets device information, returns an error if problem with obtaining
    /// that information.
    /// Returns a tuple with the blockdev's path, its size in bytes,
    /// its ownership as determined by calling determine_ownership(),
    /// and an open File handle, all of which are needed later.
    fn dev_info(dev: &Device) -> EngineResult<(PathBuf, Bytes, DevOwnership, File)> {
        let devnode = try!(dev.path().ok_or_else(|| {
            io::Error::new(ErrorKind::InvalidInput,
                           format!("could not get device node from dev {}", dev.dstr()))
        }));
        let mut f = try!(OpenOptions::new()
            .read(true)
            .write(true)
            .open(&devnode)
            .map_err(|_| {
                io::Error::new(ErrorKind::PermissionDenied,
                               format!("Could not open {}", devnode.display()))
            }));

        let dev_size = try!(blkdev_size(&f));

        let ownership = match StaticHeader::determine_ownership(&mut f) {
            Ok(ownership) => ownership,
            Err(err) => {
                let error_message = format!("{} for device {}", err, devnode.display());
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            }
        };

        Ok((devnode, dev_size, ownership, f))
    }

    /// Filter devices for admission to pool based on dev_infos.
    /// If there is an error finding out the info, return that error.
    /// Also, return an error if a device is not appropriate for this pool.
    fn filter_devs<I>(dev_infos: I,
                      pool_uuid: &PoolUuid,
                      force: bool)
                      -> EngineResult<Vec<(Device, (PathBuf, Bytes, File))>>
        where I: Iterator<Item = (Device, EngineResult<(PathBuf, Bytes, DevOwnership, File)>)>
    {
        let mut add_devs = Vec::new();
        for (dev, dev_result) in dev_infos {
            let (devnode, dev_size, ownership, f) = try!(dev_result);
            if dev_size < MIN_DEV_SIZE {
                let error_message = format!("{} too small, minimum {} bytes",
                                            devnode.display(),
                                            MIN_DEV_SIZE);
                return Err(EngineError::Engine(ErrorEnum::Invalid, error_message));
            };
            match ownership {
                DevOwnership::Unowned => add_devs.push((dev, (devnode, dev_size, f))),
                DevOwnership::Theirs => {
                    if !force {
                        let error_str = format!("First 8K of {} not zeroed", devnode.display());
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    } else {
                        add_devs.push((dev, (devnode, dev_size, f)))
                    }
                }
                DevOwnership::Ours(uuid) => {
                    if *pool_uuid != uuid {
                        let error_str = format!("Device {} already belongs to Stratis pool {}",
                                                devnode.display(),
                                                uuid);
                        return Err(EngineError::Engine(ErrorEnum::Invalid, error_str));
                    }
                }
            }
        }
        Ok(add_devs)
    }

    try!(validate_mda_size(mda_size));

    let dev_infos = devices.into_iter().map(|d: Device| (d, dev_info(&d)));

    let add_devs = try!(filter_devs(dev_infos, pool_uuid, force));

    // TODO: Fix this code.  We should deal with any number of blockdevs
    //
    if add_devs.len() < 2 {
        return Err(EngineError::Engine(ErrorEnum::Error,
                                       "Need at least 2 blockdevs to create a pool".into()));
    }

    let mut bds = Vec::new();
    for (dev, (devnode, dev_size, mut f)) in add_devs {

        let bda = try!(BDA::initialize(&mut f,
                                       pool_uuid,
                                       &Uuid::new_v4(),
                                       mda_size,
                                       dev_size.sectors()));

        let bd = BlockDev {
            dev: dev,
            devnode: devnode.clone(),
            bda: bda,
        };
        bds.push(bd);
    }
    Ok(bds)
}

#[derive(Debug)]
pub struct BlockDev {
    dev: Device,
    pub devnode: PathBuf,
    bda: BDA,
}

impl BlockDev {
    pub fn to_save(&self) -> BlockDevSave {
        BlockDevSave {
            devnode: self.devnode.clone(),
            total_size: self.size(),
        }
    }

    pub fn wipe_metadata(&self) -> EngineResult<()> {
        let mut f = try!(OpenOptions::new().write(true).open(&self.devnode));
        BDA::wipe(&mut f)
    }

    /// Get the "x:y" device string for this blockdev
    pub fn dstr(&self) -> String {
        self.dev.dstr()
    }

    pub fn save_state(&mut self, time: &Timespec, metadata: &[u8]) -> EngineResult<()> {
        let mut f = try!(OpenOptions::new().write(true).open(&self.devnode));
        self.bda.save_state(time, metadata, &mut f)
    }

    pub fn load_state(&self) -> EngineResult<Option<Vec<u8>>> {
        let mut f = try!(OpenOptions::new().read(true).open(&self.devnode));
        self.bda.load_state(&mut f)
    }

    /// List the available-for-upper-layer-use range in this blockdev.
    pub fn avail_range(&self) -> (Sectors, Sectors) {
        let start = self.bda.size();
        let size = self.size();
        // Blockdev size is at least MIN_DEV_SIZE, so this can fail only if
        // size of metadata area exceeds 1 GiB. Initial metadata area size
        // is 4 MiB.
        assert!(start <= size);
        (start, size - start)
    }

    /// The /dev/mapper/<name> device is not immediately available for use.
    /// TODO: Implement wait for event or poll.
    pub fn wait_for_dm() {
        thread::sleep(Duration::from_millis(500))
    }

    /// The device's UUID.
    pub fn uuid(&self) -> &DevUuid {
        self.bda.dev_uuid()
    }

    /// The device's pool's UUID.
    pub fn pool_uuid(&self) -> &PoolUuid {
        self.bda.pool_uuid()
    }

    /// The device's size.
    pub fn size(&self) -> Sectors {
        self.bda.dev_size()
    }

    /// Last time metadata was written to this device.
    pub fn last_update_time(&self) -> &Option<Timespec> {
        self.bda.last_update_time()
    }
}
<|endoftext|>


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

use futures::future::{err, ok};
use futures::sync::oneshot::{Receiver, Sender};
use futures::{self, Future};
use hyper::service::service_fn;
use hyper::{self, Body, Method, Request, Response, Server, StatusCode};
use tempdir::TempDir;
use tokio_threadpool::{Builder, ThreadPool};

use std::net::SocketAddr;
use std::str::FromStr;

use super::Result;
use tikv_alloc::error::ProfError;
use tikv_util::collections::HashMap;
use tikv_util::metrics::dump;
use tikv_util::timer::GLOBAL_TIMER_HANDLE;

mod profiler_guard {
    use tikv_alloc::error::ProfResult;
    use tikv_alloc::{activate_prof, deactivate_prof};

    use futures::{Future, Poll};
    use futures_locks::{Mutex, MutexFut, MutexGuard};

    lazy_static! {
        static ref PROFILER_MUTEX: Mutex<u32> = Mutex::new(0);
    }

    pub struct ProfGuard(MutexGuard<u32>);

    pub struct ProfLock(MutexFut<u32>);

    impl ProfLock {
        pub fn new() -> ProfResult<ProfLock> {
            let guard = PROFILER_MUTEX.lock();
            match activate_prof() {
                Ok(_) => Ok(ProfLock(guard)),
                Err(e) => Err(e),
            }
        }
    }

    impl Drop for ProfGuard {
        fn drop(&mut self) {
            match deactivate_prof() {
                _ => {} // TODO: handle error here
            }
        }
    }

    impl Future for ProfLock {
        type Item = ProfGuard;
        type Error = ();
        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            self.0.poll().map(|item| item.map(|guard| ProfGuard(guard)))
        }
    }
}

pub struct StatusServer {
    thread_pool: ThreadPool,
    tx: Sender<()>,
    rx: Option<Receiver<()>>,
    addr: Option<SocketAddr>,
}

impl StatusServer {
    pub fn new(status_thread_pool_size: usize) -> Self {
        let thread_pool = Builder::new()
            .pool_size(status_thread_pool_size)
            .name_prefix("status-server-")
            .after_start(|| {
                info!("Status server started");
            })
            .before_stop(|| {
                info!("stopping status server");
            })
            .build();
        let (tx, rx) = futures::sync::oneshot::channel::<()>();
        StatusServer {
            thread_pool,
            tx,
            rx: Some(rx),
            addr: None,
        }
    }

    pub fn dump_prof(seconds: u64) -> Box<dyn Future<Item = Vec<u8>, Error = ProfError> + Send> {
        let lock = match profiler_guard::ProfLock::new() {
            Err(e) => return Box::new(err(e)),
            Ok(lock) => lock,
        };
        info!("start memory profiling {} seconds", seconds);

        let timer = GLOBAL_TIMER_HANDLE.clone();
        Box::new(lock.then(move |guard| {
            timer
                .delay(std::time::Instant::now() + std::time::Duration::from_secs(seconds))
                .then(
                    move |_| -> Box<dyn Future<Item = Vec<u8>, Error = ProfError> + Send> {
                        let tmp_dir = match TempDir::new("") {
                            Ok(tmp_dir) => tmp_dir,
                            Err(e) => return Box::new(err(e.into())),
                        };
                        let os_path = tmp_dir.path().join("tikv_dump_profile").into_os_string();
                        let path = match os_path.into_string() {
                            Ok(path) => path,
                            Err(path) => return Box::new(err(ProfError::PathError(path))),
                        };

                        if let Err(e) = tikv_alloc::dump_prof(&path) {
                            return Box::new(err(e));
                        }
                        drop(guard);
                        Box::new(
                            tokio_fs::file::File::open(path)
                                .and_then(|file| {
                                    let buf: Vec<u8> = Vec::new();
                                    tokio_io::io::read_to_end(file, buf)
                                })
                                .and_then(move |(_, buf)| {
                                    drop(tmp_dir);
                                    ok(buf)
                                })
                                .map_err(|e| -> ProfError { e.into() }),
                        )
                    },
                )
        }))
    }

    pub fn dump_prof_to_resp(
        req: Request<Body>,
    ) -> Box<dyn Future<Item = Response<Body>, Error = hyper::Error> + Send> {
        let query = match req.uri().query() {
            Some(query) => query,
            None => {
                let response = Response::builder()
                    .status(StatusCode::BAD_REQUEST)
                    .body(Body::empty())
                    .unwrap();
                return Box::new(ok(response));
            }
        };
        let query_pairs: HashMap<_, _> = url::form_urlencoded::parse(query.as_bytes()).collect();
        let seconds: u64 = match query_pairs.get("seconds") {
            Some(val) => match val.parse() {
                Ok(val) => val,
                Err(_) => {
                    let response = Response::builder()
                        .status(StatusCode::BAD_REQUEST)
                        .body(Body::empty())
                        .unwrap();
                    return Box::new(ok(response));
                }
            },
            None => 10,
        };

        Box::new(
            Self::dump_prof(seconds)
                .and_then(|buf| {
                    let response = Response::builder()
                        .header("X-Content-Type-Options", "nosniff")
                        .header("Content-Disposition", "attachment; filename=\"profile\"")
                        .header("Content-Type", mime::APPLICATION_OCTET_STREAM.to_string())
                        .header("Content-Length", buf.len())
                        .body(buf.into())
                        .unwrap();
                    ok(response)
                })
                .or_else(|err| {
                    let response = Response::builder()
                        .status(StatusCode::INTERNAL_SERVER_ERROR)
                        .body(Body::from(err.to_string()))
                        .unwrap();
                    ok(response)
                }),
        )
    }

    pub fn start(&mut self, status_addr: String) -> Result<()> {
        let addr = SocketAddr::from_str(&status_addr)?;

        // TODO: support TLS for the status server.
        let builder = Server::try_bind(&addr)?;

        // Create a status service.
        let service = |req: Request<Body>| -> Box<dyn Future<Item=Response<Body>, Error=hyper::Error> + Send> {
            match (req.method(), req.uri().path()) {
                (&Method::GET, "/metrics") => {
                    let response = Response::builder().body(Body::from(dump())).unwrap();
                    Box::new(ok(response))
                }
                (&Method::GET, "/pprof/profile") => {
                    Self::dump_prof_to_resp(req)
                }
                (&Method::GET, "/status") => {
                    let response = Response::builder().body(Body::empty()).unwrap();
                    Box::new(ok(response))
                },
                _ => {
                    let response = Response::builder().status(StatusCode::NOT_FOUND).body(Body::empty()).unwrap();
                    Box::new(ok(response))
                }
            }
        };

        // Start to serve.
        let server = builder.serve(move || service_fn(service));
        self.addr = Some(server.local_addr());
        let graceful = server
            .with_graceful_shutdown(self.rx.take().unwrap())
            .map_err(|e| error!("Status server error: {:?}", e));
        self.thread_pool.spawn(graceful);
        Ok(())
    }

    pub fn stop(self) {
        let _ = self.tx.send(());
        self.thread_pool
            .shutdown_now()
            .wait()
            .unwrap_or_else(|e| error!("failed to stop the status server, error: {:?}", e));
    }

    // Return listening address, this may only be used for outer test
    // to get the real address because we may use "127.0.0.1:0"
    // in test to avoid port conflict.
    pub fn listening_addr(&self) -> SocketAddr {
        self.addr.unwrap()
    }
}

#[cfg(test)]
mod tests {
    use crate::server::status_server::StatusServer;
    use futures::future::{lazy, Future};
    use hyper::{Client, StatusCode, Uri};

    #[test]
    fn test_status_service() {
        let mut status_server = StatusServer::new(1);
        let _ = status_server.start("127.0.0.1:0".to_string());
        let client = Client::new();
        let uri = Uri::builder()
            .scheme("http")
            .authority(status_server.listening_addr().to_string().as_str())
            .path_and_query("/metrics")
            .build()
            .unwrap();

        let handle = status_server.thread_pool.spawn_handle(lazy(move || {
            client
                .get(uri)
                .map(|res| {
                    assert_eq!(res.status(), StatusCode::OK);
                })
                .map_err(|err| {
                    panic!("response status is not OK: {:?}", err);
                })
        }));
        handle.wait().unwrap();
        status_server.stop();
    }
}
<commit_msg>Add http failpoints on status_server (#4444)<commit_after>// Copyright 2018 TiKV Project Authors. Licensed under Apache-2.0.

use futures::future::{err, ok};
use futures::sync::oneshot::{Receiver, Sender};
#[cfg(not(feature = "no-fail"))]
use futures::Stream;
use futures::{self, Future};
use hyper::service::service_fn;
use hyper::{self, Body, Method, Request, Response, Server, StatusCode};
use tempdir::TempDir;
use tokio_threadpool::{Builder, ThreadPool};

use std::net::SocketAddr;
use std::str::FromStr;

use super::Result;
use tikv_alloc::error::ProfError;
use tikv_util::collections::HashMap;
use tikv_util::metrics::dump;
use tikv_util::timer::GLOBAL_TIMER_HANDLE;

mod profiler_guard {
    use tikv_alloc::error::ProfResult;
    use tikv_alloc::{activate_prof, deactivate_prof};

    use futures::{Future, Poll};
    use futures_locks::{Mutex, MutexFut, MutexGuard};

    lazy_static! {
        static ref PROFILER_MUTEX: Mutex<u32> = Mutex::new(0);
    }

    pub struct ProfGuard(MutexGuard<u32>);

    pub struct ProfLock(MutexFut<u32>);

    impl ProfLock {
        pub fn new() -> ProfResult<ProfLock> {
            let guard = PROFILER_MUTEX.lock();
            match activate_prof() {
                Ok(_) => Ok(ProfLock(guard)),
                Err(e) => Err(e),
            }
        }
    }

    impl Drop for ProfGuard {
        fn drop(&mut self) {
            match deactivate_prof() {
                _ => {} // TODO: handle error here
            }
        }
    }

    impl Future for ProfLock {
        type Item = ProfGuard;
        type Error = ();
        fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
            self.0.poll().map(|item| item.map(|guard| ProfGuard(guard)))
        }
    }
}

#[cfg(not(feature = "no-fail"))]
static MISSING_NAME: &[u8] = b"Missing param name";
#[cfg(not(feature = "no-fail"))]
static MISSING_ACTIONS: &[u8] = b"Missing param actions";
#[cfg(not(feature = "no-fail"))]
static FAIL_POINTS_REQUEST_PATH: &str = "/fail";

pub struct StatusServer {
    thread_pool: ThreadPool,
    tx: Sender<()>,
    rx: Option<Receiver<()>>,
    addr: Option<SocketAddr>,
}

impl StatusServer {
    pub fn new(status_thread_pool_size: usize) -> Self {
        let thread_pool = Builder::new()
            .pool_size(status_thread_pool_size)
            .name_prefix("status-server-")
            .after_start(|| {
                info!("Status server started");
            })
            .before_stop(|| {
                info!("stopping status server");
            })
            .build();
        let (tx, rx) = futures::sync::oneshot::channel::<()>();
        StatusServer {
            thread_pool,
            tx,
            rx: Some(rx),
            addr: None,
        }
    }

    pub fn dump_prof(seconds: u64) -> Box<dyn Future<Item = Vec<u8>, Error = ProfError> + Send> {
        let lock = match profiler_guard::ProfLock::new() {
            Err(e) => return Box::new(err(e)),
            Ok(lock) => lock,
        };
        info!("start memory profiling {} seconds", seconds);

        let timer = GLOBAL_TIMER_HANDLE.clone();
        Box::new(lock.then(move |guard| {
            timer
                .delay(std::time::Instant::now() + std::time::Duration::from_secs(seconds))
                .then(
                    move |_| -> Box<dyn Future<Item = Vec<u8>, Error = ProfError> + Send> {
                        let tmp_dir = match TempDir::new("") {
                            Ok(tmp_dir) => tmp_dir,
                            Err(e) => return Box::new(err(e.into())),
                        };
                        let os_path = tmp_dir.path().join("tikv_dump_profile").into_os_string();
                        let path = match os_path.into_string() {
                            Ok(path) => path,
                            Err(path) => return Box::new(err(ProfError::PathError(path))),
                        };

                        if let Err(e) = tikv_alloc::dump_prof(&path) {
                            return Box::new(err(e));
                        }
                        drop(guard);
                        Box::new(
                            tokio_fs::file::File::open(path)
                                .and_then(|file| {
                                    let buf: Vec<u8> = Vec::new();
                                    tokio_io::io::read_to_end(file, buf)
                                })
                                .and_then(move |(_, buf)| {
                                    drop(tmp_dir);
                                    ok(buf)
                                })
                                .map_err(|e| -> ProfError { e.into() }),
                        )
                    },
                )
        }))
    }

    pub fn dump_prof_to_resp(
        req: Request<Body>,
    ) -> Box<dyn Future<Item = Response<Body>, Error = hyper::Error> + Send> {
        let query = match req.uri().query() {
            Some(query) => query,
            None => {
                let response = Response::builder()
                    .status(StatusCode::BAD_REQUEST)
                    .body(Body::empty())
                    .unwrap();
                return Box::new(ok(response));
            }
        };
        let query_pairs: HashMap<_, _> = url::form_urlencoded::parse(query.as_bytes()).collect();
        let seconds: u64 = match query_pairs.get("seconds") {
            Some(val) => match val.parse() {
                Ok(val) => val,
                Err(_) => {
                    let response = Response::builder()
                        .status(StatusCode::BAD_REQUEST)
                        .body(Body::empty())
                        .unwrap();
                    return Box::new(ok(response));
                }
            },
            None => 10,
        };

        Box::new(
            Self::dump_prof(seconds)
                .and_then(|buf| {
                    let response = Response::builder()
                        .header("X-Content-Type-Options", "nosniff")
                        .header("Content-Disposition", "attachment; filename=\"profile\"")
                        .header("Content-Type", mime::APPLICATION_OCTET_STREAM.to_string())
                        .header("Content-Length", buf.len())
                        .body(buf.into())
                        .unwrap();
                    ok(response)
                })
                .or_else(|err| {
                    let response = Response::builder()
                        .status(StatusCode::INTERNAL_SERVER_ERROR)
                        .body(Body::from(err.to_string()))
                        .unwrap();
                    ok(response)
                }),
        )
    }

    pub fn start(&mut self, status_addr: String) -> Result<()> {
        let addr = SocketAddr::from_str(&status_addr)?;

        // TODO: support TLS for the status server.
        let builder = Server::try_bind(&addr)?;

        // Create a status service.
        let service = |req: Request<Body>| -> Box<dyn Future<Item=Response<Body>, Error=hyper::Error> + Send> {
            let path = req.uri().path().to_owned();
            let method = req.method().to_owned();

            #[cfg(not(feature = "no-fail"))]
                {
                    if path.starts_with(FAIL_POINTS_REQUEST_PATH) {
                        return handle_fail_points_request(req);
                    }
                }

            match (method, path.as_ref()) {
                (Method::GET, "/metrics") => Box::new(ok(Response::new(dump().into()))),
                (Method::GET, "/status") => Box::new(ok(Response::default())),
                (Method::GET, "/pprof/profile") => Self::dump_prof_to_resp(req),
                _ => {
                    Box::new(ok(Response::builder()
                        .status(StatusCode::NOT_FOUND)
                        .body(Body::empty())
                        .unwrap()))
                }
            }
        };

        // Start to serve.
        let server = builder.serve(move || service_fn(service));
        self.addr = Some(server.local_addr());
        let graceful = server
            .with_graceful_shutdown(self.rx.take().unwrap())
            .map_err(|e| error!("Status server error: {:?}", e));
        self.thread_pool.spawn(graceful);
        Ok(())
    }

    pub fn stop(self) {
        let _ = self.tx.send(());
        self.thread_pool
            .shutdown_now()
            .wait()
            .unwrap_or_else(|e| error!("failed to stop the status server, error: {:?}", e));
    }

    // Return listening address, this may only be used for outer test
    // to get the real address because we may use "127.0.0.1:0"
    // in test to avoid port conflict.
    pub fn listening_addr(&self) -> SocketAddr {
        self.addr.unwrap()
    }
}

// For handling fail points related requests
#[cfg(not(feature = "no-fail"))]
fn handle_fail_points_request(
    req: Request<Body>,
) -> Box<dyn Future<Item = Response<Body>, Error = hyper::Error> + Send> {
    let path = req.uri().path().to_owned();
    let method = req.method().to_owned();
    let fail_path = format!("{}/", FAIL_POINTS_REQUEST_PATH);
    let fail_path_has_sub_path: bool = path.starts_with(&fail_path);

    match (method, fail_path_has_sub_path) {
        (Method::PUT, true) => Box::new(req.into_body().concat2().map(move |chunk| {
            let (_, name) = path.split_at(fail_path.len());
            if name.is_empty() {
                return Response::builder()
                    .status(StatusCode::UNPROCESSABLE_ENTITY)
                    .body(MISSING_NAME.into())
                    .unwrap();
            };

            let actions = chunk.into_iter().collect::<Vec<u8>>();
            let actions = String::from_utf8(actions).unwrap();
            if actions.is_empty() {
                return Response::builder()
                    .status(StatusCode::UNPROCESSABLE_ENTITY)
                    .body(MISSING_ACTIONS.into())
                    .unwrap();
            };

            if let Err(e) = fail::cfg(name.to_owned(), &actions) {
                return Response::builder()
                    .status(StatusCode::BAD_REQUEST)
                    .body(e.to_string().into())
                    .unwrap();
            }
            let body = format!("Added fail point with name: {}, actions: {}", name, actions);
            Response::new(body.into())
        })),
        (Method::DELETE, true) => {
            let (_, name) = path.split_at(fail_path.len());
            if name.is_empty() {
                return Box::new(ok(Response::builder()
                    .status(StatusCode::UNPROCESSABLE_ENTITY)
                    .body(MISSING_NAME.into())
                    .unwrap()));
            };

            fail::remove(name);
            let body = format!("Deleted fail point with name: {}", name);
            Box::new(ok(Response::new(body.into())))
        }
        (Method::GET, _) => {
            // In this scope the path must be like /fail...(/...), which starts with FAIL_POINTS_REQUEST_PATH and may or may not have a sub path
            // Now we return 404 when path is neither /fail nor /fail/
            if path != FAIL_POINTS_REQUEST_PATH && path != fail_path {
                return Box::new(ok(Response::builder()
                    .status(StatusCode::NOT_FOUND)
                    .body(Body::empty())
                    .unwrap()));
            }

            // From here path is either /fail or /fail/, return lists of fail points
            let list: Vec<String> = fail::list()
                .into_iter()
                .map(move |(name, actions)| format!("{}={}", name, actions))
                .collect();
            let list = list.join("\n");
            Box::new(ok(Response::new(list.into())))
        }
        _ => Box::new(ok(Response::builder()
            .status(StatusCode::METHOD_NOT_ALLOWED)
            .body(Body::empty())
            .unwrap())),
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Mutex, MutexGuard};

    use crate::server::status_server::StatusServer;
    use futures::future::{lazy, Future};
    #[cfg(not(feature = "no-fail"))]
    use futures::Stream;
    use hyper::{Body, Client, Method, Request, StatusCode, Uri};

    lazy_static! {
        static ref LOCK: Mutex<()> = Mutex::new(());
    }

    fn setup<'a>() -> MutexGuard<'a, ()> {
        let guard = LOCK.lock().unwrap_or_else(|e| e.into_inner());
        fail::teardown();
        fail::setup();
        guard
    }

    #[test]
    fn test_status_service() {
        let mut status_server = StatusServer::new(1);
        let _ = status_server.start("127.0.0.1:0".to_string());
        let client = Client::new();
        let uri = Uri::builder()
            .scheme("http")
            .authority(status_server.listening_addr().to_string().as_str())
            .path_and_query("/metrics")
            .build()
            .unwrap();

        let handle = status_server.thread_pool.spawn_handle(lazy(move || {
            client
                .get(uri)
                .map(|res| {
                    assert_eq!(res.status(), StatusCode::OK);
                })
                .map_err(|err| {
                    panic!("response status is not OK: {:?}", err);
                })
        }));
        handle.wait().unwrap();
        status_server.stop();
    }

    #[cfg(not(feature = "no-fail"))]
    #[test]
    fn test_status_service_fail_endpoints() {
        let _gaurd = setup();
        let mut status_server = StatusServer::new(1);
        let _ = status_server.start("127.0.0.1:0".to_string());
        let client = Client::new();
        let addr = status_server.listening_addr().to_string();

        let handle = status_server.thread_pool.spawn_handle(lazy(move || {
            // test add fail point
            let uri = Uri::builder()
                .scheme("http")
                .authority(addr.as_str())
                .path_and_query("/fail/test_fail_point_name")
                .build()
                .unwrap();
            let mut req = Request::new(Body::from("panic"));
            *req.method_mut() = Method::PUT;
            *req.uri_mut() = uri.clone();

            let future_1_add_fail_point = client
                .request(req)
                .map(|res| {
                    assert_eq!(res.status(), StatusCode::OK);

                    let list: Vec<String> = fail::list()
                        .into_iter()
                        .map(move |(name, actions)| format!("{}={}", name, actions))
                        .collect();
                    assert_eq!(list.len(), 1);
                    let list = list.join(";");
                    assert_eq!("test_fail_point_name=panic", list);
                })
                .map_err(|err| {
                    panic!("response status is not OK: {:?}", err);
                });

            // test add another fail point
            let uri = Uri::builder()
                .scheme("http")
                .authority(addr.as_str())
                .path_and_query("/fail/and_another_name")
                .build()
                .unwrap();
            let mut req = Request::new(Body::from("panic"));
            *req.method_mut() = Method::PUT;
            *req.uri_mut() = uri.clone();

            let future_2_add_fail_point = client
                .request(req)
                .map(|res| {
                    assert_eq!(res.status(), StatusCode::OK);

                    let list: Vec<String> = fail::list()
                        .into_iter()
                        .map(move |(name, actions)| format!("{}={}", name, actions))
                        .collect();
                    assert_eq!(2, list.len());
                    let list = list.join(";");
                    assert!(list.contains("test_fail_point_name=panic"));
                    assert!(list.contains("and_another_name=panic"))
                })
                .map_err(|err| {
                    panic!("response status is not OK: {:?}", err);
                });

            // test list fail points
            let uri = Uri::builder()
                .scheme("http")
                .authority(addr.as_str())
                .path_and_query("/fail")
                .build()
                .unwrap();
            let mut req = Request::default();
            *req.method_mut() = Method::GET;
            *req.uri_mut() = uri.clone();

            let future_3_list_fail_points = client
                .request(req)
                .and_then(|res| {
                    assert_eq!(res.status(), StatusCode::OK);
                    res.into_body().concat2()
                })
                .map(|chunk| {
                    let body = chunk.iter().cloned().collect::<Vec<u8>>();
                    let body = String::from_utf8(body).unwrap();
                    assert!(body.contains("test_fail_point_name=panic"));
                    assert!(body.contains("and_another_name=panic"))
                })
                .map_err(|err| {
                    panic!("response status is not OK: {:?}", err);
                });

            // test delete fail point
            let uri = Uri::builder()
                .scheme("http")
                .authority(addr.as_str())
                .path_and_query("/fail/test_fail_point_name")
                .build()
                .unwrap();
            let mut req = Request::default();
            *req.method_mut() = Method::DELETE;
            *req.uri_mut() = uri.clone();

            let future_4_delete_fail_points = client
                .request(req)
                .map(|res| {
                    assert_eq!(res.status(), StatusCode::OK);

                    let list: Vec<String> = fail::list()
                        .into_iter()
                        .map(move |(name, actions)| format!("{}={}", name, actions))
                        .collect();
                    assert_eq!(1, list.len());
                    let list = list.join(";");
                    assert_eq!("and_another_name=panic", list);
                })
                .map_err(|err| {
                    panic!("response status is not OK: {:?}", err);
                });

            future_1_add_fail_point
                .then(|_| future_2_add_fail_point)
                .then(|_| future_3_list_fail_points)
                .then(|_| future_4_delete_fail_points)
        }));

        handle.wait().unwrap();
        status_server.stop();
    }

    #[cfg(not(feature = "no-fail"))]
    #[test]
    fn test_status_service_fail_endpoints_can_trigger_fails() {
        let _gaurd = setup();
        let mut status_server = StatusServer::new(1);
        let _ = status_server.start("127.0.0.1:0".to_string());
        let client = Client::new();
        let addr = status_server.listening_addr().to_string();

        let handle = status_server.thread_pool.spawn_handle(lazy(move || {
            // test add fail point
            let uri = Uri::builder()
                .scheme("http")
                .authority(addr.as_str())
                .path_and_query("/fail/a_test_fail_name_nobody_else_is_using")
                .build()
                .unwrap();
            let mut req = Request::new(Body::from("return"));
            *req.method_mut() = Method::PUT;
            *req.uri_mut() = uri.clone();

            client
                .request(req)
                .map(|res| {
                    assert_eq!(res.status(), StatusCode::OK);
                })
                .map_err(|err| {
                    panic!("response status is not OK: {:?}", err);
                })
        }));

        handle.wait().unwrap();
        status_server.stop();

        let true_only_if_fail_point_triggered = || {
            fail_point!("a_test_fail_name_nobody_else_is_using", |_| { true });
            false
        };
        assert!(true_only_if_fail_point_triggered());
    }

    #[cfg(feature = "no-fail")]
    #[test]
    fn test_status_service_fail_endpoints_should_give_404_when_feature_no_fail_exists() {
        let _gaurd = setup();
        let mut status_server = StatusServer::new(1);
        let _ = status_server.start("127.0.0.1:0".to_string());
        let client = Client::new();
        let addr = status_server.listening_addr().to_string();

        let handle = status_server.thread_pool.spawn_handle(lazy(move || {
            // test add fail point
            let uri = Uri::builder()
                .scheme("http")
                .authority(addr.as_str())
                .path_and_query("/fail/a_test_fail_name_nobody_else_is_using")
                .build()
                .unwrap();
            let mut req = Request::new(Body::from("panic"));
            *req.method_mut() = Method::PUT;
            *req.uri_mut() = uri.clone();

            client
                .request(req)
                .map(|res| {
                    assert_eq!(res.status(), StatusCode::NOT_FOUND); // with feature no-fail, this PUT endpoint should return 404
                })
                .map_err(|err| {
                    panic!("response status is not OK: {:?}", err);
                })
        }));

        handle.wait().unwrap();
        status_server.stop();
    }
}
<|endoftext|>


<commit_before>enum Foo {
    Bar = 42,
    Baz,
    Quux = 100,
}

enum Signed {
    Bar = -42,
    Baz,
    Quux = 100,
}

fn foo() -> [u8; 3] {
    [Foo::Bar as u8, Foo::Baz as u8, Foo::Quux as u8]
}

fn signed() -> [i8; 3] {
    [Signed::Bar as i8, Signed::Baz as i8, Signed::Quux as i8]
}

fn unsafe_match() -> bool {
    match unsafe { std::mem::transmute::<u8, Foo>(43) } {
        Foo::Baz => true,
        _ => false,
    }
}

fn main() {
    assert_eq!(foo(), [42, 43, 100]);
    assert_eq!(signed(), [-42, -41, 100]);
    assert!(unsafe_match());
}
<commit_msg>test more enum-int-cast code paths<commit_after>enum Foo {
    Bar = 42,
    Baz,
    Quux = 100,
}

enum Signed {
    Bar = -42,
    Baz,
    Quux = 100,
}

fn foo() -> [u8; 3] {
    let baz = Foo::Baz; // let-expansion changes the MIR significantly
    [Foo::Bar as u8, baz as u8, Foo::Quux as u8]
}

fn signed() -> [i8; 3] {
    let baz = Signed::Baz; // let-expansion changes the MIR significantly
    [Signed::Bar as i8, baz as i8, Signed::Quux as i8]
}

fn unsafe_match() -> bool {
    match unsafe { std::mem::transmute::<u8, Foo>(43) } {
        Foo::Baz => true,
        _ => false,
    }
}

fn main() {
    assert_eq!(foo(), [42, 43, 100]);
    assert_eq!(signed(), [-42, -41, 100]);
    assert!(unsafe_match());
}
<|endoftext|>


<commit_before>extern crate bytes;
extern crate fallible_iterator;
extern crate futures_cpupool;
extern crate postgres_protocol;
extern crate postgres_shared;
extern crate tokio_codec;
extern crate tokio_io;
extern crate tokio_tcp;
extern crate tokio_timer;

#[macro_use]
extern crate futures;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
#[macro_use]
extern crate state_machine_future;

#[cfg(unix)]
extern crate tokio_uds;

use futures::{Async, Future, Poll};
use std::io;
use std::sync::atomic::{AtomicUsize, Ordering};

#[doc(inline)]
pub use postgres_shared::stmt::Column;
#[doc(inline)]
pub use postgres_shared::{error, params, types};
#[doc(inline)]
pub use postgres_shared::{CancelData, Notification};

use error::Error;
use params::ConnectParams;
use types::Type;

mod proto;

static NEXT_STATEMENT_ID: AtomicUsize = AtomicUsize::new(0);

fn bad_response() -> Error {
    Error::from(io::Error::new(
        io::ErrorKind::InvalidInput,
        "the server returned an unexpected response",
    ))
}

fn disconnected() -> Error {
    Error::from(io::Error::new(
        io::ErrorKind::UnexpectedEof,
        "server disconnected",
    ))
}

pub fn connect(params: ConnectParams) -> Handshake {
    Handshake(proto::HandshakeFuture::new(params))
}

pub struct Client(proto::Client);

impl Client {
    pub fn prepare(&mut self, query: &str) -> Prepare {
        self.prepare_typed(query, &[])
    }

    pub fn prepare_typed(&mut self, query: &str, param_types: &[Type]) -> Prepare {
        let name = format!("s{}", NEXT_STATEMENT_ID.fetch_add(1, Ordering::SeqCst));
        Prepare(self.0.prepare(name, query, param_types))
    }
}

pub struct Connection(proto::Connection);

impl Connection {
    pub fn cancel_data(&self) -> CancelData {
        self.0.cancel_data()
    }

    pub fn parameter(&self, name: &str) -> Option<&str> {
        self.0.parameter(name)
    }
}

impl Future for Connection {
    type Item = ();
    type Error = Error;

    fn poll(&mut self) -> Poll<(), Error> {
        self.0.poll()
    }
}

pub struct Handshake(proto::HandshakeFuture);

impl Future for Handshake {
    type Item = (Client, Connection);
    type Error = Error;

    fn poll(&mut self) -> Poll<(Client, Connection), Error> {
        let (client, connection) = try_ready!(self.0.poll());

        Ok(Async::Ready((Client(client), Connection(connection))))
    }
}

pub struct Prepare(proto::PrepareFuture);

impl Future for Prepare {
    type Item = Statement;
    type Error = Error;

    fn poll(&mut self) -> Poll<Statement, Error> {
        let statement = try_ready!(self.0.poll());

        Ok(Async::Ready(Statement(statement)))
    }
}

pub struct Statement(proto::Statement);

impl Statement {
    pub fn params(&self) -> &[Type] {
        self.0.params()
    }

    pub fn columns(&self) -> &[Column] {
        self.0.columns()
    }
}
<commit_msg>tag futures as must use<commit_after>extern crate bytes;
extern crate fallible_iterator;
extern crate futures_cpupool;
extern crate postgres_protocol;
extern crate postgres_shared;
extern crate tokio_codec;
extern crate tokio_io;
extern crate tokio_tcp;
extern crate tokio_timer;

#[macro_use]
extern crate futures;
#[macro_use]
extern crate lazy_static;
#[macro_use]
extern crate log;
#[macro_use]
extern crate state_machine_future;

#[cfg(unix)]
extern crate tokio_uds;

use futures::{Async, Future, Poll};
use std::io;
use std::sync::atomic::{AtomicUsize, Ordering};

#[doc(inline)]
pub use postgres_shared::stmt::Column;
#[doc(inline)]
pub use postgres_shared::{error, params, types};
#[doc(inline)]
pub use postgres_shared::{CancelData, Notification};

use error::Error;
use params::ConnectParams;
use types::Type;

mod proto;

static NEXT_STATEMENT_ID: AtomicUsize = AtomicUsize::new(0);

fn bad_response() -> Error {
    Error::from(io::Error::new(
        io::ErrorKind::InvalidInput,
        "the server returned an unexpected response",
    ))
}

fn disconnected() -> Error {
    Error::from(io::Error::new(
        io::ErrorKind::UnexpectedEof,
        "server disconnected",
    ))
}

pub fn connect(params: ConnectParams) -> Handshake {
    Handshake(proto::HandshakeFuture::new(params))
}

pub struct Client(proto::Client);

impl Client {
    pub fn prepare(&mut self, query: &str) -> Prepare {
        self.prepare_typed(query, &[])
    }

    pub fn prepare_typed(&mut self, query: &str, param_types: &[Type]) -> Prepare {
        let name = format!("s{}", NEXT_STATEMENT_ID.fetch_add(1, Ordering::SeqCst));
        Prepare(self.0.prepare(name, query, param_types))
    }
}

#[must_use = "futures do nothing unless polled"]
pub struct Connection(proto::Connection);

impl Connection {
    pub fn cancel_data(&self) -> CancelData {
        self.0.cancel_data()
    }

    pub fn parameter(&self, name: &str) -> Option<&str> {
        self.0.parameter(name)
    }
}

impl Future for Connection {
    type Item = ();
    type Error = Error;

    fn poll(&mut self) -> Poll<(), Error> {
        self.0.poll()
    }
}

#[must_use = "futures do nothing unless polled"]
pub struct Handshake(proto::HandshakeFuture);

impl Future for Handshake {
    type Item = (Client, Connection);
    type Error = Error;

    fn poll(&mut self) -> Poll<(Client, Connection), Error> {
        let (client, connection) = try_ready!(self.0.poll());

        Ok(Async::Ready((Client(client), Connection(connection))))
    }
}

#[must_use = "futures do nothing unless polled"]
pub struct Prepare(proto::PrepareFuture);

impl Future for Prepare {
    type Item = Statement;
    type Error = Error;

    fn poll(&mut self) -> Poll<Statement, Error> {
        let statement = try_ready!(self.0.poll());

        Ok(Async::Ready(Statement(statement)))
    }
}

pub struct Statement(proto::Statement);

impl Statement {
    pub fn params(&self) -> &[Type] {
        self.0.params()
    }

    pub fn columns(&self) -> &[Column] {
        self.0.columns()
    }
}
<|endoftext|>