<commit_before>//! Connection pooling via r2d2.
//!
//! Note: This module requires enabling the `r2d2` feature

extern crate r2d2;

pub use self::r2d2::*;

/// A re-export of [`r2d2::Error`], which is only used by methods on [`r2d2::Pool`].
///
/// [`r2d2::Error`]: r2d2::Error
/// [`r2d2::Pool`]: r2d2::Pool
pub type PoolError = self::r2d2::Error;

use std::convert::Into;
use std::fmt;
use std::marker::PhantomData;

use crate::backend::Backend;
use crate::connection::{IterableConnection, SimpleConnection, TransactionManager};
use crate::expression::QueryMetadata;
use crate::prelude::*;
use crate::query_builder::{AsQuery, QueryFragment, QueryId};

/// An r2d2 connection manager for use with Diesel.
///
/// See the [r2d2 documentation] for usage examples.
///
/// [r2d2 documentation]: r2d2
#[derive(Clone)]
pub struct ConnectionManager<T> {
    database_url: String,
    _marker: PhantomData<T>,
}

impl<T> fmt::Debug for ConnectionManager<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ConnectionManager<{}>", std::any::type_name::<T>())
    }
}

unsafe impl<T: Send + 'static> Sync for ConnectionManager<T> {}

impl<T> ConnectionManager<T> {
    /// Returns a new connection manager,
    /// which establishes connections to the given database URL.
    pub fn new<S: Into<String>>(database_url: S) -> Self {
        ConnectionManager {
            database_url: database_url.into(),
            _marker: PhantomData,
        }
    }
}

/// The error used when managing connections with `r2d2`.
#[derive(Debug)]
pub enum Error {
    /// An error occurred establishing the connection
    ConnectionError(ConnectionError),

    /// An error occurred pinging the database
    QueryError(crate::result::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::ConnectionError(ref e) => e.fmt(f),
            Error::QueryError(ref e) => e.fmt(f),
        }
    }
}

impl ::std::error::Error for Error {}

/// A trait indicating a connection could be used inside a r2d2 pool
pub trait R2D2Connection: Connection {
    /// Check if a connection is still valid
    fn ping(&mut self) -> QueryResult<()>;
}

#[cfg(feature = "postgres")]
impl R2D2Connection for crate::pg::PgConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

#[cfg(feature = "mysql")]
impl R2D2Connection for crate::mysql::MysqlConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

#[cfg(feature = "sqlite")]
impl R2D2Connection for crate::sqlite::SqliteConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

impl<T> ManageConnection for ConnectionManager<T>
where
    T: R2D2Connection + Send + 'static,
{
    type Connection = T;
    type Error = Error;

    fn connect(&self) -> Result<T, Error> {
        T::establish(&self.database_url).map_err(Error::ConnectionError)
    }

    fn is_valid(&self, conn: &mut T) -> Result<(), Error> {
        conn.ping().map_err(Error::QueryError)
    }

    fn has_broken(&self, _conn: &mut T) -> bool {
        false
    }
}

impl<M> SimpleConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: R2D2Connection + Send + 'static,
{
    fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
        (&mut **self).batch_execute(query)
    }
}

impl<'a, DB, M> IterableConnection<'a, DB> for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection<Backend = DB>,
    DB: Backend,
{
    type Cursor = <M::Connection as IterableConnection<'a, DB>>::Cursor;
    type Row = <M::Connection as IterableConnection<'a, DB>>::Row;
}

impl<M> Connection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection + R2D2Connection + Send + 'static,
{
    type Backend = <M::Connection as Connection>::Backend;
    type TransactionManager =
        PoolTransactionManager<<M::Connection as Connection>::TransactionManager>;

    fn establish(_: &str) -> ConnectionResult<Self> {
        Err(ConnectionError::BadConnection(String::from(
            "Cannot directly establish a pooled connection",
        )))
    }

    fn execute(&mut self, query: &str) -> QueryResult<usize> {
        (&mut **self).execute(query)
    }

    fn load<'a, T>(
        &'a mut self,
        source: T,
    ) -> QueryResult<<Self as IterableConnection<'a, Self::Backend>>::Cursor>
    where
        T: AsQuery,
        T::Query: QueryFragment<Self::Backend> + QueryId,
        Self::Backend: QueryMetadata<T::SqlType>,
    {
        (&mut **self).load(source)
    }

    fn execute_returning_count<T>(&mut self, source: &T) -> QueryResult<usize>
    where
        T: QueryFragment<Self::Backend> + QueryId,
    {
        (&mut **self).execute_returning_count(source)
    }

    fn transaction_state(
        &mut self,
    ) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData {
        (&mut **self).transaction_state()
    }

    fn begin_test_transaction(&mut self) -> QueryResult<()> {
        (&mut **self).begin_test_transaction()
    }
}

#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct PoolTransactionManager<T>(std::marker::PhantomData<T>);

impl<M, T> TransactionManager<PooledConnection<M>> for PoolTransactionManager<T>
where
    M: ManageConnection,
    M::Connection: Connection<TransactionManager = T> + R2D2Connection,
    T: TransactionManager<M::Connection>,
{
    type TransactionStateData = T::TransactionStateData;

    fn begin_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::begin_transaction(&mut **conn)
    }

    fn rollback_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::rollback_transaction(&mut **conn)
    }

    fn commit_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::commit_transaction(&mut **conn)
    }

    fn get_transaction_depth(conn: &mut PooledConnection<M>) -> u32 {
        T::get_transaction_depth(&mut **conn)
    }
}

impl<M> crate::migration::MigrationConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::migration::MigrationConnection,
    Self: Connection,
{
    fn setup(&mut self) -> QueryResult<usize> {
        (&mut **self).setup()
    }
}

impl<Changes, Output, M> crate::query_dsl::UpdateAndFetchResults<Changes, Output>
    for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::query_dsl::UpdateAndFetchResults<Changes, Output>,
    Self: Connection,
{
    fn update_and_fetch(&mut self, changeset: Changes) -> QueryResult<Output> {
        (&mut **self).update_and_fetch(changeset)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::mpsc;
    use std::sync::Arc;
    use std::thread;

    use crate::r2d2::*;
    use crate::test_helpers::*;

    #[test]
    fn establish_basic_connection() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Arc::new(Pool::builder().max_size(2).build(manager).unwrap());

        let (s1, r1) = mpsc::channel();
        let (s2, r2) = mpsc::channel();

        let pool1 = Arc::clone(&pool);
        let t1 = thread::spawn(move || {
            let conn = pool1.get().unwrap();
            s1.send(()).unwrap();
            r2.recv().unwrap();
            drop(conn);
        });

        let pool2 = Arc::clone(&pool);
        let t2 = thread::spawn(move || {
            let conn = pool2.get().unwrap();
            s2.send(()).unwrap();
            r1.recv().unwrap();
            drop(conn);
        });

        t1.join().unwrap();
        t2.join().unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn is_valid() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn pooled_connection_impls_connection() {
        use crate::select;
        use crate::sql_types::Text;

        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();
        let mut conn = pool.get().unwrap();

        let query = select("foo".into_sql::<Text>());
        assert_eq!("foo", query.get_result::<String>(&mut conn).unwrap());
    }
}
<commit_msg>Fix r2d2<commit_after>//! Connection pooling via r2d2.
//!
//! Note: This module requires enabling the `r2d2` feature

extern crate r2d2;

pub use self::r2d2::*;

/// A re-export of [`r2d2::Error`], which is only used by methods on [`r2d2::Pool`].
///
/// [`r2d2::Error`]: r2d2::Error
/// [`r2d2::Pool`]: r2d2::Pool
pub type PoolError = self::r2d2::Error;

use std::convert::Into;
use std::fmt;
use std::marker::PhantomData;

use crate::backend::Backend;
use crate::connection::{ConnectionGatWorkaround, SimpleConnection, TransactionManager};
use crate::expression::QueryMetadata;
use crate::prelude::*;
use crate::query_builder::{AsQuery, QueryFragment, QueryId};

/// An r2d2 connection manager for use with Diesel.
///
/// See the [r2d2 documentation] for usage examples.
///
/// [r2d2 documentation]: r2d2
#[derive(Clone)]
pub struct ConnectionManager<T> {
    database_url: String,
    _marker: PhantomData<T>,
}

impl<T> fmt::Debug for ConnectionManager<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "ConnectionManager<{}>", std::any::type_name::<T>())
    }
}

unsafe impl<T: Send + 'static> Sync for ConnectionManager<T> {}

impl<T> ConnectionManager<T> {
    /// Returns a new connection manager,
    /// which establishes connections to the given database URL.
    pub fn new<S: Into<String>>(database_url: S) -> Self {
        ConnectionManager {
            database_url: database_url.into(),
            _marker: PhantomData,
        }
    }
}

/// The error used when managing connections with `r2d2`.
#[derive(Debug)]
pub enum Error {
    /// An error occurred establishing the connection
    ConnectionError(ConnectionError),

    /// An error occurred pinging the database
    QueryError(crate::result::Error),
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Error::ConnectionError(ref e) => e.fmt(f),
            Error::QueryError(ref e) => e.fmt(f),
        }
    }
}

impl ::std::error::Error for Error {}

/// A trait indicating a connection could be used inside a r2d2 pool
pub trait R2D2Connection: Connection {
    /// Check if a connection is still valid
    fn ping(&mut self) -> QueryResult<()>;
}

#[cfg(feature = "postgres")]
impl R2D2Connection for crate::pg::PgConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

#[cfg(feature = "mysql")]
impl R2D2Connection for crate::mysql::MysqlConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

#[cfg(feature = "sqlite")]
impl R2D2Connection for crate::sqlite::SqliteConnection {
    fn ping(&mut self) -> QueryResult<()> {
        self.execute("SELECT 1").map(|_| ())
    }
}

impl<T> ManageConnection for ConnectionManager<T>
where
    T: R2D2Connection + Send + 'static,
{
    type Connection = T;
    type Error = Error;

    fn connect(&self) -> Result<T, Error> {
        T::establish(&self.database_url).map_err(Error::ConnectionError)
    }

    fn is_valid(&self, conn: &mut T) -> Result<(), Error> {
        conn.ping().map_err(Error::QueryError)
    }

    fn has_broken(&self, _conn: &mut T) -> bool {
        false
    }
}

impl<M> SimpleConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: R2D2Connection + Send + 'static,
{
    fn batch_execute(&mut self, query: &str) -> QueryResult<()> {
        (&mut **self).batch_execute(query)
    }
}

impl<'a, DB, M> ConnectionGatWorkaround<'a, DB> for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection<Backend = DB>,
    DB: Backend,
{
    type Cursor = <M::Connection as ConnectionGatWorkaround<'a, DB>>::Cursor;
    type Row = <M::Connection as ConnectionGatWorkaround<'a, DB>>::Row;
}

impl<M> Connection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: Connection + R2D2Connection + Send + 'static,
{
    type Backend = <M::Connection as Connection>::Backend;
    type TransactionManager =
        PoolTransactionManager<<M::Connection as Connection>::TransactionManager>;

    fn establish(_: &str) -> ConnectionResult<Self> {
        Err(ConnectionError::BadConnection(String::from(
            "Cannot directly establish a pooled connection",
        )))
    }

