RUST-956 Include service_id in relevant events (#469)

Author: abr-egn

src/client/executor.rs

@@ -478,6 +478,7 @@ impl Client {
         }
 
         let connection_info = connection.info();
+        let service_id = connection.service_id();
         let request_id = crate::cmap::conn::next_request_id();
 
         if let Some(ref server_api) = self.inner.options.server_api {
@@ -509,6 +510,7 @@ impl Client {
                 command_name: raw_cmd.name.clone(),
                 request_id,
                 connection: connection_info.clone(),
+                service_id,
             };
 
             handler.handle_command_started_event(command_started_event);
@@ -600,6 +602,7 @@ impl Client {
                         failure: err.clone(),
                         request_id,
                         connection: connection_info,
+                        service_id,
                     };
 
                     handler.handle_command_failed_event(command_failed_event);
@@ -631,6 +634,7 @@ impl Client {
                         command_name: cmd_name.clone(),
                         request_id,
                         connection: connection_info,
+                        service_id,
                     };
                     handler.handle_command_succeeded_event(command_succeeded_event);
                 });

src/cmap/conn/mod.rs

@@ -167,6 +167,10 @@ impl Connection {
         }
     }
 
+    pub(crate) fn service_id(&self) -> Option<ObjectId> {
+        self.generation.service_id()
+    }
+
     pub(crate) fn address(&self) -> &ServerAddress {
         &self.address
     }

src/cmap/worker.rs

@@ -503,6 +503,7 @@ impl ConnectionPoolWorker {
             self.emit_event(|handler| {
                 let event = PoolClearedEvent {
                     address: self.address.clone(),
+                    service_id,
                 };
 
                 handler.handle_pool_cleared_event(event);

src/event/cmap.rs

@@ -5,7 +5,7 @@ use std::time::Duration;
 
 use serde::Deserialize;
 
-use crate::{bson_util, options::ServerAddress};
+use crate::{bson::oid::ObjectId, bson_util, options::ServerAddress};
 
 /// We implement `Deserialize` for all of the event types so that we can more easily parse the CMAP
 /// spec tests. However, we have no need to parse the address field from the JSON files (if it's
@@ -79,6 +79,9 @@ pub struct PoolClearedEvent {
     #[serde(default = "self::empty_address")]
     #[serde(skip)]
     pub address: ServerAddress,
+
+    /// If the connection is to a load balancer, the id of the selected backend.
+    pub service_id: Option<ObjectId>,
 }
 
 /// Event emitted when a connection pool is cleared.

src/event/command.rs

@@ -3,7 +3,11 @@
 
 use std::time::Duration;
 
-use crate::{bson::Document, cmap::ConnectionInfo, error::Error};
+use crate::{
+    bson::{oid::ObjectId, Document},
+    cmap::ConnectionInfo,
+    error::Error,
+};
 
 /// An event that triggers when a database command is initiated.
 #[derive(Clone, Debug)]
@@ -26,6 +30,9 @@ pub struct CommandStartedEvent {
 
     /// Information about the connect the command will be run on.
     pub connection: ConnectionInfo,
+
+    /// If the client connection is to a load balancer, the id of the selected backend.
+    pub service_id: Option<ObjectId>,
 }
 
 /// An event that triggers when a database command completes without an error.
@@ -48,6 +55,9 @@ pub struct CommandSucceededEvent {
 
     /// Information about the connect the command will be run on.
     pub connection: ConnectionInfo,
+
+    /// If the client connection is to a load balancer, the id of the selected backend.
+    pub service_id: Option<ObjectId>,
 }
 
 /// An event that triggers when a command failed to complete successfully.
@@ -70,6 +80,9 @@ pub struct CommandFailedEvent {
 
     /// Information about the connect the command will be run on.
     pub connection: ConnectionInfo,
+
+    /// If the client connection is to a load balancer, the id of the selected backend.
+    pub service_id: Option<ObjectId>,
 }
 
 /// Applications can implement this trait to specify custom logic to run on each command event sent

src/is_master.rs

@@ -195,6 +195,8 @@ pub(crate) struct IsMasterCommandResponse {
 
     /// The maximum number of write operations permitted in a write batch.
     pub max_write_batch_size: i64,
+
+    /// If the connection is to a load balancer, the id of the selected backend.
     pub service_id: Option<ObjectId>,
 }