<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! Execution plan for reading Parquet files

use std::fs::File;
use std::rc::Rc;
use std::sync::{Arc, Mutex};
use std::thread;

use crate::error::{ExecutionError, Result};
use crate::execution::physical_plan::common;
use crate::execution::physical_plan::{ExecutionPlan, Partition};
use arrow::datatypes::{Schema, SchemaRef};
use arrow::error::{ArrowError, Result as ArrowResult};
use arrow::record_batch::{RecordBatch, RecordBatchReader};
use parquet::file::reader::SerializedFileReader;

use crossbeam::channel::{unbounded, Receiver, RecvError, SendError, Sender};
use parquet::arrow::{ArrowReader, ParquetFileArrowReader};

/// Execution plan for scanning a Parquet file
pub struct ParquetExec {
    /// Path to directory containing partitioned Parquet files with the same schema
    filenames: Vec<String>,
    /// Schema after projection is applied
    schema: SchemaRef,
    /// Projection for which columns to load
    projection: Vec<usize>,
    /// Batch size
    batch_size: usize,
}

impl ParquetExec {
    /// Create a new Parquet reader execution plan
    pub fn try_new(
        path: &str,
        projection: Option<Vec<usize>>,
        batch_size: usize,
    ) -> Result<Self> {
        let mut filenames: Vec<String> = vec![];
        common::build_file_list(path, &mut filenames, ".parquet")?;
        if filenames.is_empty() {
            Err(ExecutionError::General("No files found".to_string()))
        } else {
            let file = File::open(&filenames[0])?;
            let file_reader = Rc::new(SerializedFileReader::new(file)?);
            let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
            let schema = arrow_reader.get_schema()?;

            let projection = match projection {
                Some(p) => p,
                None => (0..schema.fields().len()).collect(),
            };

            let projected_schema = Schema::new(
                projection
                    .iter()
                    .map(|i| schema.field(*i).clone())
                    .collect(),
            );

            Ok(Self {
                filenames,
                schema: Arc::new(projected_schema),
                projection,
                batch_size,
            })
        }
    }
}

impl ExecutionPlan for ParquetExec {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn partitions(&self) -> Result<Vec<Arc<dyn Partition>>> {
        let partitions = self
            .filenames
            .iter()
            .map(|filename| {
                Arc::new(ParquetPartition::new(
                    &filename,
                    self.projection.clone(),
                    self.schema.clone(),
                    self.batch_size,
                )) as Arc<dyn Partition>
            })
            .collect();
        Ok(partitions)
    }
}

struct ParquetPartition {
    iterator: Arc<Mutex<dyn RecordBatchReader + Send + Sync>>,
}

impl ParquetPartition {
    /// Create a new Parquet partition
    pub fn new(
        filename: &str,
        projection: Vec<usize>,
        schema: SchemaRef,
        batch_size: usize,
    ) -> Self {
        // because the parquet implementation is not thread-safe, it is necessary to execute
        // on a thread and communicate with channels
        let (request_tx, request_rx): (Sender<()>, Receiver<()>) = unbounded();
        let (response_tx, response_rx): (
            Sender<ArrowResult<Option<RecordBatch>>>,
            Receiver<ArrowResult<Option<RecordBatch>>>,
        ) = unbounded();

        let filename = filename.to_string();

        thread::spawn(move || {
            //TODO error handling, remove unwraps

            // open file
            let file = File::open(&filename).unwrap();
            match SerializedFileReader::new(file) {
                Ok(file_reader) => {
                    let file_reader = Rc::new(file_reader);

                    let mut arrow_reader = ParquetFileArrowReader::new(file_reader);

                    match arrow_reader
                        .get_record_reader_by_columns(projection, batch_size)
                    {
                        Ok(mut batch_reader) => {
                            while let Ok(_) = request_rx.recv() {
                                match batch_reader.next_batch() {
                                    Ok(Some(batch)) => {
                                        response_tx.send(Ok(Some(batch))).unwrap();
                                    }
                                    Ok(None) => {
                                        response_tx.send(Ok(None)).unwrap();
                                        break;
                                    }
                                    Err(e) => {
                                        response_tx
                                            .send(Err(ArrowError::ParquetError(format!(
                                                "{:?}",
                                                e
                                            ))))
                                            .unwrap();
                                        break;
                                    }
                                }
                            }
                        }

                        Err(e) => {
                            response_tx
                                .send(Err(ArrowError::ParquetError(format!("{:?}", e))))
                                .unwrap();
                        }
                    }
                }

                Err(e) => {
                    response_tx
                        .send(Err(ArrowError::ParquetError(format!("{:?}", e))))
                        .unwrap();
                }
            }
        });

        let iterator = Arc::new(Mutex::new(ParquetIterator {
            schema,
            request_tx,
            response_rx,
        }));

        Self { iterator }
    }
}

impl Partition for ParquetPartition {
    fn execute(&self) -> Result<Arc<Mutex<dyn RecordBatchReader + Send + Sync>>> {
        Ok(self.iterator.clone())
    }
}

struct ParquetIterator {
    schema: SchemaRef,
    request_tx: Sender<()>,
    response_rx: Receiver<ArrowResult<Option<RecordBatch>>>,
}

impl RecordBatchReader for ParquetIterator {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn next_batch(&mut self) -> ArrowResult<Option<RecordBatch>> {
        match self.request_tx.send(()) {
            Ok(_) => match self.response_rx.recv() {
                Ok(batch) => batch,
                // RecvError means receiver has exited and closed the channel
                Err(RecvError) => Ok(None),
            },
            // SendError means receiver has exited and closed the channel
            Err(SendError(())) => Ok(None),
        }
    }
}

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

    #[test]
    fn test() -> Result<()> {
        let testdata =
            env::var("PARQUET_TEST_DATA").expect("PARQUET_TEST_DATA not defined");
        let filename = format!("{}/alltypes_plain.parquet", testdata);
        let parquet_exec = ParquetExec::try_new(&filename, Some(vec![0, 1, 2]), 1024)?;
        let partitions = parquet_exec.partitions()?;
        assert_eq!(partitions.len(), 1);

        let results = partitions[0].execute()?;
        let mut results = results.lock().unwrap();
        let batch = results.next_batch()?.unwrap();

        assert_eq!(8, batch.num_rows());
        assert_eq!(3, batch.num_columns());

        let schema = batch.schema();
        let field_names: Vec<&str> =
            schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(vec!["id", "bool_col", "tinyint_col"], field_names);

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        Ok(())
    }
}
<commit_msg>ARROW-9452: [Rust] [DataFusion] Optimize ParquetScanExec<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

//! Execution plan for reading Parquet files

use std::fs::File;
use std::rc::Rc;
use std::sync::{Arc, Mutex};
use std::thread;

use crate::error::{ExecutionError, Result};
use crate::execution::physical_plan::common;
use crate::execution::physical_plan::{ExecutionPlan, Partition};
use arrow::datatypes::{Schema, SchemaRef};
use arrow::error::{ArrowError, Result as ArrowResult};
use arrow::record_batch::{RecordBatch, RecordBatchReader};
use parquet::file::reader::SerializedFileReader;

use crossbeam::channel::{bounded, Receiver, RecvError, Sender};
use parquet::arrow::{ArrowReader, ParquetFileArrowReader};

/// Execution plan for scanning a Parquet file
pub struct ParquetExec {
    /// Path to directory containing partitioned Parquet files with the same schema
    filenames: Vec<String>,
    /// Schema after projection is applied
    schema: SchemaRef,
    /// Projection for which columns to load
    projection: Vec<usize>,
    /// Batch size
    batch_size: usize,
}

impl ParquetExec {
    /// Create a new Parquet reader execution plan
    pub fn try_new(
        path: &str,
        projection: Option<Vec<usize>>,
        batch_size: usize,
    ) -> Result<Self> {
        let mut filenames: Vec<String> = vec![];
        common::build_file_list(path, &mut filenames, ".parquet")?;
        if filenames.is_empty() {
            Err(ExecutionError::General("No files found".to_string()))
        } else {
            let file = File::open(&filenames[0])?;
            let file_reader = Rc::new(SerializedFileReader::new(file)?);
            let mut arrow_reader = ParquetFileArrowReader::new(file_reader);
            let schema = arrow_reader.get_schema()?;

            let projection = match projection {
                Some(p) => p,
                None => (0..schema.fields().len()).collect(),
            };

            let projected_schema = Schema::new(
                projection
                    .iter()
                    .map(|i| schema.field(*i).clone())
                    .collect(),
            );

            Ok(Self {
                filenames,
                schema: Arc::new(projected_schema),
                projection,
                batch_size,
            })
        }
    }
}