    fn execute(&mut self, query: &str) -> QueryResult<usize> {
        (&mut **self).execute(query)
    }

    fn load<'a, T>(
        &'a mut self,
        source: T,
    ) -> QueryResult<<Self as ConnectionGatWorkaround<'a, Self::Backend>>::Cursor>
    where
        T: AsQuery,
        T::Query: QueryFragment<Self::Backend> + QueryId,
        Self::Backend: QueryMetadata<T::SqlType>,
    {
        (&mut **self).load(source)
    }

    fn execute_returning_count<T>(&mut self, source: &T) -> QueryResult<usize>
    where
        T: QueryFragment<Self::Backend> + QueryId,
    {
        (&mut **self).execute_returning_count(source)
    }

    fn transaction_state(
        &mut self,
    ) -> &mut <Self::TransactionManager as TransactionManager<Self>>::TransactionStateData {
        (&mut **self).transaction_state()
    }

    fn begin_test_transaction(&mut self) -> QueryResult<()> {
        (&mut **self).begin_test_transaction()
    }
}

#[doc(hidden)]
#[allow(missing_debug_implementations)]
pub struct PoolTransactionManager<T>(std::marker::PhantomData<T>);

impl<M, T> TransactionManager<PooledConnection<M>> for PoolTransactionManager<T>
where
    M: ManageConnection,
    M::Connection: Connection<TransactionManager = T> + R2D2Connection,
    T: TransactionManager<M::Connection>,
{
    type TransactionStateData = T::TransactionStateData;

    fn begin_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::begin_transaction(&mut **conn)
    }

    fn rollback_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::rollback_transaction(&mut **conn)
    }

    fn commit_transaction(conn: &mut PooledConnection<M>) -> QueryResult<()> {
        T::commit_transaction(&mut **conn)
    }

    fn get_transaction_depth(conn: &mut PooledConnection<M>) -> u32 {
        T::get_transaction_depth(&mut **conn)
    }
}

impl<M> crate::migration::MigrationConnection for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::migration::MigrationConnection,
    Self: Connection,
{
    fn setup(&mut self) -> QueryResult<usize> {
        (&mut **self).setup()
    }
}

impl<Changes, Output, M> crate::query_dsl::UpdateAndFetchResults<Changes, Output>
    for PooledConnection<M>
