Merge pull request #120 from ayosec/buffering

Buffered Connections

Author: vkarpov15

Cargo.lock

@@ -23,6 +23,7 @@ scan_fmt = "0.1.0"
 separator = "0.3.0"
 textnonce = "0.1.1"
 time = "0.1"
+bufstream = "0.1.1"
 
 [dev-dependencies]
 nalgebra = "0.2"

Cargo.toml

@@ -105,7 +105,7 @@ pub fn parse(address: &str) -> Result<ConnectionString> {
     // Remove scheme
     let addr = &address[URI_SCHEME.len()..];
 
-    let mut hosts: Vec<Host>;
+    let hosts: Vec<Host>;
     let mut user: Option<String> = None;
     let mut password: Option<String> = None;
     let mut database: Option<String> = Some("test".to_owned());

src/connstring.rs

@@ -250,14 +250,15 @@ impl Cursor {
                              is_cmd_cursor: bool,
                              read_pref: Option<ReadPreference>) -> Result<Cursor> {
 
+        let mut stream = stream;
         let mut socket = stream.get_socket();
         let req_id = client.get_req_id();
 
         let index = namespace.find(".").unwrap_or(namespace.len());
         let db_name = namespace[..index].to_owned();
         let coll_name = namespace[index + 1..].to_owned();
         let cmd_name = cmd_type.to_str();
-        let connstring = format!("{}", try!(socket.peer_addr()));
+        let connstring = format!("{}", try!(socket.get_ref().peer_addr()));
 
         let filter : bson::Document = match query.get("$query") {
             Some(&Bson::Document(ref doc)) => doc.clone(),
@@ -310,8 +311,8 @@ impl Cursor {
             }
         }
 
-        try_or_emit!(cmd_type, cmd_name, req_id, connstring, message.write(&mut socket), client);
-        let reply = try_or_emit!(cmd_type, cmd_name, req_id, connstring, Message::read(&mut socket),
+        try_or_emit!(cmd_type, cmd_name, req_id, connstring, message.write(socket), client);
+        let reply = try_or_emit!(cmd_type, cmd_name, req_id, connstring, Message::read(socket),
                                  client);
 
         let fin_time = time::precise_time_ns();
@@ -360,7 +361,7 @@ impl Cursor {
     }
 
     fn get_from_stream(&mut self) -> Result<()> {
-        let (stream, _, _) = try!(self.client.acquire_stream(self.read_preference.to_owned()));
+        let (mut stream, _, _) = try!(self.client.acquire_stream(self.read_preference.to_owned()));
         let mut socket = stream.get_socket();
 
         let req_id = self.client.get_req_id();
@@ -370,7 +371,7 @@ impl Cursor {
         let index = self.namespace.rfind(".").unwrap_or(self.namespace.len());
         let db_name = self.namespace[..index].to_owned();
         let cmd_name = "get_more".to_owned();
-        let connstring = format!("{}", try!(socket.peer_addr()));
+        let connstring = format!("{}", try!(socket.get_ref().peer_addr()));
 
         if self.cmd_type != CommandType::Suppressed {
             let hook_result = self.client.run_start_hooks(&CommandStarted {
@@ -386,8 +387,8 @@ impl Cursor {
             }
         }
 
-        try_or_emit!(self.cmd_type, cmd_name, req_id, connstring, get_more.write(&mut socket), self.client);
-        let reply = try!(Message::read(&mut socket));
+        try_or_emit!(self.cmd_type, cmd_name, req_id, connstring, get_more.write(socket.get_mut()), self.client);
+        let reply = try!(Message::read(socket.get_mut()));
 
         let (_, v, _) = try!(Cursor::get_bson_and_cid_from_message(reply));
         self.buffer.extend(v);

src/cursor.rs

@@ -101,6 +101,7 @@ extern crate scan_fmt;
 extern crate separator;
 extern crate textnonce;
 extern crate time;
+extern crate bufstream;
 
 pub mod db;
 pub mod coll;

src/lib.rs

@@ -4,6 +4,7 @@ use Result;
 
 use connstring::Host;
 
+use bufstream::BufStream;
 use std::net::TcpStream;
 use std::sync::{Arc, Condvar, Mutex};
 use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};
@@ -28,7 +29,7 @@ struct Pool {
     // The current number of open connections.
     pub len: Arc<AtomicUsize>,
     // The idle socket pool.
-    sockets: Vec<TcpStream>,
+    sockets: Vec<BufStream<TcpStream>>,
     // The pool iteration. When a server monitor fails to execute ismaster,
     // the connection pool is cleared and the iteration is incremented.
     iteration: usize,
@@ -39,7 +40,7 @@ struct Pool {
 pub struct PooledStream {
     // This socket option will always be Some(stream) until it is
     // returned to the pool using take().
-    socket: Option<TcpStream>,
+    socket: Option<BufStream<TcpStream>>,
     // A reference to the pool that the stream was taken from.
     pool: Arc<Mutex<Pool>>,
     // A reference to the waiting condvar associated with the pool.
@@ -50,8 +51,8 @@ pub struct PooledStream {
 
 impl PooledStream {
     /// Returns a reference to the socket.
-    pub fn get_socket<'a>(&'a self) -> &'a TcpStream {
-        self.socket.as_ref().unwrap()
+    pub fn get_socket<'a>(&'a mut self) -> &'a mut BufStream<TcpStream> {
+        self.socket.as_mut().unwrap()
     }
 }
 
@@ -150,10 +151,10 @@ impl ConnectionPool {
     }
 
     // Connects to a MongoDB server as defined by the initial configuration.
-    fn connect(&self) -> Result<TcpStream> {
+    fn connect(&self) -> Result<BufStream<TcpStream>> {
         let ref host_name = self.host.host_name;
         let port = self.host.port;
-        let stream = try!(TcpStream::connect((&host_name[..], port)));
+        let stream = BufStream::new(try!(TcpStream::connect((&host_name[..], port))));
         Ok(stream)
     }
 }

src/pool.rs

@@ -229,8 +229,8 @@ impl Message {
     /// # Return value
     ///
     /// Returns nothing on success, or an Error on failure.
-    fn write_bson_document(buffer: &mut Write,
-                           bson: &bson::Document) -> Result<()>{
+    fn write_bson_document<W: Write>(buffer: &mut W,
+                                     bson: &bson::Document) -> Result<()>{
         let mut temp_buffer = vec![];
 
         try!(bson::encode_document(&mut temp_buffer, bson));
@@ -253,9 +253,9 @@ impl Message {
     /// # Return value
     ///
     /// Returns nothing on success, or an Error on failure.
-    pub fn write_update(buffer: &mut Write, header: &Header, namespace: &str,
-                        flags: &OpUpdateFlags, selector: &bson::Document,
-                        update: &bson::Document) -> Result<()> {
+    pub fn write_update<W: Write>(buffer: &mut W, header: &Header, namespace: &str,
+                                  flags: &OpUpdateFlags, selector: &bson::Document,
+                                  update: &bson::Document) -> Result<()> {
 
         try!(header.write(buffer));
 
@@ -292,8 +292,8 @@ impl Message {
     /// # Return value
     ///
     /// Returns nothing on success, or an Error on failure.
-    fn write_insert(buffer: &mut Write, header: &Header, flags: &OpInsertFlags,
-                    namespace: &str, documents: &[bson::Document]) -> Result<()> {
+    fn write_insert<W: Write>(buffer: &mut W, header: &Header, flags: &OpInsertFlags,
+                              namespace: &str, documents: &[bson::Document]) -> Result<()> {
 
         try!(header.write(buffer));
         try!(buffer.write_i32::<LittleEndian>(flags.to_i32()));
@@ -334,11 +334,11 @@ impl Message {
     /// # Return value
     ///
     /// Returns nothing on success, or an Error on failure.
-    fn write_query(buffer: &mut Write, header: &Header,
-                   flags: &OpQueryFlags, namespace: &str,
-                   number_to_skip: i32, number_to_return: i32,
-                   query: &bson::Document,
-                   return_field_selector: &Option<bson::Document>) -> Result<()> {
+    fn write_query<W: Write>(buffer: &mut W, header: &Header,
+                             flags: &OpQueryFlags, namespace: &str,
+                             number_to_skip: i32, number_to_return: i32,
+                             query: &bson::Document,
+                             return_field_selector: &Option<bson::Document>) -> Result<()> {
 
         try!(header.write(buffer));
         try!(buffer.write_i32::<LittleEndian>(flags.to_i32()));
@@ -378,8 +378,8 @@ impl Message {
     /// # Return value
     ///
     /// Returns nothing on success, or an Error on failure.
-    pub fn write_get_more(buffer: &mut Write, header: &Header, namespace: &str,
-                          number_to_return: i32, cursor_id: i64) -> Result<()> {
+    pub fn write_get_more<W: Write>(buffer: &mut W, header: &Header, namespace: &str,
+                                    number_to_return: i32, cursor_id: i64) -> Result<()> {
 
         try!(header.write(buffer));
 
@@ -409,7 +409,7 @@ impl Message {
     /// # Return value
     ///
     /// Returns nothing on success, or an error string on failure.
-    pub fn write(&self, buffer: &mut Write) -> Result<()> {
+    pub fn write<W: Write>(&self, buffer: &mut W) -> Result<()> {
         match self {
             /// Only the server should send replies
             &Message::OpReply {..} =>
@@ -445,7 +445,7 @@ impl Message {
     /// # Return value
     ///
     /// Returns the reply message on success, or an Error on failure.
-    fn read_reply(buffer: &mut Read, header: Header) -> Result<Message> {
+    fn read_reply<R: Read>(buffer: &mut R, header: Header) -> Result<Message> {
         let mut length = header.message_length - mem::size_of::<Header>() as i32;
 
         // Read flags