meta_srv/selector/
load_based.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
15use std::collections::HashMap;
16
17use common_meta::affinity::AffinityFallback;
18use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue};
19use common_meta::peer::Peer;
20use common_telemetry::debug;
21use snafu::ResultExt;
22
23use crate::cluster::MetaPeerClientRef;
24use crate::error::{self, ListActiveDatanodesSnafu, Result};
25use crate::metasrv::SelectorContext;
26use crate::selector::common::{choose_items, filter_out_excluded_peers};
27use crate::selector::weight_compute::WeightCompute;
28use crate::selector::weighted_choose::RandomWeightedChoose;
29use crate::selector::{Selector, SelectorOptions};
30
31pub struct LoadBasedSelector<C> {
32    weight_compute: C,
33    meta_peer_client: MetaPeerClientRef,
34}
35
36impl<C> LoadBasedSelector<C> {
37    pub fn new(weight_compute: C, meta_peer_client: MetaPeerClientRef) -> Self {
38        Self {
39            weight_compute,
40            meta_peer_client,
41        }
42    }
43}
44
45impl<C> LoadBasedSelector<C>
46where
47    C: WeightCompute<Source = HashMap<DatanodeStatKey, DatanodeStatValue>>,
48{
49    async fn select_once(
50        &self,
51        ctx: &SelectorContext,
52        opts: &SelectorOptions<'_>,
53    ) -> Result<Vec<Peer>> {
54        // 1. get alive datanodes.
55        let alive_datanodes = ctx
56            .peer_discovery
57            .active_datanodes(opts.workload_filter)
58            .await
59            .context(ListActiveDatanodesSnafu)?
60            .into_iter()
61            .filter(|node| {
62                opts.affinity
63                    .is_none_or(|requirements| ctx.affinity.matches(requirements, node))
64            })
65            .map(|node| node.peer)
66            .collect::<Vec<_>>();
67
68        // 2. get stat kvs and filter out expired datanodes.
69        let stat_keys = alive_datanodes
70            .iter()
71            .map(|k| DatanodeStatKey { node_id: k.id })
72            .collect();
73        let stat_kvs = filter_out_expired_datanode(
74            self.meta_peer_client.get_dn_stat_kvs(stat_keys).await?,
75            &alive_datanodes,
76        );
77
78        // 3. compute weight array.
79        let mut weight_array = self.weight_compute.compute(&stat_kvs);
80
81        // 4. filter out excluded peers.
82        filter_out_excluded_peers(&mut weight_array, &opts.exclude_peer_ids);
83        // 5. choose peers by weight_array.
84        let mut weighted_choose = RandomWeightedChoose::new(weight_array);
85        let selected = choose_items(opts, &mut weighted_choose)?;
86
87        debug!(
88            "LoadBasedSelector select peers: {:?}, opts: {:?}.",
89            selected, opts,
90        );
91
92        Ok(selected)
93    }
94}
95
96#[async_trait::async_trait]
97impl<C> Selector for LoadBasedSelector<C>
98where
99    C: WeightCompute<Source = HashMap<DatanodeStatKey, DatanodeStatValue>>,
100{
101    type Context = SelectorContext;
102    type Output = Vec<Peer>;
103
104    async fn select(
105        &self,
106        ctx: &Self::Context,
107        mut opts: SelectorOptions<'_>,
108    ) -> Result<Self::Output> {
109        match self.select_once(ctx, &opts).await {
110            Err(error::Error::NoEnoughAvailableNode {
111                required,
112                available,
113                ..
114            }) if opts.affinity.is_some()
115                && ctx.affinity_fallback == AffinityFallback::Unconstrained
116                && available < required =>
117            {
118                opts.affinity = None;
119                self.select_once(ctx, &opts).await
120            }
121            result => result,
122        }
123    }
124}
125
126fn filter_out_expired_datanode(
127    mut stat_kvs: HashMap<DatanodeStatKey, DatanodeStatValue>,
128    datanodes: &[Peer],
129) -> HashMap<DatanodeStatKey, DatanodeStatValue> {
130    datanodes
131        .iter()
132        .filter_map(|p| stat_kvs.remove_entry(&DatanodeStatKey { node_id: p.id }))
133        .collect()
134}
135
136#[cfg(test)]
137mod tests {
138    use std::collections::HashMap;
139
140    use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue};
141    use common_meta::peer::Peer;
142
143    use crate::selector::load_based::filter_out_expired_datanode;
144
145    #[test]
146    fn test_filter_out_expired_datanode() {
147        let mut stat_kvs = HashMap::new();
148        stat_kvs.insert(
149            DatanodeStatKey { node_id: 0 },
150            DatanodeStatValue { stats: vec![] },
151        );
152        stat_kvs.insert(
153            DatanodeStatKey { node_id: 1 },
154            DatanodeStatValue { stats: vec![] },
155        );
156        stat_kvs.insert(
157            DatanodeStatKey { node_id: 2 },
158            DatanodeStatValue { stats: vec![] },
159        );
160
161        let lease_kvs = vec![Peer::new(1, "127.0.0.1:3002".to_string())];
162        let alive_stat_kvs = filter_out_expired_datanode(stat_kvs, &lease_kvs);
163
164        assert_eq!(1, alive_stat_kvs.len());
165        assert!(alive_stat_kvs.contains_key(&DatanodeStatKey { node_id: 1 }));
166    }
167}
meta_srv/selector/load_based.rs

meta_srv/selector/
load_based.rs
