meta_srv/
peer.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::HashSet;
16
17use async_trait::async_trait;
18use common_error::ext::BoxedError;
19use common_meta::error::{ExternalSnafu, Result as MetaResult};
20use common_meta::peer::{Peer, PeerAllocContext, PeerAllocator};
21use snafu::{ResultExt, ensure};
22
23use crate::discovery::utils::accept_ingest_workload;
24use crate::error::{self, Result, TooManyPartitionsSnafu};
25use crate::metasrv::{SelectTarget, SelectorContext, SelectorRef};
26use crate::selector::SelectorOptions;
27
28pub struct MetasrvPeerAllocator {
29    ctx: SelectorContext,
30    selector: SelectorRef,
31    max_items: Option<u32>,
32}
33
34impl MetasrvPeerAllocator {
35    /// Creates a new [`MetasrvPeerAllocator`] with the given [`SelectorContext`] and [`SelectorRef`].
36    pub fn new(ctx: SelectorContext, selector: SelectorRef) -> Self {
37        Self {
38            ctx,
39            selector,
40            max_items: None,
41        }
42    }
43
44    pub fn with_max_items(self, max_items: u32) -> Self {
45        Self {
46            ctx: self.ctx,
47            selector: self.selector,
48            max_items: Some(max_items),
49        }
50    }
51
52    /// Allocates a specified number (by `regions`) of [`Peer`] instances based on the number of
53    /// regions. The returned peers will have the same length as the number of regions.
54    ///
55    /// This method is mainly a wrapper around the [`SelectorRef`]::`select` method. There is
56    /// no guarantee that how the returned peers are used, like whether they are from the same
57    /// table or not. So this method isn't idempotent.
58    async fn alloc(
59        &self,
60        min_required_items: usize,
61        alloc_context: &PeerAllocContext,
62    ) -> Result<Vec<Peer>> {
63        if let Some(max_items) = self.max_items {
64            ensure!(
65                min_required_items <= max_items as usize,
66                TooManyPartitionsSnafu
67            );
68        }
69
70        let options = SelectorOptions {
71            affinity: alloc_context.affinity.as_ref(),
72            min_required_items,
73            allow_duplication: true,
74            exclude_peer_ids: HashSet::new(),
75            workload_filter: Some(accept_ingest_workload),
76        };
77        let peers = self.selector.select(&self.ctx, options).await?;
78        ensure!(
79            peers.len() >= min_required_items,
80            error::NoEnoughAvailableNodeSnafu {
81                required: min_required_items,
82                available: peers.len(),
83                select_target: SelectTarget::Datanode,
84            }
85        );
86        Ok(peers)
87    }
88}
89
90#[async_trait]
91impl PeerAllocator for MetasrvPeerAllocator {
92    async fn alloc(&self, regions: usize, ctx: &PeerAllocContext) -> MetaResult<Vec<Peer>> {
93        self.alloc(regions, ctx)
94            .await
95            .map_err(BoxedError::new)
96            .context(ExternalSnafu)
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use std::collections::HashMap;
103    use std::ptr;
104    use std::sync::Arc;
105    use std::sync::atomic::{AtomicUsize, Ordering};
106
107    use common_meta::affinity::{
108        Affinity, AffinityFallback, AffinityRef, AffinityRequirements, NoopAffinity,
109    };
110    use common_meta::cluster::NodeInfo;
111
112    use super::*;
113    use crate::selector::Selector;
114    use crate::selector::round_robin::RoundRobinSelector;
115    use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes};
116
117    struct OnlySecondNode;
118
119    impl Affinity for OnlySecondNode {
120        fn matches(&self, _: &AffinityRequirements, node: &NodeInfo) -> bool {
121            node.peer.id == 2
122        }
123    }
124
125    #[derive(Clone, Copy)]
126    enum Outcome {
127        Base,
128        BackendError,
129        Short,
130    }
131
132    struct CountingSelector {
133        base: RoundRobinSelector,
134        calls: AtomicUsize,
135        requirements_address: AtomicUsize,
136        outcome: Outcome,
137    }
138
139    #[async_trait]
140    impl Selector for CountingSelector {
141        type Context = SelectorContext;
142        type Output = Vec<Peer>;
143
144        async fn select(
145            &self,
146            ctx: &SelectorContext,
147            opts: SelectorOptions<'_>,
148        ) -> Result<Self::Output> {
149            self.calls.fetch_add(1, Ordering::Relaxed);
150            self.requirements_address.store(
151                opts.affinity
152                    .map_or(0, |requirements| ptr::from_ref(requirements) as usize),
153                Ordering::Relaxed,
154            );
155            assert!(opts.allow_duplication);
156            assert!(opts.exclude_peer_ids.is_empty());
157            assert!(opts.workload_filter.is_some());
158            assert_eq!(opts.min_required_items, 2);
159            match self.outcome {
160                Outcome::Base => self.base.select(ctx, opts).await,
161                Outcome::BackendError => error::UnexpectedSnafu {
162                    violated: "selection backend failed",
163                }
164                .fail(),
165                Outcome::Short => Ok(vec![Peer::empty(1)]),
166            }
167        }
168    }
169
170    #[tokio::test]
171    async fn test_allocator_delegates_fallback_and_validates_results() {
172        let discovery = create_meta_peer_client();
173        put_datanodes(&discovery, vec![Peer::empty(1)]).await;
174        for (has_requirements, fallback, noop, outcome, expected_calls, succeeds) in [
175            (
176                true,
177                AffinityFallback::Error,
178                false,
179                Outcome::Base,
180                1,
181                false,
182            ),
183            (
184                true,
185                AffinityFallback::Unconstrained,
186                false,
187                Outcome::Base,
188                1,
189                true,
190            ),
191            (true, AffinityFallback::Error, true, Outcome::Base, 1, true),
192            (
193                false,
194                AffinityFallback::Unconstrained,
195                false,
196                Outcome::Base,
197                1,
198                true,
199            ),
200            (
201                true,
202                AffinityFallback::Unconstrained,
203                false,
204                Outcome::BackendError,
205                1,
206                false,
207            ),
208            (
209                true,
210                AffinityFallback::Error,
211                false,
212                Outcome::Short,
213                1,
214                false,
215            ),
216            (
217                true,
218                AffinityFallback::Unconstrained,
219                false,
220                Outcome::Short,
221                1,
222                false,
223            ),
224            (
225                false,
226                AffinityFallback::Unconstrained,
227                false,
228                Outcome::Short,
229                1,
230                false,
231            ),
232        ] {
233            let selector = Arc::new(CountingSelector {
234                base: RoundRobinSelector::default(),
235                calls: AtomicUsize::new(0),
236                requirements_address: AtomicUsize::new(0),
237                outcome,
238            });
239            let affinity: AffinityRef = if noop {
240                Arc::new(NoopAffinity)
241            } else {
242                Arc::new(OnlySecondNode)
243            };
244            let allocator = MetasrvPeerAllocator::new(
245                SelectorContext {
246                    affinity_fallback: fallback,
247                    peer_discovery: Arc::new(new_node_scheduling_discovery(&discovery).1),
248                    affinity,
249                },
250                selector.clone(),
251            );
252            let alloc_context = PeerAllocContext {
253                affinity: has_requirements.then(|| {
254                    AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])).unwrap()
255                }),
256            };
257            let result = allocator.alloc(2, &alloc_context).await;
258            assert_eq!(
259                selector.requirements_address.load(Ordering::Relaxed),
260                alloc_context
261                    .affinity
262                    .as_ref()
263                    .map_or(0, |requirements| { ptr::from_ref(requirements) as usize }),
264            );
265            assert_eq!(selector.calls.load(Ordering::Relaxed), expected_calls);
266            if succeeds {
267                assert_eq!(result.unwrap(), vec![Peer::empty(1); 2]);
268            } else {
269                match outcome {
270                    Outcome::BackendError => assert!(matches!(
271                        result.unwrap_err(),
272                        error::Error::Unexpected { .. }
273                    )),
274                    _ => assert!(matches!(
275                        result.unwrap_err(),
276                        error::Error::NoEnoughAvailableNode { .. }
277                    )),
278                }
279            }
280        }
281    }
282}
meta_srv/peer.rs

meta_srv/
peer.rs