impl ExecutionPlan for ParquetExec {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn partitions(&self) -> Result<Vec<Arc<dyn Partition>>> {
        let partitions = self
            .filenames
            .iter()
            .map(|filename| {
                Arc::new(ParquetPartition::new(
                    &filename,
                    self.projection.clone(),
                    self.schema.clone(),
                    self.batch_size,
                )) as Arc<dyn Partition>
            })
            .collect();
        Ok(partitions)
    }
}

struct ParquetPartition {
    iterator: Arc<Mutex<dyn RecordBatchReader + Send + Sync>>,
}

impl ParquetPartition {
    /// Create a new Parquet partition
    pub fn new(
        filename: &str,
        projection: Vec<usize>,
        schema: SchemaRef,
        batch_size: usize,
    ) -> Self {
        // because the parquet implementation is not thread-safe, it is necessary to execute
        // on a thread and communicate with channels
        let (response_tx, response_rx): (
            Sender<ArrowResult<Option<RecordBatch>>>,
            Receiver<ArrowResult<Option<RecordBatch>>>,
        ) = bounded(2);

        let filename = filename.to_string();

        thread::spawn(move || {
            //TODO error handling, remove unwraps

            // open file
            let file = File::open(&filename).unwrap();
            match SerializedFileReader::new(file) {
                Ok(file_reader) => {
                    let file_reader = Rc::new(file_reader);

                    let mut arrow_reader = ParquetFileArrowReader::new(file_reader);

                    match arrow_reader
                        .get_record_reader_by_columns(projection, batch_size)
                    {
                        Ok(mut batch_reader) => loop {
                            match batch_reader.next_batch() {
                                Ok(Some(batch)) => {
                                    response_tx.send(Ok(Some(batch))).unwrap();
                                }
                                Ok(None) => {
                                    response_tx.send(Ok(None)).unwrap();
                                    break;
                                }
                                Err(e) => {
                                    response_tx
                                        .send(Err(ArrowError::ParquetError(format!(
                                            "{:?}",
                                            e
                                        ))))
                                        .unwrap();
                                    break;
                                }
                            }
                        },

                        Err(e) => {
                            response_tx
                                .send(Err(ArrowError::ParquetError(format!("{:?}", e))))
                                .unwrap();
                        }
                    }
                }

                Err(e) => {
                    response_tx
                        .send(Err(ArrowError::ParquetError(format!("{:?}", e))))
                        .unwrap();
                }
            }
        });

        let iterator = Arc::new(Mutex::new(ParquetIterator {
            schema,
            response_rx,
        }));

        Self { iterator }
    }
}

impl Partition for ParquetPartition {
    fn execute(&self) -> Result<Arc<Mutex<dyn RecordBatchReader + Send + Sync>>> {
        Ok(self.iterator.clone())
    }
}

struct ParquetIterator {
    schema: SchemaRef,
    response_rx: Receiver<ArrowResult<Option<RecordBatch>>>,
}

impl RecordBatchReader for ParquetIterator {
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn next_batch(&mut self) -> ArrowResult<Option<RecordBatch>> {
        match self.response_rx.recv() {
            Ok(batch) => batch,
            // RecvError means receiver has exited and closed the channel
            Err(RecvError) => Ok(None),
        }
    }
}

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

    #[test]
    fn test() -> Result<()> {
        let testdata =
            env::var("PARQUET_TEST_DATA").expect("PARQUET_TEST_DATA not defined");
        let filename = format!("{}/alltypes_plain.parquet", testdata);
        let parquet_exec = ParquetExec::try_new(&filename, Some(vec![0, 1, 2]), 1024)?;
        let partitions = parquet_exec.partitions()?;
        assert_eq!(partitions.len(), 1);

        let results = partitions[0].execute()?;
        let mut results = results.lock().unwrap();
        let batch = results.next_batch()?.unwrap();

        assert_eq!(8, batch.num_rows());
        assert_eq!(3, batch.num_columns());

        let schema = batch.schema();
        let field_names: Vec<&str> =
            schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert_eq!(vec!["id", "bool_col", "tinyint_col"], field_names);

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        let batch = results.next_batch()?;
        assert!(batch.is_none());

        Ok(())
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use ::futures::{
    task::{self, Task},
    Async::{self, NotReady, Ready},
    Future,
};

use std::{
    fmt,
    str::FromStr,
    sync::{Arc, Mutex, MutexGuard},
};

use crate::{
    conn::{pool::futures::*, Conn},
    error::*,
    opts::{Opts, PoolConstraints},
    queryable::{
        transaction::{Transaction, TransactionOptions},
        Queryable,
    },
    BoxFuture, MyFuture,
};

pub mod futures;

pub struct Inner {
    closed: bool,
    new: Vec<BoxFuture<Conn>>,
    idle: Vec<Conn>,
    disconnecting: Vec<BoxFuture<()>>,
    dropping: Vec<BoxFuture<Conn>>,
    rollback: Vec<BoxFuture<Conn>>,
    ongoing: usize,
    tasks: Vec<Task>,
}

impl Inner {
    fn conn_count(&self) -> usize {
        self.new.len()
            + self.idle.len()
            + self.disconnecting.len()
            + self.dropping.len()
            + self.rollback.len()
            + self.ongoing
    }
}

#[derive(Clone)]
/// Asynchronous pool of MySql connections.
pub struct Pool {
    opts: Opts,
    inner: Arc<Mutex<Inner>>,
    pool_constraints: PoolConstraints,
}

impl fmt::Debug for Pool {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let (new_len, idle_len, disconnecing_len, dropping_len, rollback_len, ongoing, tasks_len) =
            self.with_inner(|inner| {
                (
                    inner.new.len(),
                    inner.idle.len(),
                    inner.disconnecting.len(),
                    inner.dropping.len(),
                    inner.rollback.len(),
                    inner.ongoing,
                    inner.tasks.len(),
                )
            });
        f.debug_struct("Pool")
            .field("pool_constraints", &self.pool_constraints)
            .field("new connections count", &new_len)
            .field("idle connections count", &idle_len)
            .field("disconnecting connections count", &disconnecing_len)
            .field("dropping connections count", &dropping_len)
            .field("rollback connections count", &rollback_len)
            .field("ongoing connections count", &ongoing)
            .field("tasks count", &tasks_len)
            .finish()
    }
}

impl Pool {
    /// Creates new pool of connections.
    pub fn new<O: Into<Opts>>(opts: O) -> Pool {
        let opts = opts.into();
        let pool_constraints = opts.get_pool_constraints().clone();
        let pool = Pool {
            opts,
            inner: Arc::new(Mutex::new(Inner {
                closed: false,
                new: Vec::with_capacity(pool_constraints.min()),
                idle: Vec::new(),
                disconnecting: Vec::new(),
                dropping: Vec::new(),
                rollback: Vec::new(),
                ongoing: 0,
                tasks: Vec::new(),
            })),
            pool_constraints,
        };

        pool
    }

    /// Creates new pool of connections.
    pub fn from_url<T: AsRef<str>>(url: T) -> Result<Pool> {
        let opts = Opts::from_str(url.as_ref())?;
        Ok(Pool::new(opts))
    }

    /// Returns future that resolves to `Conn`.
    pub fn get_conn(&self) -> GetConn {
        new_get_conn(self)
    }

    /// Shortcut for `get_conn` followed by `start_transaction`.
    pub fn start_transaction(
        &self,
        options: TransactionOptions,
    ) -> impl MyFuture<Transaction<Conn>> {
        self.get_conn()
            .and_then(|conn| Queryable::start_transaction(conn, options))
    }

    /// Returns future that disconnects this pool from server and resolves to `()`.
    ///
    /// Active connections taken from this pool should be disconnected manually.
    /// Also all pending and new `GetConn`'s will resolve to error.
    pub fn disconnect(mut self) -> DisconnectPool {
        let become_closed = self.with_inner(|mut inner| {
            if !inner.closed {
                inner.closed = true;
                return true;
            } else {
                return false;
            }
        });
        if become_closed {
            while let Some(conn) = self.take_conn() {
                self.with_inner(move |mut inner| inner.disconnecting.push(conn.disconnect()));
            }
        }
        new_disconnect_pool(self)
    }

