<commit_before>//! The module contains helpers for implementing TLS-based client connections.
//!
//! Available only when the `"tls"` crate feature is enabled.
//!
//! Depends on the `openssl` crate.
//!
//! # Example
//!
//! Establishing a new client connection using the `TlsConnector` and issuing a `GET` request.
//!
//! ```no_run
//! // Remember to enable the "tls" feature for `solicit`
//! use solicit::http::client::tls::TlsConnector;
//! use solicit::client::SimpleClient;
//! use std::str;
//!
//! // Connect to an HTTP/2 aware server
//! let path = "/path/to/certs.pem";
//! let connector = TlsConnector::new("http2bin.org", &path);
//! let mut client = SimpleClient::with_connector(connector).unwrap();
//! let response = client.get(b"/get", &[]).unwrap();
//! assert_eq!(response.stream_id, 1);
//! assert_eq!(response.status_code().unwrap(), 200);
//! // Dump the headers and the response body to stdout.
//! // They are returned as raw bytes for the user to do as they please.
//! // (Note: in general directly decoding assuming a utf8 encoding might not
//! // always work -- this is meant as a simple example that shows that the
//! // response is well formed.)
//! for header in response.headers.iter() {
//!     println!("{}: {}",
//!         str::from_utf8(&header.0).unwrap(),
//!         str::from_utf8(&header.1).unwrap());
//! }
//! println!("{}", str::from_utf8(&response.body).unwrap());
//! ```

use std::convert::AsRef;
use std::net::TcpStream;
use std::path::Path;
use std::error;
use std::fmt;
use std::str;
use std::io;
use http::{HttpScheme, ALPN_PROTOCOLS};

use super::{ClientStream, write_preface, HttpConnect, HttpConnectError};

use openssl::ssl::{Ssl, SslStream, SslContext};
use openssl::ssl::{SSL_VERIFY_PEER, SSL_VERIFY_FAIL_IF_NO_PEER_CERT};
use openssl::ssl::SSL_OP_NO_COMPRESSION;
use openssl::ssl::error::SslError;
use openssl::ssl::SslMethod;

/// A struct implementing the functionality of establishing a TLS-backed TCP stream
/// that can be used by an HTTP/2 connection. Takes care to set all the TLS options
/// to those allowed by the HTTP/2 spec, as well as of the protocol negotiation.
///
/// # Example
///
/// Issue a GET request over `https` using the `TlsConnector`
///
/// ```no_run
/// use solicit::http::client::tls::TlsConnector;
/// use solicit::client::SimpleClient;
/// use std::str;
///
/// // Connect to an HTTP/2 aware server
/// let path = "/path/to/certs.pem";
/// let connector = TlsConnector::new("http2bin.org", &path);
/// let mut client = SimpleClient::with_connector(connector).unwrap();
/// let response = client.get(b"/get", &[]).unwrap();
/// assert_eq!(response.stream_id, 1);
/// assert_eq!(response.status_code().unwrap(), 200);
/// // Dump the headers and the response body to stdout.
/// // They are returned as raw bytes for the user to do as they please.
/// // (Note: in general directly decoding assuming a utf8 encoding might not
/// // always work -- this is meant as a simple example that shows that the
/// // response is well formed.)
/// for header in response.headers.iter() {
///     println!("{}: {}",
///         str::from_utf8(&header.0).unwrap(),
///         str::from_utf8(&header.1).unwrap());
/// }
/// println!("{}", str::from_utf8(&response.body).unwrap());
/// ```
pub struct TlsConnector<'a, 'ctx> {
    pub host: &'a str,
    context: Http2TlsContext<'ctx>,
}

/// A private enum that represents the two options for configuring the
/// `TlsConnector`
enum Http2TlsContext<'a> {
    /// This means that the `TlsConnector` will use the referenced `SslContext`
    /// instance when creating a new `SslStream`
    Wrapped(&'a SslContext),
    /// This means that the `TlsConnector` will create a new context with the
    /// certificates file being found at the given path.
    CertPath(&'a Path),
}

/// An enum representing possible errors that can arise when trying to
/// establish an HTTP/2 connection over TLS.
pub enum TlsConnectError {
    /// The variant corresponds to the underlying raw TCP connection returning
    /// an error.
    IoError(io::Error),
    /// The variant corresponds to the TLS negotiation returning an error.
    SslError(SslError),
    /// The variant corresponds to the case when the TLS connection is
    /// established, but the application protocol that was negotiated didn't
    /// end up being HTTP/2.
    /// It wraps the established SSL stream in order to allow the client to
    /// decide what to do with it (and the application protocol that was
    /// chosen).
    Http2NotSupported(SslStream<TcpStream>),
}

// Note: TcpStream does not implement `Debug` in 1.0.0, so deriving is not possible.
impl fmt::Debug for TlsConnectError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        // The enum variant...
        try!(write!(fmt, "TlsConnectError::{}", match *self {
            TlsConnectError::IoError(_) => "IoError",
            TlsConnectError::SslError(_) => "SslError",
            TlsConnectError::Http2NotSupported(_) => "Http2NotSupported",
        }));
        // ...and the wrapped value, except for when it's the stream.
        match *self {
            TlsConnectError::IoError(ref err) => try!(write!(fmt, "({:?})", err)),
            TlsConnectError::SslError(ref err) => try!(write!(fmt, "({:?})", err)),
            TlsConnectError::Http2NotSupported(_) => try!(write!(fmt, "(...)")),
        };

        Ok(())
    }
}

impl fmt::Display for TlsConnectError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "TLS HTTP/2 connect error: {}", (self as &error::Error).description())
    }
}

impl error::Error for TlsConnectError {
    fn description(&self) -> &str {
        match *self {
            TlsConnectError::IoError(ref err) => err.description(),
            TlsConnectError::SslError(ref err) => err.description(),
            TlsConnectError::Http2NotSupported(_) => "HTTP/2 not supported by the server",
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            TlsConnectError::IoError(ref err) => Some(err),
            TlsConnectError::SslError(ref err) => Some(err),
            TlsConnectError::Http2NotSupported(_) => None,
        }
    }
}

impl From<io::Error> for TlsConnectError {
    fn from(err: io::Error) -> TlsConnectError {
        TlsConnectError::IoError(err)
    }
}

impl From<SslError> for TlsConnectError {
    fn from(err: SslError) -> TlsConnectError {
        TlsConnectError::SslError(err)
    }
}

impl HttpConnectError for TlsConnectError {}

impl<'a, 'ctx> TlsConnector<'a, 'ctx> {
    /// Creates a new `TlsConnector` that will create a new `SslContext` before
    /// trying to establish the TLS connection. The path to the CA file that the
    /// context will use needs to be provided.
    pub fn new<P: AsRef<Path>>(host: &'a str, ca_file_path: &'ctx P) -> TlsConnector<'a, 'ctx> {
        TlsConnector {
            host: host,
            context: Http2TlsContext::CertPath(ca_file_path.as_ref()),
        }
    }

    /// Creates a new `TlsConnector` that will use the provided context to
    /// create the `SslStream` that will back the HTTP/2 connection.
    pub fn with_context(host: &'a str, context: &'ctx SslContext) -> TlsConnector<'a, 'ctx> {
        TlsConnector {
            host: host,
            context: Http2TlsContext::Wrapped(context),
        }
    }

    /// Builds up a default `SslContext` instance wth TLS settings that the
    /// HTTP/2 spec mandates. The path to the CA file needs to be provided.
    pub fn build_default_context(ca_file_path: &Path) -> Result<SslContext, TlsConnectError> {
        // HTTP/2 connections need to be on top of TLSv1.2 or newer.
        let mut context = try!(SslContext::new(SslMethod::Tlsv1_2));

        // This makes the certificate required (only VERIFY_PEER would mean optional)
        context.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
        try!(context.set_CA_file(ca_file_path));
        // Compression is not allowed by the spec
        context.set_options(SSL_OP_NO_COMPRESSION);
        // The HTTP/2 protocol identifiers are constant at the library level...
        context.set_npn_protocols(ALPN_PROTOCOLS);

        Ok(context)
    }
}

impl<'a, 'ctx> HttpConnect for TlsConnector<'a, 'ctx> {
    type Stream = SslStream<TcpStream>;
    type Err = TlsConnectError;

    fn connect(self) -> Result<ClientStream<SslStream<TcpStream>>, TlsConnectError> {
        // First, create a TCP connection to port 443
        let raw_tcp = try!(TcpStream::connect(&(self.host, 443)));
        // Now build the SSL instance, depending on which SSL context should be
        // used...
        let ssl = match self.context {
            Http2TlsContext::CertPath(path) => {
                let ctx = try!(TlsConnector::build_default_context(&path));
                try!(Ssl::new(&ctx))
            },
            Http2TlsContext::Wrapped(ctx) => try!(Ssl::new(ctx)),
        };
        // SNI must be used
        try!(ssl.set_hostname(self.host));

        // Wrap the Ssl instance into an `SslStream`
        let mut ssl_stream = try!(SslStream::new_from(ssl, raw_tcp));
        // This connector only understands HTTP/2, so if that wasn't chosen in
        // NPN, we raise an error.
        let fail = match ssl_stream.get_selected_npn_protocol() {
            None => true,
            Some(proto) => {
                // Make sure that the protocol is one of the HTTP/2 protocols.
                debug!("Selected protocol -> {:?}", str::from_utf8(proto));
                let found = ALPN_PROTOCOLS.iter().any(|&http2_proto| http2_proto == proto);

                // We fail if we don't find an HTTP/2 protcol match...
                !found
            }
        };
        if fail {
            // We need the fail flag (instead of returning from one of the match
            // arms above because we need to move the `ssl_stream` and that is
            // not possible above (since it's borrowed at that point).
            return Err(TlsConnectError::Http2NotSupported(ssl_stream));
        }

        // Now that the stream is correctly established, we write the client preface.
        try!(write_preface(&mut ssl_stream));

        // All done.
        Ok(ClientStream(ssl_stream, HttpScheme::Https, self.host.into()))
    }
}
<commit_msg>client/tls: Fix examples to use new Header struct<commit_after>//! The module contains helpers for implementing TLS-based client connections.
//!
//! Available only when the `"tls"` crate feature is enabled.
//!
//! Depends on the `openssl` crate.
//!
//! # Example
//!
//! Establishing a new client connection using the `TlsConnector` and issuing a `GET` request.
//!
//! ```no_run
//! // Remember to enable the "tls" feature for `solicit`
//! use solicit::http::client::tls::TlsConnector;
//! use solicit::client::SimpleClient;
//! use std::str;
//!
//! // Connect to an HTTP/2 aware server
//! let path = "/path/to/certs.pem";
//! let connector = TlsConnector::new("http2bin.org", &path);
//! let mut client = SimpleClient::with_connector(connector).unwrap();
//! let response = client.get(b"/get", &[]).unwrap();
//! assert_eq!(response.stream_id, 1);
//! assert_eq!(response.status_code().unwrap(), 200);
//! // Dump the headers and the response body to stdout.
//! // They are returned as raw bytes for the user to do as they please.
//! // (Note: in general directly decoding assuming a utf8 encoding might not
//! // always work -- this is meant as a simple example that shows that the
//! // response is well formed.)
//! for header in response.headers.iter() {
//!     println!("{}: {}",
//!         str::from_utf8(header.name()).unwrap(),
//!         str::from_utf8(header.value()).unwrap());
//! }
//! println!("{}", str::from_utf8(&response.body).unwrap());
//! ```

