meta_srv/discovery/
node_scheduling.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 api::v1::meta::heartbeat_request::NodeWorkloads;
16use common_error::ext::BoxedError;
17use common_meta::cluster::NodeInfo;
18use common_meta::error::Result;
19use common_meta::peer::PeerDiscovery;
20use snafu::ResultExt;
21
22use crate::cluster::MetaPeerClientRef;
23use crate::key::NodeSchedulingRole;
24use crate::node_scheduling::NodeSchedulingManagerRef;
25
26/// A [`PeerDiscovery`] decorator that stamps the scheduling state of datanodes
27/// onto [`NodeInfo::schedulable`].
28///
29/// It never hides nodes: callers that place regions must filter with
30/// `selector::common::active_schedulable_datanodes`. Keeping it additive is what allows it to be
31/// the only peer discovery in the placement path (built once by `MetasrvBuilder`), because
32/// failover still needs the full active set to compute its exclusions.
33///
34/// This is the only `PeerDiscovery` implementation over metasrv's local state: the raw
35/// [`MetaPeerClient`](crate::cluster::MetaPeerClient) only exposes inherent discovery methods, so
36/// a `PeerDiscoveryRef` cannot be built without the scheduling fence. Do not build a second
37/// instance; `MetasrvBuilder` owns the single production one and exposes it as
38/// `Metasrv::selector_ctx().peer_discovery`.
39pub struct SchedulingAwarePeerDiscovery {
40    manager: NodeSchedulingManagerRef,
41    inner: MetaPeerClientRef,
42}
43
44impl SchedulingAwarePeerDiscovery {
45    pub fn new(manager: NodeSchedulingManagerRef, inner: MetaPeerClientRef) -> Self {
46        Self { manager, inner }
47    }
48}
49
50#[async_trait::async_trait]
51impl PeerDiscovery for SchedulingAwarePeerDiscovery {
52    async fn active_frontends(&self) -> Result<Vec<NodeInfo>> {
53        self.inner.active_frontends().await
54    }
55
56    async fn active_datanodes(
57        &self,
58        filter: Option<for<'a> fn(&'a NodeWorkloads) -> bool>,
59    ) -> Result<Vec<NodeInfo>> {
60        let mut datanodes = self.inner.active_datanodes(filter).await?;
61        if datanodes.is_empty() {
62            return Ok(datanodes);
63        }
64
65        let unschedulable = self
66            .manager
67            .list_unschedulable_node_ids(NodeSchedulingRole::Datanode)
68            .await
69            .map_err(BoxedError::new)
70            .context(common_meta::error::ExternalSnafu)?;
71        for datanode in &mut datanodes {
72            datanode.schedulable = Some(!unschedulable.contains(&datanode.peer.id));
73        }
74        Ok(datanodes)
75    }
76
77    async fn active_flownodes(
78        &self,
79        filter: Option<for<'a> fn(&'a NodeWorkloads) -> bool>,
80    ) -> Result<Vec<NodeInfo>> {
81        self.inner.active_flownodes(filter).await
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use common_meta::peer::Peer;
88
89    use super::*;
90    use crate::cluster::MetaPeerClientRef;
91    use crate::test_util::{create_meta_peer_client, new_node_scheduling_discovery, put_datanodes};
92
93    async fn new_discovery() -> (
94        MetaPeerClientRef,
95        NodeSchedulingManagerRef,
96        SchedulingAwarePeerDiscovery,
97    ) {
98        let client = create_meta_peer_client();
99        put_datanodes(
100            &client,
101            vec![Peer::empty(1), Peer::empty(2), Peer::empty(3)],
102        )
103        .await;
104        let (manager, discovery) = new_node_scheduling_discovery(&client);
105        (client, manager, discovery)
106    }
107
108    #[tokio::test]
109    async fn test_active_datanodes_marks_unschedulable() {
110        let (_, manager, discovery) = new_discovery().await;
111
112        manager
113            .mark_unschedulable(NodeSchedulingRole::Datanode, 2)
114            .await
115            .unwrap();
116
117        let datanodes = discovery.active_datanodes(None).await.unwrap();
118        assert_eq!(3, datanodes.len());
119
120        let schedulable = |id: u64| {
121            datanodes
122                .iter()
123                .find(|datanode| datanode.peer.id == id)
124                .unwrap()
125                .schedulable
126        };
127        assert_eq!(Some(true), schedulable(1));
128        assert_eq!(Some(false), schedulable(2));
129        assert_eq!(Some(true), schedulable(3));
130    }
131
132    #[tokio::test]
133    async fn test_schedulable_marks_are_not_persisted() {
134        let (client, manager, discovery) = new_discovery().await;
135
136        manager
137            .mark_unschedulable(NodeSchedulingRole::Datanode, 2)
138            .await
139            .unwrap();
140
141        discovery.active_datanodes(None).await.unwrap();
142
143        let datanodes = client.active_datanodes(None).await.unwrap();
144        assert_eq!(3, datanodes.len());
145        assert!(
146            datanodes
147                .iter()
148                .all(|datanode| datanode.schedulable.is_none())
149        );
150    }
151}
meta_srv/discovery/node_scheduling.rs

meta_srv/discovery/
node_scheduling.rs
