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    pub fn new(config: ChannelConfig) -> Self {
81        Self {
82            query_channel_manager: ChannelManager::with_config(config.clone(), None),
83            control_channel_manager: ChannelManager::with_config(config, None),
84            clients: CacheBuilder::new(1024)
85                .time_to_live(Duration::from_secs(30 * 60))
86                .time_to_idle(Duration::from_secs(5 * 60))
87                .build(),
88        }
89    }
90
91    pub async fn get_client(&self, datanode: &Peer) -> Client {
92        self.clients
93            .get_with_by_ref(datanode, async move {
94                Client::with_managers_and_urls(
95                    self.query_channel_manager.clone(),
96                    self.control_channel_manager.clone(),
97                    vec![datanode.addr.clone()],
98                )
99            })
100            .await
101    }
102
103    #[cfg(feature = "testing")]
104    pub async fn insert_client(&self, datanode: Peer, client: Client) {
105        self.clients.insert(datanode, client).await
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use common_grpc::channel_manager::ChannelManager;
112    use common_meta::peer::Peer;
113
114    use super::{ChannelConfig, NodeClients};
115    use crate::Client;
116
117    const PEER_ADDR: &str = "127.0.0.1:3001";
118
119    fn assert_pool_has_one_address(manager: &ChannelManager) {
120        let mut count = 0;
121        let mut addresses = Vec::new();
122        manager.retain_channel(|addr, _| {
123            count += 1;
124            addresses.push(addr.clone());
125            true
126        });
127        assert_eq!(1, count);
128        assert_eq!([PEER_ADDR], addresses.as_slice());
129    }
130
131    #[tokio::test]
132    async fn test_node_clients_use_isolated_reused_channel_pools() {
133        let node_clients = NodeClients::new(ChannelConfig::default());
134        let peer = Peer {
135            id: 1,
136            addr: PEER_ADDR.to_string(),
137        };
138        let client = node_clients.get_client(&peer).await;
139
140        client.make_flight_client(false, false).unwrap();
141        client.make_flight_client(false, false).unwrap();
142        assert_pool_has_one_address(&node_clients.query_channel_manager);
143        let mut control_count = 0;
144        node_clients.control_channel_manager.retain_channel(|_, _| {
145            control_count += 1;
146            true
147        });
148        assert_eq!(0, control_count);
149
150        client.make_control_flight_client(false, false).unwrap();
151        client.make_control_flight_client(false, false).unwrap();
152        assert_pool_has_one_address(&node_clients.query_channel_manager);
153        assert_pool_has_one_address(&node_clients.control_channel_manager);
154
155        let manager = ChannelManager::new();
156        let legacy = Client::with_manager_and_urls(manager.clone(), [PEER_ADDR]);
157        legacy.make_flight_client(false, false).unwrap();
158        legacy.make_control_flight_client(false, false).unwrap();
159        assert_pool_has_one_address(&manager);
160    }
161}
client/client_manager.rs

client/
client_manager.rs
