implement handshake spec

Author: saghm

Cargo.toml

@@ -37,6 +37,8 @@ version = "0.6.3"
 
 [dev-dependencies]
 approx = "0.1.1"
+serde = "1.0.69"
+serde_derive = "1.0.69"
 
 [dev-dependencies.serde_json]
 features = ["preserve_order"]

src/cursor.rs

@@ -225,7 +225,7 @@ impl Cursor {
     ) -> Result<Cursor> {
 
         // Select a server stream from the topology.
-        let (stream, slave_ok, send_read_pref) = if cmd_type.is_write_command() {
+        let (mut stream, slave_ok, send_read_pref) = if cmd_type.is_write_command() {
             (try!(client.acquire_write_stream()), false, false)
         } else {
             try!(client.acquire_stream(read_pref.to_owned()))
@@ -254,7 +254,7 @@ impl Cursor {
         };
 
         Cursor::query_with_stream(
-            stream,
+            &mut stream,
             client,
             namespace,
             new_flags,
@@ -267,7 +267,7 @@ impl Cursor {
     }
 
     pub fn query_with_stream(
-        stream: PooledStream,
+        stream: &mut PooledStream,
         client: Client,
         namespace: String,
         flags: OpQueryFlags,
@@ -278,7 +278,6 @@ impl Cursor {
         read_pref: Option<ReadPreference>,
     ) -> Result<Cursor> {
 
-        let mut stream = stream;
         let socket = stream.get_socket();
         let req_id = client.get_req_id();
 

src/lib.rs

@@ -181,6 +181,8 @@ use topology::{Topology, TopologyDescription, TopologyType, DEFAULT_HEARTBEAT_FR
                DEFAULT_LOCAL_THRESHOLD_MS, DEFAULT_SERVER_SELECTION_TIMEOUT_MS};
 use topology::server::Server;
 
+pub const DRIVER_NAME: &'static str = "mongo-rust-driver-prototype";
+
 /// Interfaces with a MongoDB server or replica set.
 pub struct ClientInner {
     /// Indicates how a server should be selected for read operations.
@@ -415,11 +417,11 @@ impl ThreadedClient for Client {
         &self,
         read_preference: ReadPreference,
     ) -> Result<(PooledStream, bool, bool)> {
-        self.topology.acquire_stream(read_preference)
+        self.topology.acquire_stream(self.clone(), read_preference)
     }
 
     fn acquire_write_stream(&self) -> Result<PooledStream> {
-        self.topology.acquire_write_stream()
+        self.topology.acquire_write_stream(self.clone())
     }
 
     fn get_req_id(&self) -> i32 {

src/pool.rs

@@ -2,8 +2,13 @@
 use error::Error::{self, ArgumentError, OperationError};
 use error::Result;
 
+use Client;
+use coll::options::FindOptions;
+use command_type::CommandType;
 use connstring::Host;
+use cursor::Cursor;
 use stream::{Stream, StreamConnector};
+use wire_protocol::flags::OpQueryFlags;
 
 use bufstream::BufStream;
 use std::sync::{Arc, Condvar, Mutex};
@@ -48,6 +53,8 @@ pub struct PooledStream {
     wait_lock: Arc<Condvar>,
     // The pool iteration at the moment of extraction.
     iteration: usize,
+    // Whether the handshake occurred successfully.
+    successful_handshake: bool,
 }
 
 impl PooledStream {
@@ -59,6 +66,11 @@ impl PooledStream {
 
 impl Drop for PooledStream {
     fn drop(&mut self) {
+        // Don't add streams that couldn't successfully handshake to the pool.
+        if !self.successful_handshake {
+            return;
+        }
+
         // Attempt to lock and return the socket to the pool,
         // or give up if the pool lock has been poisoned.
         if let Ok(mut locked) = self.pool.lock() {
@@ -117,7 +129,7 @@ impl ConnectionPool {
     /// Attempts to acquire a connected socket. If none are available and
     /// the pool has not reached its maximum size, a new socket will connect.
     /// Otherwise, the function will block until a socket is returned to the pool.
-    pub fn acquire_stream(&self) -> Result<PooledStream> {
+    pub fn acquire_stream(&self, client: Client) -> Result<PooledStream> {
         let mut locked = try!(self.inner.lock());
         if locked.size == 0 {
             return Err(OperationError(String::from(
@@ -133,28 +145,32 @@ impl ConnectionPool {
                     pool: self.inner.clone(),
                     wait_lock: self.wait_lock.clone(),
                     iteration: locked.iteration,
+                    successful_handshake: true,
                 });
             }
 
             // Attempt to make a new connection
             let len = locked.len.load(Ordering::SeqCst);
             if len < locked.size {
                 let socket = try!(self.connect());
-                let _ = locked.len.fetch_add(1, Ordering::SeqCst);
-                return Ok(PooledStream {
+                let mut stream = PooledStream {
                     socket: Some(socket),
                     pool: self.inner.clone(),
                     wait_lock: self.wait_lock.clone(),
                     iteration: locked.iteration,
-                });
+                    successful_handshake: false,
+                };
+
+                self.handshake(client, &mut stream)?;
+                let _ = locked.len.fetch_add(1, Ordering::SeqCst);
+                return Ok(stream);
             }
 
             // Release lock and wait for pool to be repopulated
             locked = try!(self.wait_lock.wait(locked));
         }
     }
 
-
     // Connects to a MongoDB server as defined by the initial configuration.
     fn connect(&self) -> Result<BufStream<Stream>> {
         match self.stream_connector.connect(
@@ -165,4 +181,43 @@ impl ConnectionPool {
             Err(e) => Err(Error::from(e)),
         }
     }
-}
+
+    // This sends the client metadata to the server as described by the handshake spec.
+    //
+    // See https://github.com/mongodb/specifications/blob/master/source/mongodb-handshake/handshake.rst
+    fn handshake(&self, client: Client, stream: &mut PooledStream) -> Result<()> {
+        let mut options = FindOptions::new();
+        options.limit = Some(1);
+        options.batch_size = Some(1);
+
+        let flags = OpQueryFlags::with_find_options(&options);
+
+        Cursor::query_with_stream(
+            stream,
+            client,
+            String::from("local.$cmd"),
+            flags,
+            doc! {
+                "isMaster": 1i32,
+                "client": {
+                    "driver": {
+                        "name": ::DRIVER_NAME,
+                        "version": env!("CARGO_PKG_VERSION"),
+                    },
+                    "os": {
+                        "type": ::std::env::consts::OS,
+                        "architecture": ::std::env::consts::ARCH
+                    }
+                },
+            },
+            options,
+            CommandType::IsMaster,
+            false,
+            None,
+        )?;
+
+        stream.successful_handshake = true;
+
+        Ok(())
+    }
+}

src/topology/mod.rs

@@ -116,7 +116,7 @@ impl TopologyDescription {
     }
 
     /// Returns the nearest server stream, calculated by round trip time.
-    fn get_nearest_from_vec(&self, servers: &mut Vec<Host>) -> Result<(PooledStream, ServerType)> {
+    fn get_nearest_from_vec(&self, client: Client, servers: &mut Vec<Host>) -> Result<(PooledStream, ServerType)> {
         servers.sort_by(|a, b| {
             let mut a_rtt = i64::MAX;
             let mut b_rtt = i64::MAX;
@@ -140,7 +140,7 @@ impl TopologyDescription {
                 if let Ok(description) = server.description.read() {
                     if description.round_trip_time.is_none() {
                         break;
-                    } else if let Ok(stream) = server.acquire_stream() {
+                    } else if let Ok(stream) = server.acquire_stream(client.clone()) {
                         return Ok((stream, description.server_type));
                     }
                 }
@@ -152,13 +152,13 @@ impl TopologyDescription {
     }
 
     /// Returns a random server stream from the vector.
-    fn get_rand_from_vec(&self, servers: &mut Vec<Host>) -> Result<(PooledStream, ServerType)> {
+    fn get_rand_from_vec(&self, client: Client, servers: &mut Vec<Host>) -> Result<(PooledStream, ServerType)> {
         while !servers.is_empty() {
             let len = servers.len();
             let index = thread_rng().gen_range(0, len);
 
             if let Some(server) = self.servers.get(&servers[index]) {
-                if let Ok(stream) = server.acquire_stream() {
+                if let Ok(stream) = server.acquire_stream(client.clone()) {
                     if let Ok(description) = server.description.read() {
                         return Ok((stream, description.server_type));
                     }
@@ -174,6 +174,7 @@ impl TopologyDescription {
     /// Returns a server stream for read operations.
     pub fn acquire_stream(
         &self,
+        client: Client,
         read_preference: &ReadPreference,
     ) -> Result<(PooledStream, bool, bool)> {
         let (mut hosts, rand) = self.choose_hosts(read_preference)?;
@@ -190,7 +191,7 @@ impl TopologyDescription {
         if hosts.is_empty() && read_preference.mode == ReadMode::SecondaryPreferred {
             let mut read_pref = read_preference.clone();
             read_pref.mode = ReadMode::PrimaryPreferred;
-            return self.acquire_stream(&read_pref);
+            return self.acquire_stream(client, &read_pref);
         }
 
         // If no servers are available, request an update from all monitors.
@@ -205,9 +206,9 @@ impl TopologyDescription {
 
         // Retrieve a server stream from the list of acceptable hosts.
         let (pooled_stream, server_type) = if rand {
-            try!(self.get_rand_from_vec(&mut hosts))
+            try!(self.get_rand_from_vec(client, &mut hosts))
         } else {
-            try!(self.get_nearest_from_vec(&mut hosts))
+            try!(self.get_nearest_from_vec(client, &mut hosts))
         };
 
         // Determine how to handle server-side logic based on ReadMode and TopologyType.
@@ -251,7 +252,7 @@ impl TopologyDescription {
     }
 
     /// Returns a server stream for write operations.
-    pub fn acquire_write_stream(&self) -> Result<PooledStream> {
+    pub fn acquire_write_stream(&self, client: Client) -> Result<PooledStream> {
         let (mut hosts, rand) = self.choose_write_hosts();
 
         // If no servers are available, request an update from all monitors.
@@ -262,9 +263,9 @@ impl TopologyDescription {
         }
 
         if rand {
-            Ok(try!(self.get_rand_from_vec(&mut hosts)).0)
+            Ok(try!(self.get_rand_from_vec(client, &mut hosts)).0)
         } else {
-            Ok(try!(self.get_nearest_from_vec(&mut hosts)).0)
+            Ok(try!(self.get_nearest_from_vec(client, &mut hosts)).0)
         }
     }
 
@@ -886,6 +887,7 @@ impl Topology {
     // Private server stream acquisition helper.
     fn acquire_stream_private(
         &self,
+        client: Client,
         read_preference: Option<ReadPreference>,
         write: bool,
     ) -> Result<(PooledStream, bool, bool)> {
@@ -895,12 +897,13 @@ impl Topology {
 
         loop {
             let result = if write {
-                match self.description.read()?.acquire_write_stream() {
+                match self.description.read()?.acquire_write_stream(client.clone()) {
                     Ok(stream) => Ok((stream, false, false)),
                     Err(err) => Err(err),
                 }
             } else {
                 self.description.read()?.acquire_stream(
+                    client.clone(),
                     read_preference.as_ref().unwrap(),
                 )
             };
@@ -926,14 +929,15 @@ impl Topology {
     /// Returns a server stream for read operations.
     pub fn acquire_stream(
         &self,
+        client: Client,
         read_preference: ReadPreference,
     ) -> Result<(PooledStream, bool, bool)> {
-        self.acquire_stream_private(Some(read_preference), false)
+        self.acquire_stream_private(client, Some(read_preference), false)
     }
 
     /// Returns a server stream for write operations.
-    pub fn acquire_write_stream(&self) -> Result<PooledStream> {
-        let (stream, _, _) = try!(self.acquire_stream_private(None, true));
+    pub fn acquire_write_stream(&self, client: Client) -> Result<PooledStream> {
+        let (stream, _, _) = try!(self.acquire_stream_private(client, None, true));
         Ok(stream)
     }
 }

src/topology/monitor.rs

@@ -258,17 +258,16 @@ impl Monitor {
         options.limit = Some(1);
         options.batch_size = Some(1);
 
-
         let flags = OpQueryFlags::with_find_options(&options);
         let mut filter = bson::Document::new();
         filter.insert("isMaster", Bson::I32(1));
 
-        let stream = try!(self.personal_pool.acquire_stream());
+        let mut stream = try!(self.personal_pool.acquire_stream(self.client.clone()));
 
         let time_start = time::get_time();
 
         let cursor = try!(Cursor::query_with_stream(
-            stream,
+            &mut stream,
             self.client.clone(),
             String::from("local.$cmd"),
             flags,

src/topology/server.rs

@@ -247,8 +247,8 @@ impl Server {
     }
 
     /// Returns a server stream from the connection pool.
-    pub fn acquire_stream(&self) -> Result<PooledStream> {
-        self.pool.acquire_stream()
+    pub fn acquire_stream(&self, client: Client) -> Result<PooledStream> {
+        self.pool.acquire_stream(client)
     }
 
     /// Request an update from the monitor on the server status.

tests/client/handshake.rs

@@ -0,0 +1,46 @@
+use bson::{self, Bson};
+use mongodb::{DRIVER_NAME, Client, ThreadedClient};
+use mongodb::db::ThreadedDatabase;
+use mongodb::CommandType;
+
+#[derive(Debug, Deserialize)]
+struct Metadata {
+    #[serde(rename = "clientMetadata")]
+    pub client: ClientMetadata,
+}
+
+#[derive(Debug, Deserialize)]
+struct ClientMetadata {
+    pub driver: DriverMetadata,
+    pub os: OsMetadata,
+}
+
+#[derive(Debug, Deserialize)]
+struct DriverMetadata {
+    pub name: String,
+    pub version: String,
+}
+
+#[derive(Debug, Deserialize)]
+struct OsMetadata {
+    #[serde(rename = "type")]
+    pub os_type: String,
+    pub architecture: String,
+}
+
+#[test]
+fn metadata_sent_in_handshake() {
+    let client = Client::connect("localhost", 27017).unwrap();
+    let db = client.db("admin");
+    skip_if_db_version_below!(db, 3, 4);
+
+    let result = db.command(doc! { "currentOp" => 1 }, CommandType::Suppressed, None).unwrap();
+    let in_prog = match result.get("inprog") {
+        Some(Bson::Array(in_prog)) => in_prog,
+        _ => panic!("no `inprog` array found in response to `currentOp`"),
+    };
+
+    let metadata: Metadata = bson::from_bson(in_prog[0].clone()).unwrap();
+    assert_eq!(metadata.client.driver.name, DRIVER_NAME);
+}
+