where
    M: ManageConnection,
    M::Connection: crate::query_dsl::UpdateAndFetchResults<Changes, Output>,
    Self: Connection,
{
    fn update_and_fetch(&mut self, changeset: Changes) -> QueryResult<Output> {
        (&mut **self).update_and_fetch(changeset)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::mpsc;
    use std::sync::Arc;
    use std::thread;

    use crate::r2d2::*;
    use crate::test_helpers::*;

    #[test]
    fn establish_basic_connection() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Arc::new(Pool::builder().max_size(2).build(manager).unwrap());

        let (s1, r1) = mpsc::channel();
        let (s2, r2) = mpsc::channel();

        let pool1 = Arc::clone(&pool);
        let t1 = thread::spawn(move || {
            let conn = pool1.get().unwrap();
            s1.send(()).unwrap();
            r2.recv().unwrap();
            drop(conn);
        });

        let pool2 = Arc::clone(&pool);
        let t2 = thread::spawn(move || {
            let conn = pool2.get().unwrap();
            s2.send(()).unwrap();
            r1.recv().unwrap();
            drop(conn);
        });

        t1.join().unwrap();
        t2.join().unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn is_valid() {
        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();

        pool.get().unwrap();
    }

    #[test]
    fn pooled_connection_impls_connection() {
        use crate::select;
        use crate::sql_types::Text;

        let manager = ConnectionManager::<TestConnection>::new(database_url());
        let pool = Pool::builder()
            .max_size(1)
            .test_on_check_out(true)
            .build(manager)
            .unwrap();
        let mut conn = pool.get().unwrap();

        let query = select("foo".into_sql::<Text>());
        assert_eq!("foo", query.get_result::<String>(&mut conn).unwrap());
    }
}
<|endoftext|>


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

use std::collections::BTreeMap;
use std::sync::Arc;
pub use toml::de::Deserializer;
pub use toml::Value;

pub mod tools;

pub type CResult<T> = Result<T, String>;

pub struct BenchmarkWorkload {
    pub name: String,
    pub rate: usize,
    pub gen: Box<ProtocolGen>,
}

impl BenchmarkWorkload {
    pub fn new(name: String, rate: usize, gen: Box<ProtocolGen>) -> BenchmarkWorkload {
        BenchmarkWorkload {
            name: name,
            rate: rate,
            gen: gen,
        }
    }
}

/// Protocol specific parsing generator and workloads
pub struct ProtocolConfig {
    pub protocol: Arc<ProtocolParseFactory>,
    pub workloads: Vec<BenchmarkWorkload>,
}

#[derive(PartialEq, Debug)]
pub enum ParsedResponse {
    Error(String),
    Hit,
    Incomplete,
    Invalid,
    Miss,
    Ok,
    Unknown,
    Version(String),
}

/// Factory of protocol message buffers
pub trait ProtocolGen: Send {
    /// Generate the next buffer to send to the server
    fn generate_message(&mut self) -> Vec<u8>;

    /// The method being called on the server
    fn method(&self) -> &str;
}

/// Factory for `ProtocolParse` instances
pub trait ProtocolParseFactory: Send + Sync {
    /// Create a new protocol parser
    fn new(&self) -> Box<ProtocolParse>;

    /// Name of this protocol
    fn name(&self) -> &str;

    /// Generate some preparatory messages for the work queue
    fn prepare(&self) -> CResult<Vec<Vec<u8>>> {
        Ok(Vec::new())
    }
}

/// Protocol specific parser
pub trait ProtocolParse: Send {
    /// Parse the response buffer
    fn parse(&self, bytes: &[u8]) -> ParsedResponse;
}

/// Reusable paramter type with parser
pub trait Ptype: Sized {
    /// generate new state
    fn regen(&mut self);
    /// parse a `Ptype` from a toml tree
    fn parse(seed: usize, size: usize, num: u64, table: &BTreeMap<String, Value>) -> CResult<Self>;
}

/// `Parameter` generation style
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum Style {
    /// Values are generated based on a seed
    Static,
    /// Values are generated randomly
    Random,
}

#[derive(Clone, Debug)]
pub struct Parameter<T: Ptype> {
    /// initialization `Style` of the parameter
    pub style: Style,
    /// whether the parameter should be regenerated on each use
    pub regenerate: bool,
    /// current value of the `Parameter`
    pub value: T,
}

impl<T: Default + Ptype> Default for Parameter<T> {
    fn default() -> Parameter<T> {
        Parameter {
            style: Style::Static,
            regenerate: false,
            value: T::default(),
        }
    }
}

impl<T: Ptype> Parameter<T> {
    /// Mutate internal value if necessary
    ///
    /// If the parameter is flagged to not regenerate or was seeded, it isn't regenerated.
    pub fn regen(&mut self) {
        if self.regenerate && self.style == Style::Random {
            self.value.regen()
        }
    }
}

/// Extract a `Parameter` from the toml tree
pub fn extract_parameter<T: Ptype>(
    index: usize,
    parameter: &BTreeMap<String, Value>,
) -> CResult<Parameter<T>> {
    let style = match parameter.get("style").and_then(|k| k.as_str()) {
        Some("random") => Style::Random,
        Some("static") | None => Style::Static,
        Some(other) => return Err(format!("bad parameter style: {}", other)),
    };

    let seed = parameter
        .get("seed")
        .and_then(|k| k.as_integer())
        .map_or(index, |i| i as usize);

    let size = parameter
        .get("size")
        .and_then(|k| k.as_integer())
        .map_or(1, |i| i as usize);

    let num = parameter
        .get("num")
        .and_then(|k| k.as_integer())
        .map_or(0, |i| i as u64);

    // check that size is sufficiently large to contain num strings
    let fmtlen = format!("{}", num - 1).len();
    if  fmtlen > size {
        return Err(format!("To contain {} strings, you need to specify a \
                            size >= {}", num, fmtlen))
    }

    let regenerate = parameter
        .get("regenerate")
        .and_then(|k| k.as_bool())
        .unwrap_or(false);

    let mut value = T::parse(seed, size, num, parameter)?;

    // initialize with a random value if that is what is needed
    if style == Style::Random {
        value.regen();
    }

    Ok(Parameter {
        style: style,
        regenerate: regenerate,
        value: value,
    })
}
<commit_msg>fix input checking bug for num (#160)<commit_after>//  rpc-perf - RPC Performance Testing
//  Copyright 2015 Twitter, Inc
//
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.

use std::collections::BTreeMap;
use std::sync::Arc;
pub use toml::de::Deserializer;
pub use toml::Value;

pub mod tools;

pub type CResult<T> = Result<T, String>;

pub struct BenchmarkWorkload {
    pub name: String,
    pub rate: usize,
    pub gen: Box<ProtocolGen>,
}

impl BenchmarkWorkload {
    pub fn new(name: String, rate: usize, gen: Box<ProtocolGen>) -> BenchmarkWorkload {
        BenchmarkWorkload {
            name: name,
            rate: rate,
            gen: gen,
        }
    }
}

/// Protocol specific parsing generator and workloads
pub struct ProtocolConfig {
    pub protocol: Arc<ProtocolParseFactory>,
    pub workloads: Vec<BenchmarkWorkload>,
}

#[derive(PartialEq, Debug)]
pub enum ParsedResponse {
    Error(String),
    Hit,
    Incomplete,
    Invalid,
    Miss,
    Ok,
    Unknown,
    Version(String),
}

/// Factory of protocol message buffers
pub trait ProtocolGen: Send {
    /// Generate the next buffer to send to the server
    fn generate_message(&mut self) -> Vec<u8>;

    /// The method being called on the server
    fn method(&self) -> &str;
}

/// Factory for `ProtocolParse` instances
pub trait ProtocolParseFactory: Send + Sync {
    /// Create a new protocol parser
    fn new(&self) -> Box<ProtocolParse>;

    /// Name of this protocol
    fn name(&self) -> &str;

    /// Generate some preparatory messages for the work queue
    fn prepare(&self) -> CResult<Vec<Vec<u8>>> {
        Ok(Vec::new())
    }
}

/// Protocol specific parser
pub trait ProtocolParse: Send {
    /// Parse the response buffer
    fn parse(&self, bytes: &[u8]) -> ParsedResponse;
}

/// Reusable paramter type with parser
pub trait Ptype: Sized {
    /// generate new state
    fn regen(&mut self);
    /// parse a `Ptype` from a toml tree
    fn parse(seed: usize, size: usize, num: u64, table: &BTreeMap<String, Value>) -> CResult<Self>;
}

/// `Parameter` generation style
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum Style {
    /// Values are generated based on a seed
    Static,
    /// Values are generated randomly
    Random,
}

#[derive(Clone, Debug)]
pub struct Parameter<T: Ptype> {
    /// initialization `Style` of the parameter
    pub style: Style,
    /// whether the parameter should be regenerated on each use
    pub regenerate: bool,
    /// current value of the `Parameter`
    pub value: T,
}

impl<T: Default + Ptype> Default for Parameter<T> {
    fn default() -> Parameter<T> {
        Parameter {
            style: Style::Static,
            regenerate: false,
            value: T::default(),
        }
    }
}

impl<T: Ptype> Parameter<T> {
    /// Mutate internal value if necessary
    ///
    /// If the parameter is flagged to not regenerate or was seeded, it isn't regenerated.
    pub fn regen(&mut self) {
        if self.regenerate && self.style == Style::Random {
            self.value.regen()
        }
    }
}

/// Extract a `Parameter` from the toml tree
pub fn extract_parameter<T: Ptype>(
    index: usize,
    parameter: &BTreeMap<String, Value>,
) -> CResult<Parameter<T>> {
    let style = match parameter.get("style").and_then(|k| k.as_str()) {
        Some("random") => Style::Random,
        Some("static") | None => Style::Static,
        Some(other) => return Err(format!("bad parameter style: {}", other)),
    };

    let seed = parameter
        .get("seed")
        .and_then(|k| k.as_integer())
        .map_or(index, |i| i as usize);

    let size = parameter
        .get("size")
        .and_then(|k| k.as_integer())
        .map_or(1, |i| i as usize);

    let num = parameter
        .get("num")
        .and_then(|k| k.as_integer())
        .map_or(0, |i| i as u64);

    // check that size is sufficiently large to contain num strings
    if num > 0 {
        let fmtlen = format!("{}", num - 1).len();
        if  fmtlen > size {
            return Err(format!("To contain {} strings, you need to specify a \
                                size >= {}", num, fmtlen))
        }
    }

    let regenerate = parameter
        .get("regenerate")
        .and_then(|k| k.as_bool())
        .unwrap_or(false);

    let mut value = T::parse(seed, size, num, parameter)?;

    // initialize with a random value if that is what is needed
    if style == Style::Random {
        value.regen();
    }

    Ok(Parameter {
        style: style,
        regenerate: regenerate,
        value: value,
    })
}
<|endoftext|>


<commit_before>extern crate env_logger;
extern crate glob;
extern crate log;

extern crate protobuf_codegen_pure;

extern crate protobuf_test_common;

use std::io::Read;
use std::io::Write;

use std::env;
use std::fs;
use std::path::Path;

use protobuf_test_common::build::*;
use protobuf_test_common::print_rerun_if_changed_recursively;

fn copy_test<P1: AsRef<Path>, P2: AsRef<Path>>(src: P1, dst: P2) {
    eprintln!("copy {:?} to {:?}", src.as_ref(), dst.as_ref());
    let mut content = Vec::new();
    fs::File::open(src.as_ref())
        .expect(&format!("open {}", src.as_ref().display()))
        .read_to_end(&mut content)
        .expect(&format!("read_to_end {}", src.as_ref().display()));

    let mut write = fs::File::create(dst).expect("create");
    writeln!(write, "// @generated").expect("write");
    writeln!(write, "// copied from {}", src.as_ref().display()).expect("write");
    writeln!(write, "").expect("write");
    write.write_all(&content).expect("write_all");
    writeln!(write, "// @generated").expect("write");
    write.flush().expect("flush");
}

fn copy_from_protobuf_test(path: &str) {
    copy_test(&format!("../protobuf-test/{}", path), &format!("{}", path))
}

enum FileNameClass {
    ModRs,
    TestRs,
    Proto,
    GeneratedRs,
    Ignore,
}

fn classify_file_name(dir: &str, name: &str) -> FileNameClass {
    if name.starts_with(".") || name.ends_with(".md") || name.ends_with(".sh") {
        FileNameClass::Ignore
    } else if name.ends_with("_pb.rs") || name.ends_with("_pb_proto3.rs") {
        FileNameClass::GeneratedRs
    } else if name == "mod.rs" {
        FileNameClass::ModRs
    } else if name.ends_with(".proto") || name.ends_with(".proto3") {
        FileNameClass::Proto
    } else if name.ends_with(".rs") {
        if dir == "src/google/protobuf" {
            FileNameClass::GeneratedRs
        } else {
            FileNameClass::TestRs
        }
    } else {
        panic!("unknown test file: {}", name);
    }
}

// Copy tests from `protobuf-test` directory to the same directory here
fn copy_tests(dir: &str) {
    let src_dir = format!("../protobuf-test/{}", dir);
    for entry in fs::read_dir(&src_dir).expect(&format!("read_dir {}", src_dir)) {
        let file_name = entry.expect("entry").file_name().into_string().unwrap();

        match classify_file_name(dir, &file_name) {
            FileNameClass::ModRs | FileNameClass::Ignore | FileNameClass::GeneratedRs => {}
            FileNameClass::TestRs | FileNameClass::Proto => {
                copy_from_protobuf_test(&format!("{}/{}", dir, file_name))
            }
        }
    }
}

fn gen_in_dir(dir: &str, include_dir: &str) {
    gen_in_dir_impl(
        dir,
        |GenInDirArgs {
             out_dir,
             input,
             customize,
         }| {
            protobuf_codegen_pure::Codegen::new()
                .out_dir(out_dir)
                .inputs(input)
                .includes(&[include_dir])
                .customize(customize)
                .run()
        },
    );
}

fn generate_interop() {
    copy_from_protobuf_test("src/interop/mod.rs");
    copy_from_protobuf_test("src/interop/json.rs");

    protobuf_codegen_pure::Codegen::new()
        .out_dir("src/interop")
        .includes(&["../interop/cxx", "../proto"])
        .input("../interop/cxx/interop_pb.proto")
        .run()
        .unwrap();
}

fn generate_include_generated() {
    copy_from_protobuf_test("src/include_generated/mod.rs");

    let dir = format!("{}/include_generated", env::var("OUT_DIR").unwrap());
    if Path::new(&dir).exists() {
        fs::remove_dir_all(&dir).unwrap();
    }
    fs::create_dir(&dir).unwrap();
    protobuf_codegen_pure::Codegen::new()
        .out_dir(dir)
        .input("../protobuf-test/src/include_generated/v2.proto")
        .input("../protobuf-test/src/include_generated/v3.proto")
        .customize(Customize {
            gen_mod_rs: Some(true),
            ..Default::default()
        })
        .include("../protobuf-test/src/include_generated")
        .run()
        .unwrap();
}

fn generate_pb_rs() {
    print_rerun_if_changed_recursively("../protobuf-test");

    copy_tests("src/v2");
    gen_in_dir("src/v2", "src/v2");

    copy_tests("src/v3");
    gen_in_dir("src/v3", "src/v3");

    copy_tests("src/common/v2");
    gen_in_dir("src/common/v2", "src/common/v2");

    copy_tests_v2_v3("src/common/v2", "src/common/v3");
    gen_in_dir("src/common/v3", "src/common/v3");

    copy_tests("src/google/protobuf");
    gen_in_dir("src/google/protobuf", "src");

    generate_interop();

    generate_include_generated();
}

fn main() {
    env_logger::init();

    cfg_serde();

    clean_old_files();
    generate_pb_rs();
}
<commit_msg>Explain why generated is printed twice<commit_after>extern crate env_logger;
extern crate glob;
extern crate log;

extern crate protobuf_codegen_pure;

extern crate protobuf_test_common;

use std::io::Read;
use std::io::Write;

use std::env;
use std::fs;
use std::path::Path;

use protobuf_test_common::build::*;
use protobuf_test_common::print_rerun_if_changed_recursively;

fn copy_test<P1: AsRef<Path>, P2: AsRef<Path>>(src: P1, dst: P2) {
    eprintln!("copy {:?} to {:?}", src.as_ref(), dst.as_ref());
    let mut content = Vec::new();
    fs::File::open(src.as_ref())
        .expect(&format!("open {}", src.as_ref().display()))
        .read_to_end(&mut content)
        .expect(&format!("read_to_end {}", src.as_ref().display()));

    let mut write = fs::File::create(dst).expect("create");
    writeln!(write, "// @generated").expect("write");
    writeln!(write, "// copied from {}", src.as_ref().display()).expect("write");
    writeln!(write, "").expect("write");
    write.write_all(&content).expect("write_all");
    // Print generated twice to avoid overlooking it accidentally
    writeln!(write, "// @generated").expect("write");
    write.flush().expect("flush");
}

fn copy_from_protobuf_test(path: &str) {
    copy_test(&format!("../protobuf-test/{}", path), &format!("{}", path))
}

enum FileNameClass {
    ModRs,
    TestRs,
    Proto,
    GeneratedRs,
    Ignore,
}

fn classify_file_name(dir: &str, name: &str) -> FileNameClass {
    if name.starts_with(".") || name.ends_with(".md") || name.ends_with(".sh") {
        FileNameClass::Ignore
    } else if name.ends_with("_pb.rs") || name.ends_with("_pb_proto3.rs") {
        FileNameClass::GeneratedRs
    } else if name == "mod.rs" {
        FileNameClass::ModRs
    } else if name.ends_with(".proto") || name.ends_with(".proto3") {
        FileNameClass::Proto
    } else if name.ends_with(".rs") {
        if dir == "src/google/protobuf" {
            FileNameClass::GeneratedRs
        } else {
            FileNameClass::TestRs
        }
    } else {
        panic!("unknown test file: {}", name);
    }
}

// Copy tests from `protobuf-test` directory to the same directory here
fn copy_tests(dir: &str) {
    let src_dir = format!("../protobuf-test/{}", dir);
    for entry in fs::read_dir(&src_dir).expect(&format!("read_dir {}", src_dir)) {
        let file_name = entry.expect("entry").file_name().into_string().unwrap();

        match classify_file_name(dir, &file_name) {
            FileNameClass::ModRs | FileNameClass::Ignore | FileNameClass::GeneratedRs => {}
            FileNameClass::TestRs | FileNameClass::Proto => {
                copy_from_protobuf_test(&format!("{}/{}", dir, file_name))
            }
        }
    }
}

fn gen_in_dir(dir: &str, include_dir: &str) {
    gen_in_dir_impl(
        dir,
        |GenInDirArgs {
             out_dir,
             input,
             customize,
         }| {
            protobuf_codegen_pure::Codegen::new()
                .out_dir(out_dir)
                .inputs(input)
                .includes(&[include_dir])
                .customize(customize)
                .run()
        },
    );
}

fn generate_interop() {
    copy_from_protobuf_test("src/interop/mod.rs");
    copy_from_protobuf_test("src/interop/json.rs");

    protobuf_codegen_pure::Codegen::new()
        .out_dir("src/interop")
        .includes(&["../interop/cxx", "../proto"])
        .input("../interop/cxx/interop_pb.proto")
        .run()
        .unwrap();
}

fn generate_include_generated() {
    copy_from_protobuf_test("src/include_generated/mod.rs");

    let dir = format!("{}/include_generated", env::var("OUT_DIR").unwrap());
    if Path::new(&dir).exists() {
        fs::remove_dir_all(&dir).unwrap();
    }
    fs::create_dir(&dir).unwrap();
    protobuf_codegen_pure::Codegen::new()
        .out_dir(dir)
        .input("../protobuf-test/src/include_generated/v2.proto")
        .input("../protobuf-test/src/include_generated/v3.proto")
        .customize(Customize {
            gen_mod_rs: Some(true),
            ..Default::default()
        })
        .include("../protobuf-test/src/include_generated")
        .run()
        .unwrap();
}

fn generate_pb_rs() {
    print_rerun_if_changed_recursively("../protobuf-test");

    copy_tests("src/v2");
    gen_in_dir("src/v2", "src/v2");

    copy_tests("src/v3");
    gen_in_dir("src/v3", "src/v3");

    copy_tests("src/common/v2");
    gen_in_dir("src/common/v2", "src/common/v2");

    copy_tests_v2_v3("src/common/v2", "src/common/v3");
    gen_in_dir("src/common/v3", "src/common/v3");

    copy_tests("src/google/protobuf");
    gen_in_dir("src/google/protobuf", "src");

    generate_interop();

    generate_include_generated();
}

fn main() {
    env_logger::init();

    cfg_serde();

    clean_old_files();
    generate_pb_rs();
}
<|endoftext|>


<commit_before>//! Collection types which allows insertion of new values while shared references to its contents
//! are alive. This is done by storing each value in a stable memory location and preventing an
//! earlier inserted value to be overwritten.
use std::{
    borrow::Borrow,
    cell::RefCell,
    collections::hash_map::Entry,
    fmt,
    hash::Hash,
    iter::IntoIterator,
    mem,
    ops::{Index, IndexMut},
};

use crate::fnv::FnvMap;
use vec_map::VecMap;

// NOTE: transmute is used to circumvent the borrow checker in this module
// This is safe since the containers hold boxed values meaning allocating larger
// storage does not invalidate the references that are handed out and because values
// can only be inserted, never modified (this could be done with a Rc pointer as well but
// is not done for efficiency since it is not needed)

unsafe fn forget_lifetime<'a, 'b, T: ?Sized>(x: &'a T) -> &'b T {
    ::std::mem::transmute(x)
}

// A mapping between K and V where once a value has been inserted it cannot be changed
// Through this and the fact the all values are stored as pointers it is possible to safely
// insert new values without invalidating pointers retrieved from it
pub struct FixedMap<K, V> {
    map: RefCell<FnvMap<K, (u32, u32)>>,
    values: Buffer<V>,
}

impl<K: Eq + Hash, V> Default for FixedMap<K, V> {
    fn default() -> FixedMap<K, V> {
        FixedMap::new()
    }
}

impl<K: Eq + Hash + fmt::Debug, V: fmt::Debug> fmt::Debug for FixedMap<K, V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.map.borrow().fmt(f)
    }
}

impl<K: Eq + Hash, V> FixedMap<K, V> {
    pub fn new() -> FixedMap<K, V> {
        FixedMap {
            map: RefCell::new(FnvMap::default()),
            values: Default::default(),
        }
    }

    pub fn clear(&mut self) {
        self.map.borrow_mut().clear();
    }

    pub fn len(&self) -> usize {
        self.map.borrow().len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl<K: Eq + Hash, V> FixedMap<K, V> {
    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        match self.map.get_mut().entry(key) {
            Entry::Occupied(entry) => Some(mem::replace(&mut self.values[*entry.get()], value)),
            Entry::Vacant(entry) => {
                entry.insert(self.values.push(value));
                None
            }
        }
    }

    pub fn try_insert(&self, key: K, value: V) -> Result<(), (K, V)> {
        if self.get(&key).is_some() {
            Err((key, value))
        } else {
            let index_value = self.values.push(value);
            self.map.borrow_mut().insert(key, index_value);
            Ok(())
        }
    }

    pub fn get<Q>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: ?Sized + Eq + Hash,
    {
        self.map.borrow().get(k).map(|key| &self.values[*key])
    }

    pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        let values = &mut self.values;
        self.map.get_mut().get(k).map(move |&key| &mut values[key])
    }
}

impl<'a, Q, K, V> Index<&'a Q> for FixedMap<K, V>
where
    K: Eq + Hash + Borrow<Q>,
    Q: ?Sized + Eq + Hash,
{
    type Output = V;
    fn index(&self, index: &'a Q) -> &Self::Output {
        self.get(index).expect("Index out of bounds")
    }
}

// A mapping between K and V where once a value has been inserted it cannot be changed
// Through this and the fact the all values are stored as pointers it is possible to safely
// insert new values without invalidating pointers retrieved from it
pub struct FixedVecMap<V> {
    // Use u16 to leave space for the `Option` tag in `VecMap`
    map: RefCell<VecMap<(u16, u32)>>,
    values: Buffer<V>,
}

impl<V> Default for FixedVecMap<V> {
    fn default() -> FixedVecMap<V> {
        FixedVecMap::new()
    }
}

impl<V: fmt::Debug> fmt::Debug for FixedVecMap<V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.map.borrow().fmt(f)
    }
}