use std::convert::AsRef;
use std::net::TcpStream;
use std::path::Path;
use std::error;
use std::fmt;
use std::str;
use std::io;
use http::{HttpScheme, ALPN_PROTOCOLS};

use super::{ClientStream, write_preface, HttpConnect, HttpConnectError};

use openssl::ssl::{Ssl, SslStream, SslContext};
use openssl::ssl::{SSL_VERIFY_PEER, SSL_VERIFY_FAIL_IF_NO_PEER_CERT};
use openssl::ssl::SSL_OP_NO_COMPRESSION;
use openssl::ssl::error::SslError;
use openssl::ssl::SslMethod;

/// A struct implementing the functionality of establishing a TLS-backed TCP stream
/// that can be used by an HTTP/2 connection. Takes care to set all the TLS options
/// to those allowed by the HTTP/2 spec, as well as of the protocol negotiation.
///
/// # Example
///
/// Issue a GET request over `https` using the `TlsConnector`
///
/// ```no_run
/// use solicit::http::client::tls::TlsConnector;
/// use solicit::client::SimpleClient;
/// use std::str;
///
/// // Connect to an HTTP/2 aware server
/// let path = "/path/to/certs.pem";
/// let connector = TlsConnector::new("http2bin.org", &path);
/// let mut client = SimpleClient::with_connector(connector).unwrap();
/// let response = client.get(b"/get", &[]).unwrap();
/// assert_eq!(response.stream_id, 1);
/// assert_eq!(response.status_code().unwrap(), 200);
/// // Dump the headers and the response body to stdout.
/// // They are returned as raw bytes for the user to do as they please.
/// // (Note: in general directly decoding assuming a utf8 encoding might not
/// // always work -- this is meant as a simple example that shows that the
/// // response is well formed.)
/// for header in response.headers.iter() {
///     println!("{}: {}",
///         str::from_utf8(header.name()).unwrap(),
///         str::from_utf8(header.value()).unwrap());
/// }
/// println!("{}", str::from_utf8(&response.body).unwrap());
/// ```
pub struct TlsConnector<'a, 'ctx> {
    pub host: &'a str,
    context: Http2TlsContext<'ctx>,
}

/// A private enum that represents the two options for configuring the
/// `TlsConnector`
enum Http2TlsContext<'a> {
    /// This means that the `TlsConnector` will use the referenced `SslContext`
    /// instance when creating a new `SslStream`
    Wrapped(&'a SslContext),
    /// This means that the `TlsConnector` will create a new context with the
    /// certificates file being found at the given path.
    CertPath(&'a Path),
}

/// An enum representing possible errors that can arise when trying to
/// establish an HTTP/2 connection over TLS.
pub enum TlsConnectError {
    /// The variant corresponds to the underlying raw TCP connection returning
    /// an error.
    IoError(io::Error),
    /// The variant corresponds to the TLS negotiation returning an error.
    SslError(SslError),
    /// The variant corresponds to the case when the TLS connection is
    /// established, but the application protocol that was negotiated didn't
    /// end up being HTTP/2.
    /// It wraps the established SSL stream in order to allow the client to
    /// decide what to do with it (and the application protocol that was
    /// chosen).
    Http2NotSupported(SslStream<TcpStream>),
}

// Note: TcpStream does not implement `Debug` in 1.0.0, so deriving is not possible.
impl fmt::Debug for TlsConnectError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        // The enum variant...
        try!(write!(fmt, "TlsConnectError::{}", match *self {
            TlsConnectError::IoError(_) => "IoError",
            TlsConnectError::SslError(_) => "SslError",
            TlsConnectError::Http2NotSupported(_) => "Http2NotSupported",
        }));
        // ...and the wrapped value, except for when it's the stream.
        match *self {
            TlsConnectError::IoError(ref err) => try!(write!(fmt, "({:?})", err)),
            TlsConnectError::SslError(ref err) => try!(write!(fmt, "({:?})", err)),
            TlsConnectError::Http2NotSupported(_) => try!(write!(fmt, "(...)")),
        };

        Ok(())
    }
}

impl fmt::Display for TlsConnectError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "TLS HTTP/2 connect error: {}", (self as &error::Error).description())
    }
}

impl error::Error for TlsConnectError {
    fn description(&self) -> &str {
        match *self {
            TlsConnectError::IoError(ref err) => err.description(),
            TlsConnectError::SslError(ref err) => err.description(),
            TlsConnectError::Http2NotSupported(_) => "HTTP/2 not supported by the server",
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match *self {
            TlsConnectError::IoError(ref err) => Some(err),
            TlsConnectError::SslError(ref err) => Some(err),
            TlsConnectError::Http2NotSupported(_) => None,
        }
    }
}

impl From<io::Error> for TlsConnectError {
    fn from(err: io::Error) -> TlsConnectError {
        TlsConnectError::IoError(err)
    }
}

impl From<SslError> for TlsConnectError {
    fn from(err: SslError) -> TlsConnectError {
        TlsConnectError::SslError(err)
    }
}

impl HttpConnectError for TlsConnectError {}

impl<'a, 'ctx> TlsConnector<'a, 'ctx> {
    /// Creates a new `TlsConnector` that will create a new `SslContext` before
    /// trying to establish the TLS connection. The path to the CA file that the
    /// context will use needs to be provided.
    pub fn new<P: AsRef<Path>>(host: &'a str, ca_file_path: &'ctx P) -> TlsConnector<'a, 'ctx> {
        TlsConnector {
            host: host,
            context: Http2TlsContext::CertPath(ca_file_path.as_ref()),
        }
    }

    /// Creates a new `TlsConnector` that will use the provided context to
    /// create the `SslStream` that will back the HTTP/2 connection.
    pub fn with_context(host: &'a str, context: &'ctx SslContext) -> TlsConnector<'a, 'ctx> {
        TlsConnector {
            host: host,
            context: Http2TlsContext::Wrapped(context),
        }
    }

    /// Builds up a default `SslContext` instance wth TLS settings that the
    /// HTTP/2 spec mandates. The path to the CA file needs to be provided.
    pub fn build_default_context(ca_file_path: &Path) -> Result<SslContext, TlsConnectError> {
        // HTTP/2 connections need to be on top of TLSv1.2 or newer.
        let mut context = try!(SslContext::new(SslMethod::Tlsv1_2));

        // This makes the certificate required (only VERIFY_PEER would mean optional)
        context.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
        try!(context.set_CA_file(ca_file_path));
        // Compression is not allowed by the spec
        context.set_options(SSL_OP_NO_COMPRESSION);
        // The HTTP/2 protocol identifiers are constant at the library level...
        context.set_npn_protocols(ALPN_PROTOCOLS);

        Ok(context)
    }
}

impl<'a, 'ctx> HttpConnect for TlsConnector<'a, 'ctx> {
    type Stream = SslStream<TcpStream>;
    type Err = TlsConnectError;

    fn connect(self) -> Result<ClientStream<SslStream<TcpStream>>, TlsConnectError> {
        // First, create a TCP connection to port 443
        let raw_tcp = try!(TcpStream::connect(&(self.host, 443)));
        // Now build the SSL instance, depending on which SSL context should be
        // used...
        let ssl = match self.context {
            Http2TlsContext::CertPath(path) => {
                let ctx = try!(TlsConnector::build_default_context(&path));
                try!(Ssl::new(&ctx))
            },
            Http2TlsContext::Wrapped(ctx) => try!(Ssl::new(ctx)),
        };
        // SNI must be used
        try!(ssl.set_hostname(self.host));

        // Wrap the Ssl instance into an `SslStream`
        let mut ssl_stream = try!(SslStream::new_from(ssl, raw_tcp));
        // This connector only understands HTTP/2, so if that wasn't chosen in
        // NPN, we raise an error.
        let fail = match ssl_stream.get_selected_npn_protocol() {
            None => true,
            Some(proto) => {
                // Make sure that the protocol is one of the HTTP/2 protocols.
                debug!("Selected protocol -> {:?}", str::from_utf8(proto));
                let found = ALPN_PROTOCOLS.iter().any(|&http2_proto| http2_proto == proto);

                // We fail if we don't find an HTTP/2 protcol match...
                !found
            }
        };
        if fail {
            // We need the fail flag (instead of returning from one of the match
            // arms above because we need to move the `ssl_stream` and that is
            // not possible above (since it's borrowed at that point).
            return Err(TlsConnectError::Http2NotSupported(ssl_stream));
        }

        // Now that the stream is correctly established, we write the client preface.
        try!(write_preface(&mut ssl_stream));

        // All done.
        Ok(ClientStream(ssl_stream, HttpScheme::Https, self.host.into()))
    }
}
<|endoftext|>


<commit_before>extern mod extra;

use std::comm::SharedChan;
use std::io::{Listener, Acceptor, Writer};
use std::io::net::ip::SocketAddr;
use std::io::io_error;
use extra::time::precise_time_ns;

use std::io::net::tcp::TcpListener;

use buffer::BufferedStream;

pub use self::request::{RequestBuffer, Request};
pub use self::response::ResponseWriter;

pub mod request;
pub mod response;

pub trait Server: Send + Clone {
	fn handle_request(&self, request: &Request, response: &mut ResponseWriter) -> ();

