switch to callback

Author: abr-egn

src/client.rs

@@ -35,6 +35,7 @@ use crate::{
     error::{Error, ErrorKind, Result},
     event::command::CommandEvent,
     id_set::IdSet,
+    operation::OverrideCriteriaFn,
     options::{ClientOptions, DatabaseOptions, ReadPreference, SelectionCriteria, ServerAddress},
     sdam::{
         server_selection::{self, attempt_to_select_server},
@@ -447,7 +448,7 @@ impl Client {
         criteria: Option<&SelectionCriteria>,
     ) -> Result<ServerAddress> {
         let (server, _) = self
-            .select_server(criteria, "Test select server", None, false)
+            .select_server(criteria, "Test select server", None, |_, _| None)
             .await?;
         Ok(server.address.clone())
     }
@@ -460,7 +461,7 @@ impl Client {
         #[allow(unused_variables)] // we only use the operation_name for tracing.
         operation_name: &str,
         deprioritized: Option<&ServerAddress>,
-        is_out_or_merge: bool,
+        override_criteria: OverrideCriteriaFn,
     ) -> Result<(SelectedServer, SelectionCriteria)> {
         let criteria =
             criteria.unwrap_or(&SelectionCriteria::ReadPreference(ReadPreference::Primary));
@@ -489,15 +490,14 @@ impl Client {
         let mut watcher = self.inner.topology.watch();
         loop {
             let state = watcher.observe_latest();
-            let override_criteria;
-            let effective_criteria = if let Some(oc) =
-                Self::override_criteria(criteria, &state.description, is_out_or_merge)
-            {
-                override_criteria = oc;
-                &override_criteria
-            } else {
-                criteria
-            };
+            let override_slot;
+            let effective_criteria =
+                if let Some(oc) = override_criteria(criteria, &state.description) {
+                    override_slot = oc;
+                    &override_slot
+                } else {
+                    criteria
+                };
             let result = server_selection::attempt_to_select_server(
                 effective_criteria,
                 &state.description,
@@ -549,38 +549,6 @@ impl Client {
         }
     }
 
-    /// Check to see if selection criteria need to be overridden.  Currently only required for
-    /// aggregate operations with $merge/$out stages.
-    fn override_criteria(
-        criteria: &SelectionCriteria,
-        desc: &crate::sdam::TopologyDescription,
-        is_out_or_merge: bool,
-    ) -> Option<SelectionCriteria> {
-        if is_out_or_merge {
-            eprintln!("aggregate: checking override");
-        }
-        if !is_out_or_merge
-            || criteria == &SelectionCriteria::ReadPreference(ReadPreference::Primary)
-            || desc.topology_type() == crate::TopologyType::LoadBalanced
-        {
-            if is_out_or_merge {
-                eprintln!("aggregate: skipping override");
-            }
-            return None;
-        }
-        for server in desc.servers.values() {
-            let _ = dbg!(server.hello_response());
-            if let Ok(Some(v)) = server.max_wire_version() {
-                static SERVER_5_0_0_WIRE_VERSION: i32 = 13;
-                if v < SERVER_5_0_0_WIRE_VERSION {
-                    eprintln!("aggregate: overriding criteria");
-                    return Some(SelectionCriteria::ReadPreference(ReadPreference::Primary));
-                }
-            }
-        }
-        return None;
-    }
-
     #[cfg(all(test, feature = "dns-resolver"))]
     pub(crate) fn get_hosts(&self) -> Vec<String> {
         let watcher = self.inner.topology.watch();

src/client/executor.rs

@@ -326,7 +326,7 @@ impl Client {
                     selection_criteria,
                     op.name(),
                     retry.as_ref().map(|r| &r.first_server),
-                    op.is_out_or_merge(),
+                    op.override_criteria(),
                 )
                 .await
             {
@@ -863,7 +863,7 @@ impl Client {
                 || server_type.is_data_bearing()
         }));
         let _ = self
-            .select_server(Some(&criteria), operation_name, None, false)
+            .select_server(Some(&criteria), operation_name, None, |_, _| None)
             .await?;
         Ok(())
     }

src/operation.rs

@@ -76,6 +76,7 @@ pub(crate) use update::{Update, UpdateOrReplace};
 
 const SERVER_4_2_0_WIRE_VERSION: i32 = 8;
 const SERVER_4_4_0_WIRE_VERSION: i32 = 9;
+const SERVER_5_0_0_WIRE_VERSION: i32 = 13;
 const SERVER_8_0_0_WIRE_VERSION: i32 = 25;
 // The maximum number of bytes that may be included in a write payload when auto-encryption is
 // enabled.
@@ -148,14 +149,18 @@ pub(crate) trait Operation {
     /// Updates this operation as needed for a retry.
     fn update_for_retry(&mut self);
 
-    /// Returns whether this is a $out or $merge aggregation operation.
-    fn is_out_or_merge(&self) -> bool;
+    /// Returns a function handle to potentially override selection criteria based on server
+    /// topology.
+    fn override_criteria(&self) -> OverrideCriteriaFn;
 
     fn pinned_connection(&self) -> Option<&PinnedConnectionHandle>;
 
     fn name(&self) -> &str;
 }
 
+pub(crate) type OverrideCriteriaFn =
+    fn(&SelectionCriteria, &crate::sdam::TopologyDescription) -> Option<SelectionCriteria>;
+
 // A mirror of the `Operation` trait, with default behavior where appropriate.  Should only be
 // implemented by operation types that do not delegate to other operations.
 pub(crate) trait OperationWithDefaults: Send + Sync {
@@ -238,9 +243,10 @@ pub(crate) trait OperationWithDefaults: Send + Sync {
     /// Updates this operation as needed for a retry.
     fn update_for_retry(&mut self) {}
 
-    /// Returns whether this is a $out or $merge aggregation operation.
-    fn is_out_or_merge(&self) -> bool {
-        false
+    /// Returns a function handle to potentially override selection criteria based on server
+    /// topology.
+    fn override_criteria(&self) -> OverrideCriteriaFn {
+        |_, _| None
     }
 
     fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> {
@@ -295,8 +301,8 @@ where
     fn update_for_retry(&mut self) {
         self.update_for_retry()
     }
-    fn is_out_or_merge(&self) -> bool {
-        self.is_out_or_merge()
+    fn override_criteria(&self) -> OverrideCriteriaFn {
+        self.override_criteria()
     }
     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,
 };
 
@@ -135,6 +135,30 @@ impl OperationWithDefaults for Aggregate {
         }
     }
 
+    fn override_criteria(&self) -> super::OverrideCriteriaFn {
+        if !self.is_out_or_merge() {
+            return |_, _| None;
+        }
+        |criteria, topology| {
+            if criteria == &SelectionCriteria::ReadPreference(ReadPreference::Primary)
+                || topology.topology_type() == crate::TopologyType::LoadBalanced
+            {
+                return None;
+            }
+            for server in topology.servers.values() {
+                if let Ok(Some(v)) = server.max_wire_version() {
+                    if v < super::SERVER_5_0_0_WIRE_VERSION {
+                        return Some(SelectionCriteria::ReadPreference(ReadPreference::Primary));
+                    }
+                }
+            }
+            None
+        }
+    }
+}
+
+impl Aggregate {
+    /// Returns whether this is a $out or $merge aggregation operation.
     fn is_out_or_merge(&self) -> bool {
         self.pipeline
             .last()

src/operation/raw_output.rs

@@ -68,8 +68,8 @@ impl<Op: Operation> Operation for RawOutput<Op> {
         self.0.update_for_retry()
     }
 
-    fn is_out_or_merge(&self) -> bool {
-        self.0.is_out_or_merge()
+    fn override_criteria(&self) -> super::OverrideCriteriaFn {
+        self.0.override_criteria()
     }
 
     fn pinned_connection(&self) -> Option<&crate::cmap::conn::PinnedConnectionHandle> {

src/operation/run_cursor_command.rs

@@ -79,8 +79,8 @@ impl Operation for RunCursorCommand<'_> {
         self.run_command.update_for_retry()
     }
 
-    fn is_out_or_merge(&self) -> bool {
-        self.run_command.is_out_or_merge()
+    fn override_criteria(&self) -> super::OverrideCriteriaFn {
+        self.run_command.override_criteria()
     }
 
     fn pinned_connection(&self) -> Option<&PinnedConnectionHandle> {

src/sdam/description/server.rs

@@ -386,13 +386,6 @@ impl ServerDescription {
         Ok(me)
     }
 
-    pub(crate) fn hello_response(&self) -> Result<Option<&HelloCommandResponse>> {
-        self.reply
-            .as_ref()
-            .map_err(Clone::clone)
-            .map(|o| o.as_ref().map(|r| &r.command_response))
-    }
-
     pub(crate) fn last_write_date(&self) -> Result<Option<DateTime>> {
         match self.reply {
             Ok(None) => Ok(None),