impl<V> FixedVecMap<V> {
    pub fn new() -> FixedVecMap<V> {
        FixedVecMap {
            map: Default::default(),
            values: Default::default(),
        }
    }

    pub fn clear(&mut self) {
        self.map.borrow_mut().clear();
    }

    pub fn len(&self) -> usize {
        self.map.borrow().len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl<V> FixedVecMap<V> {
    pub fn insert(&mut self, key: usize, value: V) -> Option<V> {
        use vec_map::Entry;
        match self.map.get_mut().entry(key) {
            Entry::Occupied(entry) => Some(mem::replace(&mut self.values[*entry.get()], value)),
            Entry::Vacant(entry) => {
                let (i, j) = self.values.push(value);
                entry.insert((i as u16, j));
                None
            }
        }
    }

    pub fn try_insert(&self, key: usize, value: V) -> Result<(), (usize, V)> {
        if self.get(key).is_some() {
            Err((key, value))
        } else {
            let (i, j) = self.values.push(value);
            self.map.borrow_mut().insert(key, (i as u16, j));
            Ok(())
        }
    }

    pub fn get(&self, k: usize) -> Option<&V> {
        self.map.borrow().get(k).map(|key| &self.values[*key])
    }

    pub fn get_mut(&mut self, k: usize) -> Option<&mut V> {
        let values = &mut self.values;
        self.map.get_mut().get(k).map(move |&key| &mut values[key])
    }

    pub fn remove(&mut self, k: usize) -> Option<V> {
        let values = self.values.values.get_mut();
        self.map.get_mut().remove(k).and_then(|(i, j)| {
            if values.len() == i as usize + 1 && values[i as usize].len() == j as usize + 1 {
                let x = values[i as usize].pop();
                if values[i as usize].is_empty() {
                    values.pop();
                }
                x
            } else {
                unimplemented!()
            }
        })
    }

    pub fn truncate(&mut self, index: usize) {
        self.map.get_mut().retain(|i, _| i < index);
        self.values.truncate(index);
    }

    pub fn drain<'a>(&'a mut self) -> impl Iterator<Item = V> + 'a {
        self.map.get_mut().clear();
        self.values.drain()
    }
}

impl<V> Index<usize> for FixedVecMap<V> {
    type Output = V;
    fn index(&self, index: usize) -> &Self::Output {
        self.get(index).expect("Index out of bounds")
    }
}

#[derive(Debug)]
pub struct FixedVec<T> {
    vec: RefCell<Vec<(u32, u32)>>,
    values: Buffer<T>,
}

impl<T> Default for FixedVec<T> {
    fn default() -> Self {
        FixedVec::new()
    }
}

impl<T> FixedVec<T> {
    pub fn new() -> FixedVec<T> {
        FixedVec {
            vec: RefCell::new(Vec::new()),
            values: Default::default(),
        }
    }

    pub fn clear(&mut self) {
        self.vec.borrow_mut().clear();
    }

    pub fn push(&self, value: T) {
        let key = self.values.push(value);
        self.vec.borrow_mut().push(key);
    }

    pub fn extend<I: IntoIterator<Item = T>>(&self, iter: I) {
        for item in iter {
            self.push(item);
        }
    }

    pub fn len(&self) -> usize {
        self.vec.borrow().len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn iter(&self) -> impl Iterator<Item = &T> {
        (0..).scan((), move |_, i| self.get(i))
    }

    pub fn get(&self, index: usize) -> Option<&T> {
        if index < self.len() {
            Some(&self[index])
        } else {
            None
        }
    }

    pub fn truncate(&mut self, index: usize) {
        self.vec.get_mut().truncate(index);
        self.values.truncate(index)
    }

    pub fn pop(&mut self) -> Option<T> {
        if self.vec.get_mut().pop().is_some() {
            self.values.pop()
        } else {
            None
        }
    }
}

impl<T> Index<usize> for FixedVec<T> {
    type Output = T;
    fn index(&self, index: usize) -> &T {
        let vec = self.vec.borrow();
        &self.values[vec[index]]
    }
}

impl<T> IndexMut<usize> for FixedVec<T> {
    fn index_mut(&mut self, index: usize) -> &mut T {
        let vec = self.vec.borrow();
        &mut self.values[vec[index]]
    }
}

#[derive(Debug)]
struct Buffer<T> {
    values: RefCell<Vec<Vec<T>>>,
}

impl<T> Default for Buffer<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Buffer<T> {
    pub fn new() -> Self {
        Self {
            values: Default::default(),
        }
    }

    fn total_len(&self) -> usize {
        self.values
            .borrow()
            .iter()
            .map(|vec| vec.len())
            .sum::<usize>()
    }

    fn push(&self, value: T) -> (u32, u32) {
        let mut values = self.values.borrow_mut();
        let cap = match values.last() {
            Some(vec) => {
                if vec.len() == vec.capacity() {
                    Some(vec.capacity() * 2)
                } else {
                    None
                }
            }
            None => Some(4),
        };
        if let Some(cap) = cap {
            values.push(Vec::with_capacity(cap));
        }

        let i = values.len() as u32 - 1;
        let inner = values.last_mut().unwrap();
        let j = inner.len() as u32;
        inner.push(value);
        (i, j)
    }

    fn truncate(&mut self, index: usize) {
        let mut left = self.total_len() - index;
        let values = self.values.get_mut();
        while left != 0 {
            let inner = values.last_mut().expect("Not values left");
            if inner.len() <= left {
                left -= inner.len();
                values.pop();
            } else {
                let i = inner.len() - left;
                inner.truncate(i);
                left = 0;
            }
        }
    }

    fn drain<'a>(&'a mut self) -> impl Iterator<Item = T> + 'a {
        self.values.get_mut().drain(..).flat_map(|vec| vec)
    }

    fn pop(&mut self) -> Option<T> {
        let values = self.values.get_mut();
        let out = values.last_mut().and_then(|vec| vec.pop());
        if out.is_some() {
            out
        } else {
            values.pop();
            values.last_mut().and_then(|vec| vec.pop())
        }
    }
}

impl<T> Index<(u32, u32)> for Buffer<T> {
    type Output = T;
    fn index(&self, (i, j): (u32, u32)) -> &T {
        unsafe {
            forget_lifetime(
                &self
                    .values
                    .borrow()
                    .get(i as usize)
                    .and_then(|v| v.get(j as usize))
                    .unwrap_or_else(|| panic!("Index out of bounds: {:?}", (i, j))),
            )
        }
    }
}

impl<T> IndexMut<(u32, u32)> for Buffer<T> {
    fn index_mut(&mut self, (i, j): (u32, u32)) -> &mut T {
        self.values
            .get_mut()
            .get_mut(i as usize)
            .and_then(|v| v.get_mut(j as usize))
            .unwrap_or_else(|| panic!("Index out of bounds: {:?}", (i, j)))
    }
}