	// XXX: this could also be implemented on the serve methods
	fn get_config(&self) -> Config;

	/**
	 * Attempt to bind to the address and port and start serving forever.
	 *
	 * This will only return if the initial connection fails or something else blows up.
	 */
    fn serve_forever(self) {
        let config = self.get_config();
        debug!("About to bind to {:?}", config.bind_address);
        match TcpListener::bind(config.bind_address).listen() {
            None => {
                error!("bind or listen failed :-(");
                return;
            },
            Some(ref mut acceptor) => {
                debug!("listening");
                let (perf_po, perf_ch) = stream();
                let perf_ch = SharedChan::new(perf_ch);
                do spawn {
                    perf_dumper(perf_po);
                }
                loop {
                    // OK, we're sort of shadowing an IoError here. Perhaps this should be done in a
                    // separate task so that it can safely fail...
                    let mut error = None;
                    let optstream = io_error::cond.trap(|e| {
                        error = Some(e);
                    }).inside(|| {
                        acceptor.accept()
                    });

                    let time_start = precise_time_ns();
                    if optstream.is_none() {
                        debug!("accept failed: {:?}", error);
                        // Question: is this the correct thing to do? We should probably be more
                        // intelligent, for there are some accept failures that are likely to be
                        // permanent, such that continuing would be a very bad idea, such as
                        // ENOBUFS/ENOMEM; and some where it should just be ignored, e.g.
                        // ECONNABORTED. TODO.
                        continue;
                    }
                    let child_perf_ch = perf_ch.clone();
                    let child_self = self.clone();
                    do spawn {
                        let mut time_start = time_start;
                        let mut stream = BufferedStream::new(optstream.unwrap());
                        debug!("accepted connection, got {:?}", stream);
                        loop {  // A keep-alive loop, condition at end
                            let time_spawned = precise_time_ns();
                            let (request, err_status) = Request::load(&mut stream);
                            let time_request_made = precise_time_ns();
                            let mut response = ~ResponseWriter::new(&mut stream, request);
                            let time_response_made = precise_time_ns();
                            match err_status {
                                Ok(()) => {
                                    child_self.handle_request(request, response);
                                    // Ensure that we actually do send a response:
                                    response.try_write_headers();
                                },
                                Err(status) => {
                                    // Uh oh, it's a response that I as a server cannot cope with.
                                    // No good user-agent should have caused this, so for the moment
                                    // at least I am content to send no body in the response.
                                    response.status = status;
                                    response.headers.content_length = Some(0);
                                    response.write_headers();
                                },
                            }
                            // Ensure the request is flushed, any Transfer-Encoding completed, etc.
                            response.finish_response();
                            let time_finished = precise_time_ns();
                            child_perf_ch.send((time_start, time_spawned, time_request_made, time_response_made, time_finished));

                            // Subsequent requests on this connection have no spawn time
                            time_start = time_finished;

                            if request.close_connection {
                                break;
                            }
                        }
                    }
                }
            }
        }
    }
}

/// The necessary configuration for an HTTP server.
///
/// At present, only the IP address and port to bind to are needed, but it's possible that other
/// options may turn up later.
pub struct Config {
	bind_address: SocketAddr,
}

static PERF_DUMP_FREQUENCY : u64 = 10_000;

/// Simple function to dump out perf stats every `PERF_DUMP_FREQUENCY` requests
fn perf_dumper(perf_po: Port<(u64, u64, u64, u64, u64)>) {
    // Total durations
    let mut td_spawn = 0u64;
    let mut td_request = 0u64;
    let mut td_response = 0u64;
    let mut td_handle = 0u64;
    let mut td_total = 0u64;
    let mut i = 0u64;
    loop {
        let data = perf_po.recv();
        let (start, spawned, request_made, response_made, finished) = data;
        td_spawn += spawned - start;
        td_request += request_made - spawned;
        td_response += response_made - request_made;
        td_handle += finished - response_made;
        td_total += finished - start;
        i += 1;
        if i % PERF_DUMP_FREQUENCY == 0 {
            println("");
            println!("{} requests made thus far. Current means:", i);
            println!("- Total:               100%, {:12f}",
                     td_total as f64 / i as f64);
            println!("- Spawn:               {:3f}%, {:12f}",
                     100f64 * td_spawn as f64 / td_total as f64,
                     td_spawn as f64 / i as f64);
            println!("- Load request:        {:3f}%, {:12f}",
                     100f64 * td_request as f64 / td_total as f64,
                     td_request as f64 / i as f64);
            println!("- Initialise response: {:3f}%, {:12f}",
                     100f64 * td_response as f64 / td_total as f64,
                     td_response as f64 / i as f64);
            println!("- Handle:              {:3f}%, {:12f}",
                     100f64 * td_handle as f64 / td_total as f64,
                     td_handle as f64 / i as f64);
        }
    }
}
<commit_msg>Fix some rightward drift.<commit_after>extern mod extra;

use std::comm::SharedChan;
use std::io::{Listener, Acceptor, Writer};
use std::io::net::ip::SocketAddr;
use std::io::io_error;
use extra::time::precise_time_ns;

use std::io::net::tcp::TcpListener;

use buffer::BufferedStream;

pub use self::request::{RequestBuffer, Request};
pub use self::response::ResponseWriter;

pub mod request;
pub mod response;

pub trait Server: Send + Clone {
	fn handle_request(&self, request: &Request, response: &mut ResponseWriter) -> ();

	// XXX: this could also be implemented on the serve methods
	fn get_config(&self) -> Config;

