common_grpc/
channel_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::path::Path;
16use std::sync::Arc;
17use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
18use std::time::Duration;
19
20use common_base::readable_size::ReadableSize;
21use common_telemetry::info;
22use dashmap::DashMap;
23use dashmap::mapref::entry::Entry;
24use lazy_static::lazy_static;
25use serde::{Deserialize, Serialize};
26use snafu::ResultExt;
27use tokio_util::sync::CancellationToken;
28use tonic::transport::{
29    Certificate, Channel as InnerChannel, ClientTlsConfig, Endpoint, Identity, Uri,
30};
31use tower::Service;
32
33use crate::error::{CreateChannelSnafu, InvalidConfigFilePathSnafu, Result};
34use crate::reloadable_tls::{ReloadableTlsConfig, TlsConfigLoader, maybe_watch_tls_config};
35
36const RECYCLE_CHANNEL_INTERVAL_SECS: u64 = 60;
37pub const DEFAULT_GRPC_REQUEST_TIMEOUT_SECS: u64 = 10;
38pub const DEFAULT_GRPC_CONNECT_TIMEOUT_SECS: u64 = 1;
39pub const DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE: ReadableSize = ReadableSize::mb(512);
40pub const DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE: ReadableSize = ReadableSize::mb(512);
41
42lazy_static! {
43    static ref ID: AtomicU64 = AtomicU64::new(0);
44}
45
46#[derive(Clone, Debug, Default)]
47pub struct ChannelManager {
48    inner: Arc<Inner>,
49}
50
51#[derive(Debug)]
52struct Inner {
53    id: u64,
54    config: ChannelConfig,
55    reloadable_client_tls_config: Option<Arc<ReloadableClientTlsConfig>>,
56    pool: Arc<Pool>,
57    channel_recycle_started: AtomicBool,
58    cancel: CancellationToken,
59}
60
61impl Default for Inner {
62    fn default() -> Self {
63        Self::with_config(ChannelConfig::default())
64    }
65}
66
67impl Drop for Inner {
68    fn drop(&mut self) {
69        // Cancel the channel recycle task.
70        self.cancel.cancel();
71    }
72}
73
74impl Inner {
75    fn with_config(config: ChannelConfig) -> Self {
76        let id = ID.fetch_add(1, Ordering::Relaxed);
77        let pool = Arc::new(Pool::default());
78        let cancel = CancellationToken::new();
79
80        Self {
81            id,
82            config,
83            reloadable_client_tls_config: None,
84            pool,
85            channel_recycle_started: AtomicBool::new(false),
86            cancel,
87        }
88    }
89}
90
91impl ChannelManager {
92    pub fn new() -> Self {
93        Default::default()
94    }
95
96    /// Create a ChannelManager with configuration and optional TLS config
97    ///
98    /// Use [`load_client_tls_config`] to create TLS configuration from `ClientTlsOption`.
99    /// The TLS config supports both static (watch disabled) and dynamic reloading (watch enabled).
100    /// If you want to use dynamic reloading, please **manually** invoke [`maybe_watch_client_tls_config`] after this method.
101    pub fn with_config(
102        config: ChannelConfig,
103        reloadable_tls_config: Option<Arc<ReloadableClientTlsConfig>>,
104    ) -> Self {
105        let mut inner = Inner::with_config(config.clone());
106        inner.reloadable_client_tls_config = reloadable_tls_config;
107        Self {
108            inner: Arc::new(inner),
109        }
110    }
111
112    pub fn config(&self) -> &ChannelConfig {
113        &self.inner.config
114    }
115
116    fn pool(&self) -> &Arc<Pool> {
117        &self.inner.pool
118    }
119
120    pub fn get(&self, addr: impl AsRef<str>) -> Result<InnerChannel> {
121        self.trigger_channel_recycling();
122
123        let addr = addr.as_ref();
124        // It will acquire the read lock.
125        if let Some(inner_ch) = self.pool().get(addr) {
126            return Ok(inner_ch);
127        }
128
129        // It will acquire the write lock.
130        let entry = match self.pool().entry(addr.to_string()) {
131            Entry::Occupied(entry) => {
132                entry.get().increase_access();
133                entry.into_ref()
134            }
135            Entry::Vacant(entry) => {
136                let endpoint = self.build_endpoint(addr)?;
137                let inner_channel = endpoint.connect_lazy();
138
139                let channel = Channel {
140                    channel: inner_channel,
141                    access: AtomicUsize::new(1),
142                    use_default_connector: true,
143                };
144                entry.insert(channel)
145            }
146        };
147        Ok(entry.channel.clone())
148    }
149
150    pub fn reset_with_connector<C>(
151        &self,
152        addr: impl AsRef<str>,
153        connector: C,
154    ) -> Result<InnerChannel>
155    where
156        C: Service<Uri> + Send + 'static,
157        C::Response: hyper::rt::Read + hyper::rt::Write + Send + Unpin,
158        C::Future: Send + 'static,
159        Box<dyn std::error::Error + Send + Sync>: From<C::Error> + Send + 'static,
160    {
161        let addr = addr.as_ref();
162        let endpoint = self.build_endpoint(addr)?;
163        let inner_channel = endpoint.connect_with_connector_lazy(connector);
164        let channel = Channel {
165            channel: inner_channel.clone(),
166            access: AtomicUsize::new(1),
167            use_default_connector: false,
168        };
169        self.pool().put(addr, channel);
170
171        Ok(inner_channel)
172    }
173
174    pub fn retain_channel<F>(&self, f: F)
175    where
176        F: FnMut(&String, &mut Channel) -> bool,
177    {
178        self.pool().retain_channel(f);
179    }
180
181    /// Clear all channels to force reconnection.
182    /// This should be called when TLS configuration changes to ensure new connections use updated certificates.
183    pub fn clear_all_channels(&self) {
184        self.pool().retain_channel(|_, _| false);
185    }
186
187    fn build_endpoint(&self, addr: &str) -> Result<Endpoint> {
188        // Get the latest TLS config from reloadable config (which handles both static and dynamic cases)
189        let tls_config = self
190            .inner
191            .reloadable_client_tls_config
192            .as_ref()
193            .and_then(|c| c.get_config());
194
195        let http_prefix = if tls_config.is_some() {
196            "https"
197        } else {
198            "http"
199        };
200
201        let mut endpoint = Endpoint::new(format!("{http_prefix}://{addr}"))
202            .context(CreateChannelSnafu { addr })?;
203
204        if let Some(dur) = self.config().timeout {
205            endpoint = endpoint.timeout(dur);
206        }
207        if let Some(dur) = self.config().connect_timeout {
208            endpoint = endpoint.connect_timeout(dur);
209        }
210        if let Some(limit) = self.config().concurrency_limit {
211            endpoint = endpoint.concurrency_limit(limit);
212        }
213        if let Some((limit, dur)) = self.config().rate_limit {
214            endpoint = endpoint.rate_limit(limit, dur);
215        }
216        if let Some(size) = self.config().initial_stream_window_size {
217            endpoint = endpoint.initial_stream_window_size(size);
218        }
219        if let Some(size) = self.config().initial_connection_window_size {
220            endpoint = endpoint.initial_connection_window_size(size);
221        }
222        if let Some(dur) = self.config().http2_keep_alive_interval {
223            endpoint = endpoint.http2_keep_alive_interval(dur);
224        }
225        if let Some(dur) = self.config().http2_keep_alive_timeout {
226            endpoint = endpoint.keep_alive_timeout(dur);
227        }
228        if let Some(enabled) = self.config().http2_keep_alive_while_idle {
229            endpoint = endpoint.keep_alive_while_idle(enabled);
230        }
231        if let Some(enabled) = self.config().http2_adaptive_window {
232            endpoint = endpoint.http2_adaptive_window(enabled);
233        }
234        if let Some(tls_config) = tls_config {
235            endpoint = endpoint
236                .tls_config(tls_config)
237                .context(CreateChannelSnafu { addr })?;
238        }
239
240        endpoint = endpoint
241            .tcp_keepalive(self.config().tcp_keepalive)
242            .tcp_nodelay(self.config().tcp_nodelay);
243
244        Ok(endpoint)
245    }
246
247    fn trigger_channel_recycling(&self) {
248        if self
249            .inner
250            .channel_recycle_started
251            .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
252            .is_err()
253        {
254            return;
255        }
256
257        let pool = self.pool().clone();
258        let cancel = self.inner.cancel.clone();
259        let id = self.inner.id;
260        let _handle = common_runtime::spawn_global(async move {
261            recycle_channel_in_loop(pool, id, cancel, RECYCLE_CHANNEL_INTERVAL_SECS).await;
262        });
263        info!(
264            "ChannelManager: {}, channel recycle is started, running in the background!",
265            self.inner.id
266        );
267    }
268}
269
270fn load_tls_config(tls_option: Option<&ClientTlsOption>) -> Result<Option<ClientTlsConfig>> {
271    let path_config = match tls_option {
272        Some(path_config) if path_config.enabled => path_config,
273        _ => return Ok(None),
274    };
275
276    let mut tls_config = ClientTlsConfig::new();
277
278    if let Some(server_ca) = &path_config.server_ca_cert_path {
279        let server_root_ca_cert =
280            std::fs::read_to_string(server_ca).context(InvalidConfigFilePathSnafu)?;
281        let server_root_ca_cert = Certificate::from_pem(server_root_ca_cert);
282        tls_config = tls_config.ca_certificate(server_root_ca_cert);
283    }
284
285    if let (Some(client_cert_path), Some(client_key_path)) =
286        (&path_config.client_cert_path, &path_config.client_key_path)
287    {
288        let client_cert =
289            std::fs::read_to_string(client_cert_path).context(InvalidConfigFilePathSnafu)?;
290        let client_key =
291            std::fs::read_to_string(client_key_path).context(InvalidConfigFilePathSnafu)?;
292        let client_identity = Identity::from_pem(client_cert, client_key);
293        tls_config = tls_config.identity(client_identity);
294    }
295    Ok(Some(tls_config))
296}
297
298impl TlsConfigLoader<ClientTlsConfig> for ClientTlsOption {
299    type Error = crate::error::Error;
300
301    fn load(&self) -> Result<Option<ClientTlsConfig>> {
302        load_tls_config(Some(self))
303    }
304
305    fn watch_paths(&self) -> Vec<&Path> {
306        let mut paths = Vec::new();
307        if let Some(cert_path) = &self.client_cert_path {
308            paths.push(Path::new(cert_path.as_str()));
309        }
310        if let Some(key_path) = &self.client_key_path {
311            paths.push(Path::new(key_path.as_str()));
312        }
313        if let Some(ca_path) = &self.server_ca_cert_path {
314            paths.push(Path::new(ca_path.as_str()));
315        }
316        paths
317    }
318
319    fn watch_enabled(&self) -> bool {
320        self.enabled && self.watch
321    }
322}
323
324/// Type alias for client-side reloadable TLS config
325pub type ReloadableClientTlsConfig = ReloadableTlsConfig<ClientTlsConfig, ClientTlsOption>;
326
327/// Load client TLS configuration from `ClientTlsOption` and return a `ReloadableClientTlsConfig`.
328/// This is the primary way to create TLS configuration for the ChannelManager.
329pub fn load_client_tls_config(
330    tls_option: Option<ClientTlsOption>,
331) -> Result<Option<Arc<ReloadableClientTlsConfig>>> {
332    match tls_option {
333        Some(option) if option.enabled => {
334            let reloadable = ReloadableClientTlsConfig::try_new(option)?;
335            Ok(Some(Arc::new(reloadable)))
336        }
337        _ => Ok(None),
338    }
339}
340
341pub fn maybe_watch_client_tls_config(
342    client_tls_config: Arc<ReloadableClientTlsConfig>,
343    channel_manager: ChannelManager,
344) -> Result<()> {
345    maybe_watch_tls_config(client_tls_config, move || {
346        // Clear all existing channels to force reconnection with new certificates
347        channel_manager.clear_all_channels();
348        info!("Cleared all existing channels to use new TLS certificates.");
349    })
350}
351
352#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Default)]
353pub struct ClientTlsOption {
354    /// Whether to enable TLS for client.
355    pub enabled: bool,
356    pub server_ca_cert_path: Option<String>,
357    pub client_cert_path: Option<String>,
358    pub client_key_path: Option<String>,
359    #[serde(default)]
360    pub watch: bool,
361}
362
363#[derive(Clone, Debug, PartialEq, Eq)]
364pub struct ChannelConfig {
365    pub timeout: Option<Duration>,
366    pub connect_timeout: Option<Duration>,
367    pub concurrency_limit: Option<usize>,
368    pub rate_limit: Option<(u64, Duration)>,
369    pub initial_stream_window_size: Option<u32>,
370    pub initial_connection_window_size: Option<u32>,
371    pub http2_keep_alive_interval: Option<Duration>,
372    pub http2_keep_alive_timeout: Option<Duration>,
373    pub http2_keep_alive_while_idle: Option<bool>,
374    pub http2_adaptive_window: Option<bool>,
375    pub tcp_keepalive: Option<Duration>,
376    pub tcp_nodelay: bool,
377    pub client_tls: Option<ClientTlsOption>,
378    // Max gRPC receiving(decoding) message size
379    pub max_recv_message_size: ReadableSize,
380    // Max gRPC sending(encoding) message size
381    pub max_send_message_size: ReadableSize,
382    pub send_compression: bool,
383    pub accept_compression: bool,
384}
385
386/// Configures a generated tonic client with compression and message size limits.
387#[macro_export]
388macro_rules! configure_tonic_client {
389    ($client:expr, $channel_manager:expr $(,)?) => {{
390        let channel_manager = &$channel_manager;
391        let config = channel_manager.config();
392        $client
393            .accept_compressed(::tonic::codec::CompressionEncoding::Gzip)
394            .accept_compressed(::tonic::codec::CompressionEncoding::Zstd)
395            .send_compressed(::tonic::codec::CompressionEncoding::Zstd)
396            .max_decoding_message_size(config.max_recv_message_size.as_bytes() as usize)
397            .max_encoding_message_size(config.max_send_message_size.as_bytes() as usize)
398    }};
399}
400
401impl Default for ChannelConfig {
402    fn default() -> Self {
403        Self {
404            timeout: Some(Duration::from_secs(DEFAULT_GRPC_REQUEST_TIMEOUT_SECS)),
405            connect_timeout: Some(Duration::from_secs(DEFAULT_GRPC_CONNECT_TIMEOUT_SECS)),
406            concurrency_limit: None,
407            rate_limit: None,
408            initial_stream_window_size: None,
409            initial_connection_window_size: None,
410            http2_keep_alive_interval: Some(Duration::from_secs(30)),
411            http2_keep_alive_timeout: None,
412            http2_keep_alive_while_idle: Some(true),
413            http2_adaptive_window: None,
414            tcp_keepalive: None,
415            tcp_nodelay: true,
416            client_tls: None,
417            max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE,
418            max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE,
419            send_compression: false,
420            accept_compression: false,
421        }
422    }
423}
424
425impl ChannelConfig {
426    pub fn new() -> Self {
427        Default::default()
428    }
429
430    /// A timeout to each request.
431    pub fn timeout(mut self, timeout: Option<Duration>) -> Self {
432        self.timeout = timeout;
433        self
434    }
435
436    /// A timeout to connecting to the uri.
437    ///
438    /// Defaults to no timeout.
439    pub fn connect_timeout(mut self, timeout: Duration) -> Self {
440        self.connect_timeout = Some(timeout);
441        self
442    }
443
444    /// A concurrency limit to each request.
445    pub fn concurrency_limit(mut self, limit: usize) -> Self {
446        self.concurrency_limit = Some(limit);
447        self
448    }
449
450    /// A rate limit to each request.
451    pub fn rate_limit(mut self, limit: u64, duration: Duration) -> Self {
452        self.rate_limit = Some((limit, duration));
453        self
454    }
455
456    /// Sets the SETTINGS_INITIAL_WINDOW_SIZE option for HTTP2 stream-level flow control.
457    /// Default is 65,535
458    pub fn initial_stream_window_size(mut self, size: u32) -> Self {
459        self.initial_stream_window_size = Some(size);
460        self
461    }
462
463    /// Sets the max connection-level flow control for HTTP2
464    ///
465    /// Default is 65,535
466    pub fn initial_connection_window_size(mut self, size: u32) -> Self {
467        self.initial_connection_window_size = Some(size);
468        self
469    }
470
471    /// Set http2 KEEP_ALIVE_INTERVAL. Uses hyper’s default otherwise.
472    pub fn http2_keep_alive_interval(mut self, duration: Duration) -> Self {
473        self.http2_keep_alive_interval = Some(duration);
474        self
475    }
476
477    /// Set http2 KEEP_ALIVE_TIMEOUT. Uses hyper’s default otherwise.
478    pub fn http2_keep_alive_timeout(mut self, duration: Duration) -> Self {
479        self.http2_keep_alive_timeout = Some(duration);
480        self
481    }
482
483    /// Set http2 KEEP_ALIVE_WHILE_IDLE. Uses hyper’s default otherwise.
484    pub fn http2_keep_alive_while_idle(mut self, enabled: bool) -> Self {
485        self.http2_keep_alive_while_idle = Some(enabled);
486        self
487    }
488
489    /// Sets whether to use an adaptive flow control. Uses hyper’s default otherwise.
490    pub fn http2_adaptive_window(mut self, enabled: bool) -> Self {
491        self.http2_adaptive_window = Some(enabled);
492        self
493    }
494
495    /// Set whether TCP keepalive messages are enabled on accepted connections.
496    ///
497    /// If None is specified, keepalive is disabled, otherwise the duration specified
498    /// will be the time to remain idle before sending TCP keepalive probes.
499    ///
500    /// Default is no keepalive (None)
501    pub fn tcp_keepalive(mut self, duration: Duration) -> Self {
502        self.tcp_keepalive = Some(duration);
503        self
504    }
505
506    /// Set the value of TCP_NODELAY option for accepted connections.
507    ///
508    /// Enabled by default.
509    pub fn tcp_nodelay(mut self, enabled: bool) -> Self {
510        self.tcp_nodelay = enabled;
511        self
512    }
513
514    /// Set the value of tls client auth.
515    ///
516    /// Disabled by default.
517    pub fn client_tls_config(mut self, client_tls_option: ClientTlsOption) -> Self {
518        self.client_tls = Some(client_tls_option);
519        self
520    }
521}
522
523#[derive(Debug)]
524pub struct Channel {
525    channel: InnerChannel,
526    access: AtomicUsize,
527    use_default_connector: bool,
528}
529
530impl Channel {
531    #[inline]
532    pub fn access(&self) -> usize {
533        self.access.load(Ordering::Relaxed)
534    }
535
536    #[inline]
537    pub fn use_default_connector(&self) -> bool {
538        self.use_default_connector
539    }
540
541    #[inline]
542    pub fn increase_access(&self) {
543        let _ = self.access.fetch_add(1, Ordering::Relaxed);
544    }
545}
546
547#[derive(Debug, Default)]
548struct Pool {
549    channels: DashMap<String, Channel>,
550}
551
552impl Pool {
553    fn get(&self, addr: &str) -> Option<InnerChannel> {
554        let channel = self.channels.get(addr);
555        channel.map(|ch| {
556            ch.increase_access();
557            ch.channel.clone()
558        })
559    }
560
561    fn entry(&self, addr: String) -> Entry<'_, String, Channel> {
562        self.channels.entry(addr)
563    }
564
565    #[cfg(test)]
566    fn get_access(&self, addr: &str) -> Option<usize> {
567        let channel = self.channels.get(addr);
568        channel.map(|ch| ch.access())
569    }
570
571    fn put(&self, addr: &str, channel: Channel) {
572        let _ = self.channels.insert(addr.to_string(), channel);
573    }
574
575    fn retain_channel<F>(&self, f: F)
576    where
577        F: FnMut(&String, &mut Channel) -> bool,
578    {
579        self.channels.retain(f);
580    }
581}
582
583async fn recycle_channel_in_loop(
584    pool: Arc<Pool>,
585    id: u64,
586    cancel: CancellationToken,
587    interval_secs: u64,
588) {
589    let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));
590
591    loop {
592        tokio::select! {
593            _ = cancel.cancelled() => {
594                info!("Stop channel recycle, ChannelManager id: {}", id);
595                break;
596            },
597            _ = interval.tick() => {}
598        }
599
600        pool.retain_channel(|_, c| c.access.swap(0, Ordering::Relaxed) != 0)
601    }
602}
603
604#[cfg(test)]
605mod tests {
606    use tower::service_fn;
607
608    use super::*;
609
610    #[should_panic]
611    #[test]
612    fn test_invalid_addr() {
613        let mgr = ChannelManager::default();
614        let addr = "http://test";
615
616        let _ = mgr.get(addr).unwrap();
617    }
618
619    #[tokio::test]
620    async fn test_access_count() {
621        let mgr = ChannelManager::new();
622        // Do not start recycle
623        mgr.inner
624            .channel_recycle_started
625            .store(true, Ordering::Relaxed);
626        let mgr = Arc::new(mgr);
627        let addr = "test_uri";
628
629        let mut joins = Vec::with_capacity(10);
630        for _ in 0..10 {
631            let mgr_clone = mgr.clone();
632            let join = tokio::spawn(async move {
633                for _ in 0..100 {
634                    let _ = mgr_clone.get(addr);
635                }
636            });
637            joins.push(join);
638        }
639        for join in joins {
640            join.await.unwrap();
641        }
642
643        assert_eq!(1000, mgr.pool().get_access(addr).unwrap());
644
645        mgr.pool()
646            .retain_channel(|_, c| c.access.swap(0, Ordering::Relaxed) != 0);
647
648        assert_eq!(0, mgr.pool().get_access(addr).unwrap());
649    }
650
651    #[test]
652    fn test_config() {
653        let default_cfg = ChannelConfig::new();
654        assert_eq!(
655            ChannelConfig {
656                timeout: Some(Duration::from_secs(DEFAULT_GRPC_REQUEST_TIMEOUT_SECS)),
657                connect_timeout: Some(Duration::from_secs(DEFAULT_GRPC_CONNECT_TIMEOUT_SECS)),
658                concurrency_limit: None,
659                rate_limit: None,
660                initial_stream_window_size: None,
661                initial_connection_window_size: None,
662                http2_keep_alive_interval: Some(Duration::from_secs(30)),
663                http2_keep_alive_timeout: None,
664                http2_keep_alive_while_idle: Some(true),
665                http2_adaptive_window: None,
666                tcp_keepalive: None,
667                tcp_nodelay: true,
668                client_tls: None,
669                max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE,
670                max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE,
671                send_compression: false,
672                accept_compression: false,
673            },
674            default_cfg
675        );
676
677        let cfg = default_cfg
678            .timeout(Some(Duration::from_secs(3)))
679            .connect_timeout(Duration::from_secs(5))
680            .concurrency_limit(6)
681            .rate_limit(5, Duration::from_secs(1))
682            .initial_stream_window_size(10)
683            .initial_connection_window_size(20)
684            .http2_keep_alive_interval(Duration::from_secs(1))
685            .http2_keep_alive_timeout(Duration::from_secs(3))
686            .http2_keep_alive_while_idle(true)
687            .http2_adaptive_window(true)
688            .tcp_keepalive(Duration::from_secs(2))
689            .tcp_nodelay(false)
690            .client_tls_config(ClientTlsOption {
691                enabled: true,
692                server_ca_cert_path: Some("some_server_path".to_string()),
693                client_cert_path: Some("some_cert_path".to_string()),
694                client_key_path: Some("some_key_path".to_string()),
695                watch: false,
696            });
697
698        assert_eq!(
699            ChannelConfig {
700                timeout: Some(Duration::from_secs(3)),
701                connect_timeout: Some(Duration::from_secs(5)),
702                concurrency_limit: Some(6),
703                rate_limit: Some((5, Duration::from_secs(1))),
704                initial_stream_window_size: Some(10),
705                initial_connection_window_size: Some(20),
706                http2_keep_alive_interval: Some(Duration::from_secs(1)),
707                http2_keep_alive_timeout: Some(Duration::from_secs(3)),
708                http2_keep_alive_while_idle: Some(true),
709                http2_adaptive_window: Some(true),
710                tcp_keepalive: Some(Duration::from_secs(2)),
711                tcp_nodelay: false,
712                client_tls: Some(ClientTlsOption {
713                    enabled: true,
714                    server_ca_cert_path: Some("some_server_path".to_string()),
715                    client_cert_path: Some("some_cert_path".to_string()),
716                    client_key_path: Some("some_key_path".to_string()),
717                    watch: false,
718                }),
719                max_recv_message_size: DEFAULT_MAX_GRPC_RECV_MESSAGE_SIZE,
720                max_send_message_size: DEFAULT_MAX_GRPC_SEND_MESSAGE_SIZE,
721                send_compression: false,
722                accept_compression: false,
723            },
724            cfg
725        );
726    }
727
728    #[derive(Default)]
729    struct FakeTonicClient {
730        accepted_gzip: bool,
731        accepted_zstd: bool,
732        sent_gzip: bool,
733        sent_zstd: bool,
734        max_decoding_message_size: usize,
735        max_encoding_message_size: usize,
736    }
737
738    impl FakeTonicClient {
739        fn accept_compressed(mut self, encoding: ::tonic::codec::CompressionEncoding) -> Self {
740            match encoding {
741                ::tonic::codec::CompressionEncoding::Gzip => self.accepted_gzip = true,
742                ::tonic::codec::CompressionEncoding::Zstd => self.accepted_zstd = true,
743                _ => unreachable!(),
744            }
745            self
746        }
747
748        fn send_compressed(mut self, encoding: ::tonic::codec::CompressionEncoding) -> Self {
749            match encoding {
750                ::tonic::codec::CompressionEncoding::Gzip => self.sent_gzip = true,
751                ::tonic::codec::CompressionEncoding::Zstd => self.sent_zstd = true,
752                _ => unreachable!(),
753            }
754            self
755        }
756
757        fn max_decoding_message_size(mut self, size: usize) -> Self {
758            self.max_decoding_message_size = size;
759            self
760        }
761
762        fn max_encoding_message_size(mut self, size: usize) -> Self {
763            self.max_encoding_message_size = size;
764            self
765        }
766    }
767
768    #[test]
769    fn test_configure_tonic_client() {
770        let recv_message_size = ReadableSize::mb(2);
771        let send_message_size = ReadableSize::mb(3);
772        let channel_manager = ChannelManager::with_config(
773            ChannelConfig {
774                max_recv_message_size: recv_message_size,
775                max_send_message_size: send_message_size,
776                ..ChannelConfig::new()
777            },
778            None,
779        );
780
781        let client = crate::configure_tonic_client!(FakeTonicClient::default(), channel_manager,);
782
783        assert!(client.accepted_gzip);
784        assert!(client.accepted_zstd);
785        assert!(!client.sent_gzip);
786        assert!(client.sent_zstd);
787        assert_eq!(
788            recv_message_size.as_bytes() as usize,
789            client.max_decoding_message_size
790        );
791        assert_eq!(
792            send_message_size.as_bytes() as usize,
793            client.max_encoding_message_size
794        );
795    }
796
797    #[test]
798    fn test_build_endpoint() {
799        let config = ChannelConfig::new()
800            .timeout(Some(Duration::from_secs(3)))
801            .connect_timeout(Duration::from_secs(5))
802            .concurrency_limit(6)
803            .rate_limit(5, Duration::from_secs(1))
804            .initial_stream_window_size(10)
805            .initial_connection_window_size(20)
806            .http2_keep_alive_interval(Duration::from_secs(1))
807            .http2_keep_alive_timeout(Duration::from_secs(3))
808            .http2_keep_alive_while_idle(true)
809            .http2_adaptive_window(true)
810            .tcp_keepalive(Duration::from_secs(2))
811            .tcp_nodelay(true);
812        let mgr = ChannelManager::with_config(config, None);
813
814        let res = mgr.build_endpoint("test_addr");
815
816        let _ = res.unwrap();
817    }
818
819    #[tokio::test]
820    async fn test_channel_with_connector() {
821        let mgr = ChannelManager::new();
822
823        let addr = "test_addr";
824        let res = mgr.get(addr);
825        let _ = res.unwrap();
826
827        mgr.retain_channel(|addr, channel| {
828            assert_eq!("test_addr", addr);
829            assert!(channel.use_default_connector());
830            true
831        });
832
833        let (client, _) = tokio::io::duplex(1024);
834        let mut client = Some(hyper_util::rt::TokioIo::new(client));
835        let res = mgr.reset_with_connector(
836            addr,
837            service_fn(move |_| {
838                let client = client.take().unwrap();
839                async move { Ok::<_, std::io::Error>(client) }
840            }),
841        );
842
843        let _ = res.unwrap();
844
845        mgr.retain_channel(|addr, channel| {
846            assert_eq!("test_addr", addr);
847            assert!(!channel.use_default_connector());
848            true
849        });
850    }
851
852    #[tokio::test]
853    async fn test_pool_release_with_channel_recycle() {
854        let mgr = ChannelManager::new();
855
856        let pool_holder = mgr.pool().clone();
857
858        // start channel recycle task
859        let addr = "test_addr";
860        let _ = mgr.get(addr);
861
862        let mgr_clone_1 = mgr.clone();
863        let mgr_clone_2 = mgr.clone();
864        assert_eq!(3, Arc::strong_count(mgr.pool()));
865
866        drop(mgr_clone_1);
867        drop(mgr_clone_2);
868        assert_eq!(3, Arc::strong_count(mgr.pool()));
869
870        drop(mgr);
871
872        // wait for the channel recycle task to finish
873        tokio::time::sleep(Duration::from_millis(10)).await;
874
875        assert_eq!(1, Arc::strong_count(&pool_holder));
876    }
877
878    #[tokio::test]
879    async fn test_pool_release_without_channel_recycle() {
880        let mgr = ChannelManager::new();
881
882        let pool_holder = mgr.pool().clone();
883
884        let mgr_clone_1 = mgr.clone();
885        let mgr_clone_2 = mgr.clone();
886        assert_eq!(2, Arc::strong_count(mgr.pool()));
887
888        drop(mgr_clone_1);
889        drop(mgr_clone_2);
890        assert_eq!(2, Arc::strong_count(mgr.pool()));
891
892        drop(mgr);
893
894        assert_eq!(1, Arc::strong_count(&pool_holder));
895    }
896}
common_grpc/channel_manager.rs

common_grpc/
channel_manager.rs
