meta_srv/selector/
round_robin.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::sync::atomic::AtomicUsize;
16
17use common_meta::affinity::AffinityFallback;
18use common_meta::peer::Peer;
19use snafu::{ResultExt, ensure};
20
21use crate::error::{self, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result};
22use crate::metasrv::{SelectTarget, SelectorContext};
23use crate::selector::{Selector, SelectorOptions, common};
24
25/// Round-robin selector that returns the next peer in the list in sequence.
26/// Datanodes are ordered by their node_id.
27///
28/// This selector is useful when you want to distribute the load evenly across
29/// all datanodes. But **it's not recommended** to use this selector in serious
30/// production environments because it doesn't take into account the load of
31/// each datanode.
32pub struct RoundRobinSelector {
33    select_target: SelectTarget,
34    counter: AtomicUsize,
35}
36
37impl Default for RoundRobinSelector {
38    fn default() -> Self {
39        Self {
40            select_target: SelectTarget::Datanode,
41            counter: AtomicUsize::new(0),
42        }
43    }
44}
45
46impl RoundRobinSelector {
47    pub fn new(select_target: SelectTarget) -> Self {
48        Self {
49            select_target,
50            ..Default::default()
51        }
52    }
53
54    async fn get_peers(
55        &self,
56        opts: &SelectorOptions<'_>,
57        ctx: &SelectorContext,
58    ) -> Result<Vec<Peer>> {
59        let mut peers = match self.select_target {
60            SelectTarget::Datanode => {
61                // 1. get alive, schedulable datanodes.
62                let alive_datanodes =
63                    common::active_schedulable_datanodes(ctx, opts.workload_filter).await?;
64
65                // 2. filter by affinity, then drop the excluded datanodes.
66                alive_datanodes
67                    .into_iter()
68                    .filter(|node| {
69                        opts.affinity
70                            .is_none_or(|requirements| ctx.affinity.matches(requirements, node))
71                    })
72                    .map(|node| node.peer)
73                    .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id))
74                    .collect::<Vec<_>>()
75            }
76            SelectTarget::Flownode => ctx
77                .peer_discovery
78                .active_flownodes(opts.workload_filter)
79                .await
80                .context(ListActiveFlownodesSnafu)?
81                .into_iter()
82                .filter(|node| {
83                    opts.affinity
84                        .is_none_or(|requirements| ctx.affinity.matches(requirements, node))
85                })
86                .map(|node| node.peer)
87                .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id))
88                .collect::<Vec<_>>(),
89        };
90
91        ensure!(
92            !peers.is_empty() && (opts.allow_duplication || peers.len() >= opts.min_required_items),
93            NoEnoughAvailableNodeSnafu {
94                required: opts.min_required_items,
95                available: peers.len(),
96                select_target: self.select_target
97            }
98        );
99
100        // 3. sort by node id
101        peers.sort_by_key(|p| p.id);
102
103        Ok(peers)
104    }
105
106    async fn select_once(
107        &self,
108        ctx: &SelectorContext,
109        opts: &SelectorOptions<'_>,
110    ) -> Result<Vec<Peer>> {
111        let peers = self.get_peers(opts, ctx).await?;
112        // choose peers
113        let mut selected = Vec::with_capacity(opts.min_required_items);
114        for _ in 0..opts.min_required_items {
115            let idx = self
116                .counter
117                .fetch_add(1, std::sync::atomic::Ordering::Relaxed)
118                % peers.len();
119            selected.push(peers[idx].clone());
120        }
121
122        Ok(selected)
123    }
124}
125
126#[async_trait::async_trait]
127impl Selector for RoundRobinSelector {
128    type Context = SelectorContext;
129    type Output = Vec<Peer>;
130
131    async fn select(
132        &self,
133        ctx: &Self::Context,
134        mut opts: SelectorOptions<'_>,
135    ) -> Result<Self::Output> {
136        match self.select_once(ctx, &opts).await {
137            Err(error::Error::NoEnoughAvailableNode {
138                required,
139                available,
140                ..
141            }) if opts.affinity.is_some()
142                && ctx.affinity_fallback == AffinityFallback::Unconstrained
143                && available < required =>
144            {
145                opts.affinity = None;
146                self.select_once(ctx, &opts).await
147            }
148            result => result,
149        }
150    }
151}
152
153#[cfg(test)]
154mod test {
155    use std::collections::HashSet;
156    use std::sync::Arc;
157
158    use common_meta::affinity::NoopAffinity;
159
160    use super::*;
161    use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes};
162
163    #[tokio::test]
164    async fn test_round_robin_selector() {
165        let selector = RoundRobinSelector::default();
166        let meta_peer_client = create_meta_peer_client();
167        let ctx = SelectorContext {
168            affinity_fallback: Default::default(),
169            affinity: Arc::new(NoopAffinity),
170            peer_discovery: Arc::new(new_node_scheduling_discovery(&meta_peer_client).1),
171        };
172        // add three nodes
173        let peer1 = Peer {
174            id: 2,
175            addr: "node1".to_string(),
176        };
177        let peer2 = Peer {
178            id: 5,
179            addr: "node2".to_string(),
180        };
181        let peer3 = Peer {
182            id: 8,
183            addr: "node3".to_string(),
184        };
185        let peers = vec![peer1.clone(), peer2.clone(), peer3.clone()];
186        put_datanodes(&meta_peer_client, peers).await;
187
188        let peers = selector
189            .select(
190                &ctx,
191                SelectorOptions {
192                    affinity: None,
193                    min_required_items: 4,
194                    allow_duplication: true,
195                    exclude_peer_ids: HashSet::new(),
196                    workload_filter: None,
197                },
198            )
199            .await
200            .unwrap();
201        assert_eq!(peers.len(), 4);
202        assert_eq!(
203            peers,
204            vec![peer1.clone(), peer2.clone(), peer3.clone(), peer1.clone()]
205        );
206
207        let peers = selector
208            .select(
209                &ctx,
210                SelectorOptions {
211                    affinity: None,
212                    min_required_items: 2,
213                    allow_duplication: true,
214                    exclude_peer_ids: HashSet::new(),
215                    workload_filter: None,
216                },
217            )
218            .await
219            .unwrap();
220        assert_eq!(peers.len(), 2);
221        assert_eq!(peers, vec![peer2.clone(), peer3.clone()]);
222    }
223
224    #[tokio::test]
225    async fn test_round_robin_selector_with_exclude_peer_ids() {
226        let selector = RoundRobinSelector::new(SelectTarget::Datanode);
227        let meta_peer_client = create_meta_peer_client();
228        let ctx = SelectorContext {
229            affinity_fallback: Default::default(),
230            affinity: Arc::new(NoopAffinity),
231            peer_discovery: Arc::new(new_node_scheduling_discovery(&meta_peer_client).1),
232        };
233        // add three nodes
234        let peer1 = Peer {
235            id: 2,
236            addr: "node1".to_string(),
237        };
238        let peer2 = Peer {
239            id: 5,
240            addr: "node2".to_string(),
241        };
242        let peer3 = Peer {
243            id: 8,
244            addr: "node3".to_string(),
245        };
246        put_datanodes(
247            &meta_peer_client,
248            vec![peer1.clone(), peer2.clone(), peer3.clone()],
249        )
250        .await;
251
252        let peers = selector
253            .select(
254                &ctx,
255                SelectorOptions {
256                    affinity: None,
257                    min_required_items: 1,
258                    allow_duplication: true,
259                    exclude_peer_ids: HashSet::from([2, 5]),
260                    workload_filter: None,
261                },
262            )
263            .await
264            .unwrap();
265        assert_eq!(peers.len(), 1);
266        assert_eq!(peers, vec![peer3.clone()]);
267    }
268}
meta_srv/selector/round_robin.rs

meta_srv/selector/
round_robin.rs