	/**
	 * Attempt to bind to the address and port and start serving forever.
	 *
	 * This will only return if the initial connection fails or something else blows up.
	 */
    fn serve_forever(self) {
        let config = self.get_config();
        debug!("About to bind to {:?}", config.bind_address);
        let mut acceptor = match TcpListener::bind(config.bind_address).listen() {
            None => {
                error!("bind or listen failed :-(");
                return;
            },
            Some(acceptor) => acceptor,
        };
        debug!("listening");
        let (perf_po, perf_ch) = stream();
        let perf_ch = SharedChan::new(perf_ch);
        do spawn {
            perf_dumper(perf_po);
        }
        loop {
            // OK, we're sort of shadowing an IoError here. Perhaps this should be done in a
            // separate task so that it can safely fail...
            let mut error = None;
            let optstream = io_error::cond.trap(|e| {
                error = Some(e);
            }).inside(|| {
                acceptor.accept()
            });

            let time_start = precise_time_ns();
            if optstream.is_none() {
                debug!("accept failed: {:?}", error);
                // Question: is this the correct thing to do? We should probably be more
                // intelligent, for there are some accept failures that are likely to be
                // permanent, such that continuing would be a very bad idea, such as
                // ENOBUFS/ENOMEM; and some where it should just be ignored, e.g.
                // ECONNABORTED. TODO.
                continue;
            }
            let child_perf_ch = perf_ch.clone();
            let child_self = self.clone();
            do spawn {
                let mut time_start = time_start;
                let mut stream = BufferedStream::new(optstream.unwrap());
                debug!("accepted connection, got {:?}", stream);
                loop {  // A keep-alive loop, condition at end
                    let time_spawned = precise_time_ns();
                    let (request, err_status) = Request::load(&mut stream);
                    let time_request_made = precise_time_ns();
                    let mut response = ~ResponseWriter::new(&mut stream, request);
                    let time_response_made = precise_time_ns();
                    match err_status {
                        Ok(()) => {
                            child_self.handle_request(request, response);
                            // Ensure that we actually do send a response:
                            response.try_write_headers();
                        },
                        Err(status) => {
                            // Uh oh, it's a response that I as a server cannot cope with.
                            // No good user-agent should have caused this, so for the moment
                            // at least I am content to send no body in the response.
                            response.status = status;
                            response.headers.content_length = Some(0);
                            response.write_headers();
                        },
                    }
                    // Ensure the request is flushed, any Transfer-Encoding completed, etc.
                    response.finish_response();
                    let time_finished = precise_time_ns();
                    child_perf_ch.send((time_start, time_spawned, time_request_made, time_response_made, time_finished));

                    // Subsequent requests on this connection have no spawn time
                    time_start = time_finished;

                    if request.close_connection {
                        break;
                    }
                }
            }
        }
    }
}

/// The necessary configuration for an HTTP server.
///
/// At present, only the IP address and port to bind to are needed, but it's possible that other
/// options may turn up later.
pub struct Config {
	bind_address: SocketAddr,
}

static PERF_DUMP_FREQUENCY : u64 = 10_000;

/// Simple function to dump out perf stats every `PERF_DUMP_FREQUENCY` requests
fn perf_dumper(perf_po: Port<(u64, u64, u64, u64, u64)>) {
    // Total durations
    let mut td_spawn = 0u64;
    let mut td_request = 0u64;
    let mut td_response = 0u64;
    let mut td_handle = 0u64;
    let mut td_total = 0u64;
    let mut i = 0u64;
    loop {
        let data = perf_po.recv();
        let (start, spawned, request_made, response_made, finished) = data;
        td_spawn += spawned - start;
        td_request += request_made - spawned;
        td_response += response_made - request_made;
        td_handle += finished - response_made;
        td_total += finished - start;
        i += 1;
        if i % PERF_DUMP_FREQUENCY == 0 {
            println("");
            println!("{} requests made thus far. Current means:", i);
            println!("- Total:               100%, {:12f}",
                     td_total as f64 / i as f64);
            println!("- Spawn:               {:3f}%, {:12f}",
                     100f64 * td_spawn as f64 / td_total as f64,
                     td_spawn as f64 / i as f64);
            println!("- Load request:        {:3f}%, {:12f}",
                     100f64 * td_request as f64 / td_total as f64,
                     td_request as f64 / i as f64);
            println!("- Initialise response: {:3f}%, {:12f}",
                     100f64 * td_response as f64 / td_total as f64,
                     td_response as f64 / i as f64);
            println!("- Handle:              {:3f}%, {:12f}",
                     100f64 * td_handle as f64 / td_total as f64,
                     td_handle as f64 / i as f64);
        }
    }
}
<|endoftext|>


<commit_before>extern crate websocket;
extern crate rustc_serialize;

mod emulator;

use std::thread;

use websocket::{Server, Message, Sender, Receiver};
use websocket::header::WebSocketProtocol;

fn main() {
    let server = Server::bind("127.0.0.1:2794").unwrap();

    for connection in server {
        // Spawn a new thread for each connection.
        thread::spawn(move || {
            let request = connection.unwrap().read_request().unwrap(); // Get the request
            let headers = request.headers.clone(); // Keep the headers so we can check them

            request.validate().unwrap(); // Validate the request

            let mut response = request.accept(); // Form a response

            if let Some(&WebSocketProtocol(ref protocols)) = headers.get() {
                if protocols.contains(&("rust-websocket".to_string())) {
                    // We have a protocol we want to use
                    response.headers.set(WebSocketProtocol(vec!["rust-websocket".to_string()]));
                }
            }

            let mut client = response.send().unwrap(); // Send the response

            let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

            println!("Connection from {}", ip);

            let emulator_instance = emulator::Emulator {
                data_memory: emulator::AvrDataMemory {
                    registers: vec![0,2,3],
                    io: vec![],
                    ram: vec![]
                }
            };

            let encoded = Message::Text(emulator::serialize(&emulator_instance));
            client.send_message(encoded).unwrap();

            let (mut sender, mut receiver) = client.split();

            for message in receiver.incoming_messages() {
                let message = message.unwrap();

                match message {
                    Message::Close(_) => {
                        let message = Message::Close(None);
                        sender.send_message(message).unwrap();
                        println!("Client {} disconnected", ip);
                        return;
                    }
                    Message::Ping(data) => {
                        let message = Message::Pong(data);
                        sender.send_message(message).unwrap();
                    }
                    Message::Text(data) => {
                        let new_emulator = emulator::perform_instruction(&emulator_instance, data);
                        let encoded = Message::Text(emulator::serialize(&emulator_instance));
                        sender.send_message(encoded).unwrap();
                    }
                    _ => {
                        // this is where we are echoing back the message sent to us
                        let message = Message::Text("Should not be here".to_string());
                        sender.send_message(message).unwrap();
                    }
                }
            }
        });
    }
}
<commit_msg>Return new emulator state in websocket response<commit_after>extern crate websocket;
extern crate rustc_serialize;

mod emulator;

use std::thread;

use websocket::{Server, Message, Sender, Receiver};
use websocket::header::WebSocketProtocol;

fn main() {
    let server = Server::bind("127.0.0.1:2794").unwrap();

    for connection in server {
        // Spawn a new thread for each connection.
        thread::spawn(move || {
            let request = connection.unwrap().read_request().unwrap(); // Get the request
            let headers = request.headers.clone(); // Keep the headers so we can check them

            request.validate().unwrap(); // Validate the request

            let mut response = request.accept(); // Form a response

            if let Some(&WebSocketProtocol(ref protocols)) = headers.get() {
                if protocols.contains(&("rust-websocket".to_string())) {
                    // We have a protocol we want to use
                    response.headers.set(WebSocketProtocol(vec!["rust-websocket".to_string()]));
                }
            }

            let mut client = response.send().unwrap(); // Send the response

            let ip = client.get_mut_sender()
            .get_mut()
            .peer_addr()
            .unwrap();

            println!("Connection from {}", ip);

            let emulator_instance = emulator::Emulator {
                data_memory: emulator::AvrDataMemory {
                    registers: vec![0,2,3],
                    io: vec![],
                    ram: vec![]
                }
            };

            let encoded = Message::Text(emulator::serialize(&emulator_instance));
            client.send_message(encoded).unwrap();

            let (mut sender, mut receiver) = client.split();

            for message in receiver.incoming_messages() {
                let message = message.unwrap();

                match message {
                    Message::Close(_) => {
                        let message = Message::Close(None);
                        sender.send_message(message).unwrap();
                        println!("Client {} disconnected", ip);
                        return;
                    }
                    Message::Ping(data) => {
                        let message = Message::Pong(data);
                        sender.send_message(message).unwrap();
                    }
                    Message::Text(data) => {
                        let new_emulator = emulator::perform_instruction(&emulator_instance, data);
                        let encoded = Message::Text(emulator::serialize(&new_emulator));
                        sender.send_message(encoded).unwrap();
                    }
                    _ => {
                        // this is where we are echoing back the message sent to us
                        let message = Message::Text("Should not be here".to_string());
                        sender.send_message(message).unwrap();
                    }
                }
            }
        });
    }
}
<|endoftext|>


<commit_before>use schema::Schema;
use schema::Document;
use indexer::SegmentSerializer;
use core::Index;
use core::SerializableSegment;
use core::Segment;
use std::thread::JoinHandle;
use indexer::SegmentWriter;
use std::clone::Clone;
use std::io;
use std::path::Path;
use std::thread;
use std::collections::HashSet;
use indexer::merger::IndexMerger;
use core::SegmentId;
use datastruct::stacker::Heap;
use std::mem::swap;
use chan;
use directory::WritePtr;

use Result;
use Error;

// Size of the margin for the heap. A segment is closed when the remaining memory
// in the heap goes below MARGIN_IN_BYTES.
pub const MARGIN_IN_BYTES: u32 = 10_000_000u32;

// We impose the memory per thread to be at least 30 MB.
pub const HEAP_SIZE_LIMIT: u32 = MARGIN_IN_BYTES * 3u32;

// Add document will block if the number of docs waiting in the queue to be indexed reaches PIPELINE_MAX_SIZE_IN_DOCS
const PIPELINE_MAX_SIZE_IN_DOCS: usize = 10_000;

pub const LOCKFILE_NAME: &'static str = ".tantivy-indexer.lock";

type DocumentSender = chan::Sender<Document>;
type DocumentReceiver = chan::Receiver<Document>;

type NewSegmentSender = chan::Sender<Result<(SegmentId, usize)>>;
type NewSegmentReceiver = chan::Receiver<Result<(SegmentId, usize)>>;

/// `IndexWriter` is the user entry-point to add document to an index.
///
/// It manages a small number of indexing thread, as well as a shared
/// indexing queue.
/// Each indexing thread builds its own independant `Segment`, via
/// a `SegmentWriter` object.
pub struct IndexWriter {
	index: Index,
	heap_size_in_bytes_per_thread: usize,
	workers_join_handle: Vec<JoinHandle<()>>,
	segment_ready_sender: NewSegmentSender,
	segment_ready_receiver: NewSegmentReceiver,
	document_receiver: DocumentReceiver,
	document_sender: DocumentSender,
	num_threads: usize,
	  docstamp: u64,
}

fn index_documents(heap: &mut Heap,
				   segment: Segment,
				   schema: &Schema,
				   document_iterator: &mut Iterator<Item=Document>) -> Result<usize> {
	heap.clear();
	let mut segment_writer = try!(SegmentWriter::for_segment(heap, segment, &schema));
	for doc in document_iterator {
		try!(segment_writer.add_document(&doc, &schema));
		if segment_writer.is_buffer_full() {
			info!("Buffer limit reached, flushing segment with maxdoc={}.", segment_writer.max_doc());
			break;
		}
	}
	let num_docs = segment_writer.max_doc() as usize;
	try!(segment_writer.finalize());
	Ok(num_docs)
}

impl Drop for IndexWriter {
    fn drop(&mut self) {
        let lockfile_path = Path::new(LOCKFILE_NAME);
        match self.index.directory_mut().delete(lockfile_path) {
            Ok(_) => (),
            Err(_) => ()
        }
    }
}
impl IndexWriter {

	/// Spawns a new worker thread for indexing.
	/// The thread consumes documents from the pipeline.
	///
	fn add_indexing_worker(&mut self,) -> Result<()> {
		let index = self.index.clone();
		let schema = self.index.schema();
		let segment_ready_sender_clone = self.segment_ready_sender.clone();
		let document_receiver_clone = self.document_receiver.clone();

		let mut heap = Heap::with_capacity(self.heap_size_in_bytes_per_thread);
		let join_handle: JoinHandle<()> = thread::spawn(move || {
			loop {
				let segment = index.new_segment();
				let segment_id = segment.id();
				let mut document_iterator = document_receiver_clone
					.clone()
					.into_iter()
					.peekable();
				// the peeking here is to avoid
				// creating a new segment's files
				// if no document are available.
				if document_iterator.peek().is_some() {
					let index_result = index_documents(&mut heap, segment, &schema, &mut document_iterator)
						.map(|num_docs| (segment_id, num_docs));
					segment_ready_sender_clone.send(index_result);
				}
				else {
					return;
				}
			}
		});
		self.workers_join_handle.push(join_handle);
		Ok(())
	}