impl<T> Index<(u16, u32)> for Buffer<T> {
    type Output = T;
    fn index(&self, (i, j): (u16, u32)) -> &T {
        &self[(i as u32, j)]
    }
}

impl<T> IndexMut<(u16, u32)> for Buffer<T> {
    fn index_mut(&mut self, (i, j): (u16, u32)) -> &mut T {
        &mut self[(i as u32, j)]
    }
}

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

    #[test]
    fn truncate_buffer() {
        let mut buffer = Buffer::new();
        for i in 0..10 {
            buffer.push(i);
        }
        buffer.truncate(6);
        assert_eq!(buffer.total_len(), 6);
        buffer.truncate(2);
        assert_eq!(buffer.total_len(), 2);
    }
}
<commit_msg>perf: Re-use the value buffers used in fixed structures<commit_after>//! Collection types which allows insertion of new values while shared references to its contents
//! are alive. This is done by storing each value in a stable memory location and preventing an
//! earlier inserted value to be overwritten.
use std::{
    borrow::Borrow,
    cell::RefCell,
    collections::hash_map::Entry,
    fmt,
    hash::Hash,
    iter::IntoIterator,
    mem,
    ops::{Index, IndexMut},
};

use crate::fnv::FnvMap;
use vec_map::VecMap;

// NOTE: transmute is used to circumvent the borrow checker in this module
// This is safe since the containers hold boxed values meaning allocating larger
// storage does not invalidate the references that are handed out and because values
// can only be inserted, never modified (this could be done with a Rc pointer as well but
// is not done for efficiency since it is not needed)

unsafe fn forget_lifetime<'a, 'b, T: ?Sized>(x: &'a T) -> &'b T {
    ::std::mem::transmute(x)
}

// A mapping between K and V where once a value has been inserted it cannot be changed
// Through this and the fact the all values are stored as pointers it is possible to safely
// insert new values without invalidating pointers retrieved from it
pub struct FixedMap<K, V> {
    map: RefCell<FnvMap<K, (u32, u32)>>,
    values: Buffer<V>,
}

impl<K: Eq + Hash, V> Default for FixedMap<K, V> {
    fn default() -> FixedMap<K, V> {
        FixedMap::new()
    }
}

impl<K: Eq + Hash + fmt::Debug, V: fmt::Debug> fmt::Debug for FixedMap<K, V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.map.borrow().fmt(f)
    }
}

impl<K: Eq + Hash, V> FixedMap<K, V> {
    pub fn new() -> FixedMap<K, V> {
        FixedMap {
            map: RefCell::new(FnvMap::default()),
            values: Default::default(),
        }
    }

    pub fn clear(&mut self) {
        self.map.borrow_mut().clear();
    }

    pub fn len(&self) -> usize {
        self.map.borrow().len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl<K: Eq + Hash, V> FixedMap<K, V> {
    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
        match self.map.get_mut().entry(key) {
            Entry::Occupied(entry) => Some(mem::replace(&mut self.values[*entry.get()], value)),
            Entry::Vacant(entry) => {
                entry.insert(self.values.push(value));
                None
            }
        }
    }

    pub fn try_insert(&self, key: K, value: V) -> Result<(), (K, V)> {
        if self.get(&key).is_some() {
            Err((key, value))
        } else {
            let index_value = self.values.push(value);
            self.map.borrow_mut().insert(key, index_value);
            Ok(())
        }
    }

    pub fn get<Q>(&self, k: &Q) -> Option<&V>
    where
        K: Borrow<Q>,
        Q: ?Sized + Eq + Hash,
    {
        self.map.borrow().get(k).map(|key| &self.values[*key])
    }

    pub fn get_mut<Q>(&mut self, k: &Q) -> Option<&mut V>
    where
        K: Borrow<Q>,
        Q: Eq + Hash,
    {
        let values = &mut self.values;
        self.map.get_mut().get(k).map(move |&key| &mut values[key])
    }
}

impl<'a, Q, K, V> Index<&'a Q> for FixedMap<K, V>
where
    K: Eq + Hash + Borrow<Q>,
    Q: ?Sized + Eq + Hash,
{
    type Output = V;
    fn index(&self, index: &'a Q) -> &Self::Output {
        self.get(index).expect("Index out of bounds")
    }
}

// A mapping between K and V where once a value has been inserted it cannot be changed
// Through this and the fact the all values are stored as pointers it is possible to safely
// insert new values without invalidating pointers retrieved from it
pub struct FixedVecMap<V> {
    // Use u16 to leave space for the `Option` tag in `VecMap`
    map: RefCell<VecMap<(u16, u32)>>,
    values: Buffer<V>,
}

impl<V> Default for FixedVecMap<V> {
    fn default() -> FixedVecMap<V> {
        FixedVecMap::new()
    }
}

impl<V: fmt::Debug> fmt::Debug for FixedVecMap<V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.map.borrow().fmt(f)
    }
}

impl<V> FixedVecMap<V> {
    pub fn new() -> FixedVecMap<V> {
        FixedVecMap {
            map: Default::default(),
            values: Default::default(),
        }
    }

    pub fn clear(&mut self) {
        self.map.borrow_mut().clear();
    }

    pub fn len(&self) -> usize {
        self.map.borrow().len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

impl<V> FixedVecMap<V> {
    pub fn insert(&mut self, key: usize, value: V) -> Option<V> {
        use vec_map::Entry;
        match self.map.get_mut().entry(key) {
            Entry::Occupied(entry) => Some(mem::replace(&mut self.values[*entry.get()], value)),
            Entry::Vacant(entry) => {
                let (i, j) = self.values.push(value);
                entry.insert((i as u16, j));
                None
            }
        }
    }

    pub fn try_insert(&self, key: usize, value: V) -> Result<(), (usize, V)> {
        if self.get(key).is_some() {
            Err((key, value))
        } else {
            let (i, j) = self.values.push(value);
            self.map.borrow_mut().insert(key, (i as u16, j));
            Ok(())
        }
    }

    pub fn get(&self, k: usize) -> Option<&V> {
        self.map.borrow().get(k).map(|key| &self.values[*key])
    }

    pub fn get_mut(&mut self, k: usize) -> Option<&mut V> {
        let values = &mut self.values;
        self.map.get_mut().get(k).map(move |&key| &mut values[key])
    }

    pub fn truncate(&mut self, index: usize) {
        self.map.get_mut().retain(|i, _| i < index);
        self.values.truncate(index);
    }

    pub fn drain<'a>(&'a mut self) -> impl Iterator<Item = V> + 'a {
        self.map.get_mut().clear();
        self.values.drain()
    }
}

impl<V> Index<usize> for FixedVecMap<V> {
    type Output = V;
    fn index(&self, index: usize) -> &Self::Output {
        self.get(index).expect("Index out of bounds")
    }
}

#[derive(Debug)]
pub struct FixedVec<T> {
    vec: RefCell<Vec<(u32, u32)>>,
    values: Buffer<T>,
}

impl<T> Default for FixedVec<T> {
    fn default() -> Self {
        FixedVec::new()
    }
}

impl<T> FixedVec<T> {
    pub fn new() -> FixedVec<T> {
        FixedVec {
            vec: RefCell::new(Vec::new()),
            values: Default::default(),
        }
    }

    pub fn clear(&mut self) {
        self.vec.borrow_mut().clear();
    }

    pub fn push(&self, value: T) {
        let key = self.values.push(value);
        self.vec.borrow_mut().push(key);
    }

    pub fn extend<I: IntoIterator<Item = T>>(&self, iter: I) {
        for item in iter {
            self.push(item);
        }
    }

    pub fn len(&self) -> usize {
        self.vec.borrow().len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn iter(&self) -> impl Iterator<Item = &T> {
        (0..).scan((), move |_, i| self.get(i))
    }

    pub fn get(&self, index: usize) -> Option<&T> {
        if index < self.len() {
            Some(&self[index])
        } else {
            None
        }
    }

    pub fn truncate(&mut self, index: usize) {
        self.vec.get_mut().truncate(index);
        self.values.truncate(index)
    }

    pub fn pop(&mut self) -> Option<T> {
        if self.vec.get_mut().pop().is_some() {
            self.values.pop()
        } else {
            None
        }
    }
}

impl<T> Index<usize> for FixedVec<T> {
    type Output = T;
    fn index(&self, index: usize) -> &T {
        let vec = self.vec.borrow();
        &self.values[vec[index]]
    }
}

impl<T> IndexMut<usize> for FixedVec<T> {
    fn index_mut(&mut self, index: usize) -> &mut T {
        let vec = self.vec.borrow();
        &mut self.values[vec[index]]
    }
}

#[derive(Debug)]
struct BufferInner<T> {
    values: Vec<Vec<T>>,
    current: usize,
}

#[derive(Debug)]
struct Buffer<T> {
    values: RefCell<BufferInner<T>>,
}

impl<T> Default for BufferInner<T> {
    fn default() -> Self {
        BufferInner {
            values: Vec::new(),
            current: 0,
        }
    }
}

impl<T> Default for Buffer<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> Buffer<T> {
    pub fn new() -> Self {
        Self {
            values: Default::default(),
        }
    }

    fn total_len(&self) -> usize {
        let values = self.values.borrow();
        values.values.iter().map(|vec| vec.len()).sum::<usize>()
    }

    fn push(&self, value: T) -> (u32, u32) {
        let mut values = self.values.borrow_mut();
        let cap = match values.current().map(|vec| (vec.len(), vec.capacity())) {
            Some((len, capacity)) => {
                if len == capacity {
                    values.current += 1;
                    if values.current + 1 < values.values.len() {
                        None
                    } else {
                        Some(capacity * 2)
                    }
                } else {
                    None
                }
            }
            None => Some(4),
        };
        if let Some(cap) = cap {
            values.values.push(Vec::with_capacity(cap));
        }

        let i = values.current as u32;
        let inner = values.current_mut().unwrap();
        let j = inner.len() as u32;
        inner.push(value);
        (i, j)
    }

    fn truncate(&mut self, index: usize) {
        let mut left = self.total_len() - index;
        let values = self.values.get_mut();
        while left != 0 {
            let inner = values.current_mut().expect("No values left");
            if inner.len() < left {
                left -= inner.len();
                inner.clear();
                values.current -= 1;
            } else {
                let i = inner.len() - left;
                inner.truncate(i);
                left = 0;
            }
        }
    }

    fn drain<'a>(&'a mut self) -> impl Iterator<Item = T> + 'a {
        let values = self.values.get_mut();
        values.current = 0;
        values.values.iter_mut().flat_map(|vec| vec.drain(..))
    }

    fn pop(&mut self) -> Option<T> {
        let values = self.values.get_mut();
        let out = values.current_mut().and_then(|vec| vec.pop());
        if out.is_some() {
            out
        } else {
            if values.current != 0 {
                values.current -= 1;
            }
            values.current_mut().and_then(|vec| vec.pop())
        }
    }
}

impl<T> BufferInner<T> {
    fn current(&self) -> Option<&Vec<T>> {
        self.values.get(self.current)
    }
    fn current_mut(&mut self) -> Option<&mut Vec<T>> {
        self.values.get_mut(self.current)
    }
}