    /// Returns true if futures is in queue.
    fn in_queue(&self) -> bool {
        self.with_inner(|inner| {
            let count = inner.new.len()
                + inner.disconnecting.len()
                + inner.dropping.len()
                + inner.rollback.len();
            count > 0
        })
    }

    /// A way to take connection from a pool.
    fn take_conn(&mut self) -> Option<Conn> {
        if self.in_queue() {
            // Do not return connection until queue is empty
            return None;
        }
        self.with_inner(|mut inner| {
            while let Some(mut conn) = inner.idle.pop() {
                if conn.expired() {
                    inner.disconnecting.push(conn.disconnect());
                } else {
                    conn.inner.pool = Some(self.clone());
                    inner.ongoing += 1;
                    return Some(conn);
                }
            }
            None
        })
    }

    /// A way to return connection taken from a pool.
    fn return_conn(&mut self, conn: Conn) {
        let min = self.pool_constraints.min();

        self.with_inner(|mut inner| {
            if inner.closed {
                return;
            }

            if conn.inner.has_result.is_some() {
                inner.dropping.push(Box::new(conn.drop_result()));
            } else if conn.inner.in_transaction {
                inner.rollback.push(Box::new(conn.rollback_transaction()));
            } else {
                if inner.idle.len() >= min {
                    inner.disconnecting.push(conn.disconnect());
                } else {
                    inner.ongoing -= 1;
                    inner.idle.push(conn);
                }
            }

            while let Some(task) = inner.tasks.pop() {
                task.notify()
            }
        });
    }

    fn with_inner<F, T>(&self, fun: F) -> T
    where
        F: FnOnce(MutexGuard<'_, Inner>) -> T,
        T: 'static,
    {
        fun(self.inner.lock().unwrap())
    }

    /// Will manage lifetime of futures stored in a pool.
    fn handle_futures(&mut self) -> Result<()> {
        if !self.in_queue() {
            // There is no futures in queue
            return Ok(());
        }

        let mut handled = false;

        let mut returned_conns: Vec<Conn> = vec![];

        self.with_inner(|mut inner| {
            macro_rules! handle {
                ($vec:ident { $($p:pat => $b:block,)+ }) => ({
                    let len = inner.$vec.len();
                    let mut done_fut_idxs = Vec::new();
                    for i in 0..len {
                        let result = inner.$vec.get_mut(i).unwrap().poll();
                        match result {
                            Ok(Ready(_)) | Err(_) => done_fut_idxs.push(i),
                            _ => (),
                        }

                        let out: Result<()> = match result {
                            $($p => $b),+
                            _ => {
                                Ok(())
                            }
                        };

                        match out {
                            Err(err) => {
                                // early return in case of error
                                while let Some(i) = done_fut_idxs.pop() {
                                    inner.$vec.swap_remove(i);
                                }
                                return Err(err)
                            }
                            _ => (),
                        }
                    }

                    while let Some(i) = done_fut_idxs.pop() {
                        inner.$vec.swap_remove(i);
                    }
                });
            }

            // Handle closing connections.
            handle!(disconnecting {
                Ok(Ready(_)) => {
                    handled = true;
                    Ok(())
                },
                Err(_) => { Ok(()) },
            });

            // Handle dirty connections.
            handle!(dropping {
                Ok(Ready(conn)) => {
                    if inner.closed {
                        inner.disconnecting.push(conn.disconnect());
                    } else {
                        returned_conns.push(conn);
                    }
                    handled = true;
                    Ok(())
                },
                Err(_) => { Ok(()) },
            });

            // Handle in-transaction connections
            handle!(rollback {
                Ok(Ready(conn)) => {
                    if inner.closed {
                        inner.disconnecting.push(conn.disconnect());
                    } else {
                        returned_conns.push(conn);
                    }
                    handled = true;
                    Ok(())
                },
                Err(_) => { Ok(()) },
            });

            // Handle connecting connections.
            handle!(new {
                Ok(Ready(conn)) => {
                    if inner.closed {
                        inner.disconnecting.push(conn.disconnect());
                    } else {
                        inner.ongoing += 1;
                        returned_conns.push(conn);
                    }
                    handled = true;
                    Ok(())
                },
                Err(err) => {
                    if !inner.closed {
                        Err(err)
                    } else {
                        Ok(())
                    }
                },
            });

            Ok(())
        })?;

        for conn in returned_conns {
            self.return_conn(conn);
        }

        if handled {
            self.handle_futures()
        } else {
            Ok(())
        }
    }

    /// Will poll pool for connection.
    fn poll(&mut self) -> Result<Async<Conn>> {
        if self.with_inner(|inner| inner.closed) {
            return Err(DriverError::PoolDisconnected.into());
        }

        self.handle_futures()?;

        match self.take_conn() {
            Some(conn) => Ok(Ready(conn)),
            None => {
                let new_conn_created = self.with_inner(|mut inner| {
                    if inner.new.len() == 0 && inner.conn_count() < self.pool_constraints.max() {
                        let new_conn = Conn::new(self.opts.clone());
                        inner.new.push(Box::new(new_conn));
                        true
                    } else {
                        inner.tasks.push(task::current());
                        false
                    }
                });
                if new_conn_created {
                    self.poll()
                } else {
                    Ok(NotReady)
                }
            }
        }
    }
}

impl Drop for Conn {
    fn drop(&mut self) {
        if let Some(mut pool) = self.inner.pool.take() {
            let conn = self.take();
            if conn.inner.stream.is_some() {
                pool.return_conn(conn)
            } // drop incomplete connection
        }
    }
}

#[cfg(test)]
mod test {
    use futures::Future;

    use crate::{
        conn::pool::Pool, queryable::Queryable, test_misc::DATABASE_URL, TransactionOptions,
    };

    /// Same as `tokio::run`, but will panic if future panics and will return the result
    /// of future execution.
    fn run<F, T, U>(future: F) -> Result<T, U>
    where
        F: Future<Item = T, Error = U> + Send + 'static,
        T: Send + 'static,
        U: Send + 'static,
    {
        let mut runtime = tokio::runtime::Runtime::new().unwrap();
        let result = runtime.block_on(future);
        runtime.shutdown_on_idle().wait().unwrap();
        result
    }

    #[test]
    fn should_connect() {
        let pool = Pool::new(&**DATABASE_URL);
        let fut = pool
            .get_conn()
            .and_then(|conn| conn.ping().map(|_| ()))
            .and_then(|_| pool.disconnect());

        run(fut).unwrap();
    }

    #[test]
    fn should_start_transaction() {
        let pool = Pool::new(format!("{}?pool_min=1&pool_max=1", &**DATABASE_URL));
        let fut = pool
            .get_conn()
            .and_then(|conn| conn.drop_query("CREATE TABLE IF NOT EXISTS tmp(id int)"))
            .and_then({
                let pool = pool.clone();
                move |_| pool.start_transaction(TransactionOptions::default())
            })
            .and_then(|transaction| {
                transaction.batch_exec("INSERT INTO tmp (id) VALUES (?)", vec![(1,), (2,)])
            })
            .and_then(|transaction| transaction.prep_exec("SELECT * FROM tmp", ()))
            .map(|_| ())
            .and_then({
                let pool = pool.clone();
                move |_| pool.get_conn()
            })
            .and_then(|conn| conn.first("SELECT COUNT(*) FROM tmp"))
            .and_then(|(_, row_opt)| {
                assert_eq!(row_opt, Some((0u8,)));
                pool.get_conn()
                    .and_then(|conn| conn.drop_query("DROP TABLE tmp"))
                    .and_then(move |_| pool.disconnect())
            });

        run(fut).unwrap();
    }