	/// Open a new index writer
	///
	/// num_threads tells the number of indexing worker that
	/// should work at the same time.
	pub fn open(index: &Index,
				num_threads: usize,
				heap_size_in_bytes_per_thread: usize) -> Result<IndexWriter> {
		if heap_size_in_bytes_per_thread <= HEAP_SIZE_LIMIT as usize {
			panic!(format!("The heap size per thread needs to be at least {}.", HEAP_SIZE_LIMIT));
		}

      let mut cloned_index = index.clone();
      let lockfile_path = Path::new(LOCKFILE_NAME);

      try!(cloned_index.directory_mut().open_write(lockfile_path));
		let (document_sender, document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (segment_ready_sender, segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		let mut index_writer = IndexWriter {
			heap_size_in_bytes_per_thread: heap_size_in_bytes_per_thread,
			index: cloned_index,
			segment_ready_receiver: segment_ready_receiver,
			segment_ready_sender: segment_ready_sender,
			document_receiver: document_receiver,
			document_sender: document_sender,
			workers_join_handle: Vec::new(),
			num_threads: num_threads,
			docstamp: try!(index.docstamp()),
		};
		try!(index_writer.start_workers());
		Ok(index_writer)
	}

	fn start_workers(&mut self,) -> Result<()> {
		for _ in 0 .. self.num_threads {
			try!(self.add_indexing_worker());
		}
		Ok(())
	}

	/// Merges a given list of segments
	pub fn merge(&mut self, segments: &[Segment]) -> Result<()> {
		let schema = self.index.schema();
		let merger = try!(IndexMerger::open(schema, segments));
		let mut merged_segment = self.index.new_segment();
		let segment_serializer = try!(SegmentSerializer::for_segment(&mut merged_segment));
		try!(merger.write(segment_serializer));
		let merged_segment_ids: HashSet<SegmentId> = segments.iter().map(|segment| segment.id()).collect();
		try!(self.index.publish_merge_segment(merged_segment_ids, merged_segment.id()));
		Ok(())
	}

	/// Closes the current document channel send.
	/// and replace all the channels by new ones.
	///
	/// The current workers will keep on indexing
	/// the pending document and stop
	/// when no documents are remaining.
	///
	/// Returns the former segment_ready channel.
	fn recreate_channels(&mut self,) -> (DocumentReceiver, chan::Receiver<Result<(SegmentId, usize)>>) {
		let (mut document_sender, mut document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (mut segment_ready_sender, mut segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		swap(&mut self.document_sender, &mut document_sender);
		swap(&mut self.document_receiver, &mut document_receiver);
		swap(&mut self.segment_ready_sender, &mut segment_ready_sender);
		swap(&mut self.segment_ready_receiver, &mut segment_ready_receiver);
		(document_receiver, segment_ready_receiver)
	}

	/// Rollback to the last commit
	///
	/// This cancels all of the update that
	/// happened before after the last commit.
	/// After calling rollback, the index is in the same
	/// state as it was after the last commit.
	///
	/// The docstamp at the last commit is returned.
	pub fn rollback(&mut self,) -> Result<u64> {

		// we cannot drop segment ready receiver yet
		// as it would block the workers.
		let (document_receiver, mut _segment_ready_receiver) = self.recreate_channels();

		// consumes the document receiver pipeline
		// worker don't need to index the pending documents.
		for _ in document_receiver {};

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		// wait for all the worker to finish their work
		// (it should be fast since we consumed all pending documents)
		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		// reset the docstamp to what it was before
		self.docstamp = try!(self.index.docstamp());
		Ok(self.docstamp)
	}

	/// Commits all of the pending changes
	///
	/// A call to commit blocks.
	/// After it returns, all of the document that
	/// were added since the last commit are published
	/// and persisted.
	///
	/// In case of a crash or an hardware failure (as
	/// long as the hard disk is spared), it will be possible
	/// to resume indexing from this point.
	///
	/// Commit returns the `docstamp` of the last document
	/// that made it in the commit.
	///
	pub fn commit(&mut self,) -> Result<u64> {

		let (document_receiver, segment_ready_receiver) = self.recreate_channels();
		drop(document_receiver);

		// Docstamp of the last document in this commit.
		let commit_docstamp = self.docstamp;

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		let segment_ids_and_size: Vec<(SegmentId, usize)> = try!(
			segment_ready_receiver
				.into_iter()
				.collect()
		);

		let segment_ids: Vec<SegmentId> = segment_ids_and_size
			.iter()
			.map(|&(segment_id, _num_docs)| segment_id)
			.collect();

		try!(self.index.publish_segments(&segment_ids, commit_docstamp));

		Ok(commit_docstamp)
	}

	/// Adds a document.
	///
	/// If the indexing pipeline is full, this call may block.
	///
	/// The docstamp is an increasing `u64` that can
	/// be used by the client to align commits with its own
	/// document queue.
	///
	/// Currently it represents the number of documents that
	/// have been added since the creation of the index.
	pub fn add_document(&mut self, doc: Document) -> io::Result<u64> {
		self.document_sender.send(doc);
		self.docstamp += 1;
		Ok(self.docstamp)
	}

}

#[cfg(test)]
mod tests {

	use schema::{self, Document};
	use Index;
	use Term;

	#[test]
	fn test_commit_and_rollback() {
		let mut schema_builder = schema::SchemaBuilder::default();
		let text_field = schema_builder.add_text_field("text", schema::TEXT);
		let index = Index::create_in_ram(schema_builder.build());

		let num_docs_containing = |s: &str| {
			let searcher = index.searcher();
			let term_a = Term::from_field_text(text_field, s);
			searcher.doc_freq(&term_a)
		};

		{
			// writing the segment
			let mut index_writer = index.writer_with_num_threads(3, 40_000_000).unwrap();
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "a");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.rollback().unwrap(), 0u64);
			assert_eq!(num_docs_containing("a"), 0);

			{
				let mut doc = Document::default();
				doc.add_text(text_field, "b");
				index_writer.add_document(doc).unwrap();
			}
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "c");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.commit().unwrap(), 2u64);

			assert_eq!(num_docs_containing("a"), 0);
			assert_eq!(num_docs_containing("b"), 1);
			assert_eq!(num_docs_containing("c"), 1);
		}
		index.searcher();
	}

}
<commit_msg>added tests for correct handling of lockfiles<commit_after>use schema::Schema;
use schema::Document;
use indexer::SegmentSerializer;
use core::Index;
use core::SerializableSegment;
use core::Segment;
use std::thread::JoinHandle;
use indexer::SegmentWriter;
use std::clone::Clone;
use std::io;
use std::path::Path;
use std::thread;
use std::collections::HashSet;
use indexer::merger::IndexMerger;
use core::SegmentId;
use datastruct::stacker::Heap;
use std::mem::swap;
use chan;
use directory::WritePtr;

use Result;
use Error;

// Size of the margin for the heap. A segment is closed when the remaining memory
// in the heap goes below MARGIN_IN_BYTES.
pub const MARGIN_IN_BYTES: u32 = 10_000_000u32;

// We impose the memory per thread to be at least 30 MB.
pub const HEAP_SIZE_LIMIT: u32 = MARGIN_IN_BYTES * 3u32;

// Add document will block if the number of docs waiting in the queue to be indexed reaches PIPELINE_MAX_SIZE_IN_DOCS
const PIPELINE_MAX_SIZE_IN_DOCS: usize = 10_000;

pub const LOCKFILE_NAME: &'static str = ".tantivy-indexer.lock";

type DocumentSender = chan::Sender<Document>;
type DocumentReceiver = chan::Receiver<Document>;

type NewSegmentSender = chan::Sender<Result<(SegmentId, usize)>>;
type NewSegmentReceiver = chan::Receiver<Result<(SegmentId, usize)>>;

/// `IndexWriter` is the user entry-point to add document to an index.
///
/// It manages a small number of indexing thread, as well as a shared
/// indexing queue.
/// Each indexing thread builds its own independant `Segment`, via
/// a `SegmentWriter` object.
pub struct IndexWriter {
	index: Index,
	heap_size_in_bytes_per_thread: usize,
	workers_join_handle: Vec<JoinHandle<()>>,
	segment_ready_sender: NewSegmentSender,
	segment_ready_receiver: NewSegmentReceiver,
	document_receiver: DocumentReceiver,
	document_sender: DocumentSender,
	num_threads: usize,
	  docstamp: u64,
}

fn index_documents(heap: &mut Heap,
				   segment: Segment,
				   schema: &Schema,
				   document_iterator: &mut Iterator<Item=Document>) -> Result<usize> {
	heap.clear();
	let mut segment_writer = try!(SegmentWriter::for_segment(heap, segment, &schema));
	for doc in document_iterator {
		try!(segment_writer.add_document(&doc, &schema));
		if segment_writer.is_buffer_full() {
			info!("Buffer limit reached, flushing segment with maxdoc={}.", segment_writer.max_doc());
			break;
		}
	}
	let num_docs = segment_writer.max_doc() as usize;
	try!(segment_writer.finalize());
	Ok(num_docs)
}

impl Drop for IndexWriter {
    fn drop(&mut self) {
        let lockfile_path = Path::new(LOCKFILE_NAME);
        match self.index.directory_mut().delete(lockfile_path) {
            Ok(_) => (),
            Err(_) => ()
        }
    }
}
impl IndexWriter {

	/// Spawns a new worker thread for indexing.
	/// The thread consumes documents from the pipeline.
	///
	fn add_indexing_worker(&mut self,) -> Result<()> {
		let index = self.index.clone();
		let schema = self.index.schema();
		let segment_ready_sender_clone = self.segment_ready_sender.clone();
		let document_receiver_clone = self.document_receiver.clone();

		let mut heap = Heap::with_capacity(self.heap_size_in_bytes_per_thread);
		let join_handle: JoinHandle<()> = thread::spawn(move || {
			loop {
				let segment = index.new_segment();
				let segment_id = segment.id();
				let mut document_iterator = document_receiver_clone
					.clone()
					.into_iter()
					.peekable();
				// the peeking here is to avoid
				// creating a new segment's files
				// if no document are available.
				if document_iterator.peek().is_some() {
					let index_result = index_documents(&mut heap, segment, &schema, &mut document_iterator)
						.map(|num_docs| (segment_id, num_docs));
					segment_ready_sender_clone.send(index_result);
				}
				else {
					return;
				}
			}
		});
		self.workers_join_handle.push(join_handle);
		Ok(())
	}

