use std::{collections::HashMap, sync::Arc, time::Duration};
use derivative::Derivative;
use serde::{de::Error as SerdeError, Deserialize, Deserializer, Serialize};
use typed_builder::TypedBuilder;
use crate::{
bson::doc,
error::{ErrorKind, Result},
options::ServerAddress,
sdam::public::ServerInfo,
serde_util,
};
#[derive(Clone, Derivative, derive_more::Display)]
#[derivative(Debug)]
#[non_exhaustive]
pub enum SelectionCriteria {
#[display(fmt = "ReadPreference {}", _0)]
ReadPreference(ReadPreference),
#[display(fmt = "Custom predicate")]
Predicate(#[derivative(Debug = "ignore")] Predicate),
}
impl PartialEq for SelectionCriteria {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(Self::ReadPreference(r1), Self::ReadPreference(r2)) => r1 == r2,
_ => false,
}
}
}
impl From<ReadPreference> for SelectionCriteria {
fn from(read_pref: ReadPreference) -> Self {
Self::ReadPreference(read_pref)
}
}
impl SelectionCriteria {
pub(crate) fn as_read_pref(&self) -> Option<&ReadPreference> {
match self {
Self::ReadPreference(ref read_pref) => Some(read_pref),
Self::Predicate(..) => None,
}
}
#[cfg(test)]
pub(crate) fn as_predicate(&self) -> Option<&Predicate> {
match self {
Self::Predicate(ref p) => Some(p),
_ => None,
}
}
#[cfg(test)]
pub(crate) fn is_read_pref_primary(&self) -> bool {
matches!(self, Self::ReadPreference(ReadPreference::Primary))
}
pub(crate) fn from_address(address: ServerAddress) -> Self {
SelectionCriteria::Predicate(Arc::new(move |server| server.address() == &address))
}
#[cfg(test)]
pub(crate) fn serialize_for_client_options<S>(
selection_criteria: &Option<SelectionCriteria>,
serializer: S,
) -> std::result::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
match selection_criteria {
Some(SelectionCriteria::ReadPreference(pref)) => {
ReadPreference::serialize_for_client_options(pref, serializer)
}
_ => serializer.serialize_none(),
}
}
}
impl<'de> Deserialize<'de> for SelectionCriteria {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: Deserializer<'de>,
{
Ok(SelectionCriteria::ReadPreference(
ReadPreference::deserialize(deserializer)?,
))
}
}
pub type Predicate = Arc<dyn Send + Sync + Fn(&ServerInfo) -> bool>;
#[allow(missing_docs)]
#[derive(Clone, Debug, PartialEq)]
pub enum ReadPreference {
Primary,
Secondary { options: ReadPreferenceOptions },
PrimaryPreferred { options: ReadPreferenceOptions },
SecondaryPreferred { options: ReadPreferenceOptions },
Nearest { options: ReadPreferenceOptions },
}
impl std::fmt::Display for ReadPreference {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{{ Mode: ")?;
let opts_ref = match self {
ReadPreference::Primary => {
write!(f, "Primary")?;
None
}
ReadPreference::Secondary { options } => {
write!(f, "Secondary")?;
Some(options)
}
ReadPreference::PrimaryPreferred { options } => {
write!(f, "PrimaryPreferred")?;
Some(options)
}
ReadPreference::SecondaryPreferred { options } => {
write!(f, "SecondaryPreferred")?;
Some(options)
}
ReadPreference::Nearest { options } => {
write!(f, "Nearest")?;
Some(options)
}
};
if let Some(opts) = opts_ref {
if !opts.is_default() {
if let Some(ref tag_sets) = opts.tag_sets {
write!(f, ", Tag Sets: {:?}", tag_sets)?;
}
if let Some(ref max_staleness) = opts.max_staleness {
write!(f, ", Max Staleness: {:?}", max_staleness)?;
}
if let Some(ref hedge) = opts.hedge {
write!(f, ", Hedge: {}", hedge.enabled)?;
}
}
}
write!(f, " }}")
}
}
impl<'de> Deserialize<'de> for ReadPreference {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Serialize, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ReadPreferenceHelper {
mode: String,
#[serde(flatten)]
options: ReadPreferenceOptions,
}
let preference = ReadPreferenceHelper::deserialize(deserializer)?;
match preference.mode.to_ascii_lowercase().as_str() {
"primary" => {
if !preference.options.is_default() {
return Err(D::Error::custom(&format!(
"no options can be specified with read preference mode = primary, but got \
{:?}",
preference.options
)));
}
Ok(ReadPreference::Primary)
}
"secondary" => Ok(ReadPreference::Secondary {
options: preference.options,
}),
"primarypreferred" => Ok(ReadPreference::PrimaryPreferred {
options: preference.options,
}),
"secondarypreferred" => Ok(ReadPreference::SecondaryPreferred {
options: preference.options,
}),
"nearest" => Ok(ReadPreference::Nearest {
options: preference.options,
}),
other => Err(D::Error::custom(format!(
"Unknown read preference mode: {}",
other
))),
}
}
}
impl Serialize for ReadPreference {
fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
#[serde_with::skip_serializing_none]
#[derive(Serialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct ReadPreferenceHelper<'a> {
mode: &'static str,
#[serde(flatten)]
options: Option<&'a ReadPreferenceOptions>,
}
let helper = match self {
ReadPreference::Primary => ReadPreferenceHelper {
mode: "primary",
options: None,
},
ReadPreference::PrimaryPreferred { options } => ReadPreferenceHelper {
mode: "primaryPreferred",
options: Some(options),
},
ReadPreference::Secondary { options } => ReadPreferenceHelper {
mode: "secondary",
options: Some(options),
},
ReadPreference::SecondaryPreferred { options } => ReadPreferenceHelper {
mode: "secondaryPreferred",
options: Some(options),
},
ReadPreference::Nearest { options } => ReadPreferenceHelper {
mode: "nearest",
options: Some(options),
},
};
helper.serialize(serializer)
}
}
#[serde_with::skip_serializing_none]
#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, TypedBuilder)]
#[builder(field_defaults(default, setter(into)))]
#[serde(rename_all = "camelCase")]
#[non_exhaustive]
pub struct ReadPreferenceOptions {
pub tag_sets: Option<Vec<TagSet>>,
#[serde(
rename = "maxStalenessSeconds",
default,
with = "serde_util::duration_option_as_int_seconds"
)]
pub max_staleness: Option<Duration>,
pub hedge: Option<HedgedReadOptions>,
}
impl ReadPreferenceOptions {
pub(crate) fn is_default(&self) -> bool {
self.hedge.is_none()
&& self.max_staleness.is_none()
&& self
.tag_sets
.as_ref()
.map(|ts| ts.is_empty() || ts[..] == [HashMap::default()])
.unwrap_or(true)
}
}
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, TypedBuilder)]
#[non_exhaustive]
pub struct HedgedReadOptions {
pub enabled: bool,
}
impl HedgedReadOptions {
pub fn with_enabled(enabled: bool) -> Self {
Self { enabled }
}
}
impl ReadPreference {
pub(crate) fn max_staleness(&self) -> Option<Duration> {
match self {
ReadPreference::Primary => None,
ReadPreference::Secondary { ref options }
| ReadPreference::PrimaryPreferred { ref options }
| ReadPreference::SecondaryPreferred { ref options }
| ReadPreference::Nearest { ref options } => options.max_staleness,
}
}
pub(crate) fn with_tags(mut self, tag_sets: Vec<TagSet>) -> Result<Self> {
let options = match self {
ReadPreference::Primary => {
return Err(ErrorKind::InvalidArgument {
message: "read preference tags can only be specified when a non-primary mode \
is specified"
.to_string(),
}
.into());
}
ReadPreference::Secondary { ref mut options } => options,
ReadPreference::PrimaryPreferred { ref mut options } => options,
ReadPreference::SecondaryPreferred { ref mut options } => options,
ReadPreference::Nearest { ref mut options } => options,
};
options.tag_sets = Some(tag_sets);
Ok(self)
}
pub(crate) fn with_max_staleness(mut self, max_staleness: Duration) -> Result<Self> {
let options = match self {
ReadPreference::Primary => {
return Err(ErrorKind::InvalidArgument {
message: "max staleness can only be specified when a non-primary mode is \
specified"
.to_string(),
}
.into());
}
ReadPreference::Secondary { ref mut options } => options,
ReadPreference::PrimaryPreferred { ref mut options } => options,
ReadPreference::SecondaryPreferred { ref mut options } => options,
ReadPreference::Nearest { ref mut options } => options,
};
options.max_staleness = Some(max_staleness);
Ok(self)
}
#[cfg(test)]
pub(crate) fn serialize_for_client_options<S>(
read_preference: &ReadPreference,
serializer: S,
) -> std::result::Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
#[derive(serde::Serialize)]
struct ReadPreferenceHelper<'a> {
readpreference: &'a str,
readpreferencetags: Option<&'a Vec<HashMap<String, String>>>,
#[serde(serialize_with = "serde_util::duration_option_as_int_seconds::serialize")]
maxstalenessseconds: Option<Duration>,
}
let state = match read_preference {
ReadPreference::Primary => ReadPreferenceHelper {
readpreference: "primary",
readpreferencetags: None,
maxstalenessseconds: None,
},
ReadPreference::PrimaryPreferred { options } => ReadPreferenceHelper {
readpreference: "primaryPreferred",
readpreferencetags: options.tag_sets.as_ref(),
maxstalenessseconds: options.max_staleness,
},
ReadPreference::Secondary { options } => ReadPreferenceHelper {
readpreference: "secondary",
readpreferencetags: options.tag_sets.as_ref(),
maxstalenessseconds: options.max_staleness,
},
ReadPreference::SecondaryPreferred { options } => ReadPreferenceHelper {
readpreference: "secondaryPreferred",
readpreferencetags: options.tag_sets.as_ref(),
maxstalenessseconds: options.max_staleness,
},
ReadPreference::Nearest { options } => ReadPreferenceHelper {
readpreference: "nearest",
readpreferencetags: options.tag_sets.as_ref(),
maxstalenessseconds: options.max_staleness,
},
};
state.serialize(serializer)
}
}
pub type TagSet = HashMap<String, String>;
#[cfg(test)]
mod test {
use super::{HedgedReadOptions, ReadPreference, ReadPreferenceOptions};
use crate::bson::doc;
#[test]
fn hedged_read_included_in_document() {
let options = ReadPreferenceOptions::builder()
.hedge(HedgedReadOptions { enabled: true })
.build();
let read_pref = ReadPreference::Secondary { options };
let doc = bson::to_document(&read_pref).unwrap();
assert_eq!(
doc,
doc! { "mode": "secondary", "hedge": { "enabled": true } }
);
}
}