mongodb/
selection_criteria.rs

1use std::{collections::HashMap, sync::Arc, time::Duration};
2
3use derive_where::derive_where;
4use serde::{de::Error as SerdeError, Deserialize, Deserializer, Serialize};
5use typed_builder::TypedBuilder;
6
7use crate::{
8    bson::doc,
9    error::{ErrorKind, Result},
10    options::ServerAddress,
11    sdam::public::ServerInfo,
12    serde_util,
13};
14
15/// Describes which servers are suitable for a given operation.
16#[derive(Clone, derive_more::Display)]
17#[derive_where(Debug)]
18#[non_exhaustive]
19pub enum SelectionCriteria {
20    /// A read preference that describes the suitable servers based on the server type, max
21    /// staleness, and server tags.
22    ///
23    /// See the documentation [here](https://www.mongodb.com/docs/manual/core/read-preference/) for more details.
24    #[display("ReadPreference {_0}")]
25    ReadPreference(ReadPreference),
26
27    /// A predicate used to filter servers that are considered suitable. A `server` will be
28    /// considered suitable by a `predicate` if `predicate(server)` returns true.
29    #[display("Custom predicate")]
30    Predicate(#[derive_where(skip)] Predicate),
31}
32
33impl PartialEq for SelectionCriteria {
34    fn eq(&self, other: &Self) -> bool {
35        match (self, other) {
36            (Self::ReadPreference(r1), Self::ReadPreference(r2)) => r1 == r2,
37            _ => false,
38        }
39    }
40}
41
42impl From<ReadPreference> for SelectionCriteria {
43    fn from(read_pref: ReadPreference) -> Self {
44        Self::ReadPreference(read_pref)
45    }
46}
47
48impl SelectionCriteria {
49    pub(crate) fn as_read_pref(&self) -> Option<&ReadPreference> {
50        match self {
51            Self::ReadPreference(ref read_pref) => Some(read_pref),
52            Self::Predicate(..) => None,
53        }
54    }
55
56    pub(crate) fn from_address(address: ServerAddress) -> Self {
57        SelectionCriteria::Predicate(Arc::new(move |server| server.address() == &address))
58    }
59
60    #[cfg(test)]
61    pub(crate) fn serialize_for_client_options<S>(
62        selection_criteria: &Option<SelectionCriteria>,
63        serializer: S,
64    ) -> std::result::Result<S::Ok, S::Error>
65    where
66        S: serde::Serializer,
67    {
68        match selection_criteria {
69            Some(SelectionCriteria::ReadPreference(read_preference)) => {
70                read_preference.serialize(serializer)
71            }
72            _ => serializer.serialize_none(),
73        }
74    }
75}
76
77impl<'de> Deserialize<'de> for SelectionCriteria {
78    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
79    where
80        D: Deserializer<'de>,
81    {
82        Ok(SelectionCriteria::ReadPreference(
83            ReadPreference::deserialize(deserializer)?,
84        ))
85    }
86}
87
88/// A predicate used to filter servers that are considered suitable.
89pub type Predicate = Arc<dyn Send + Sync + Fn(&ServerInfo) -> bool>;
90
91/// Specifies how the driver should route a read operation to members of a replica set.
92///
93/// If applicable, `tag_sets` can be used to target specific nodes in a replica set, and
94/// `max_staleness` specifies the maximum lag behind the primary that a secondary can be to remain
95/// eligible for the operation. The max staleness value maps to the `maxStalenessSeconds` MongoDB
96/// option and will be sent to the server as an integer number of seconds.
97///
98/// See the [MongoDB docs](https://www.mongodb.com/docs/manual/core/read-preference) for more details.
99#[allow(missing_docs)]
100#[derive(Clone, Debug, PartialEq)]
101#[non_exhaustive]
102pub enum ReadPreference {
103    /// Only route this operation to the primary.
104    Primary,
105
106    /// Only route this operation to a secondary.
107    Secondary {
108        options: Option<ReadPreferenceOptions>,
109    },
110
111    /// Route this operation to the primary if it's available, but fall back to the secondaries if
112    /// not.
113    PrimaryPreferred {
114        options: Option<ReadPreferenceOptions>,
115    },
116
117    /// Route this operation to a secondary if one is available, but fall back to the primary if
118    /// not.
119    SecondaryPreferred {
120        options: Option<ReadPreferenceOptions>,
121    },
122
123    /// Route this operation to the node with the least network latency regardless of whether it's
124    /// the primary or a secondary.
125    Nearest {
126        options: Option<ReadPreferenceOptions>,
127    },
128}
129
130impl std::fmt::Display for ReadPreference {
131    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
132        let mut mode = self.mode().to_string();
133        mode[0..1].make_ascii_uppercase();
134        write!(f, "{{ Mode: {mode}")?;
135
136        if let Some(options) = self.options() {
137            if let Some(ref tag_sets) = options.tag_sets {
138                write!(f, ", Tag Sets: {tag_sets:?}")?;
139            }
140            if let Some(ref max_staleness) = options.max_staleness {
141                write!(f, ", Max Staleness: {max_staleness:?}")?;
142            }
143            #[allow(deprecated)]
144            if let Some(ref hedge) = options.hedge {
145                write!(f, ", Hedge: {}", hedge.enabled)?;
146            }
147        }
148
149        write!(f, " }}")
150    }
151}
152
153impl<'de> Deserialize<'de> for ReadPreference {
154    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
155    where
156        D: Deserializer<'de>,
157    {
158        #[derive(Serialize, Deserialize)]
159        #[serde(rename_all = "camelCase", deny_unknown_fields)]
160        struct ReadPreferenceHelper {
161            mode: String,
162            #[serde(flatten)]
163            options: ReadPreferenceOptions,
164        }
165        let helper = ReadPreferenceHelper::deserialize(deserializer)?;
166        match helper.mode.to_ascii_lowercase().as_str() {
167            "primary" => {
168                if !helper.options.is_default() {
169                    return Err(D::Error::custom(format!(
170                        "cannot specify options for primary read preference, got {:?}",
171                        helper.options
172                    )));
173                }
174                Ok(ReadPreference::Primary)
175            }
176            "secondary" => Ok(ReadPreference::Secondary {
177                options: Some(helper.options),
178            }),
179            "primarypreferred" => Ok(ReadPreference::PrimaryPreferred {
180                options: Some(helper.options),
181            }),
182            "secondarypreferred" => Ok(ReadPreference::SecondaryPreferred {
183                options: Some(helper.options),
184            }),
185            "nearest" => Ok(ReadPreference::Nearest {
186                options: Some(helper.options),
187            }),
188            other => Err(D::Error::custom(format!(
189                "Unknown read preference mode: {other}"
190            ))),
191        }
192    }
193}
194
195impl Serialize for ReadPreference {
196    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
197    where
198        S: serde::Serializer,
199    {
200        #[serde_with::skip_serializing_none]
201        #[derive(Serialize)]
202        #[serde(rename_all = "camelCase")]
203        struct ReadPreferenceHelper<'a> {
204            mode: &'static str,
205            #[serde(flatten)]
206            options: Option<&'a ReadPreferenceOptions>,
207        }
208
209        let helper = ReadPreferenceHelper {
210            mode: self.mode(),
211            options: self.options(),
212        };
213        helper.serialize(serializer)
214    }
215}
216
217/// Specifies read preference options for non-primary read preferences.
218#[allow(deprecated)]
219#[serde_with::skip_serializing_none]
220#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, TypedBuilder)]
221#[builder(field_defaults(default, setter(into)))]
222#[serde(rename_all = "camelCase")]
223#[non_exhaustive]
224pub struct ReadPreferenceOptions {
225    /// Specifies which replica set members should be considered for operations. Each tag set will
226    /// be checked in order until one or more servers is found with each tag in the set.
227    #[serde(alias = "tag_sets")]
228    pub tag_sets: Option<Vec<TagSet>>,
229
230    /// Specifies the maximum amount of lag behind the primary that a secondary can be to be
231    /// considered for the given operation. Any secondaries lagging behind more than
232    /// `max_staleness` will not be considered for the operation.
233    ///
234    /// `max_staleness` must be at least 90 seconds. If a `max_staleness` less than 90 seconds is
235    /// specified for an operation, the operation will return an error.
236    #[serde(
237        rename = "maxStalenessSeconds",
238        default,
239        with = "serde_util::duration_option_as_int_seconds"
240    )]
241    pub max_staleness: Option<Duration>,
242
243    /// Specifies hedging behavior for reads. These options only apply to sharded clusters on
244    /// servers that are at least version 4.4. Note that hedged reads are automatically enabled for
245    /// read preference mode "nearest" on server versions less than 8.0.
246    ///
247    /// See the [MongoDB docs](https://www.mongodb.com/docs/manual/core/read-preference-hedge-option/) for more details.
248    #[deprecated(
249        note = "hedged reads are deprecated as of MongoDB 8.0 and will be removed in a future \
250                server version"
251    )]
252    pub hedge: Option<HedgedReadOptions>,
253}
254
255impl ReadPreferenceOptions {
256    pub(crate) fn is_default(&self) -> bool {
257        #[allow(deprecated)]
258        let hedge = self.hedge.is_some();
259        !hedge
260            && self.max_staleness.is_none()
261            && self
262                .tag_sets
263                .as_ref()
264                .map(|ts| ts.is_empty() || ts[..] == [HashMap::default()])
265                .unwrap_or(true)
266    }
267}
268
269/// Specifies hedging behavior for reads.
270///
271/// See the [MongoDB docs](https://www.mongodb.com/docs/manual/core/read-preference-hedge-option/) for more details.
272#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, TypedBuilder)]
273#[builder(field_defaults(default, setter(into)))]
274#[non_exhaustive]
275pub struct HedgedReadOptions {
276    /// Whether or not to allow reads from a sharded cluster to be "hedged" across two replica
277    /// set members per shard, with the results from the first response received back from either
278    /// being returned.
279    pub enabled: bool,
280}
281
282impl ReadPreference {
283    pub(crate) fn mode(&self) -> &'static str {
284        match self {
285            Self::Primary => "primary",
286            Self::Secondary { .. } => "secondary",
287            Self::PrimaryPreferred { .. } => "primaryPreferred",
288            Self::SecondaryPreferred { .. } => "secondaryPreferred",
289            Self::Nearest { .. } => "nearest",
290        }
291    }
292
293    pub(crate) fn options(&self) -> Option<&ReadPreferenceOptions> {
294        match self {
295            Self::Primary => None,
296            Self::Secondary { options }
297            | Self::PrimaryPreferred { options }
298            | Self::SecondaryPreferred { options }
299            | Self::Nearest { options } => options.as_ref(),
300        }
301    }
302
303    pub(crate) fn max_staleness(&self) -> Option<Duration> {
304        self.options().and_then(|options| options.max_staleness)
305    }
306
307    pub(crate) fn tag_sets(&self) -> Option<&Vec<TagSet>> {
308        self.options().and_then(|options| options.tag_sets.as_ref())
309    }
310
311    pub(crate) fn with_tags(mut self, tag_sets: Vec<TagSet>) -> Result<Self> {
312        let options = match self {
313            Self::Primary => {
314                return Err(ErrorKind::InvalidArgument {
315                    message: "read preference tags can only be specified when a non-primary mode \
316                              is specified"
317                        .to_string(),
318                }
319                .into());
320            }
321            Self::Secondary { ref mut options } => options,
322            Self::PrimaryPreferred { ref mut options } => options,
323            Self::SecondaryPreferred { ref mut options } => options,
324            Self::Nearest { ref mut options } => options,
325        };
326
327        options.get_or_insert_with(Default::default).tag_sets = Some(tag_sets);
328
329        Ok(self)
330    }
331
332    pub(crate) fn with_max_staleness(mut self, max_staleness: Duration) -> Result<Self> {
333        let options = match self {
334            ReadPreference::Primary => {
335                return Err(ErrorKind::InvalidArgument {
336                    message: "max staleness can only be specified when a non-primary mode is \
337                              specified"
338                        .to_string(),
339                }
340                .into());
341            }
342            ReadPreference::Secondary { ref mut options } => options,
343            ReadPreference::PrimaryPreferred { ref mut options } => options,
344            ReadPreference::SecondaryPreferred { ref mut options } => options,
345            ReadPreference::Nearest { ref mut options } => options,
346        };
347
348        options.get_or_insert_with(Default::default).max_staleness = Some(max_staleness);
349
350        Ok(self)
351    }
352}
353
354/// A read preference tag set. See the documentation [here](https://www.mongodb.com/docs/manual/tutorial/configure-replica-set-tag-sets/) for more details.
355pub type TagSet = HashMap<String, String>;
356
357#[cfg(test)]
358mod test {
359    use super::{HedgedReadOptions, ReadPreference, ReadPreferenceOptions};
360    use crate::bson::doc;
361
362    #[test]
363    fn hedged_read_included_in_document() {
364        #[allow(deprecated)]
365        let options = Some(
366            ReadPreferenceOptions::builder()
367                .hedge(HedgedReadOptions { enabled: true })
368                .build(),
369        );
370
371        let read_pref = ReadPreference::Secondary { options };
372        let doc = crate::bson_compat::serialize_to_document(&read_pref).unwrap();
373
374        assert_eq!(
375            doc,
376            doc! { "mode": "secondary", "hedge": { "enabled": true } }
377        );
378    }
379}
mongodb/selection_criteria.rs

mongodb/
selection_criteria.rs