	/// Open a new index writer
	///
	/// num_threads tells the number of indexing worker that
	/// should work at the same time.
	pub fn open(index: &Index,
				num_threads: usize,
				heap_size_in_bytes_per_thread: usize) -> Result<IndexWriter> {
		if heap_size_in_bytes_per_thread <= HEAP_SIZE_LIMIT as usize {
			panic!(format!("The heap size per thread needs to be at least {}.", HEAP_SIZE_LIMIT));
		}

      let mut cloned_index = index.clone();
      let lockfile_path = Path::new(LOCKFILE_NAME);

      try!(cloned_index.directory_mut().open_write(lockfile_path));
		let (document_sender, document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (segment_ready_sender, segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		let mut index_writer = IndexWriter {
			heap_size_in_bytes_per_thread: heap_size_in_bytes_per_thread,
			index: cloned_index,
			segment_ready_receiver: segment_ready_receiver,
			segment_ready_sender: segment_ready_sender,
			document_receiver: document_receiver,
			document_sender: document_sender,
			workers_join_handle: Vec::new(),
			num_threads: num_threads,
			docstamp: try!(index.docstamp()),
		};
		try!(index_writer.start_workers());
		Ok(index_writer)
	}

	fn start_workers(&mut self,) -> Result<()> {
		for _ in 0 .. self.num_threads {
			try!(self.add_indexing_worker());
		}
		Ok(())
	}

	/// Merges a given list of segments
	pub fn merge(&mut self, segments: &[Segment]) -> Result<()> {
		let schema = self.index.schema();
		let merger = try!(IndexMerger::open(schema, segments));
		let mut merged_segment = self.index.new_segment();
		let segment_serializer = try!(SegmentSerializer::for_segment(&mut merged_segment));
		try!(merger.write(segment_serializer));
		let merged_segment_ids: HashSet<SegmentId> = segments.iter().map(|segment| segment.id()).collect();
		try!(self.index.publish_merge_segment(merged_segment_ids, merged_segment.id()));
		Ok(())
	}

	/// Closes the current document channel send.
	/// and replace all the channels by new ones.
	///
	/// The current workers will keep on indexing
	/// the pending document and stop
	/// when no documents are remaining.
	///
	/// Returns the former segment_ready channel.
	fn recreate_channels(&mut self,) -> (DocumentReceiver, chan::Receiver<Result<(SegmentId, usize)>>) {
		let (mut document_sender, mut document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (mut segment_ready_sender, mut segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		swap(&mut self.document_sender, &mut document_sender);
		swap(&mut self.document_receiver, &mut document_receiver);
		swap(&mut self.segment_ready_sender, &mut segment_ready_sender);
		swap(&mut self.segment_ready_receiver, &mut segment_ready_receiver);
		(document_receiver, segment_ready_receiver)
	}

	/// Rollback to the last commit
	///
	/// This cancels all of the update that
	/// happened before after the last commit.
	/// After calling rollback, the index is in the same
	/// state as it was after the last commit.
	///
	/// The docstamp at the last commit is returned.
	pub fn rollback(&mut self,) -> Result<u64> {

		// we cannot drop segment ready receiver yet
		// as it would block the workers.
		let (document_receiver, mut _segment_ready_receiver) = self.recreate_channels();

		// consumes the document receiver pipeline
		// worker don't need to index the pending documents.
		for _ in document_receiver {};

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		// wait for all the worker to finish their work
		// (it should be fast since we consumed all pending documents)
		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		// reset the docstamp to what it was before
		self.docstamp = try!(self.index.docstamp());
		Ok(self.docstamp)
	}

	/// Commits all of the pending changes
	///
	/// A call to commit blocks.
	/// After it returns, all of the document that
	/// were added since the last commit are published
	/// and persisted.
	///
	/// In case of a crash or an hardware failure (as
	/// long as the hard disk is spared), it will be possible
	/// to resume indexing from this point.
	///
	/// Commit returns the `docstamp` of the last document
	/// that made it in the commit.
	///
	pub fn commit(&mut self,) -> Result<u64> {

		let (document_receiver, segment_ready_receiver) = self.recreate_channels();
		drop(document_receiver);

		// Docstamp of the last document in this commit.
		let commit_docstamp = self.docstamp;

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		let segment_ids_and_size: Vec<(SegmentId, usize)> = try!(
			segment_ready_receiver
				.into_iter()
				.collect()
		);

		let segment_ids: Vec<SegmentId> = segment_ids_and_size
			.iter()
			.map(|&(segment_id, _num_docs)| segment_id)
			.collect();

		try!(self.index.publish_segments(&segment_ids, commit_docstamp));

		Ok(commit_docstamp)
	}

	/// Adds a document.
	///
	/// If the indexing pipeline is full, this call may block.
	///
	/// The docstamp is an increasing `u64` that can
	/// be used by the client to align commits with its own
	/// document queue.
	///
	/// Currently it represents the number of documents that
	/// have been added since the creation of the index.
	pub fn add_document(&mut self, doc: Document) -> io::Result<u64> {
		self.document_sender.send(doc);
		self.docstamp += 1;
		Ok(self.docstamp)
	}

}

#[cfg(test)]
mod tests {

	use schema::{self, Document};
	use Index;
	use Term;

    #[test]
    #[should_panic]
    fn test_lockfile_stops_duplicates() {

		    let mut schema_builder = schema::SchemaBuilder::default();
		    let text_field = schema_builder.add_text_field("text", schema::TEXT);
		    let index = Index::create_in_ram(schema_builder.build());
			  let index_writer = index.writer(40_000_000).unwrap();
        let index_writer_two = index.writer(40_000_000).unwrap();
    }

    #[test]
    fn test_lockfile_released_on_drop() {
		    let mut schema_builder = schema::SchemaBuilder::default();
		    let text_field = schema_builder.add_text_field("text", schema::TEXT);
		    let index = Index::create_in_ram(schema_builder.build());
			  {
            let index_writer = index.writer(40_000_000).unwrap();
        }

        let index_writer_two = index.writer(40_000_000).unwrap();
    }
	#[test]
	fn test_commit_and_rollback() {
		let mut schema_builder = schema::SchemaBuilder::default();
		let text_field = schema_builder.add_text_field("text", schema::TEXT);
		let index = Index::create_in_ram(schema_builder.build());

		let num_docs_containing = |s: &str| {
			let searcher = index.searcher();
			let term_a = Term::from_field_text(text_field, s);
			searcher.doc_freq(&term_a)
		};

		{
			// writing the segment
			let mut index_writer = index.writer_with_num_threads(3, 40_000_000).unwrap();
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "a");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.rollback().unwrap(), 0u64);
			assert_eq!(num_docs_containing("a"), 0);

			{
				let mut doc = Document::default();
				doc.add_text(text_field, "b");
				index_writer.add_document(doc).unwrap();
			}
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "c");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.commit().unwrap(), 2u64);

			assert_eq!(num_docs_containing("a"), 0);
			assert_eq!(num_docs_containing("b"), 1);
			assert_eq!(num_docs_containing("c"), 1);
		}
		index.searcher();
	}

}
<|endoftext|>


<commit_before>use schema::Schema;
use schema::Document;
use indexer::SegmentSerializer;
use core::Index;
use core::SerializableSegment;
use core::Segment;
use std::thread::JoinHandle;
use indexer::SegmentWriter;
use std::clone::Clone;
use std::io;
use std::path::Path;
use std::thread;
use std::collections::HashSet;
use indexer::merger::IndexMerger;
use core::SegmentId;
use datastruct::stacker::Heap;
use std::mem::swap;
use chan;
use directory::WritePtr;

use Result;
use Error;

// Size of the margin for the heap. A segment is closed when the remaining memory
// in the heap goes below MARGIN_IN_BYTES.
pub const MARGIN_IN_BYTES: u32 = 10_000_000u32;

// We impose the memory per thread to be at least 30 MB.
pub const HEAP_SIZE_LIMIT: u32 = MARGIN_IN_BYTES * 3u32;

// Add document will block if the number of docs waiting in the queue to be indexed reaches PIPELINE_MAX_SIZE_IN_DOCS
const PIPELINE_MAX_SIZE_IN_DOCS: usize = 10_000;

pub const LOCKFILE_NAME: &'static str = ".tantivy-indexer.lock";

type DocumentSender = chan::Sender<Document>;
type DocumentReceiver = chan::Receiver<Document>;

type NewSegmentSender = chan::Sender<Result<(SegmentId, usize)>>;
type NewSegmentReceiver = chan::Receiver<Result<(SegmentId, usize)>>;

/// `IndexWriter` is the user entry-point to add document to an index.
///
/// It manages a small number of indexing thread, as well as a shared
/// indexing queue.
/// Each indexing thread builds its own independant `Segment`, via
/// a `SegmentWriter` object.
pub struct IndexWriter {
	index: Index,
	heap_size_in_bytes_per_thread: usize,
	workers_join_handle: Vec<JoinHandle<()>>,
	segment_ready_sender: NewSegmentSender,
	segment_ready_receiver: NewSegmentReceiver,
	document_receiver: DocumentReceiver,
	document_sender: DocumentSender,
	num_threads: usize,
	  docstamp: u64,
    lockfile: WritePtr
}

fn index_documents(heap: &mut Heap,
				   segment: Segment,
				   schema: &Schema,
				   document_iterator: &mut Iterator<Item=Document>) -> Result<usize> {
	heap.clear();
	let mut segment_writer = try!(SegmentWriter::for_segment(heap, segment, &schema));
	for doc in document_iterator {
		try!(segment_writer.add_document(&doc, &schema));
		if segment_writer.is_buffer_full() {
			info!("Buffer limit reached, flushing segment with maxdoc={}.", segment_writer.max_doc());
			break;
		}
	}
	let num_docs = segment_writer.max_doc() as usize;
	try!(segment_writer.finalize());
	Ok(num_docs)
}

impl Drop for IndexWriter {
    fn drop(&mut self) {
        let lockfile_path = Path::new(LOCKFILE_NAME);
        match self.index.directory_mut().delete(lockfile_path) {
            Ok(_) => (),
            Err(_) => ()
        }
    }
}
impl IndexWriter {

	/// Spawns a new worker thread for indexing.
	/// The thread consumes documents from the pipeline.
	///
	fn add_indexing_worker(&mut self,) -> Result<()> {
		let index = self.index.clone();
		let schema = self.index.schema();
		let segment_ready_sender_clone = self.segment_ready_sender.clone();
		let document_receiver_clone = self.document_receiver.clone();

		let mut heap = Heap::with_capacity(self.heap_size_in_bytes_per_thread);
		let join_handle: JoinHandle<()> = thread::spawn(move || {
			loop {
				let segment = index.new_segment();
				let segment_id = segment.id();
				let mut document_iterator = document_receiver_clone
					.clone()
					.into_iter()
					.peekable();
				// the peeking here is to avoid
				// creating a new segment's files
				// if no document are available.
				if document_iterator.peek().is_some() {
					let index_result = index_documents(&mut heap, segment, &schema, &mut document_iterator)
						.map(|num_docs| (segment_id, num_docs));
					segment_ready_sender_clone.send(index_result);
				}
				else {
					return;
				}
			}
		});
		self.workers_join_handle.push(join_handle);
		Ok(())
	}

