Skip to main content

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}