    #[test]
    fn should_hold_bounds() {
        let pool = Pool::new(format!("{}?pool_min=1&pool_max=2", &**DATABASE_URL));
        let pool_clone = pool.clone();
        let fut = pool
            .get_conn()
            .join(pool.get_conn())
            .and_then(move |(mut conn1, conn2)| {
                let new_conn = pool_clone.get_conn();
                conn1.inner.pool.as_mut().unwrap().handle_futures().unwrap();
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.new.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.idle.len()),
                    0
                );
                assert_eq!(
                    conn2
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.disconnecting.len()),
                    0
                );
                assert_eq!(
                    conn2
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.dropping.len()),
                    0
                );
                new_conn
            })
            .and_then(|conn1| {
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.new.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.idle.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.disconnecting.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.dropping.len()),
                    0
                );
                Ok(())
            })
            .and_then(|_| {
                assert_eq!(pool.with_inner(|inner| inner.new.len()), 0);
                assert_eq!(pool.with_inner(|inner| inner.idle.len()), 1);
                assert_eq!(pool.with_inner(|inner| inner.disconnecting.len()), 0);
                assert_eq!(pool.with_inner(|inner| inner.dropping.len()), 0);
                pool.disconnect()
            });

        run(fut).unwrap();
    }

    #[cfg(feature = "nightly")]
    mod bench {
        use futures::Future;

        use crate::{conn::pool::Pool, queryable::Queryable, test_misc::DATABASE_URL};

        #[bench]
        fn connect(bencher: &mut test::Bencher) {
            let mut runtime = tokio::executor::current_thread::CurrentThread::new();
            let pool = Pool::new(&**DATABASE_URL);

            bencher.iter(|| {
                let fut = pool.get_conn().and_then(|conn| conn.ping());
                runtime.block_on(fut).unwrap();
            });

            runtime.block_on(pool.disconnect()).unwrap();
        }
    }
}
<commit_msg>Fix pool's ongoing count accumulating over time even if connections are dropped (#44)<commit_after>// Copyright (c) 2016 Anatoly Ikorsky
//
// Licensed under the Apache License, Version 2.0
// <LICENSE-APACHE or http://www.apache.org/licenses/LICENSE-2.0> or the MIT
// license <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. All files in the project carrying such notice may not be copied,
// modified, or distributed except according to those terms.

use ::futures::{
    task::{self, Task},
    Async::{self, NotReady, Ready},
    Future,
};

use std::{
    fmt,
    str::FromStr,
    sync::{Arc, Mutex, MutexGuard},
};

use crate::{
    conn::{pool::futures::*, Conn},
    error::*,
    opts::{Opts, PoolConstraints},
    queryable::{
        transaction::{Transaction, TransactionOptions},
        Queryable,
    },
    BoxFuture, MyFuture,
};

pub mod futures;

pub struct Inner {
    closed: bool,
    new: Vec<BoxFuture<Conn>>,
    idle: Vec<Conn>,
    disconnecting: Vec<BoxFuture<()>>,
    dropping: Vec<BoxFuture<Conn>>,
    rollback: Vec<BoxFuture<Conn>>,
    ongoing: usize,
    tasks: Vec<Task>,
}

impl Inner {
    fn conn_count(&self) -> usize {
        self.new.len()
            + self.idle.len()
            + self.disconnecting.len()
            + self.dropping.len()
            + self.rollback.len()
            + self.ongoing
    }
}

#[derive(Clone)]
/// Asynchronous pool of MySql connections.
pub struct Pool {
    opts: Opts,
    inner: Arc<Mutex<Inner>>,
    pool_constraints: PoolConstraints,
}

impl fmt::Debug for Pool {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let (new_len, idle_len, disconnecing_len, dropping_len, rollback_len, ongoing, tasks_len) =
            self.with_inner(|inner| {
                (
                    inner.new.len(),
                    inner.idle.len(),
                    inner.disconnecting.len(),
                    inner.dropping.len(),
                    inner.rollback.len(),
                    inner.ongoing,
                    inner.tasks.len(),
                )
            });
        f.debug_struct("Pool")
            .field("pool_constraints", &self.pool_constraints)
            .field("new connections count", &new_len)
            .field("idle connections count", &idle_len)
            .field("disconnecting connections count", &disconnecing_len)
            .field("dropping connections count", &dropping_len)
            .field("rollback connections count", &rollback_len)
            .field("ongoing connections count", &ongoing)
            .field("tasks count", &tasks_len)
            .finish()
    }
}

impl Pool {
    /// Creates new pool of connections.
    pub fn new<O: Into<Opts>>(opts: O) -> Pool {
        let opts = opts.into();
        let pool_constraints = opts.get_pool_constraints().clone();
        let pool = Pool {
            opts,
            inner: Arc::new(Mutex::new(Inner {
                closed: false,
                new: Vec::with_capacity(pool_constraints.min()),
                idle: Vec::new(),
                disconnecting: Vec::new(),
                dropping: Vec::new(),
                rollback: Vec::new(),
                ongoing: 0,
                tasks: Vec::new(),
            })),
            pool_constraints,
        };

        pool
    }

    /// Creates new pool of connections.
    pub fn from_url<T: AsRef<str>>(url: T) -> Result<Pool> {
        let opts = Opts::from_str(url.as_ref())?;
        Ok(Pool::new(opts))
    }

    /// Returns future that resolves to `Conn`.
    pub fn get_conn(&self) -> GetConn {
        new_get_conn(self)
    }

    /// Shortcut for `get_conn` followed by `start_transaction`.
    pub fn start_transaction(
        &self,
        options: TransactionOptions,
    ) -> impl MyFuture<Transaction<Conn>> {
        self.get_conn()
            .and_then(|conn| Queryable::start_transaction(conn, options))
    }

    /// Returns future that disconnects this pool from server and resolves to `()`.
    ///
    /// Active connections taken from this pool should be disconnected manually.
    /// Also all pending and new `GetConn`'s will resolve to error.
    pub fn disconnect(mut self) -> DisconnectPool {
        let become_closed = self.with_inner(|mut inner| {
            if !inner.closed {
                inner.closed = true;
                return true;
            } else {
                return false;
            }
        });
        if become_closed {
            while let Some(conn) = self.take_conn() {
                self.with_inner(move |mut inner| inner.disconnecting.push(conn.disconnect()));
            }
        }
        new_disconnect_pool(self)
    }

    /// Returns true if futures is in queue.
    fn in_queue(&self) -> bool {
        self.with_inner(|inner| {
            let count = inner.new.len()
                + inner.disconnecting.len()
                + inner.dropping.len()
                + inner.rollback.len();
            count > 0
        })
    }

    /// A way to take connection from a pool.
    fn take_conn(&mut self) -> Option<Conn> {
        if self.in_queue() {
            // Do not return connection until queue is empty
            return None;
        }
        self.with_inner(|mut inner| {
            while let Some(mut conn) = inner.idle.pop() {
                if conn.expired() {
                    inner.disconnecting.push(conn.disconnect());
                } else {
                    conn.inner.pool = Some(self.clone());
                    inner.ongoing += 1;
                    return Some(conn);
                }
            }
            None
        })
    }

    /// A way to return connection taken from a pool.
    fn return_conn(&mut self, conn: Conn) {
        let min = self.pool_constraints.min();

        self.with_inner(|mut inner| {
            if inner.closed {
                return;
            }

            if conn.inner.stream.is_none() {
                // drop incomplete connection
                inner.ongoing -= 1;
                return;
            }

            if conn.inner.has_result.is_some() {
                inner.dropping.push(Box::new(conn.drop_result()));
            } else if conn.inner.in_transaction {
                inner.rollback.push(Box::new(conn.rollback_transaction()));
            } else {
                inner.ongoing -= 1;

                if inner.idle.len() >= min {
                    inner.disconnecting.push(conn.disconnect());
                } else {
                    inner.idle.push(conn);
                }
            }

            while let Some(task) = inner.tasks.pop() {
                task.notify()
            }
        });
    }