	/// Open a new index writer
	///
	/// num_threads tells the number of indexing worker that
	/// should work at the same time.
	pub fn open(index: &Index,
				num_threads: usize,
				heap_size_in_bytes_per_thread: usize) -> Result<IndexWriter> {
		if heap_size_in_bytes_per_thread <= HEAP_SIZE_LIMIT as usize {
			panic!(format!("The heap size per thread needs to be at least {}.", HEAP_SIZE_LIMIT));
		}

      let mut cloned_index = index.clone();
      let lockfile_path = Path::new(LOCKFILE_NAME);

      let lf = try!(cloned_index.directory_mut().open_write(lockfile_path));
		let (document_sender, document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (segment_ready_sender, segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		let mut index_writer = IndexWriter {
			heap_size_in_bytes_per_thread: heap_size_in_bytes_per_thread,
			index: cloned_index,
			segment_ready_receiver: segment_ready_receiver,
			segment_ready_sender: segment_ready_sender,
			document_receiver: document_receiver,
			document_sender: document_sender,
			workers_join_handle: Vec::new(),
			num_threads: num_threads,
			docstamp: try!(index.docstamp()),
        lockfile: lf
		};
		try!(index_writer.start_workers());
		Ok(index_writer)
	}

	fn start_workers(&mut self,) -> Result<()> {
		for _ in 0 .. self.num_threads {
			try!(self.add_indexing_worker());
		}
		Ok(())
	}

	/// Merges a given list of segments
	pub fn merge(&mut self, segments: &[Segment]) -> Result<()> {
		let schema = self.index.schema();
		let merger = try!(IndexMerger::open(schema, segments));
		let mut merged_segment = self.index.new_segment();
		let segment_serializer = try!(SegmentSerializer::for_segment(&mut merged_segment));
		try!(merger.write(segment_serializer));
		let merged_segment_ids: HashSet<SegmentId> = segments.iter().map(|segment| segment.id()).collect();
		try!(self.index.publish_merge_segment(merged_segment_ids, merged_segment.id()));
		Ok(())
	}

	/// Closes the current document channel send.
	/// and replace all the channels by new ones.
	///
	/// The current workers will keep on indexing
	/// the pending document and stop
	/// when no documents are remaining.
	///
	/// Returns the former segment_ready channel.
	fn recreate_channels(&mut self,) -> (DocumentReceiver, chan::Receiver<Result<(SegmentId, usize)>>) {
		let (mut document_sender, mut document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (mut segment_ready_sender, mut segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		swap(&mut self.document_sender, &mut document_sender);
		swap(&mut self.document_receiver, &mut document_receiver);
		swap(&mut self.segment_ready_sender, &mut segment_ready_sender);
		swap(&mut self.segment_ready_receiver, &mut segment_ready_receiver);
		(document_receiver, segment_ready_receiver)
	}

	/// Rollback to the last commit
	///
	/// This cancels all of the update that
	/// happened before after the last commit.
	/// After calling rollback, the index is in the same
	/// state as it was after the last commit.
	///
	/// The docstamp at the last commit is returned.
	pub fn rollback(&mut self,) -> Result<u64> {

		// we cannot drop segment ready receiver yet
		// as it would block the workers.
		let (document_receiver, mut _segment_ready_receiver) = self.recreate_channels();

		// consumes the document receiver pipeline
		// worker don't need to index the pending documents.
		for _ in document_receiver {};

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		// wait for all the worker to finish their work
		// (it should be fast since we consumed all pending documents)
		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		// reset the docstamp to what it was before
		self.docstamp = try!(self.index.docstamp());
		Ok(self.docstamp)
	}

	/// Commits all of the pending changes
	///
	/// A call to commit blocks.
	/// After it returns, all of the document that
	/// were added since the last commit are published
	/// and persisted.
	///
	/// In case of a crash or an hardware failure (as
	/// long as the hard disk is spared), it will be possible
	/// to resume indexing from this point.
	///
	/// Commit returns the `docstamp` of the last document
	/// that made it in the commit.
	///
	pub fn commit(&mut self,) -> Result<u64> {

		let (document_receiver, segment_ready_receiver) = self.recreate_channels();
		drop(document_receiver);

		// Docstamp of the last document in this commit.
		let commit_docstamp = self.docstamp;

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		let segment_ids_and_size: Vec<(SegmentId, usize)> = try!(
			segment_ready_receiver
				.into_iter()
				.collect()
		);

		let segment_ids: Vec<SegmentId> = segment_ids_and_size
			.iter()
			.map(|&(segment_id, _num_docs)| segment_id)
			.collect();

		try!(self.index.publish_segments(&segment_ids, commit_docstamp));

		Ok(commit_docstamp)
	}

	/// Adds a document.
	///
	/// If the indexing pipeline is full, this call may block.
	///
	/// The docstamp is an increasing `u64` that can
	/// be used by the client to align commits with its own
	/// document queue.
	///
	/// Currently it represents the number of documents that
	/// have been added since the creation of the index.
	pub fn add_document(&mut self, doc: Document) -> io::Result<u64> {
		self.document_sender.send(doc);
		self.docstamp += 1;
		Ok(self.docstamp)
	}

}

#[cfg(test)]
mod tests {

	use schema::{self, Document};
	use Index;
	use Term;

	#[test]
	fn test_commit_and_rollback() {
		let mut schema_builder = schema::SchemaBuilder::default();
		let text_field = schema_builder.add_text_field("text", schema::TEXT);
		let index = Index::create_in_ram(schema_builder.build());

		let num_docs_containing = |s: &str| {
			let searcher = index.searcher();
			let term_a = Term::from_field_text(text_field, s);
			searcher.doc_freq(&term_a)
		};

		{
			// writing the segment
			let mut index_writer = index.writer_with_num_threads(3, 40_000_000).unwrap();
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "a");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.rollback().unwrap(), 0u64);
			assert_eq!(num_docs_containing("a"), 0);

			{
				let mut doc = Document::default();
				doc.add_text(text_field, "b");
				index_writer.add_document(doc).unwrap();
			}
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "c");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.commit().unwrap(), 2u64);

			assert_eq!(num_docs_containing("a"), 0);
			assert_eq!(num_docs_containing("b"), 1);
			assert_eq!(num_docs_containing("c"), 1);
		}
		index.searcher();
	}

}
<commit_msg>don't actually store a pointer to the lockfile, there's no need<commit_after>use schema::Schema;
use schema::Document;
use indexer::SegmentSerializer;
use core::Index;
use core::SerializableSegment;
use core::Segment;
use std::thread::JoinHandle;
use indexer::SegmentWriter;
use std::clone::Clone;
use std::io;
use std::path::Path;
use std::thread;
use std::collections::HashSet;
use indexer::merger::IndexMerger;
use core::SegmentId;
use datastruct::stacker::Heap;
use std::mem::swap;
use chan;
use directory::WritePtr;

use Result;
use Error;

// Size of the margin for the heap. A segment is closed when the remaining memory
// in the heap goes below MARGIN_IN_BYTES.
pub const MARGIN_IN_BYTES: u32 = 10_000_000u32;

// We impose the memory per thread to be at least 30 MB.
pub const HEAP_SIZE_LIMIT: u32 = MARGIN_IN_BYTES * 3u32;

// Add document will block if the number of docs waiting in the queue to be indexed reaches PIPELINE_MAX_SIZE_IN_DOCS
const PIPELINE_MAX_SIZE_IN_DOCS: usize = 10_000;

pub const LOCKFILE_NAME: &'static str = ".tantivy-indexer.lock";

type DocumentSender = chan::Sender<Document>;
type DocumentReceiver = chan::Receiver<Document>;

type NewSegmentSender = chan::Sender<Result<(SegmentId, usize)>>;
type NewSegmentReceiver = chan::Receiver<Result<(SegmentId, usize)>>;

/// `IndexWriter` is the user entry-point to add document to an index.
///
/// It manages a small number of indexing thread, as well as a shared
/// indexing queue.
/// Each indexing thread builds its own independant `Segment`, via
/// a `SegmentWriter` object.
pub struct IndexWriter {
	index: Index,
	heap_size_in_bytes_per_thread: usize,
	workers_join_handle: Vec<JoinHandle<()>>,
	segment_ready_sender: NewSegmentSender,
	segment_ready_receiver: NewSegmentReceiver,
	document_receiver: DocumentReceiver,
	document_sender: DocumentSender,
	num_threads: usize,
	  docstamp: u64,
}

fn index_documents(heap: &mut Heap,
				   segment: Segment,
				   schema: &Schema,
				   document_iterator: &mut Iterator<Item=Document>) -> Result<usize> {
	heap.clear();
	let mut segment_writer = try!(SegmentWriter::for_segment(heap, segment, &schema));
	for doc in document_iterator {
		try!(segment_writer.add_document(&doc, &schema));
		if segment_writer.is_buffer_full() {
			info!("Buffer limit reached, flushing segment with maxdoc={}.", segment_writer.max_doc());
			break;
		}
	}
	let num_docs = segment_writer.max_doc() as usize;
	try!(segment_writer.finalize());
	Ok(num_docs)
}

impl Drop for IndexWriter {
    fn drop(&mut self) {
        let lockfile_path = Path::new(LOCKFILE_NAME);
        match self.index.directory_mut().delete(lockfile_path) {
            Ok(_) => (),
            Err(_) => ()
        }
    }
}
impl IndexWriter {

	/// Spawns a new worker thread for indexing.
	/// The thread consumes documents from the pipeline.
	///
	fn add_indexing_worker(&mut self,) -> Result<()> {
		let index = self.index.clone();
		let schema = self.index.schema();
		let segment_ready_sender_clone = self.segment_ready_sender.clone();
		let document_receiver_clone = self.document_receiver.clone();

		let mut heap = Heap::with_capacity(self.heap_size_in_bytes_per_thread);
		let join_handle: JoinHandle<()> = thread::spawn(move || {
			loop {
				let segment = index.new_segment();
				let segment_id = segment.id();
				let mut document_iterator = document_receiver_clone
					.clone()
					.into_iter()
					.peekable();
				// the peeking here is to avoid
				// creating a new segment's files
				// if no document are available.
				if document_iterator.peek().is_some() {
					let index_result = index_documents(&mut heap, segment, &schema, &mut document_iterator)
						.map(|num_docs| (segment_id, num_docs));
					segment_ready_sender_clone.send(index_result);
				}
				else {
					return;
				}
			}
		});
		self.workers_join_handle.push(join_handle);
		Ok(())
	}