impl<T> Index<(u32, u32)> for Buffer<T> {
    type Output = T;
    fn index(&self, (i, j): (u32, u32)) -> &T {
        unsafe {
            forget_lifetime(
                &self
                    .values
                    .borrow()
                    .values
                    .get(i as usize)
                    .and_then(|v| v.get(j as usize))
                    .unwrap_or_else(|| panic!("Index out of bounds: {:?}", (i, j))),
            )
        }
    }
}

impl<T> IndexMut<(u32, u32)> for Buffer<T> {
    fn index_mut(&mut self, (i, j): (u32, u32)) -> &mut T {
        self.values
            .get_mut()
            .values
            .get_mut(i as usize)
            .and_then(|v| v.get_mut(j as usize))
            .unwrap_or_else(|| panic!("Index out of bounds: {:?}", (i, j)))
    }
}

impl<T> Index<(u16, u32)> for Buffer<T> {
    type Output = T;
    fn index(&self, (i, j): (u16, u32)) -> &T {
        &self[(i as u32, j)]
    }
}

impl<T> IndexMut<(u16, u32)> for Buffer<T> {
    fn index_mut(&mut self, (i, j): (u16, u32)) -> &mut T {
        &mut self[(i as u32, j)]
    }
}

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

    #[test]
    fn truncate_buffer() {
        let mut buffer = Buffer::new();
        for i in 0..10 {
            buffer.push(i);
        }
        buffer.truncate(6);
        assert_eq!(buffer.total_len(), 6);
        buffer.truncate(2);
        assert_eq!(buffer.total_len(), 2);
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

//! Tests client/server control flow.

#![feature(generators, await_macro, async_await, futures_api)]

use futures::{
    compat::{Executor01CompatExt, Future01CompatExt},
    prelude::*,
    stream,
};
use log::{info, trace};
use rand::distributions::{Distribution, Normal};
use rpc::{client, context, server::Server};
use std::{
    io,
    time::{Duration, Instant, SystemTime},
};
use tokio::timer::Delay;

pub trait AsDuration {
    /// Delay of 0 if self is in the past
    fn as_duration(&self) -> Duration;
}

impl AsDuration for SystemTime {
    fn as_duration(&self) -> Duration {
        self.duration_since(SystemTime::now()).unwrap_or_default()
    }
}

async fn run() -> io::Result<()> {
    let listener = tarpc_bincode_transport::listen(&"0.0.0.0:0".parse().unwrap())?;
    let addr = listener.local_addr();
    let server = Server::<String, String>::default()
        .incoming(listener)
        .take(1)
        .for_each(async move |channel| {
            let channel = if let Ok(channel) = channel {
                channel
            } else {
                return;
            };
            let client_addr = *channel.client_addr();
            let handler = channel.respond_with(move |ctx, request| {
                // Sleep for a time sampled from a normal distribution with:
                // - mean: 1/2 the deadline.
                // - std dev: 1/2 the deadline.
                let deadline: Duration = ctx.deadline.as_duration();
                let deadline_millis = deadline.as_secs() * 1000 + deadline.subsec_millis() as u64;
                let distribution =
                    Normal::new(deadline_millis as f64 / 2., deadline_millis as f64 / 2.);
                let delay_millis = distribution.sample(&mut rand::thread_rng()).max(0.);
                let delay = Duration::from_millis(delay_millis as u64);

                trace!(
                    "[{}/{}] Responding to request in {:?}.",
                    ctx.trace_id(),
                    client_addr,
                    delay,
                );

                let wait = Delay::new(Instant::now() + delay).compat();
                async move {
                    await!(wait).unwrap();
                    Ok(request)
                }
            });
            tokio_executor::spawn(handler.unit_error().boxed().compat());
        });

    tokio_executor::spawn(server.unit_error().boxed().compat());

    let conn = await!(tarpc_bincode_transport::connect(&addr))?;
    let client = await!(client::new::<String, String, _>(
        client::Config::default(),
        conn
    ))?;

    // Proxy service
    let listener = tarpc_bincode_transport::listen(&"0.0.0.0:0".parse().unwrap())?;
    let addr = listener.local_addr();
    let proxy_server = Server::<String, String>::default()
        .incoming(listener)
        .take(1)
        .for_each(move |channel| {
            let client = client.clone();
            async move {
                let channel = if let Ok(channel) = channel {
                    channel
                } else {
                    return;
                };
                let client_addr = *channel.client_addr();
                let handler = channel.respond_with(move |ctx, request| {
                    trace!("[{}/{}] Proxying request.", ctx.trace_id(), client_addr);
                    let mut client = client.clone();
                    async move { await!(client.call(ctx, request)) }
                });
                tokio_executor::spawn(handler.unit_error().boxed().compat());
            }
        });

    tokio_executor::spawn(proxy_server.unit_error().boxed().compat());

    let mut config = client::Config::default();
    config.max_in_flight_requests = 10;
    config.pending_request_buffer = 10;

    let client = await!(client::new::<String, String, _>(
        config,
        await!(tarpc_bincode_transport::connect(&addr))?
    ))?;

    // Make 3 speculative requests, returning only the quickest.
    let mut clients: Vec<_> = (1..=3u32).map(|_| client.clone()).collect();
    let mut requests = vec![];
    for client in &mut clients {
        let mut ctx = context::current();
        ctx.deadline = SystemTime::now() + Duration::from_millis(200);
        let trace_id = *ctx.trace_id();
        let response = client.call(ctx, "ping".into());
        requests.push(response.map(move |r| (trace_id, r)));
    }
    let (fastest_response, _) = await!(stream::futures_unordered(requests).into_future());
    let (trace_id, resp) = fastest_response.unwrap();
    info!("[{}] fastest_response = {:?}", trace_id, resp);

    Ok::<_, io::Error>(())
}

#[test]
fn cancel_slower() -> io::Result<()> {
    env_logger::init();
    rpc::init(tokio::executor::DefaultExecutor::current().compat());

    tokio::run(run().boxed().map_err(|e| panic!(e)).compat());
    Ok(())
}
<commit_msg>Fix test client breakage by 9100ea46f997f24d4bc8c1764d0fe3ff8226ad2a<commit_after>// Copyright 2018 Google LLC
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

//! Tests client/server control flow.

#![feature(generators, await_macro, async_await, futures_api)]

use futures::{
    compat::{Executor01CompatExt, Future01CompatExt},
    prelude::*,
    stream::FuturesUnordered,
};
use log::{info, trace};
use rand::distributions::{Distribution, Normal};
use rpc::{client, context, server::Server};
use std::{
    io,
    time::{Duration, Instant, SystemTime},
};
use tokio::timer::Delay;

pub trait AsDuration {
    /// Delay of 0 if self is in the past
    fn as_duration(&self) -> Duration;
}

impl AsDuration for SystemTime {
    fn as_duration(&self) -> Duration {
        self.duration_since(SystemTime::now()).unwrap_or_default()
    }
}

async fn run() -> io::Result<()> {
    let listener = tarpc_bincode_transport::listen(&"0.0.0.0:0".parse().unwrap())?;
    let addr = listener.local_addr();
    let server = Server::<String, String>::default()
        .incoming(listener)
        .take(1)
        .for_each(async move |channel| {
            let channel = if let Ok(channel) = channel {
                channel
            } else {
                return;
            };
            let client_addr = *channel.client_addr();
            let handler = channel.respond_with(move |ctx, request| {
                // Sleep for a time sampled from a normal distribution with:
                // - mean: 1/2 the deadline.
                // - std dev: 1/2 the deadline.
                let deadline: Duration = ctx.deadline.as_duration();
                let deadline_millis = deadline.as_secs() * 1000 + deadline.subsec_millis() as u64;
                let distribution =
                    Normal::new(deadline_millis as f64 / 2., deadline_millis as f64 / 2.);
                let delay_millis = distribution.sample(&mut rand::thread_rng()).max(0.);
                let delay = Duration::from_millis(delay_millis as u64);

                trace!(
                    "[{}/{}] Responding to request in {:?}.",
                    ctx.trace_id(),
                    client_addr,
                    delay,
                );

                let wait = Delay::new(Instant::now() + delay).compat();
                async move {
                    await!(wait).unwrap();
                    Ok(request)
                }
            });
            tokio_executor::spawn(handler.unit_error().boxed().compat());
        });

    tokio_executor::spawn(server.unit_error().boxed().compat());

    let conn = await!(tarpc_bincode_transport::connect(&addr))?;
    let client = await!(client::new::<String, String, _>(
        client::Config::default(),
        conn
    ))?;

    // Proxy service
    let listener = tarpc_bincode_transport::listen(&"0.0.0.0:0".parse().unwrap())?;
    let addr = listener.local_addr();
    let proxy_server = Server::<String, String>::default()
        .incoming(listener)
        .take(1)
        .for_each(move |channel| {
            let client = client.clone();
            async move {
                let channel = if let Ok(channel) = channel {
                    channel
                } else {
                    return;
                };
                let client_addr = *channel.client_addr();
                let handler = channel.respond_with(move |ctx, request| {
                    trace!("[{}/{}] Proxying request.", ctx.trace_id(), client_addr);
                    let mut client = client.clone();
                    async move { await!(client.call(ctx, request)) }
                });
                tokio_executor::spawn(handler.unit_error().boxed().compat());
            }
        });

    tokio_executor::spawn(proxy_server.unit_error().boxed().compat());

    let mut config = client::Config::default();
    config.max_in_flight_requests = 10;
    config.pending_request_buffer = 10;

    let client = await!(client::new::<String, String, _>(
        config,
        await!(tarpc_bincode_transport::connect(&addr))?
    ))?;

    // Make 3 speculative requests, returning only the quickest.
    let mut clients: Vec<_> = (1..=3u32).map(|_| client.clone()).collect();
    let mut requests = vec![];
    for client in &mut clients {
        let mut ctx = context::current();
        ctx.deadline = SystemTime::now() + Duration::from_millis(200);
        let trace_id = *ctx.trace_id();
        let response = client.call(ctx, "ping".into());
        requests.push(response.map(move |r| (trace_id, r)));
    }
    let (fastest_response, _) = await!(requests
        .into_iter()
        .collect::<FuturesUnordered<_>>()
        .into_future());
    let (trace_id, resp) = fastest_response.unwrap();
    info!("[{}] fastest_response = {:?}", trace_id, resp);

    Ok::<_, io::Error>(())
}

#[test]
fn cancel_slower() -> io::Result<()> {
    env_logger::init();
    rpc::init(tokio::executor::DefaultExecutor::current().compat());

    tokio::run(run().boxed().map_err(|e| panic!(e)).compat());
    Ok(())
}
<|endoftext|>


<commit_before>extern crate rsoundio;

use std::thread;
use std::time::Duration;

#[test]
fn test_outstream() {
    let sio = rsoundio::SoundIo::new();
    sio.connect_backend(rsoundio::SioBackend::Dummy).unwrap();
    let current_backend = sio.current_backend().unwrap();
    assert_eq!(current_backend, rsoundio::SioBackend::Dummy);
    sio.flush_events();
    let dev_idx = sio.default_output_device_index().unwrap();
    let dev = sio.output_device(dev_idx).unwrap();
    let mut stream = dev.create_outstream().unwrap();
    stream.open().unwrap();
    assert!(stream.sample_rate() > 0);
    let layout = stream.layout();
    assert_eq!(layout.channel_count(), 2);
    let cb = |out: rsoundio::OutStream, min_frame_count: i32, max_frame_count: i32| {
        let l: Vec<f32> = (0..max_frame_count as usize)
                              .map(|i| {
                                  match i % 2 == 0 {
                                      true => -1.0,
                                      false => 1.0,
                                  }
                              })
                              .collect();
        let r = l.clone();
        let frames = vec![l, r];
        let frames_written = out.write_stream_f32(min_frame_count, &frames).unwrap();
        assert!(frames_written > 0);
        assert!(out.latency().is_ok());
    };
    stream.register_write_callback(Box::new(cb));
    let ucb = |_: rsoundio::OutStream| println!("Underflow!");
    stream.register_underflow_callback(Box::new(ucb));
    let ecb = |_: rsoundio::OutStream, err: rsoundio::SioError| println!("Error: {}", err);
    stream.register_error_callback(Box::new(ecb));
    stream.start().unwrap();
    thread::sleep(Duration::new(1, 0));
    assert!(stream.pause().is_none());
    thread::sleep(Duration::new(1, 0));
    assert!(stream.unpause().is_none());
    assert!(stream.clear_buffer().is_none());
    thread::sleep(Duration::new(1, 0));
    stream.destroy();
    dev.dec_ref();
}
<commit_msg>Fix test by removing unnecessary device unref<commit_after>extern crate rsoundio;

use std::thread;
use std::time::Duration;

#[test]
fn test_outstream() {
    let sio = rsoundio::SoundIo::new();
    sio.connect_backend(rsoundio::SioBackend::Dummy).unwrap();
    let current_backend = sio.current_backend().unwrap();
    assert_eq!(current_backend, rsoundio::SioBackend::Dummy);
    sio.flush_events();
    let dev_idx = sio.default_output_device_index().unwrap();
    let dev = sio.output_device(dev_idx).unwrap();
    let mut stream = dev.create_outstream().unwrap();
    stream.open().unwrap();
    assert!(stream.sample_rate() > 0);
    let layout = stream.layout();
    assert_eq!(layout.channel_count(), 2);
    let cb = |out: rsoundio::OutStream, min_frame_count: i32, max_frame_count: i32| {
        let l: Vec<f32> = (0..max_frame_count as usize)
                              .map(|i| {
                                  match i % 2 == 0 {
                                      true => -1.0,
                                      false => 1.0,
                                  }
                              })
                              .collect();
        let r = l.clone();
        let frames = vec![l, r];
        let frames_written = out.write_stream_f32(min_frame_count, &frames).unwrap();
        assert!(frames_written > 0);
        assert!(out.latency().is_ok());
    };
    stream.register_write_callback(Box::new(cb));
    let ucb = |_: rsoundio::OutStream| println!("Underflow!");
    stream.register_underflow_callback(Box::new(ucb));
    let ecb = |_: rsoundio::OutStream, err: rsoundio::SioError| println!("Error: {}", err);
    stream.register_error_callback(Box::new(ecb));
    stream.start().unwrap();
    thread::sleep(Duration::new(1, 0));
    assert!(stream.pause().is_none());
    thread::sleep(Duration::new(1, 0));
    assert!(stream.unpause().is_none());
    assert!(stream.clear_buffer().is_none());
    thread::sleep(Duration::new(1, 0));
    stream.destroy();
}
<|endoftext|>


<commit_before>extern crate rust_libindy_wrapper as indy;
#[macro_use] extern crate serde_json;
#[macro_use] extern crate serde_derive;
extern crate rmp_serde;
extern crate byteorder;

use indy::did::Did;
use indy::wallet::Wallet;

use indy::ErrorCode;

use std::path::Path;
use std::panic;
use std::sync::mpsc::channel;
use std::time::Duration;

mod utils;

use utils::{export_config_json, export_path};
use utils::constants::DEFAULT_CREDENTIALS;
use utils::file::TempDir;
use utils::rand;

const VALID_TIMEOUT: Duration = Duration::from_secs(5);
const INVALID_TIMEOUT: Duration = Duration::from_micros(1);

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

    #[test]
    fn create_delete_wallet_works() {
        let wallet_name = r#"{"id":"create_delete_wallet_works"}"#;
        match Wallet::create(wallet_name, DEFAULT_CREDENTIALS) {
            Ok(..) => assert!(Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).is_ok()),
            Err(e) => match e {
                ErrorCode::WalletAlreadyExistsError => {
                    //This is ok, just delete
                    assert!(Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).is_ok())
                }
                _ => {
                    panic!("{:#?}", e)
                }
            }
        }
    }

    #[test]
    fn open_close_wallet_works() {
        let wallet_name = r#"{"id":"open_wallet_works"}"#;
        let open_closure = || {
            match Wallet::open(wallet_name, DEFAULT_CREDENTIALS) {
                Ok(handle) => {
                    Wallet::close(handle).unwrap();
                    Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).unwrap();
                },
                Err(e) => {
                    Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).unwrap();
                    panic!("{:#?}", e);
                }
            }
        };

        match Wallet::create(wallet_name, DEFAULT_CREDENTIALS) {
            Err(e) => match e {
                ErrorCode::WalletAlreadyExistsError => {
                    open_closure()
                }
                _ => panic!("{:#?}", e)
            }
            _ => open_closure()
        };
    }

    #[test]
    fn export_import_wallet_works() {
        let wallet_name = r#"{"id":"export_import_wallet_works"}"#;

        let open_closure = || {
            match Wallet::open(wallet_name, DEFAULT_CREDENTIALS) {
                Ok(handle) => {
                    Did::new(handle, "{}").unwrap();

                    Wallet::export(handle, &export_config_json(wallet_name)).unwrap();

                    assert!(Path::new(&export_path(wallet_name)).exists());

                    Wallet::close(handle).unwrap();
                    Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).unwrap();
                },
                Err(e) => {
                    Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).unwrap();
                    panic!("{:#?}", e);
                }
            }
        };

        match Wallet::create(wallet_name, DEFAULT_CREDENTIALS) {
            Err(e) => match e {
                ErrorCode::WalletAlreadyExistsError => {
                    open_closure()
                }
                _ => panic!("{:#?}", e)
            }
            _ => open_closure()
        };
    }
}