    fn with_inner<F, T>(&self, fun: F) -> T
    where
        F: FnOnce(MutexGuard<'_, Inner>) -> T,
        T: 'static,
    {
        fun(self.inner.lock().unwrap())
    }

    /// Will manage lifetime of futures stored in a pool.
    fn handle_futures(&mut self) -> Result<()> {
        if !self.in_queue() {
            // There is no futures in queue
            return Ok(());
        }

        let mut handled = false;

        let mut returned_conns: Vec<Conn> = vec![];

        self.with_inner(|mut inner| {
            macro_rules! handle {
                ($vec:ident { $($p:pat => $b:block,)+ }) => ({
                    let len = inner.$vec.len();
                    let mut done_fut_idxs = Vec::new();
                    for i in 0..len {
                        let result = inner.$vec.get_mut(i).unwrap().poll();
                        match result {
                            Ok(Ready(_)) | Err(_) => done_fut_idxs.push(i),
                            _ => (),
                        }

                        let out: Result<()> = match result {
                            $($p => $b),+
                            _ => {
                                Ok(())
                            }
                        };

                        match out {
                            Err(err) => {
                                // early return in case of error
                                while let Some(i) = done_fut_idxs.pop() {
                                    inner.$vec.swap_remove(i);
                                }
                                return Err(err)
                            }
                            _ => (),
                        }
                    }

                    while let Some(i) = done_fut_idxs.pop() {
                        inner.$vec.swap_remove(i);
                    }
                });
            }

            // Handle closing connections.
            handle!(disconnecting {
                Ok(Ready(_)) => {
                    handled = true;
                    Ok(())
                },
                Err(_) => { Ok(()) },
            });

            // Handle dirty connections.
            handle!(dropping {
                Ok(Ready(conn)) => {
                    if inner.closed {
                        inner.disconnecting.push(conn.disconnect());
                    } else {
                        returned_conns.push(conn);
                    }
                    handled = true;
                    Ok(())
                },
                Err(_) => { Ok(()) },
            });

            // Handle in-transaction connections
            handle!(rollback {
                Ok(Ready(conn)) => {
                    if inner.closed {
                        inner.disconnecting.push(conn.disconnect());
                    } else {
                        returned_conns.push(conn);
                    }
                    handled = true;
                    Ok(())
                },
                Err(_) => { Ok(()) },
            });

            // Handle connecting connections.
            handle!(new {
                Ok(Ready(conn)) => {
                    if inner.closed {
                        inner.disconnecting.push(conn.disconnect());
                    } else {
                        inner.ongoing += 1;
                        returned_conns.push(conn);
                    }
                    handled = true;
                    Ok(())
                },
                Err(err) => {
                    if !inner.closed {
                        Err(err)
                    } else {
                        Ok(())
                    }
                },
            });

            Ok(())
        })?;

        for conn in returned_conns {
            self.return_conn(conn);
        }

        if handled {
            self.handle_futures()
        } else {
            Ok(())
        }
    }

    /// Will poll pool for connection.
    fn poll(&mut self) -> Result<Async<Conn>> {
        if self.with_inner(|inner| inner.closed) {
            return Err(DriverError::PoolDisconnected.into());
        }

        self.handle_futures()?;

        match self.take_conn() {
            Some(conn) => Ok(Ready(conn)),
            None => {
                let new_conn_created = self.with_inner(|mut inner| {
                    if inner.new.len() == 0 && inner.conn_count() < self.pool_constraints.max() {
                        let new_conn = Conn::new(self.opts.clone());
                        inner.new.push(Box::new(new_conn));
                        true
                    } else {
                        inner.tasks.push(task::current());
                        false
                    }
                });
                if new_conn_created {
                    self.poll()
                } else {
                    Ok(NotReady)
                }
            }
        }
    }
}

impl Drop for Conn {
    fn drop(&mut self) {
        if let Some(mut pool) = self.inner.pool.take() {
            pool.return_conn(self.take());
        }
    }
}

#[cfg(test)]
mod test {
    use futures::Future;

    use crate::{
        conn::pool::Pool, queryable::Queryable, test_misc::DATABASE_URL, TransactionOptions,
    };

    /// Same as `tokio::run`, but will panic if future panics and will return the result
    /// of future execution.
    fn run<F, T, U>(future: F) -> Result<T, U>
    where
        F: Future<Item = T, Error = U> + Send + 'static,
        T: Send + 'static,
        U: Send + 'static,
    {
        let mut runtime = tokio::runtime::Runtime::new().unwrap();
        let result = runtime.block_on(future);
        runtime.shutdown_on_idle().wait().unwrap();
        result
    }

    #[test]
    fn should_connect() {
        let pool = Pool::new(&**DATABASE_URL);
        let fut = pool
            .get_conn()
            .and_then(|conn| conn.ping().map(|_| ()))
            .and_then(|_| pool.disconnect());

        run(fut).unwrap();
    }

    #[test]
    fn should_start_transaction() {
        let pool = Pool::new(format!("{}?pool_min=1&pool_max=1", &**DATABASE_URL));
        let fut = pool
            .get_conn()
            .and_then(|conn| conn.drop_query("CREATE TABLE IF NOT EXISTS tmp(id int)"))
            .and_then({
                let pool = pool.clone();
                move |_| pool.start_transaction(TransactionOptions::default())
            })
            .and_then(|transaction| {
                transaction.batch_exec("INSERT INTO tmp (id) VALUES (?)", vec![(1,), (2,)])
            })
            .and_then(|transaction| transaction.prep_exec("SELECT * FROM tmp", ()))
            .map(|_| ())
            .and_then({
                let pool = pool.clone();
                move |_| pool.get_conn()
            })
            .and_then(|conn| conn.first("SELECT COUNT(*) FROM tmp"))
            .and_then(|(_, row_opt)| {
                assert_eq!(row_opt, Some((0u8,)));
                pool.get_conn()
                    .and_then(|conn| conn.drop_query("DROP TABLE tmp"))
                    .and_then(move |_| pool.disconnect())
            });

        run(fut).unwrap();
    }

    #[test]
    fn should_hold_bounds() {
        let pool = Pool::new(format!("{}?pool_min=1&pool_max=2", &**DATABASE_URL));
        let pool_clone = pool.clone();
        let fut = pool
            .get_conn()
            .join(pool.get_conn())
            .and_then(move |(mut conn1, conn2)| {
                let new_conn = pool_clone.get_conn();
                conn1.inner.pool.as_mut().unwrap().handle_futures().unwrap();
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.new.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.idle.len()),
                    0
                );
                assert_eq!(
                    conn2
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.disconnecting.len()),
                    0
                );
                assert_eq!(
                    conn2
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.dropping.len()),
                    0
                );
                new_conn
            })
            .and_then(|conn1| {
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.new.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.idle.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.disconnecting.len()),
                    0
                );
                assert_eq!(
                    conn1
                        .inner
                        .pool
                        .as_ref()
                        .unwrap()
                        .with_inner(|inner| inner.dropping.len()),
                    0
                );
                Ok(())
            })
            .and_then(|_| {
                assert_eq!(pool.with_inner(|inner| inner.new.len()), 0);
                assert_eq!(pool.with_inner(|inner| inner.idle.len()), 1);
                assert_eq!(pool.with_inner(|inner| inner.disconnecting.len()), 0);
                assert_eq!(pool.with_inner(|inner| inner.dropping.len()), 0);
                pool.disconnect()
            });

        run(fut).unwrap();
    }

    #[cfg(feature = "nightly")]
    mod bench {
        use futures::Future;

        use crate::{conn::pool::Pool, queryable::Queryable, test_misc::DATABASE_URL};

        #[bench]
        fn connect(bencher: &mut test::Bencher) {
            let mut runtime = tokio::executor::current_thread::CurrentThread::new();
            let pool = Pool::new(&**DATABASE_URL);

            bencher.iter(|| {
                let fut = pool.get_conn().and_then(|conn| conn.ping());
                runtime.block_on(fut).unwrap();
            });

            runtime.block_on(pool.disconnect()).unwrap();
        }
    }
}
<|endoftext|>


<commit_before>/*!
`SHA-256`.

There has been considerable degradation of public confidence in the
security conjectures for many hash functions, including `SHA-256`.
However, for the moment, there do not appear to be alternatives that
inspire satisfactory levels of confidence. One can hope that NIST's
SHA-3 competition will improve the situation.
*/
#[cfg(test)]
extern crate rustc_serialize;
use ffi::{crypto_hash_sha256, crypto_hash_sha256_BYTES};

hash_module!(crypto_hash_sha256,
             crypto_hash_sha256_BYTES,
             64);

#[test]
fn test_vector_1() {
    // hash of empty string
    let x = [];
    let h_expected = [0xe3, 0xb0, 0xc4, 0x42, 0x98, 0xfc, 0x1c, 0x14,
                      0x9a, 0xfb, 0xf4, 0xc8, 0x99, 0x6f, 0xb9, 0x24,
                      0x27, 0xae, 0x41, 0xe4, 0x64, 0x9b, 0x93, 0x4c,
                      0xa4, 0x95, 0x99, 0x1b, 0x78, 0x52, 0xb8, 0x55];
    let Digest(h) = hash(&x);
    assert!(h == h_expected);
}

