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::{HashMap, 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/// A selector that aware of region statistics
42///
43/// It selects the best destination peer for a list of regions.
44/// The selection is based on the region statistics, such as the region leader's write throughput.
45#[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    /// Minimum number of selected results.
62    pub min_required_items: usize,
63    /// Whether duplicates are allowed in the selected result, default false.
64    pub allow_duplication: bool,
65    /// The peers to exclude from the selection.
66    pub exclude_peer_ids: HashSet<u64>,
67    /// The filter to select the peers based on their workloads.
68    pub workload_filter: Option<fn(&NodeWorkloads) -> bool>,
69    /// Additional placement context, such as the frontend identity that originated the request.
70    pub extensions: HashMap<String, String>,
71}
72
73impl Default for SelectorOptions {
74    fn default() -> Self {
75        Self {
76            min_required_items: 1,
77            allow_duplication: false,
78            exclude_peer_ids: HashSet::new(),
79            workload_filter: None,
80            extensions: HashMap::new(),
81        }
82    }
83}
84
85/// [`SelectorType`] refers to the load balancer used when creating tables.
86#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default, AsRefStr)]
87#[serde(try_from = "String")]
88pub enum SelectorType {
89    /// The current load balancing is based on the number of regions on each datanode node;
90    /// the more regions, the higher the load (it may be changed to Capacity Units(CU)
91    /// calculation in the future).
92    LoadBased,
93    /// This one randomly selects from all available (in lease) nodes. Its characteristic
94    /// is simplicity and fast.
95    LeaseBased,
96    /// This one selects the node in a round-robin way.
97    /// In most cases, it's recommended and is the default option. If you're unsure which
98    /// to choose, using it is usually correct.
99    #[default]
100    RoundRobin,
101}
102
103impl TryFrom<&str> for SelectorType {
104    type Error = error::Error;
105
106    fn try_from(value: &str) -> Result<Self> {
107        match value {
108            "load_based" | "LoadBased" => Ok(SelectorType::LoadBased),
109            "lease_based" | "LeaseBased" => Ok(SelectorType::LeaseBased),
110            "round_robin" | "RoundRobin" => Ok(SelectorType::RoundRobin),
111            other => error::UnsupportedSelectorTypeSnafu {
112                selector_type: other,
113            }
114            .fail(),
115        }
116    }
117}
118
119impl TryFrom<String> for SelectorType {
120    type Error = error::Error;
121
122    fn try_from(value: String) -> Result<Self> {
123        SelectorType::try_from(value.as_str())
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::SelectorType;
130    use crate::error::Result;
131
132    #[test]
133    fn test_default_selector_type() {
134        assert_eq!(SelectorType::RoundRobin, SelectorType::default());
135    }
136
137    #[test]
138    fn test_convert_str_to_selector_type() {
139        let lease_based = "lease_based";
140        let selector_type = lease_based.try_into().unwrap();
141        assert_eq!(SelectorType::LeaseBased, selector_type);
142        let lease_based = "LeaseBased";
143        let selector_type = lease_based.try_into().unwrap();
144        assert_eq!(SelectorType::LeaseBased, selector_type);
145
146        let load_based = "load_based";
147        let selector_type = load_based.try_into().unwrap();
148        assert_eq!(SelectorType::LoadBased, selector_type);
149        let load_based = "LoadBased";
150        let selector_type = load_based.try_into().unwrap();
151        assert_eq!(SelectorType::LoadBased, selector_type);
152
153        let unknown = "unknown";
154        let selector_type: Result<SelectorType> = unknown.try_into();
155        assert!(selector_type.is_err());
156    }
157}
meta_srv/selector.rs

meta_srv/
selector.rs
