meta_srv/
selector.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15pub 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::{BTreeSet, HashMap, HashSet};
24
25use api::v1::meta::heartbeat_request::NodeWorkloads;
26use common_meta::affinity::AffinityRequirements;
27use serde::{Deserialize, Serialize};
28use store_api::storage::RegionId;
29use strum::AsRefStr;
30
31use crate::error;
32use crate::error::Result;
33
34#[async_trait::async_trait]
35pub trait Selector: Send + Sync {
36    type Context;
37    type Output;
38
39    async fn select(&self, ctx: &Self::Context, opts: SelectorOptions<'_>) -> Result<Self::Output>;
40}
41
42/// A selector that aware of region statistics
43///
44/// It selects the best destination peer for a list of regions.
45/// The selection is based on the region statistics, such as the region leader's write throughput.
46///
47/// A region is never placed on `exclude_peer_ids`, nor on its own peers in
48/// `exclude_peer_ids_by_region`. It returns the regions it placed, possibly none: a region
49/// without a remaining candidate is left out. Errors are for other failures.
50#[async_trait::async_trait]
51pub trait RegionStatAwareSelector: Send + Sync {
52    type Context;
53    type Output;
54
55    async fn select(
56        &self,
57        ctx: &Self::Context,
58        from_peer_id: u64,
59        region_ids: &[RegionId],
60        exclude_peer_ids: HashSet<u64>,
61        exclude_peer_ids_by_region: &HashMap<RegionId, BTreeSet<u64>>,
62        affinity: Option<&AffinityRequirements>,
63    ) -> Result<Self::Output>;
64}
65
66#[derive(Debug, Clone)]
67pub struct SelectorOptions<'a> {
68    /// Resolved requirements used to filter candidate nodes.
69    pub affinity: Option<&'a AffinityRequirements>,
70    /// Minimum number of selected results.
71    pub min_required_items: usize,
72    /// Whether duplicates are allowed in the selected result, default false.
73    pub allow_duplication: bool,
74    /// The peers to exclude from the selection.
75    pub exclude_peer_ids: HashSet<u64>,
76    /// The filter to select the peers based on their workloads.
77    pub workload_filter: Option<fn(&NodeWorkloads) -> bool>,
78}
79
80impl Default for SelectorOptions<'_> {
81    fn default() -> Self {
82        Self {
83            affinity: None,
84            min_required_items: 1,
85            allow_duplication: false,
86            exclude_peer_ids: HashSet::new(),
87            workload_filter: None,
88        }
89    }
90}
91
92/// [`SelectorType`] refers to the load balancer used when creating tables.
93#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default, AsRefStr)]
94#[serde(try_from = "String")]
95pub enum SelectorType {
96    /// The current load balancing is based on the number of regions on each datanode node;
97    /// the more regions, the higher the load (it may be changed to Capacity Units(CU)
98    /// calculation in the future).
99    LoadBased,
100    /// This one randomly selects from all available (in lease) nodes. Its characteristic
101    /// is simplicity and fast.
102    LeaseBased,
103    /// This one selects the node in a round-robin way.
104    /// In most cases, it's recommended and is the default option. If you're unsure which
105    /// to choose, using it is usually correct.
106    #[default]
107    RoundRobin,
108}
109
110impl TryFrom<&str> for SelectorType {
111    type Error = error::Error;
112
113    fn try_from(value: &str) -> Result<Self> {
114        match value {
115            "load_based" | "LoadBased" => Ok(SelectorType::LoadBased),
116            "lease_based" | "LeaseBased" => Ok(SelectorType::LeaseBased),
117            "round_robin" | "RoundRobin" => Ok(SelectorType::RoundRobin),
118            other => error::UnsupportedSelectorTypeSnafu {
119                selector_type: other,
120            }
121            .fail(),
122        }
123    }
124}
125
126impl TryFrom<String> for SelectorType {
127    type Error = error::Error;
128
129    fn try_from(value: String) -> Result<Self> {
130        SelectorType::try_from(value.as_str())
131    }
132}
133
134#[cfg(test)]
135mod tests {
136    use super::SelectorType;
137    use crate::error::Result;
138
139    #[test]
140    fn test_default_selector_type() {
141        assert_eq!(SelectorType::RoundRobin, SelectorType::default());
142    }
143
144    #[test]
145    fn test_convert_str_to_selector_type() {
146        let lease_based = "lease_based";
147        let selector_type = lease_based.try_into().unwrap();
148        assert_eq!(SelectorType::LeaseBased, selector_type);
149        let lease_based = "LeaseBased";
150        let selector_type = lease_based.try_into().unwrap();
151        assert_eq!(SelectorType::LeaseBased, selector_type);
152
153        let load_based = "load_based";
154        let selector_type = load_based.try_into().unwrap();
155        assert_eq!(SelectorType::LoadBased, selector_type);
156        let load_based = "LoadBased";
157        let selector_type = load_based.try_into().unwrap();
158        assert_eq!(SelectorType::LoadBased, selector_type);
159
160        let unknown = "unknown";
161        let selector_type: Result<SelectorType> = unknown.try_into();
162        assert!(selector_type.is_err());
163    }
164}
meta_srv/selector.rs

meta_srv/
selector.rs