#[test]
fn test_vector_2() {
    // The quick brown fox jumps over the lazy dog
    let x = [0x54, 0x68, 0x65, 0x20, 0x71, 0x75, 0x69, 0x63,
             0x6b, 0x20, 0x62, 0x72, 0x6f, 0x77, 0x6e, 0x20,
             0x66, 0x6f, 0x78, 0x20, 0x6a, 0x75, 0x6d, 0x70,
             0x73, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x20, 0x74,
             0x68, 0x65, 0x20, 0x6c, 0x61, 0x7a, 0x79, 0x20,
             0x64, 0x6f, 0x67];
    let h_expected = [0xd7, 0xa8, 0xfb, 0xb3, 0x07, 0xd7, 0x80, 0x94,
                      0x69, 0xca, 0x9a, 0xbc, 0xb0, 0x08, 0x2e, 0x4f,
                      0x8d, 0x56, 0x51, 0xe4, 0x6d, 0x3c, 0xdb, 0x76,
                      0x2d, 0x02, 0xd0, 0xbf, 0x37, 0xc9, 0xe5, 0x92];
    let Digest(h) = hash(&x);
    assert!(h == h_expected);
}

#[cfg(test)]
fn test_nist_vector(filename: &str) {
    use self::rustc_serialize::hex::{FromHex};
    use std::fs::File;
    use std::io::{BufRead, BufReader};

    let mut r = BufReader::new(File::open(filename).unwrap());
    let mut line = String::new();
    loop {
        line.clear();
        r.read_line(&mut line).unwrap();
        if line.len() == 0 {
            break;
        }
        let starts_with_len = line.starts_with("Len = ");
        if  starts_with_len {
            let len: usize = line[6..].trim().parse().unwrap();
            line.clear();
            r.read_line(&mut line).unwrap();
            let rawmsg = line[6..].from_hex().unwrap();
            let msg = &rawmsg[..len/8];
            line.clear();
            r.read_line(&mut line).unwrap();
            let md = line[5..].from_hex().unwrap();
            let Digest(digest) = hash(msg);
            assert!(digest == md);
       }
    }
}

#[test]
fn test_vectors_nist_short() {
    test_nist_vector("testvectors/SHA256ShortMsg.rsp");
}

#[test]
fn test_vectors_nist_long() {
    test_nist_vector("testvectors/SHA256LongMsg.rsp");
}
<commit_msg>Fixes for latest rust-nightly<commit_after>/*!
`SHA-256`.

There has been considerable degradation of public confidence in the
security conjectures for many hash functions, including `SHA-256`.
However, for the moment, there do not appear to be alternatives that
inspire satisfactory levels of confidence. One can hope that NIST's
SHA-3 competition will improve the situation.
*/
#[cfg(test)]
extern crate rustc_serialize;
use ffi::{crypto_hash_sha256, crypto_hash_sha256_BYTES};

hash_module!(crypto_hash_sha256,
             crypto_hash_sha256_BYTES,
             64);

#[test]
fn test_vector_1() {
    // hash of empty string
    let x = [];
    let h_expected = [0xe3, 0xb0, 0xc4, 0x42, 0x98, 0xfc, 0x1c, 0x14,
                      0x9a, 0xfb, 0xf4, 0xc8, 0x99, 0x6f, 0xb9, 0x24,
                      0x27, 0xae, 0x41, 0xe4, 0x64, 0x9b, 0x93, 0x4c,
                      0xa4, 0x95, 0x99, 0x1b, 0x78, 0x52, 0xb8, 0x55];
    let Digest(h) = hash(&x);
    assert!(h == h_expected);
}

#[test]
fn test_vector_2() {
    // The quick brown fox jumps over the lazy dog
    let x = [0x54, 0x68, 0x65, 0x20, 0x71, 0x75, 0x69, 0x63,
             0x6b, 0x20, 0x62, 0x72, 0x6f, 0x77, 0x6e, 0x20,
             0x66, 0x6f, 0x78, 0x20, 0x6a, 0x75, 0x6d, 0x70,
             0x73, 0x20, 0x6f, 0x76, 0x65, 0x72, 0x20, 0x74,
             0x68, 0x65, 0x20, 0x6c, 0x61, 0x7a, 0x79, 0x20,
             0x64, 0x6f, 0x67];
    let h_expected = [0xd7, 0xa8, 0xfb, 0xb3, 0x07, 0xd7, 0x80, 0x94,
                      0x69, 0xca, 0x9a, 0xbc, 0xb0, 0x08, 0x2e, 0x4f,
                      0x8d, 0x56, 0x51, 0xe4, 0x6d, 0x3c, 0xdb, 0x76,
                      0x2d, 0x02, 0xd0, 0xbf, 0x37, 0xc9, 0xe5, 0x92];
    let Digest(h) = hash(&x);
    assert!(h == h_expected);
}

#[cfg(test)]
fn test_nist_vector(filename: &str) {
    use self::rustc_serialize::hex::{FromHex};
    use std::fs::File;
    use std::io::{BufRead, BufReader};

    let mut r = BufReader::new(File::open(filename).unwrap());
    let mut line = String::new();
    loop {
        line.clear();
        r.read_line(&mut line).unwrap();
        if line.len() == 0 {
            break;
        }
        let starts_with_len = line.starts_with("Len = ");
        if  starts_with_len {
            let len: usize = line[6..].trim().parse().unwrap();
            line.clear();
            r.read_line(&mut line).unwrap();
            let rawmsg = line[6..].from_hex().unwrap();
            let msg = &rawmsg[..len/8];
            line.clear();
            r.read_line(&mut line).unwrap();
            let md = line[5..].from_hex().unwrap();
            let Digest(digest) = hash(msg);
            assert!(&digest[..] == &md[..]);
       }
    }
}

#[test]
fn test_vectors_nist_short() {
    test_nist_vector("testvectors/SHA256ShortMsg.rsp");
}

#[test]
fn test_vectors_nist_long() {
    test_nist_vector("testvectors/SHA256LongMsg.rsp");
}
<|endoftext|>


<commit_before>extern crate herder;
extern crate serde_json;

use herder::Mastodon;
use herder::oauth::{CreateApp, OAuthApp};

use std::sync::{Arc, Mutex};

fn main() {
    // # Ways to create an app
    //
    // let social_app = CreateApp {
    //     client_name: String::from("herder"),
    //     redirect_uris: String::from("urn:ietf:wg:oauth:2.0:oob"),
    //     scopes: String::from("read write follow")
    // };
    //
    // let social_app = CreateApp::new("herder-app", "urn:ietf:wg:oauth:2.0:oob", "read write follow");
    //
    // let social_app = CreateApp::default();
    //
    let api_url = "https://mastodon.cloud/api/v1/apps";
    let out = Arc::new(Mutex::new(Vec::new()));

    let social_app = CreateApp::new("herder-app", "urn:ietf:wg:oauth:2.0:oob", "read write follow");

    {
        social_app.register_app(api_url, out.clone()).unwrap();

        let result = out.lock().unwrap();
        let herder_app: OAuthApp = serde_json::from_slice(&result).unwrap();

        println!("\n{:?}\n", &herder_app);
        //println!("\n{}\n", serde_json::to_string(&herder_app).unwrap());
    }

    // TODO !!! after successfully getting the herder_app registered,
    // we still need to show the user the link where they can authorize
    // our app. After that, we're good to go using our API!
    //
    // We haven't used herder::Mastodon yet!
    let mastodon = Mastodon(String::from(api_url));
    println!("Mastodon: {:?}", mastodon);
}

<commit_msg>oops... use 'localhost' on examples<commit_after>extern crate herder;
extern crate serde_json;

use herder::Mastodon;
use herder::oauth::{CreateApp, OAuthApp};

use std::sync::{Arc, Mutex};

fn main() {
    // # Ways to create an app
    //
    // let social_app = CreateApp {
    //     client_name: String::from("herder"),
    //     redirect_uris: String::from("urn:ietf:wg:oauth:2.0:oob"),
    //     scopes: String::from("read write follow")
    // };
    //
    // let social_app = CreateApp::new("herder-app", "urn:ietf:wg:oauth:2.0:oob", "read write follow");
    //
    // let social_app = CreateApp::default();
    //
    let api_url = "http://localhost:3000/api/v1/apps";
    let out = Arc::new(Mutex::new(Vec::new()));

    let social_app = CreateApp::new("herder-app", "urn:ietf:wg:oauth:2.0:oob", "read write follow");

    {
        social_app.register_app(api_url, out.clone()).unwrap();

        let result = out.lock().unwrap();
        let herder_app: OAuthApp = serde_json::from_slice(&result).unwrap();

        println!("\n{:?}\n", &herder_app);
        //println!("\n{}\n", serde_json::to_string(&herder_app).unwrap());
    }

    // TODO !!! after successfully getting the herder_app registered,
    // we still need to show the user the link where they can authorize
    // our app. After that, we're good to go using our API!
    //
    // We haven't used herder::Mastodon yet!
    let mastodon = Mastodon(String::from(api_url));
    println!("Mastodon: {:?}", mastodon);
}

