common_meta/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::HashMap;
16use std::sync::Arc;
17
18pub use api::v1::meta::Peer;
19use api::v1::meta::heartbeat_request::NodeWorkloads;
20
21use crate::cluster::NodeInfo;
22use crate::error::Error;
23use crate::{DatanodeId, FlownodeId};
24
25/// Shared reference to a [`Peer`].
26pub type PeerRef = Arc<Peer>;
27
28/// [`PeerResolver`] provides methods to look up peer information by node ID.
29///
30/// This trait allows resolving both datanode and flownode peers, regardless of their current activity status.
31/// Returned peers may be inactive (i.e., not currently alive in the cluster).
32#[async_trait::async_trait]
33pub trait PeerResolver: Send + Sync {
34 /// Looks up a datanode peer by its ID.
35 ///
36 /// Returns `Some(Peer)` if the datanode exists, or `None` if not found.
37 /// The returned peer may be inactive.
38 async fn datanode(&self, id: DatanodeId) -> Result<Option<Peer>, Error>;
39
40 /// Looks up a flownode peer by its ID.
41 ///
42 /// Returns `Some(Peer)` if the flownode exists, or `None` if not found.
43 /// The returned peer may be inactive.
44 async fn flownode(&self, id: FlownodeId) -> Result<Option<Peer>, Error>;
45}
46
47pub type PeerResolverRef = Arc<dyn PeerResolver>;
48
49/// [`PeerDiscovery`] is a service for discovering active peers in the cluster.
50#[async_trait::async_trait]
51pub trait PeerDiscovery: Send + Sync {
52 /// Returns all currently active frontend nodes.
53 ///
54 /// A frontend is considered active if it has reported a heartbeat within the most recent heartbeat interval,
55 /// as determined by the in-memory backend.
56 async fn active_frontends(&self) -> Result<Vec<NodeInfo>, Error>;
57
58 /// Returns all currently active datanodes, optionally filtered by a predicate on their workloads.
59 ///
60 /// A datanode is considered active if it has reported a heartbeat within the most recent heartbeat interval,
61 /// as determined by the in-memory backend.
62 /// The optional `filter` allows further selection based on the node's workloads.
63 async fn active_datanodes(
64 &self,
65 filter: Option<for<'a> fn(&'a NodeWorkloads) -> bool>,
66 ) -> Result<Vec<NodeInfo>, Error>;
67
68 /// Returns all currently active flownodes, optionally filtered by a predicate on their workloads.
69 ///
70 /// A flownode is considered active if it has reported a heartbeat within the most recent heartbeat interval,
71 /// as determined by the in-memory backend.
72 /// The optional `filter` allows further selection based on the node's workloads.
73 async fn active_flownodes(
74 &self,
75 filter: Option<for<'a> fn(&'a NodeWorkloads) -> bool>,
76 ) -> Result<Vec<NodeInfo>, Error>;
77}
78
79pub type PeerDiscoveryRef = Arc<dyn PeerDiscovery>;
80
81#[derive(Debug, Clone, Default)]
82pub struct PeerAllocContext {
83 pub extensions: HashMap<String, String>,
84}
85
86/// [`PeerAllocator`] allocates [`Peer`]s for creating region or flow.
87#[async_trait::async_trait]
88pub trait PeerAllocator: Send + Sync {
89 async fn alloc(&self, num: usize, ctx: &PeerAllocContext) -> Result<Vec<Peer>, Error>;
90}
91
92pub type PeerAllocatorRef = Arc<dyn PeerAllocator>;
93
94#[async_trait::async_trait]
95impl<T: PeerAllocator + ?Sized> PeerAllocator for Arc<T> {
96 async fn alloc(&self, num: usize, ctx: &PeerAllocContext) -> Result<Vec<Peer>, Error> {
97 T::alloc(self, num, ctx).await
98 }
99}
100
101pub struct NoopPeerAllocator;
102
103#[async_trait::async_trait]
104impl PeerAllocator for NoopPeerAllocator {
105 async fn alloc(&self, num: usize, _ctx: &PeerAllocContext) -> Result<Vec<Peer>, Error> {
106 Ok(vec![Peer::default(); num])
107 }
108}