common_meta/
affinity.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
15//! Resolve reference attributes into requirements, then evaluate candidates.
16
17use std::collections::HashMap;
18use std::sync::Arc;
19
20use crate::DatanodeId;
21use crate::cluster::NodeInfo;
22use crate::error::Result;
23use crate::peer::PeerAllocContext;
24use crate::rpc::ddl::QueryContext;
25
26/// Named attributes used to express affinity requirements.
27pub type AffinityAttributes = HashMap<String, String>;
28
29/// Identifies the node whose attributes define affinity requirements.
30#[derive(Debug, Clone, PartialEq, Eq, Hash)]
31pub enum AffinityReference {
32    /// Uses the datanode's last-known attributes, even if its lease has expired.
33    Datanode(DatanodeId),
34    /// Uses the frontend identified by its advertised address.
35    Frontend(String),
36}
37
38/// Behavior when affinity cannot be satisfied.
39#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
40pub enum AffinityFallback {
41    /// Returns an error instead of selecting without affinity.
42    #[default]
43    Error,
44    /// Allows selection among eligible nodes without affinity.
45    Unconstrained,
46}
47
48/// Builds allocation inputs from the query that initiated an operation.
49#[async_trait::async_trait]
50pub trait PeerAllocContextProvider: Send + Sync {
51    /// Resolves query-specific allocation inputs using metadata, without selecting peers.
52    async fn build(&self, query_context: &QueryContext) -> Result<PeerAllocContext>;
53}
54
55/// Shared allocation-context provider.
56pub type PeerAllocContextProviderRef = Arc<dyn PeerAllocContextProvider>;
57
58/// Nonempty attributes a candidate must match.
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub struct AffinityRequirements {
61    attributes: AffinityAttributes,
62}
63
64impl AffinityRequirements {
65    /// Returns no requirements when the supplied attributes are empty.
66    pub fn new(attributes: AffinityAttributes) -> Option<Self> {
67        (!attributes.is_empty()).then_some(Self { attributes })
68    }
69
70    /// Returns the attributes that candidates are evaluated against.
71    pub fn attributes(&self) -> &AffinityAttributes {
72        &self.attributes
73    }
74}
75
76/// Resolves references using current metadata, without selecting candidates.
77#[async_trait::async_trait]
78pub trait AffinityResolver: Send + Sync {
79    /// Returns one result per reference, in order.
80    /// None means requirements are unavailable; Err means resolution failed.
81    async fn resolve(
82        &self,
83        references: &[AffinityReference],
84    ) -> Result<Vec<Option<AffinityRequirements>>>;
85}
86
87/// Shared requirement resolver.
88pub type AffinityResolverRef = Arc<dyn AffinityResolver>;
89
90/// Evaluates candidate nodes without I/O or scheduling decisions.
91pub trait Affinity: Send + Sync {
92    /// Evaluates a node snapshot; implementations choose which fields to inspect.
93    fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool;
94}
95
96/// Accepts every candidate without imposing affinity constraints.
97pub struct NoopAffinity;
98
99impl Affinity for NoopAffinity {
100    fn matches(&self, _: &AffinityRequirements, _: &NodeInfo) -> bool {
101        true
102    }
103}
104
105/// Shared affinity evaluator.
106pub type AffinityRef = Arc<dyn Affinity>;
107
108#[cfg(test)]
109mod tests {
110    use std::sync::RwLock;
111
112    use super::*;
113    use crate::cluster::{DatanodeStatus, NodeStatus};
114    use crate::peer::Peer;
115
116    struct NodeAffinityResolver {
117        attributes: RwLock<HashMap<AffinityReference, AffinityAttributes>>,
118    }
119
120    #[async_trait::async_trait]
121    impl AffinityResolver for NodeAffinityResolver {
122        async fn resolve(
123            &self,
124            references: &[AffinityReference],
125        ) -> Result<Vec<Option<AffinityRequirements>>> {
126            let snapshot = self.attributes.read().unwrap();
127            Ok(references
128                .iter()
129                .map(|reference| {
130                    snapshot
131                        .get(reference)
132                        .cloned()
133                        .and_then(AffinityRequirements::new)
134                })
135                .collect())
136        }
137    }
138
139    struct ExactAffinity;
140
141    impl Affinity for ExactAffinity {
142        fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool {
143            requirements
144                .attributes()
145                .iter()
146                .all(|(key, expected)| candidate.env_vars.get(key) == Some(expected))
147        }
148    }
149
150    fn zone(name: &str) -> AffinityAttributes {
151        HashMap::from([("AZ".to_string(), name.to_string())])
152    }
153
154    fn node(id: DatanodeId, env_vars: AffinityAttributes) -> NodeInfo {
155        NodeInfo {
156            peer: Peer::empty(id),
157            last_activity_ts: 0,
158            status: NodeStatus::Datanode(DatanodeStatus {
159                rcus: 0,
160                wcus: 0,
161                leader_regions: 0,
162                follower_regions: 0,
163                workloads: Default::default(),
164            }),
165            version: String::new(),
166            git_commit: String::new(),
167            start_time_ms: 0,
168            total_cpu_millicores: 0,
169            total_memory_bytes: 0,
170            cpu_usage_millicores: 0,
171            memory_usage_bytes: 0,
172            hostname: String::new(),
173            env_vars,
174        }
175    }
176
177    #[tokio::test]
178    async fn test_resolve_reference_and_match_candidates() {
179        let resolver = NodeAffinityResolver {
180            attributes: RwLock::new(HashMap::from([(AffinityReference::Datanode(1), zone("a"))])),
181        };
182        let references = [AffinityReference::Datanode(1)];
183        let resolved = resolver.resolve(&references).await.unwrap();
184        assert_eq!(resolved, vec![AffinityRequirements::new(zone("a"))]);
185        let requirements = resolved[0].as_ref().unwrap();
186        let candidates = [
187            node(2, zone("a")),
188            node(3, zone("b")),
189            node(4, HashMap::new()),
190        ];
191        assert_eq!(
192            candidates
193                .iter()
194                .map(|node| (node.peer.id, ExactAffinity.matches(requirements, node)))
195                .collect::<Vec<_>>(),
196            vec![(2, true), (3, false), (4, false),]
197        );
198
199        assert!(
200            candidates
201                .iter()
202                .all(|node| NoopAffinity.matches(requirements, node))
203        );
204
205        // A new allocation attempt resolves current attributes rather than a saved requirement.
206        resolver
207            .attributes
208            .write()
209            .unwrap()
210            .insert(AffinityReference::Datanode(1), zone("b"));
211        let resolved = resolver.resolve(&references).await.unwrap();
212        let requirements = resolved[0].as_ref().unwrap();
213        let matched: Vec<_> = candidates
214            .iter()
215            .filter(|node| ExactAffinity.matches(requirements, node))
216            .map(|node| node.peer.id)
217            .collect();
218        assert_eq!(matched, vec![3]);
219    }
220
221    #[tokio::test]
222    async fn test_resolution_preserves_order_and_unavailable_requirements() {
223        let resolver = NodeAffinityResolver {
224            attributes: RwLock::new(HashMap::from([
225                (AffinityReference::Datanode(1), zone("a")),
226                (AffinityReference::Datanode(2), HashMap::new()),
227                (
228                    AffinityReference::Frontend("frontend:4001".into()),
229                    zone("b"),
230                ),
231            ])),
232        };
233        let references = [
234            AffinityReference::Frontend("missing:4001".into()),
235            AffinityReference::Datanode(1),
236            AffinityReference::Frontend("frontend:4001".into()),
237            AffinityReference::Datanode(2),
238            AffinityReference::Datanode(3),
239            AffinityReference::Frontend("frontend:4001".into()),
240        ];
241        assert_eq!(
242            resolver.resolve(&references).await.unwrap(),
243            vec![
244                None,
245                AffinityRequirements::new(zone("a")),
246                AffinityRequirements::new(zone("b")),
247                None,
248                None,
249                AffinityRequirements::new(zone("b")),
250            ]
251        );
252        assert!(AffinityRequirements::new(HashMap::new()).is_none());
253    }
254}
common_meta/affinity.rs

common_meta/
affinity.rs