#[cfg(test)]
mod test_wallet_create {
    use super::*;
    const CREDENTIALS: &str = r#"{"key":"9DXvkIMD7iSgD&RT$XYjHo0t"}"#;

    fn wallet_config(storage_type: Option<&str>) -> String {
        let name = rand::random_string(20);

        if let Some(storage) = storage_type {
            json!({
                "id": name,
                "storage_type": storage
            }).to_string()
        } else {
            json!({"id": name}).to_string()
        }
    }

    #[test]
    fn create_default_wallet() {
        let config = wallet_config(Some("default"));

        let result = Wallet::create(&config, CREDENTIALS);

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    #[test]
    fn create_default_wallet_custom_path() {
        let dir = TempDir::new(None).unwrap();
        let config = json!({
            "id": rand::random_string(20),
            "storage_type": "default",
            "storage_config": {
                "path": dir.as_ref().to_str()
            }
        }).to_string();

        let result = Wallet::create(&config, CREDENTIALS);

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    // #[test]
    // fn create_wallet_custom_storage_type() {
    //     unimplemented!();
    // }

    #[test]
    fn create_wallet_unknown_storage_type() {
        let config = wallet_config(Some("unknown"));

        let result = Wallet::create(&config, CREDENTIALS);

        assert_eq!(ErrorCode::WalletUnknownTypeError, result.unwrap_err());
    }

    #[test]
    fn create_wallet_empty_storage_type() {
        let config = wallet_config(None);

        let result = Wallet::create(&config, CREDENTIALS);

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    #[test]
    fn create_wallet_without_key() {
        let config = wallet_config(None);
        let credentials = "{}";

        let result = Wallet::create(&config, credentials);

        assert_eq!(ErrorCode::CommonInvalidStructure, result.unwrap_err());
    }

    #[test]
    fn create_wallet_without_encryption() {
        let config = wallet_config(None);
        let credentials = json!({"key": ""}).to_string();

        let result = Wallet::create(&config, &credentials);

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, &credentials).unwrap();
    }

    #[test]
    fn create_default_wallet_async() {
        let (sender, receiver) = channel();
        let config = wallet_config(Some("default"));

        Wallet::create_async(
            &config,
            CREDENTIALS,
            move |ec| sender.send(ec).unwrap()
        );

        let ec = receiver.recv_timeout(VALID_TIMEOUT).unwrap();

        assert_eq!(ErrorCode::Success, ec);

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    #[test]
    fn create_wallet_unknown_storage_type_async() {
        let (sender, receiver) = channel();
        let config = wallet_config(Some("unknown"));

        Wallet::create_async(
            &config,
            CREDENTIALS,
            move |ec| sender.send(ec).unwrap()
        );

        let ec = receiver.recv_timeout(VALID_TIMEOUT).unwrap();

        assert_eq!(ErrorCode::WalletUnknownTypeError, ec);
    }

    #[test]
    fn create_default_wallet_timeout() {
        let config = wallet_config(Some("default"));

        let result = Wallet::create_timeout(
            &config,
            CREDENTIALS,
            VALID_TIMEOUT
        );

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    #[test]
    fn create_wallet_unknown_storage_type_timeout() {
        let config = wallet_config(Some("unknown"));

        let result = Wallet::create_timeout(
            &config,
            CREDENTIALS,
            VALID_TIMEOUT
        );

        assert_eq!(ErrorCode::WalletUnknownTypeError, result.unwrap_err());
    }

    #[test]
    fn create_wallet_timeout_timeouts() {
        let config = wallet_config(Some("unknown"));

        let result = Wallet::create_timeout(
            &config,
            CREDENTIALS,
            INVALID_TIMEOUT
        );

        assert_eq!(ErrorCode::CommonIOError, result.unwrap_err());
    }
}
<commit_msg>TOK-424 Add `Wallet::delete` tests.<commit_after>extern crate rust_libindy_wrapper as indy;
#[macro_use] extern crate serde_json;
#[macro_use] extern crate serde_derive;
extern crate rmp_serde;
extern crate byteorder;

use indy::did::Did;
use indy::wallet::Wallet;

use indy::ErrorCode;

use std::path::Path;
use std::panic;
use std::sync::mpsc::channel;
use std::time::Duration;

mod utils;

use utils::{export_config_json, export_path};
use utils::constants::DEFAULT_CREDENTIALS;
use utils::file::TempDir;
use utils::rand;

const VALID_TIMEOUT: Duration = Duration::from_secs(5);
const INVALID_TIMEOUT: Duration = Duration::from_micros(1);

mod wallet_config {
    use super::*;

    #[inline]
    pub fn new() -> String {
        json!({
            "id": rand::random_string(20)
        }).to_string()
    }

    #[inline]
    pub fn with_storage(storage: &str) -> String {
        json!({
            "id": rand::random_string(20),
            "storage_type": storage,
        }).to_string()
    }

    #[inline]
    pub fn with_custom_path<P: AsRef<Path>>(path: P) -> String {
        json!({
            "id": rand::random_string(20),
            "storage_type": "default",
            "storage_config": {
                "path": path.as_ref().to_str()
            }
        }).to_string()
    }
}

#[cfg(test)]
mod wallet_register {

}

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

    #[test]
    fn create_delete_wallet_works() {
        let wallet_name = r#"{"id":"create_delete_wallet_works"}"#;
        match Wallet::create(wallet_name, DEFAULT_CREDENTIALS) {
            Ok(..) => assert!(Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).is_ok()),
            Err(e) => match e {
                ErrorCode::WalletAlreadyExistsError => {
                    //This is ok, just delete
                    assert!(Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).is_ok())
                }
                _ => {
                    panic!("{:#?}", e)
                }
            }
        }
    }

    #[test]
    fn open_close_wallet_works() {
        let wallet_name = r#"{"id":"open_wallet_works"}"#;
        let open_closure = || {
            match Wallet::open(wallet_name, DEFAULT_CREDENTIALS) {
                Ok(handle) => {
                    Wallet::close(handle).unwrap();
                    Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).unwrap();
                },
                Err(e) => {
                    Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).unwrap();
                    panic!("{:#?}", e);
                }
            }
        };

