split topology update from retry

Author: abr-egn

src/client/executor.rs

@@ -309,6 +309,8 @@ impl Client {
         let mut retry: Option<ExecutionRetry> = None;
         let mut implicit_session: Option<ClientSession> = None;
         loop {
+            op.update_for_topology(&self.inner.topology.description());
+
             if retry.is_some() {
                 op.update_for_retry();
             }

src/operation.rs

@@ -52,6 +52,7 @@ use crate::{
         WriteFailure,
     },
     options::WriteConcern,
+    sdam::TopologyDescription,
     selection_criteria::SelectionCriteria,
     BoxFuture,
     ClientSession,
@@ -148,6 +149,9 @@ pub(crate) trait Operation {
     /// Updates this operation as needed for a retry.
     fn update_for_retry(&mut self);
 
+    /// Updates this operation based on server topology.
+    fn update_for_topology(&mut self, _topology: &TopologyDescription);
+
     fn pinned_connection(&self) -> Option<&PinnedConnectionHandle>;
 
     fn name(&self) -> &str;
@@ -235,6 +239,9 @@ pub(crate) trait OperationWithDefaults: Send + Sync {
     /// Updates this operation as needed for a retry.
     fn update_for_retry(&mut self) {}
 
+    /// Updates this operation based on server topology.
+    fn update_for_topology(&mut self, _topology: &TopologyDescription) {}
+
     fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> {
         None
     }
@@ -287,6 +294,9 @@ where
     fn update_for_retry(&mut self) {
         self.update_for_retry()
     }
+    fn update_for_topology(&mut self, topology: &TopologyDescription) {
+        self.update_for_topology(topology)
+    }
     fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> {
         self.pinned_connection()
     }

src/operation/aggregate.rs

@@ -7,7 +7,7 @@ use crate::{
     cursor::CursorSpecification,
     error::Result,
     operation::{append_options, remove_empty_write_concern, Retryability},
-    options::{AggregateOptions, SelectionCriteria, WriteConcern},
+    options::{AggregateOptions, ReadPreference, SelectionCriteria, WriteConcern},
     Namespace,
 };
 
@@ -109,6 +109,14 @@ impl OperationWithDefaults for Aggregate {
         ))
     }
 
+    fn update_for_topology(&mut self, topology: &crate::sdam::TopologyDescription) {
+        match self.selection_criteria() {
+            None => return,
+            Some(SelectionCriteria::ReadPreference(ReadPreference::Primary)) => return,
+            _ => (),
+        }
+    }
+
     fn selection_criteria(&self) -> Option<&SelectionCriteria> {
         self.options
             .as_ref()

src/operation/raw_output.rs

@@ -68,6 +68,10 @@ impl<Op: Operation> Operation for RawOutput<Op> {
         self.0.update_for_retry()
     }
 
+    fn update_for_topology(&mut self, topology: &crate::sdam::TopologyDescription) {
+        self.0.update_for_topology(topology)
+    }
+
     fn pinned_connection(&self) -> Option<&crate::cmap::conn::PinnedConnectionHandle> {
         self.0.pinned_connection()
     }

src/operation/run_cursor_command.rs

@@ -79,6 +79,10 @@ impl Operation for RunCursorCommand<'_> {
         self.run_command.update_for_retry()
     }
 
+    fn update_for_topology(&mut self, topology: &crate::sdam::TopologyDescription) {
+        self.run_command.update_for_topology(topology)
+    }
+
     fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> {
         self.run_command.pinned_connection()
     }

src/sdam/topology.rs

@@ -229,7 +229,6 @@ impl Topology {
         self.watcher.peek_latest().servers()
     }
 
-    #[cfg(test)]
     pub(crate) fn description(&self) -> TopologyDescription {
         self.watcher.peek_latest().description.clone()
     }