1pub mod common;
16pub mod lease_based;
17pub mod load_based;
18pub mod round_robin;
19#[cfg(test)]
20pub(crate) mod test_utils;
21pub mod weight_compute;
22pub mod weighted_choose;
23use std::collections::HashSet;
24
25use api::v1::meta::heartbeat_request::NodeWorkloads;
26use serde::{Deserialize, Serialize};
27use store_api::storage::RegionId;
28use strum::AsRefStr;
29
30use crate::error;
31use crate::error::Result;
32
33#[async_trait::async_trait]
34pub trait Selector: Send + Sync {
35 type Context;
36 type Output;
37
38 async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result<Self::Output>;
39}
40
41#[async_trait::async_trait]
46pub trait RegionStatAwareSelector: Send + Sync {
47 type Context;
48 type Output;
49
50 async fn select(
51 &self,
52 ctx: &Self::Context,
53 from_peer_id: u64,
54 region_ids: &[RegionId],
55 exclude_peer_ids: HashSet<u64>,
56 ) -> Result<Self::Output>;
57}
58
59#[derive(Debug)]
60pub struct SelectorOptions {
61 pub min_required_items: usize,
63 pub allow_duplication: bool,
65 pub exclude_peer_ids: HashSet<u64>,
67 pub workload_filter: Option<fn(&NodeWorkloads) -> bool>,
69}
70
71impl Default for SelectorOptions {
72 fn default() -> Self {
73 Self {
74 min_required_items: 1,
75 allow_duplication: false,
76 exclude_peer_ids: HashSet::new(),
77 workload_filter: None,
78 }
79 }
80}
81
82#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default, AsRefStr)]
84#[serde(try_from = "String")]
85pub enum SelectorType {
86 LoadBased,
90 LeaseBased,
93 #[default]
97 RoundRobin,
98}
99
100impl TryFrom<&str> for SelectorType {
101 type Error = error::Error;
102
103 fn try_from(value: &str) -> Result<Self> {
104 match value {
105 "load_based" | "LoadBased" => Ok(SelectorType::LoadBased),
106 "lease_based" | "LeaseBased" => Ok(SelectorType::LeaseBased),
107 "round_robin" | "RoundRobin" => Ok(SelectorType::RoundRobin),
108 other => error::UnsupportedSelectorTypeSnafu {
109 selector_type: other,
110 }
111 .fail(),
112 }
113 }
114}
115
116impl TryFrom<String> for SelectorType {
117 type Error = error::Error;
118
119 fn try_from(value: String) -> Result<Self> {
120 SelectorType::try_from(value.as_str())
121 }
122}
123
124#[cfg(test)]
125mod tests {
126 use super::SelectorType;
127 use crate::error::Result;
128
129 #[test]
130 fn test_default_selector_type() {
131 assert_eq!(SelectorType::RoundRobin, SelectorType::default());
132 }
133
134 #[test]
135 fn test_convert_str_to_selector_type() {
136 let lease_based = "lease_based";
137 let selector_type = lease_based.try_into().unwrap();
138 assert_eq!(SelectorType::LeaseBased, selector_type);
139 let lease_based = "LeaseBased";
140 let selector_type = lease_based.try_into().unwrap();
141 assert_eq!(SelectorType::LeaseBased, selector_type);
142
143 let load_based = "load_based";
144 let selector_type = load_based.try_into().unwrap();
145 assert_eq!(SelectorType::LoadBased, selector_type);
146 let load_based = "LoadBased";
147 let selector_type = load_based.try_into().unwrap();
148 assert_eq!(SelectorType::LoadBased, selector_type);
149
150 let unknown = "unknown";
151 let selector_type: Result<SelectorType> = unknown.try_into();
152 assert!(selector_type.is_err());
153 }
154}