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::sync::Arc;
16
17pub use api::v1::meta::Peer;
18use api::v1::meta::heartbeat_request::NodeWorkloads;
19
20use crate::affinity::AffinityRequirements;
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    /// Resolved requirements passed to node selection.
84    pub affinity: Option<AffinityRequirements>,
85}
86
87/// [`PeerAllocator`] allocates [`Peer`]s for creating region or flow.
88#[async_trait::async_trait]
89pub trait PeerAllocator: Send + Sync {
90    async fn alloc(&self, num: usize, ctx: &PeerAllocContext) -> Result<Vec<Peer>, Error>;
91}
92
93pub type PeerAllocatorRef = Arc<dyn PeerAllocator>;
94
95#[async_trait::async_trait]
96impl<T: PeerAllocator + ?Sized> PeerAllocator for Arc<T> {
97    async fn alloc(&self, num: usize, ctx: &PeerAllocContext) -> Result<Vec<Peer>, Error> {
98        T::alloc(self, num, ctx).await
99    }
100}
101
102pub struct NoopPeerAllocator;
103
104#[async_trait::async_trait]
105impl PeerAllocator for NoopPeerAllocator {
106    async fn alloc(&self, num: usize, _ctx: &PeerAllocContext) -> Result<Vec<Peer>, Error> {
107        Ok(vec![Peer::default(); num])
108    }
109}