<|endoftext|>


<commit_before>#![deny(warnings)]
extern crate futures;
extern crate hyper;
extern crate pretty_env_logger;

use futures::future::FutureResult;

use hyper::{Get, Post, StatusCode};
use hyper::header::ContentLength;
use hyper::server::{Http, Service, Request, Response};

static INDEX: &'static [u8] = b"Try POST /echo";

#[derive(Clone, Copy)]
struct Echo;

impl Service for Echo {
    type Request = Request;
    type Response = Response;
    type Error = hyper::Error;
    type Future = FutureResult<Response, hyper::Error>;

    fn call(&self, req: Request) -> Self::Future {
        futures::future::ok(match (req.method(), req.path()) {
            (&Get, "/") | (&Get, "/echo") => {
                Response::new()
                    .with_header(ContentLength(INDEX.len() as u64))
                    .with_body(INDEX)
            },
            (&Post, "/echo") => {
                let mut res = Response::new();
                if let Some(len) = req.headers().get::<ContentLength>() {
                    res.headers_mut().set(len.clone());
                }
                res.with_body(req.body())
            },
            _ => {
                Response::new()
                    .with_status(StatusCode::NotFound)
            }
        })
    }

}

fn main() {
    pretty_env_logger::init().unwrap();
    let addr = "127.0.0.1:1337".parse().unwrap();

    let server = Http::new().bind(&addr, || Ok(Echo)).unwrap();
    println!("Listening on http://{} with 1 thread.", server.local_addr().unwrap());
    server.run().unwrap();
}
<commit_msg>docs(examples): remove unnecessary Clone and Copy from Echo<commit_after>#![deny(warnings)]
extern crate futures;
extern crate hyper;
extern crate pretty_env_logger;

use futures::future::FutureResult;

use hyper::{Get, Post, StatusCode};
use hyper::header::ContentLength;
use hyper::server::{Http, Service, Request, Response};

static INDEX: &'static [u8] = b"Try POST /echo";

struct Echo;

impl Service for Echo {
    type Request = Request;
    type Response = Response;
    type Error = hyper::Error;
    type Future = FutureResult<Response, hyper::Error>;

    fn call(&self, req: Request) -> Self::Future {
        futures::future::ok(match (req.method(), req.path()) {
            (&Get, "/") | (&Get, "/echo") => {
                Response::new()
                    .with_header(ContentLength(INDEX.len() as u64))
                    .with_body(INDEX)
            },
            (&Post, "/echo") => {
                let mut res = Response::new();
                if let Some(len) = req.headers().get::<ContentLength>() {
                    res.headers_mut().set(len.clone());
                }
                res.with_body(req.body())
            },
            _ => {
                Response::new()
                    .with_status(StatusCode::NotFound)
            }
        })
    }

}

fn main() {
    pretty_env_logger::init().unwrap();
    let addr = "127.0.0.1:1337".parse().unwrap();

    let server = Http::new().bind(&addr, || Ok(Echo)).unwrap();
    println!("Listening on http://{} with 1 thread.", server.local_addr().unwrap());
    server.run().unwrap();
}
<|endoftext|>


<commit_before>use std::{
    ops::Deref,
    sync::{atomic::AtomicU64, Arc, RwLock},
};

use pagecache::FastMap8;

use super::*;

/// The `sled` embedded database!
#[derive(Clone)]
pub struct Db {
    context: Context,
    default: Arc<Tree>,
    tenants: Arc<RwLock<FastMap8<Vec<u8>, Arc<Tree>>>>,
}

unsafe impl Send for Db {}

unsafe impl Sync for Db {}

impl Deref for Db {
    type Target = Tree;

    fn deref(&self) -> &Tree {
        &self.default
    }
}

impl std::fmt::Debug for Db {
    fn fmt(
        &self,
        f: &mut std::fmt::Formatter<'_>,
    ) -> std::result::Result<(), std::fmt::Error> {
        let tenants = self.tenants.read().unwrap();
        for (name, tree) in tenants.iter() {
            write!(f, "tree: {:?} contents: {:?}", name, tree)?;
        }
        Ok(())
    }
}

impl Db {
    /// Load existing or create a new `Db` with a default configuration.
    pub fn start_default<P: AsRef<std::path::Path>>(path: P) -> Result<Db> {
        let config = ConfigBuilder::new().path(path).build();
        Self::start(config)
    }

    /// Load existing or create a new `Db`.
    pub fn start(config: Config) -> Result<Db> {
        let _measure = Measure::new(&M.tree_start);

        let context = Context::start(config)?;

        let flusher_pagecache = context.pagecache.clone();
        let flusher = context.flush_every_ms.map(move |fem| {
            flusher::Flusher::new(
                "log flusher".to_owned(),
                flusher_pagecache,
                fem,
            )
        });
        *context._flusher.lock().unwrap() = flusher;

        // create or open the default tree
        let tx = context.pagecache.begin()?;
        let default = Arc::new(meta::open_tree(
            context.clone(),
            DEFAULT_TREE_ID.to_vec(),
            &tx,
        )?);

        let ret = Db {
            context: context.clone(),
            default,
            tenants: Arc::new(RwLock::new(FastMap8::default())),
        };

        let mut tenants = ret.tenants.write().unwrap();

        for (id, root) in context.pagecache.meta(&tx)?.tenants().into_iter() {
            let tree = Tree {
                tree_id: id.clone(),
                subscriptions: Arc::new(Subscriptions::default()),
                context: context.clone(),
                root: Arc::new(AtomicU64::new(root)),
            };
            tenants.insert(id, Arc::new(tree));
        }

        drop(tenants);

        #[cfg(feature = "event_log")]
        context.event_log.meta_after_restart(
            context
                .pagecache
                .meta(&tx)
                .expect("should be able to get meta under test")
                .clone(),
        );

        Ok(ret)
    }

    /// Open or create a new disk-backed Tree with its own keyspace,
    /// accessible from the `Db` via the provided identifier.
    pub fn open_tree<V: AsRef<[u8]>>(&self, name: V) -> Result<Arc<Tree>> {
        let name = name.as_ref();
        let tenants = self.tenants.read().unwrap();
        if let Some(tree) = tenants.get(name) {
            return Ok(tree.clone());
        }
        drop(tenants);

        let tx = self.context.pagecache.begin()?;

        let mut tenants = self.tenants.write().unwrap();
        let tree = Arc::new(meta::open_tree(
            self.context.clone(),
            name.to_vec(),
            &tx,
        )?);
        tenants.insert(name.to_vec(), tree.clone());
        drop(tenants);
        Ok(tree)
    }

    /// Remove a disk-backed collection.
    pub fn drop_tree(&self, name: &[u8]) -> Result<bool> {
        if name == DEFAULT_TREE_ID {
            return Err(Error::Unsupported(
                "cannot remove the core structures".into(),
            ));
        }
        trace!("dropping tree {:?}", name,);

        let mut tenants = self.tenants.write().unwrap();

        let tree = if let Some(tree) = tenants.remove(&*name) {
            tree
        } else {
            return Ok(false);
        };

        let tx = self.context.pagecache.begin()?;

        let mut root_id =
            Some(self.context.pagecache.meta_pid_for_name(&name, &tx)?);

        let mut leftmost_chain: Vec<PageId> = vec![root_id.unwrap()];
        let mut cursor = root_id.unwrap();
        loop {
            let view_opt = self.view_for_pid(cursor, &tx)?;
            let view = if let Some(view) = view_opt {
                view
            } else {
                break;
            };

            let node = view.compact(&self.context);

            if let Some(index) = node.data.index_ref() {
                let leftmost_child = index[0].1;
                leftmost_chain.push(leftmost_child);
                cursor = leftmost_child;
            } else {
                break;
            }
        }

        loop {
            let res = self.context.pagecache.cas_root_in_meta(
                name.to_vec(),
                root_id,
                None,
                &tx,
            )?;

            if let Err(actual_root) = res {
                root_id = actual_root;
            } else {
                break;
            }
        }

        tree.root
            .store(u64::max_value(), std::sync::atomic::Ordering::SeqCst);

        // drop writer lock
        drop(tenants);

        tree.gc_pages(leftmost_chain)?;

        tx.flush();

        Ok(true)
    }