        match Wallet::create(wallet_name, DEFAULT_CREDENTIALS) {
            Err(e) => match e {
                ErrorCode::WalletAlreadyExistsError => {
                    open_closure()
                }
                _ => panic!("{:#?}", e)
            }
            _ => open_closure()
        };
    }

    #[test]
    fn export_import_wallet_works() {
        let wallet_name = r#"{"id":"export_import_wallet_works"}"#;

        let open_closure = || {
            match Wallet::open(wallet_name, DEFAULT_CREDENTIALS) {
                Ok(handle) => {
                    Did::new(handle, "{}").unwrap();

                    Wallet::export(handle, &export_config_json(wallet_name)).unwrap();

                    assert!(Path::new(&export_path(wallet_name)).exists());

                    Wallet::close(handle).unwrap();
                    Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).unwrap();
                },
                Err(e) => {
                    Wallet::delete(wallet_name, DEFAULT_CREDENTIALS).unwrap();
                    panic!("{:#?}", e);
                }
            }
        };

        match Wallet::create(wallet_name, DEFAULT_CREDENTIALS) {
            Err(e) => match e {
                ErrorCode::WalletAlreadyExistsError => {
                    open_closure()
                }
                _ => panic!("{:#?}", e)
            }
            _ => open_closure()
        };
    }
}

#[cfg(test)]
mod test_wallet_create {
    use super::*;
    const CREDENTIALS: &str = r#"{"key":"9DXvkIMD7iSgD&RT$XYjHo0t"}"#;

    #[test]
    fn create_default_wallet() {
        let config = wallet_config::with_storage("default");

        let result = Wallet::create(&config, CREDENTIALS);

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    #[test]
    fn create_default_wallet_custom_path() {
        let dir = TempDir::new(None).unwrap();
        let config = wallet_config::with_custom_path(&dir);

        let result = Wallet::create(&config, CREDENTIALS);

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    // #[test]
    // fn create_wallet_custom_storage_type() {
    //     unimplemented!();
    // }

    #[test]
    fn create_wallet_unknown_storage_type() {
        let config = wallet_config::with_storage("unknown");

        let result = Wallet::create(&config, CREDENTIALS);

        assert_eq!(ErrorCode::WalletUnknownTypeError, result.unwrap_err());
    }

    #[test]
    fn create_wallet_empty_storage_type() {
        let config = wallet_config::new();

        let result = Wallet::create(&config, CREDENTIALS);

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    #[test]
    fn create_wallet_without_key() {
        let config = wallet_config::new();
        let credentials = "{}";

        let result = Wallet::create(&config, credentials);

        assert_eq!(ErrorCode::CommonInvalidStructure, result.unwrap_err());
    }

    #[test]
    fn create_wallet_without_encryption() {
        let config = wallet_config::new();
        let credentials = json!({"key": ""}).to_string();

        let result = Wallet::create(&config, &credentials);

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, &credentials).unwrap();
    }

    #[test]
    fn create_default_wallet_async() {
        let (sender, receiver) = channel();
        let config = wallet_config::with_storage("default");

        Wallet::create_async(
            &config,
            CREDENTIALS,
            move |ec| sender.send(ec).unwrap()
        );

        let ec = receiver.recv_timeout(VALID_TIMEOUT).unwrap();

        assert_eq!(ErrorCode::Success, ec);

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    #[test]
    fn create_wallet_unknown_storage_type_async() {
        let (sender, receiver) = channel();
        let config = wallet_config::with_storage("unknown");

        Wallet::create_async(
            &config,
            CREDENTIALS,
            move |ec| sender.send(ec).unwrap()
        );

        let ec = receiver.recv_timeout(VALID_TIMEOUT).unwrap();

        assert_eq!(ErrorCode::WalletUnknownTypeError, ec);
    }

    #[test]
    fn create_default_wallet_timeout() {
        let config = wallet_config::with_storage("default");

        let result = Wallet::create_timeout(
            &config,
            CREDENTIALS,
            VALID_TIMEOUT
        );

        assert_eq!((), result.unwrap());

        Wallet::delete(&config, CREDENTIALS).unwrap();
    }

    #[test]
    fn create_wallet_unknown_storage_type_timeout() {
        let config = wallet_config::with_storage("unknown");

        let result = Wallet::create_timeout(
            &config,
            CREDENTIALS,
            VALID_TIMEOUT
        );

        assert_eq!(ErrorCode::WalletUnknownTypeError, result.unwrap_err());
    }

    #[test]
    fn create_wallet_timeout_timeouts() {
        let config = wallet_config::with_storage("unknown");

        let result = Wallet::create_timeout(
            &config,
            CREDENTIALS,
            INVALID_TIMEOUT
        );

        assert_eq!(ErrorCode::CommonIOError, result.unwrap_err());
    }
}

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

    #[inline]
    fn assert_wallet_deleted(config: &str, credentials: &str) {
        let result = Wallet::open(config, credentials);
        assert_eq!(ErrorCode::WalletNotFoundError, result.unwrap_err());
    }

    #[test]
    fn delete_wallet() {
        let config = wallet_config::new();
        Wallet::create(&config, DEFAULT_CREDENTIALS).unwrap();

        let result = Wallet::delete(&config, DEFAULT_CREDENTIALS);

        assert_eq!((), result.unwrap());
        assert_wallet_deleted(&config, DEFAULT_CREDENTIALS);
    }

    #[test]
    fn delete_wallet_custom_path() {
        let dir = TempDir::new(None).unwrap();
        let config = wallet_config::with_custom_path(&dir);
        Wallet::create(&config, DEFAULT_CREDENTIALS).unwrap();

        let result = Wallet::delete(&config, DEFAULT_CREDENTIALS);

        assert_eq!((), result.unwrap());
        assert_wallet_deleted(&config, DEFAULT_CREDENTIALS);
    }

    #[test]
    fn delete_wallet_closed() {
        let config = wallet_config::new();

        Wallet::create(&config, DEFAULT_CREDENTIALS).unwrap();
        let handle = Wallet::open(&config, DEFAULT_CREDENTIALS).unwrap();
        Wallet::close(handle).unwrap();

        let result = Wallet::delete(&config, DEFAULT_CREDENTIALS);

        assert_eq!((), result.unwrap());
        assert_wallet_deleted(&config, DEFAULT_CREDENTIALS);
    }

    #[test]
    fn delete_wallet_opened() {
        let config = wallet_config::new();

        Wallet::create(&config, DEFAULT_CREDENTIALS).unwrap();
        let handle = Wallet::open(&config, DEFAULT_CREDENTIALS).unwrap();

        let result = Wallet::delete(&config, DEFAULT_CREDENTIALS);

        assert_eq!(ErrorCode::CommonInvalidState, result.unwrap_err());

        Wallet::close(handle).unwrap();
        Wallet::delete(&config, DEFAULT_CREDENTIALS).unwrap();
    }

    // #[test]
    // fn delete_registered_wallet() {
    //     unimplemented!();
    // }

    #[test]
    fn delete_wallet_repeated_command() {
        let config = wallet_config::new();
        Wallet::create(&config, DEFAULT_CREDENTIALS).unwrap();
        Wallet::delete(&config, DEFAULT_CREDENTIALS).unwrap();

        let result = Wallet::delete(&config, DEFAULT_CREDENTIALS);

        assert_eq!(ErrorCode::WalletNotFoundError, result.unwrap_err());
    }

    #[test]
    fn delete_wallet_invalid_credentials() {
        let config = wallet_config::new();
        Wallet::create(&config, DEFAULT_CREDENTIALS).unwrap();

        let result = Wallet::delete(&config, r#"{"key": "badkey"}"#);

        assert_eq!(ErrorCode::WalletAccessFailed, result.unwrap_err());

        Wallet::delete(&config, DEFAULT_CREDENTIALS).unwrap();
    }

    #[test]
    fn delete_wallet_uncreated() {
        let config = wallet_config::new();

        let result = Wallet::delete(&config, DEFAULT_CREDENTIALS);

        assert_eq!(ErrorCode::WalletNotFoundError, result.unwrap_err());
    }

    #[test]
    fn delete_wallet_async() {
        let (sender, receiver) = channel();
        let config = wallet_config::new();
        Wallet::create(&config, DEFAULT_CREDENTIALS);

        Wallet::delete_async(
            &config,
            DEFAULT_CREDENTIALS,
            move |ec| sender.send(ec).unwrap()
        );

        let ec = receiver.recv().unwrap();

        assert_eq!(ErrorCode::Success, ec);
        assert_wallet_deleted(&config, DEFAULT_CREDENTIALS);
    }

    #[test]
    fn delete_wallet_uncreated_async() {
        let (sender, receiver) = channel();
        let config = wallet_config::new();

        Wallet::delete_async(
            &config,
            DEFAULT_CREDENTIALS,
            move |ec| sender.send(ec).unwrap()
        );

        let ec = receiver.recv().unwrap();

        assert_eq!(ErrorCode::WalletNotFoundError, ec);
    }

    #[test]
    fn delete_wallet_timeout() {
        let config = wallet_config::new();
        Wallet::create(&config, DEFAULT_CREDENTIALS).unwrap();

        let result = Wallet::delete_timeout(
            &config,
            DEFAULT_CREDENTIALS,
            VALID_TIMEOUT,
        );

        assert_eq!((), result.unwrap());
        assert_wallet_deleted(&config, DEFAULT_CREDENTIALS);
    }

    #[test]
    fn delete_wallet_uncreated_timeout() {
        let config = wallet_config::new();

        let result = Wallet::delete_timeout(
            &config,
            DEFAULT_CREDENTIALS,
            VALID_TIMEOUT,
        );

        assert_eq!(ErrorCode::WalletNotFoundError, result.unwrap_err());
    }

    #[test]
    fn delete_wallet_timeout_timeouts() {
        let config = wallet_config::new();

        let result = Wallet::delete_timeout(
            &config,
            DEFAULT_CREDENTIALS,
            INVALID_TIMEOUT,
        );

        assert_eq!(ErrorCode::CommonIOError, result.unwrap_err());
    }
}<|endoftext|>