RUST-185: Implement command monitoring (#57)

Author: shsaskin

src/bson_util.rs

@@ -1,7 +1,7 @@
 use std::time::Duration;
 
 use bson::{bson, doc, oid::ObjectId, Bson, Document};
-use serde::{ser, Serialize, Serializer};
+use serde::{ser, Deserialize, Deserializer, Serialize, Serializer};
 
 use crate::error::{ErrorKind, Result};
 
@@ -69,6 +69,16 @@ pub(crate) fn serialize_duration_as_i64_millis<S: Serializer>(
     }
 }
 
+pub(crate) fn deserialize_duration_from_u64_millis<'de, D>(
+    deserializer: D,
+) -> std::result::Result<Option<Duration>, D::Error>
+where
+    D: Deserializer<'de>,
+{
+    let millis = Option::<u64>::deserialize(deserializer)?;
+    Ok(millis.map(Duration::from_millis))
+}
+
 #[allow(clippy::trivially_copy_pass_by_ref)]
 pub(crate) fn serialize_u32_as_i32<S: Serializer>(
     val: &Option<u32>,

src/client/auth/scram.rs

@@ -98,7 +98,7 @@ impl ScramVersion {
             client_first.to_command(self),
         );
 
-        let server_first_response = conn.send_command(command)?;
+        let server_first_response = conn.send_command(command, None)?;
         let server_first = ServerFirst::parse(server_first_response.raw_response)?;
         server_first.validate(nonce.as_str())?;
 
@@ -135,7 +135,7 @@ impl ScramVersion {
             client_final.to_command(),
         );
 
-        let server_final_response = conn.send_command(command)?;
+        let server_final_response = conn.send_command(command, None)?;
         let server_final = ServerFinal::parse(server_final_response.raw_response)?;
         server_final.validate(salted_password.as_slice(), &client_final, self)?;
 
@@ -149,7 +149,7 @@ impl ScramVersion {
         };
         let command = Command::new("saslContinue".into(), source.into(), noop);
 
-        let server_noop_response = conn.send_command(command)?;
+        let server_noop_response = conn.send_command(command, None)?;
 
         if server_noop_response
             .raw_response

src/client/executor.rs

@@ -1,11 +1,35 @@
 use super::Client;
+
+use std::collections::HashSet;
+
+use bson::Document;
+use lazy_static::lazy_static;
+use time::PreciseTime;
+
 use crate::{
     cmap::Connection,
     error::Result,
+    event::command::{CommandFailedEvent, CommandStartedEvent, CommandSucceededEvent},
     operation::Operation,
     sdam::{update_topology, ServerDescription},
 };
 
+lazy_static! {
+    static ref REDACTED_COMMANDS: HashSet<&'static str> = {
+        let mut hash_set = HashSet::new();
+        hash_set.insert("authenticate");
+        hash_set.insert("saslstart");
+        hash_set.insert("saslcontinue");
+        hash_set.insert("getnonce");
+        hash_set.insert("createuser");
+        hash_set.insert("updateuser");
+        hash_set.insert("copydbgetnonce");
+        hash_set.insert("copydbsaslstart");
+        hash_set.insert("copydb");
+        hash_set
+    };
+}
+
 impl Client {
     /// Executes an operation and returns the connection used to do so along with the result of the
     /// operation. This will be used primarily for the opening of exhaust cursors.
@@ -83,13 +107,75 @@ impl Client {
         connection: &mut Connection,
     ) -> Result<T::O> {
         let mut cmd = op.build(connection.stream_description()?)?;
-
         self.topology()
             .read()
             .unwrap()
             .update_command_with_read_pref(connection.address(), &mut cmd, op.selection_criteria());
 
-        let response = connection.send_command(cmd)?;
-        op.handle_response(response)
+        let connection_info = connection.info();
+        let request_id = crate::cmap::conn::next_request_id();
+
+        let should_redact = REDACTED_COMMANDS.contains(cmd.name.to_lowercase().as_str());
+
+        let command_body = if should_redact {
+            Document::new()
+        } else {
+            cmd.body.clone()
+        };
+        let command_started_event = CommandStartedEvent {
+            command: command_body,
+            db: cmd.target_db.clone(),
+            command_name: cmd.name.clone(),
+            request_id,
+            connection: connection_info.clone(),
+        };
+
+        self.send_command_started_event(command_started_event);
+
+        let start_time = PreciseTime::now();
+
+        let response_result =
+            connection
+                .send_command(cmd.clone(), request_id)
+                .and_then(|response| {
+                    if !op.handles_command_errors() {
+                        response.validate()?;
+                    }
+                    Ok(response)
+                });
+
+        let end_time = PreciseTime::now();
+        let duration = start_time.to(end_time).to_std()?;
+
+        match response_result {
+            Err(error) => {
+                let command_failed_event = CommandFailedEvent {
+                    duration,
+                    command_name: cmd.name.clone(),
+                    failure: error.clone(),
+                    request_id,
+                    connection: connection_info,
+                };
+                self.send_command_failed_event(command_failed_event);
+                Err(error)
+            }
+            Ok(response) => {
+                let reply = if should_redact {
+                    Document::new()
+                } else {
+                    response.raw_response.clone()
+                };
+
+                let command_succeeded_event = CommandSucceededEvent {
+                    duration,
+                    reply,
+                    command_name: cmd.name.clone(),
+                    request_id,
+                    connection: connection_info,
+                };
+                self.send_command_succeeded_event(command_succeeded_event);
+                op.handle_response(response)
+            }
+        }
     }
 }