    /// Returns the trees names saved in this Db.
    pub fn tree_names(&self) -> Vec<Vec<u8>> {
        let tenants = self.tenants.read().unwrap();
        tenants.iter().map(|(name, _)| name.clone()).collect()
    }

    /// Returns `true` if the database was
    /// recovered from a previous process.
    /// Note that database state is only
    /// guaranteed to be present up to the
    /// last call to `flush`! Otherwise state
    /// is synced to disk periodically if the
    /// `sync_every_ms` configuration option
    /// is set to `Some(number_of_ms_between_syncs)`
    /// or if the IO buffer gets filled to
    /// capacity before being rotated.
    pub fn was_recovered(&self) -> bool {
        self.context.was_recovered()
    }

    /// Generate a monotonic ID. Not guaranteed to be
    /// contiguous. Written to disk every `idgen_persist_interval`
    /// operations, followed by a blocking flush. During recovery, we
    /// take the last recovered generated ID and add 2x
    /// the `idgen_persist_interval` to it. While persisting, if the
    /// previous persisted counter wasn't synced to disk yet, we will do
    /// a blocking flush to fsync the latest counter, ensuring
    /// that we will never give out the same counter twice.
    pub fn generate_id(&self) -> Result<u64> {
        self.context.generate_id()
    }
}
<commit_msg>Improve Debug impl on Db<commit_after>use std::{
    ops::Deref,
    sync::{atomic::AtomicU64, Arc, RwLock},
};

use pagecache::FastMap8;

use super::*;

/// The `sled` embedded database!
#[derive(Clone)]
pub struct Db {
    context: Context,
    default: Arc<Tree>,
    tenants: Arc<RwLock<FastMap8<Vec<u8>, Arc<Tree>>>>,
}

unsafe impl Send for Db {}

unsafe impl Sync for Db {}

impl Deref for Db {
    type Target = Tree;

    fn deref(&self) -> &Tree {
        &self.default
    }
}

impl std::fmt::Debug for Db {
    fn fmt(
        &self,
        f: &mut std::fmt::Formatter<'_>,
    ) -> std::result::Result<(), std::fmt::Error> {
        let tenants = self.tenants.read().unwrap();
        for (raw_name, tree) in tenants.iter() {
            let name = std::str::from_utf8(&raw_name)
                .map(String::from)
                .unwrap_or_else(|_| format!("{:?}", raw_name));
            write!(f, "tree: {:?} contents: {:?}", name, tree)?;
        }
        Ok(())
    }
}

impl Db {
    /// Load existing or create a new `Db` with a default configuration.
    pub fn start_default<P: AsRef<std::path::Path>>(path: P) -> Result<Db> {
        let config = ConfigBuilder::new().path(path).build();
        Self::start(config)
    }

    /// Load existing or create a new `Db`.
    pub fn start(config: Config) -> Result<Db> {
        let _measure = Measure::new(&M.tree_start);

        let context = Context::start(config)?;

        let flusher_pagecache = context.pagecache.clone();
        let flusher = context.flush_every_ms.map(move |fem| {
            flusher::Flusher::new(
                "log flusher".to_owned(),
                flusher_pagecache,
                fem,
            )
        });
        *context._flusher.lock().unwrap() = flusher;

        // create or open the default tree
        let tx = context.pagecache.begin()?;
        let default = Arc::new(meta::open_tree(
            context.clone(),
            DEFAULT_TREE_ID.to_vec(),
            &tx,
        )?);

        let ret = Db {
            context: context.clone(),
            default,
            tenants: Arc::new(RwLock::new(FastMap8::default())),
        };

        let mut tenants = ret.tenants.write().unwrap();

        for (id, root) in context.pagecache.meta(&tx)?.tenants().into_iter() {
            let tree = Tree {
                tree_id: id.clone(),
                subscriptions: Arc::new(Subscriptions::default()),
                context: context.clone(),
                root: Arc::new(AtomicU64::new(root)),
            };
            tenants.insert(id, Arc::new(tree));
        }

        drop(tenants);

        #[cfg(feature = "event_log")]
        context.event_log.meta_after_restart(
            context
                .pagecache
                .meta(&tx)
                .expect("should be able to get meta under test")
                .clone(),
        );

        Ok(ret)
    }

    /// Open or create a new disk-backed Tree with its own keyspace,
    /// accessible from the `Db` via the provided identifier.
    pub fn open_tree<V: AsRef<[u8]>>(&self, name: V) -> Result<Arc<Tree>> {
        let name = name.as_ref();
        let tenants = self.tenants.read().unwrap();
        if let Some(tree) = tenants.get(name) {
            return Ok(tree.clone());
        }
        drop(tenants);

        let tx = self.context.pagecache.begin()?;

        let mut tenants = self.tenants.write().unwrap();
        let tree = Arc::new(meta::open_tree(
            self.context.clone(),
            name.to_vec(),
            &tx,
        )?);
        tenants.insert(name.to_vec(), tree.clone());
        drop(tenants);
        Ok(tree)
    }

    /// Remove a disk-backed collection.
    pub fn drop_tree(&self, name: &[u8]) -> Result<bool> {
        if name == DEFAULT_TREE_ID {
            return Err(Error::Unsupported(
                "cannot remove the core structures".into(),
            ));
        }
        trace!("dropping tree {:?}", name,);

        let mut tenants = self.tenants.write().unwrap();

        let tree = if let Some(tree) = tenants.remove(&*name) {
            tree
        } else {
            return Ok(false);
        };

        let tx = self.context.pagecache.begin()?;

        let mut root_id =
            Some(self.context.pagecache.meta_pid_for_name(&name, &tx)?);

        let mut leftmost_chain: Vec<PageId> = vec![root_id.unwrap()];
        let mut cursor = root_id.unwrap();
        loop {
            let view_opt = self.view_for_pid(cursor, &tx)?;
            let view = if let Some(view) = view_opt {
                view
            } else {
                break;
            };

            let node = view.compact(&self.context);

            if let Some(index) = node.data.index_ref() {
                let leftmost_child = index[0].1;
                leftmost_chain.push(leftmost_child);
                cursor = leftmost_child;
            } else {
                break;
            }
        }

        loop {
            let res = self.context.pagecache.cas_root_in_meta(
                name.to_vec(),
                root_id,
                None,
                &tx,
            )?;

            if let Err(actual_root) = res {
                root_id = actual_root;
            } else {
                break;
            }
        }

        tree.root
            .store(u64::max_value(), std::sync::atomic::Ordering::SeqCst);

        // drop writer lock
        drop(tenants);

        tree.gc_pages(leftmost_chain)?;

        tx.flush();

        Ok(true)
    }

    /// Returns the trees names saved in this Db.
    pub fn tree_names(&self) -> Vec<Vec<u8>> {
        let tenants = self.tenants.read().unwrap();
        tenants.iter().map(|(name, _)| name.clone()).collect()
    }

    /// Returns `true` if the database was
    /// recovered from a previous process.
    /// Note that database state is only
    /// guaranteed to be present up to the
    /// last call to `flush`! Otherwise state
    /// is synced to disk periodically if the
    /// `sync_every_ms` configuration option
    /// is set to `Some(number_of_ms_between_syncs)`
    /// or if the IO buffer gets filled to
    /// capacity before being rotated.
    pub fn was_recovered(&self) -> bool {
        self.context.was_recovered()
    }

    /// Generate a monotonic ID. Not guaranteed to be
    /// contiguous. Written to disk every `idgen_persist_interval`
    /// operations, followed by a blocking flush. During recovery, we
    /// take the last recovered generated ID and add 2x
    /// the `idgen_persist_interval` to it. While persisting, if the
    /// previous persisted counter wasn't synced to disk yet, we will do
    /// a blocking flush to fsync the latest counter, ensuring
    /// that we will never give out the same counter twice.
    pub fn generate_id(&self) -> Result<u64> {
        self.context.generate_id()
    }
}
<|endoftext|>


<commit_before><commit_msg>Fix several bugs in civilization reading<commit_after><|endoftext|>