client/
client_manager.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::fmt::{Debug, Formatter};
16use std::sync::Arc;
17use std::time::Duration;
18
19use common_grpc::channel_manager::{ChannelConfig, ChannelManager};
20use common_meta::node_manager::{DatanodeManager, DatanodeRef, FlownodeManager, FlownodeRef};
21use common_meta::peer::Peer;
22use moka::future::{Cache, CacheBuilder};
23
24use crate::Client;
25use crate::flow::FlowRequester;
26use crate::region::RegionRequester;
27
28pub struct NodeClients {
29    // Keep the channel managers independent by lane:
30    // query = Flight DoGet/distributed reads; control/mutation = health, unary RPCs,
31    // gateway, and Flight DoPut.
32    query_channel_manager: ChannelManager,
33    control_channel_manager: ChannelManager,
34    clients: Cache<Peer, Client>,
35}
36
37impl Default for NodeClients {
38    fn default() -> Self {
39        Self::new(ChannelConfig::new())
40    }
41}
42
43impl Debug for NodeClients {
44    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("NodeClients")
46            .field("query_channel_manager", &self.query_channel_manager)
47            .field("control_channel_manager", &self.control_channel_manager)
48            .finish()
49    }
50}
51
52#[async_trait::async_trait]
53impl DatanodeManager for NodeClients {
54    async fn datanode(&self, datanode: &Peer) -> DatanodeRef {
55        let client = self.get_client(datanode).await;
56
57        let ChannelConfig {
58            send_compression,
59            accept_compression,
60            ..
61        } = self.control_channel_manager.config();
62        Arc::new(RegionRequester::new(
63            client,
64            *send_compression,
65            *accept_compression,
66        ))
67    }
68}
69
70#[async_trait::async_trait]
71impl FlownodeManager for NodeClients {
72    async fn flownode(&self, flownode: &Peer) -> FlownodeRef {
73        let client = self.get_client(flownode).await;
74
75        Arc::new(FlowRequester::new(client))
76    }
77}
78
79impl NodeClients {
80    /// Creates node clients with independent query and control channel pools.
81    pub fn new(config: ChannelConfig) -> Self {
82        Self {
83            query_channel_manager: ChannelManager::with_config(config.clone(), None),
84            control_channel_manager: ChannelManager::with_config(config, None),
85            clients: CacheBuilder::new(1024)
86                .time_to_live(Duration::from_secs(30 * 60))
87                .time_to_idle(Duration::from_secs(5 * 60))
88                .build(),
89        }
90    }
91
92    pub async fn get_client(&self, datanode: &Peer) -> Client {
93        self.clients
94            .get_with_by_ref(datanode, async move {
95                Client::with_query_and_control_managers(
96                    self.query_channel_manager.clone(),
97                    self.control_channel_manager.clone(),
98                    vec![datanode.addr.clone()],
99                )
100            })
101            .await
102    }
103
104    #[cfg(feature = "testing")]
105    pub async fn insert_client(&self, datanode: Peer, client: Client) {
106        self.clients.insert(datanode, client).await
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use common_grpc::channel_manager::ChannelManager;
113    use common_meta::peer::Peer;
114
115    use super::{ChannelConfig, NodeClients};
116    use crate::Client;
117
118    const PEER_ADDR: &str = "127.0.0.1:3001";
119
120    fn assert_pool_has_one_address(manager: &ChannelManager) {
121        let mut count = 0;
122        let mut addresses = Vec::new();
123        manager.retain_channel(|addr, _| {
124            count += 1;
125            addresses.push(addr.clone());
126            true
127        });
128        assert_eq!(1, count);
129        assert_eq!([PEER_ADDR], addresses.as_slice());
130    }
131
132    #[tokio::test]
133    async fn test_node_clients_use_isolated_reused_channel_pools() {
134        let node_clients = NodeClients::new(ChannelConfig::default());
135        let peer = Peer {
136            id: 1,
137            addr: PEER_ADDR.to_string(),
138        };
139        let client = node_clients.get_client(&peer).await;
140
141        client.make_flight_client(false, false).unwrap();
142        client.make_flight_client(false, false).unwrap();
143        assert_pool_has_one_address(&node_clients.query_channel_manager);
144        let mut control_count = 0;
145        node_clients.control_channel_manager.retain_channel(|_, _| {
146            control_count += 1;
147            true
148        });
149        assert_eq!(0, control_count);
150
151        client.make_control_flight_client(false, false).unwrap();
152        client.make_control_flight_client(false, false).unwrap();
153        assert_pool_has_one_address(&node_clients.query_channel_manager);
154        assert_pool_has_one_address(&node_clients.control_channel_manager);
155
156        let manager = ChannelManager::new();
157        #[allow(deprecated)]
158        let legacy = Client::with_manager_and_urls(manager.clone(), [PEER_ADDR]);
159        legacy.make_flight_client(false, false).unwrap();
160        legacy.make_control_flight_client(false, false).unwrap();
161        assert_pool_has_one_address(&manager);
162    }
163}
client/client_manager.rs

client/
client_manager.rs