	/// Open a new index writer
	///
	/// num_threads tells the number of indexing worker that
	/// should work at the same time.
	pub fn open(index: &Index,
				num_threads: usize,
				heap_size_in_bytes_per_thread: usize) -> Result<IndexWriter> {
		if heap_size_in_bytes_per_thread <= HEAP_SIZE_LIMIT as usize {
			panic!(format!("The heap size per thread needs to be at least {}.", HEAP_SIZE_LIMIT));
		}

      let mut cloned_index = index.clone();
      let lockfile_path = Path::new(LOCKFILE_NAME);

      try!(cloned_index.directory_mut().open_write(lockfile_path));
		let (document_sender, document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (segment_ready_sender, segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		let mut index_writer = IndexWriter {
			heap_size_in_bytes_per_thread: heap_size_in_bytes_per_thread,
			index: cloned_index,
			segment_ready_receiver: segment_ready_receiver,
			segment_ready_sender: segment_ready_sender,
			document_receiver: document_receiver,
			document_sender: document_sender,
			workers_join_handle: Vec::new(),
			num_threads: num_threads,
			docstamp: try!(index.docstamp()),
		};
		try!(index_writer.start_workers());
		Ok(index_writer)
	}

	fn start_workers(&mut self,) -> Result<()> {
		for _ in 0 .. self.num_threads {
			try!(self.add_indexing_worker());
		}
		Ok(())
	}

	/// Merges a given list of segments
	pub fn merge(&mut self, segments: &[Segment]) -> Result<()> {
		let schema = self.index.schema();
		let merger = try!(IndexMerger::open(schema, segments));
		let mut merged_segment = self.index.new_segment();
		let segment_serializer = try!(SegmentSerializer::for_segment(&mut merged_segment));
		try!(merger.write(segment_serializer));
		let merged_segment_ids: HashSet<SegmentId> = segments.iter().map(|segment| segment.id()).collect();
		try!(self.index.publish_merge_segment(merged_segment_ids, merged_segment.id()));
		Ok(())
	}

	/// Closes the current document channel send.
	/// and replace all the channels by new ones.
	///
	/// The current workers will keep on indexing
	/// the pending document and stop
	/// when no documents are remaining.
	///
	/// Returns the former segment_ready channel.
	fn recreate_channels(&mut self,) -> (DocumentReceiver, chan::Receiver<Result<(SegmentId, usize)>>) {
		let (mut document_sender, mut document_receiver): (DocumentSender, DocumentReceiver) = chan::sync(PIPELINE_MAX_SIZE_IN_DOCS);
		let (mut segment_ready_sender, mut segment_ready_receiver): (NewSegmentSender, NewSegmentReceiver) = chan::async();
		swap(&mut self.document_sender, &mut document_sender);
		swap(&mut self.document_receiver, &mut document_receiver);
		swap(&mut self.segment_ready_sender, &mut segment_ready_sender);
		swap(&mut self.segment_ready_receiver, &mut segment_ready_receiver);
		(document_receiver, segment_ready_receiver)
	}

	/// Rollback to the last commit
	///
	/// This cancels all of the update that
	/// happened before after the last commit.
	/// After calling rollback, the index is in the same
	/// state as it was after the last commit.
	///
	/// The docstamp at the last commit is returned.
	pub fn rollback(&mut self,) -> Result<u64> {

		// we cannot drop segment ready receiver yet
		// as it would block the workers.
		let (document_receiver, mut _segment_ready_receiver) = self.recreate_channels();

		// consumes the document receiver pipeline
		// worker don't need to index the pending documents.
		for _ in document_receiver {};

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		// wait for all the worker to finish their work
		// (it should be fast since we consumed all pending documents)
		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		// reset the docstamp to what it was before
		self.docstamp = try!(self.index.docstamp());
		Ok(self.docstamp)
	}

	/// Commits all of the pending changes
	///
	/// A call to commit blocks.
	/// After it returns, all of the document that
	/// were added since the last commit are published
	/// and persisted.
	///
	/// In case of a crash or an hardware failure (as
	/// long as the hard disk is spared), it will be possible
	/// to resume indexing from this point.
	///
	/// Commit returns the `docstamp` of the last document
	/// that made it in the commit.
	///
	pub fn commit(&mut self,) -> Result<u64> {

		let (document_receiver, segment_ready_receiver) = self.recreate_channels();
		drop(document_receiver);

		// Docstamp of the last document in this commit.
		let commit_docstamp = self.docstamp;

		let mut former_workers_join_handle = Vec::new();
		swap(&mut former_workers_join_handle, &mut self.workers_join_handle);

		for worker_handle in former_workers_join_handle {
			try!(worker_handle
				.join()
				.map_err(|e| Error::ErrorInThread(format!("{:?}", e)))
			);
			// add a new worker for the next generation.
			try!(self.add_indexing_worker());
		}

		let segment_ids_and_size: Vec<(SegmentId, usize)> = try!(
			segment_ready_receiver
				.into_iter()
				.collect()
		);

		let segment_ids: Vec<SegmentId> = segment_ids_and_size
			.iter()
			.map(|&(segment_id, _num_docs)| segment_id)
			.collect();

		try!(self.index.publish_segments(&segment_ids, commit_docstamp));

		Ok(commit_docstamp)
	}

	/// Adds a document.
	///
	/// If the indexing pipeline is full, this call may block.
	///
	/// The docstamp is an increasing `u64` that can
	/// be used by the client to align commits with its own
	/// document queue.
	///
	/// Currently it represents the number of documents that
	/// have been added since the creation of the index.
	pub fn add_document(&mut self, doc: Document) -> io::Result<u64> {
		self.document_sender.send(doc);
		self.docstamp += 1;
		Ok(self.docstamp)
	}

}

#[cfg(test)]
mod tests {

	use schema::{self, Document};
	use Index;
	use Term;

	#[test]
	fn test_commit_and_rollback() {
		let mut schema_builder = schema::SchemaBuilder::default();
		let text_field = schema_builder.add_text_field("text", schema::TEXT);
		let index = Index::create_in_ram(schema_builder.build());

		let num_docs_containing = |s: &str| {
			let searcher = index.searcher();
			let term_a = Term::from_field_text(text_field, s);
			searcher.doc_freq(&term_a)
		};

		{
			// writing the segment
			let mut index_writer = index.writer_with_num_threads(3, 40_000_000).unwrap();
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "a");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.rollback().unwrap(), 0u64);
			assert_eq!(num_docs_containing("a"), 0);

			{
				let mut doc = Document::default();
				doc.add_text(text_field, "b");
				index_writer.add_document(doc).unwrap();
			}
			{
				let mut doc = Document::default();
				doc.add_text(text_field, "c");
				index_writer.add_document(doc).unwrap();
			}
			assert_eq!(index_writer.commit().unwrap(), 2u64);

			assert_eq!(num_docs_containing("a"), 0);
			assert_eq!(num_docs_containing("b"), 1);
			assert_eq!(num_docs_containing("c"), 1);
		}
		index.searcher();
	}

}
<|endoftext|>


<commit_before>use std::mem::MaybeUninit;
use std::ptr::NonNull;

#[test]
fn unitialized_zero_size_box() {
    assert_eq!(
        &*Box::<()>::new_uninit() as *const _,
        NonNull::<MaybeUninit<()>>::dangling().as_ptr(),
    );
    assert_eq!(
        Box::<[()]>::new_uninit_slice(4).as_ptr(),
        NonNull::<MaybeUninit<()>>::dangling().as_ptr(),
    );
    assert_eq!(
        Box::<[String]>::new_uninit_slice(0).as_ptr(),
        NonNull::<MaybeUninit<String>>::dangling().as_ptr(),
    );
}
<commit_msg>Add tests<commit_after>use std::mem::MaybeUninit;
use std::ptr::NonNull;

#[test]
fn unitialized_zero_size_box() {
    assert_eq!(
        &*Box::<()>::new_uninit() as *const _,
        NonNull::<MaybeUninit<()>>::dangling().as_ptr(),
    );
    assert_eq!(
        Box::<[()]>::new_uninit_slice(4).as_ptr(),
        NonNull::<MaybeUninit<()>>::dangling().as_ptr(),
    );
    assert_eq!(
        Box::<[String]>::new_uninit_slice(0).as_ptr(),
        NonNull::<MaybeUninit<String>>::dangling().as_ptr(),
    );
}

#[derive(Clone, PartialEq, Eq, Debug)]
struct Dummy {
    _data: u8,
}

#[test]
fn box_clone_and_clone_from_equivalence() {
    for size in (0..8).map(|i| 2usize.pow(i)) {
        let control = vec![Dummy { _data: 42 }; size].into_boxed_slice();
        let clone = control.clone();
        let mut copy = vec![Dummy { _data: 84 }; size].into_boxed_slice();
        copy.clone_from(&control);
        assert_eq!(control, clone);
        assert_eq!(control, copy);
    }
}

#[test]
fn box_clone_from_ptr_stability() {
    for size in (0..8).map(|i| 2usize.pow(i)) {
        let control = vec![Dummy { _data: 42 }; size].into_boxed_slice();
        let mut copy = vec![Dummy { _data: 84 }; size].into_boxed_slice();
        let copy_raw = copy.as_ptr() as usize;
        copy.clone_from(&control);
        assert_eq!(copy.as_ptr() as usize, copy_raw);
    }

    for size in (0..8).map(|i| 2usize.pow(i)) {
        let control = vec![Dummy { _data: 42 }; size].into_boxed_slice();
        let mut copy = vec![Dummy { _data: 84 }; size + 1].into_boxed_slice();
        let copy_raw = copy.as_ptr() as usize;
        copy.clone_from(&control);
        assert_ne!(copy.as_ptr() as usize, copy_raw);
    }
}
<|endoftext|>