src/cmap/conn/command.rs

@@ -104,10 +104,7 @@ impl CommandResponse {
     }
 
     /// Deserialize the body of the response.
-    /// If this response corresponds to a command failure, an appropriate CommandError result will
-    /// be returned.
     pub(crate) fn body<T: DeserializeOwned>(&self) -> Result<T> {
-        self.validate()?;
         match bson::from_bson(Bson::Document(self.raw_response.clone())) {
             Ok(body) => Ok(body),
             Err(e) => Err(ErrorKind::ResponseError {

src/cmap/conn/mod.rs

@@ -27,6 +27,7 @@ use crate::{
 };
 pub(crate) use command::{Command, CommandResponse};
 pub(crate) use stream_description::StreamDescription;
+pub(crate) use wire::next_request_id;
 
 /// User-facing information about a connection to the database.
 #[derive(Clone, Debug)]
@@ -88,6 +89,13 @@ impl Connection {
         Ok(conn)
     }
 
+    pub(crate) fn info(&self) -> ConnectionInfo {
+        ConnectionInfo {
+            id: self.id,
+            address: self.address.clone(),
+        }
+    }
+
     pub(crate) fn address(&self) -> &StreamAddress {
         &self.address
     }
@@ -189,8 +197,12 @@ impl Connection {
     /// An `Ok(...)` result simply means the server received the command and that the driver
     /// driver received the response; it does not imply anything about the success of the command
     /// itself.
-    pub(crate) fn send_command(&mut self, command: Command) -> Result<CommandResponse> {
-        let message = Message::from_command(command);
+    pub(crate) fn send_command(
+        &mut self,
+        command: Command,
+        request_id: impl Into<Option<i32>>,
+    ) -> Result<CommandResponse> {
+        let message = Message::with_command(command, request_id.into());
         message.write_to(&mut self.stream)?;
 
         let response_message = Message::read_from(&mut self.stream)?;

src/cmap/conn/wire/message.rs

@@ -20,13 +20,14 @@ pub(crate) struct Message {
     pub(crate) flags: MessageFlags,
     pub(crate) sections: Vec<MessageSection>,
     pub(crate) checksum: Option<u32>,
+    pub(crate) request_id: Option<i32>,
 }
 
 impl Message {
     /// Creates a `Message` from a given `Command`.
     ///
     /// Note that `response_to` will need to be set manually.
-    pub(crate) fn from_command(mut command: Command) -> Self {
+    pub(crate) fn with_command(mut command: Command, request_id: Option<i32>) -> Self {
         command.body.insert("$db", command.target_db);
 
         if let Some(read_pref) = command.read_pref {
@@ -40,6 +41,7 @@ impl Message {
             flags: MessageFlags::empty(),
             sections: vec![MessageSection::Document(command.body)],
             checksum: None,
+            request_id,
         }
     }
 
@@ -110,6 +112,7 @@ impl Message {
             flags,
             sections,
             checksum,
+            request_id: None,
         })
     }
 
@@ -132,7 +135,7 @@ impl Message {
 
         let header = Header {
             length: total_length as i32,
-            request_id: super::util::next_request_id(),
+            request_id: self.request_id.unwrap_or_else(super::util::next_request_id),
             response_to: self.response_to,
             op_code: OpCode::Message,
         };

src/cmap/conn/wire/mod.rs

@@ -4,4 +4,4 @@ mod message;
 mod test;
 mod util;
 
-pub(super) use message::Message;
+pub(crate) use self::{message::Message, util::next_request_id};

src/cmap/conn/wire/test.rs

@@ -23,6 +23,7 @@ fn basic() {
             doc! { "isMaster": 1, "$db": "admin" },
         )],
         checksum: None,
+        request_id: None,
     };
 
     let StreamAddress { ref hostname, port } = CLIENT_OPTIONS.hosts[0];

src/cmap/conn/wire/util.rs

@@ -8,7 +8,7 @@ use crate::error::Result;
 use lazy_static::lazy_static;
 
 /// Closure to obtain a new, unique request ID.
-pub(super) fn next_request_id() -> i32 {
+pub(crate) fn next_request_id() -> i32 {
     lazy_static! {
         static ref REQUEST_ID: AtomicI32 = AtomicI32::new(0);
     }

src/cmap/establish/handshake.rs

@@ -83,7 +83,7 @@ impl Handshaker {
 
     /// Handshakes a connection.
     pub(super) fn handshake(&self, conn: &mut Connection) -> Result<()> {
-        let response = conn.send_command(self.command.clone())?;
+        let response = conn.send_command(self.command.clone(), None)?;
         let command_response = response.body()?;
 
         // TODO RUST-192: Calculate round trip time.