iota_network_stack/
client.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use std::{
6    collections::HashMap,
7    fmt,
8    future::Future,
9    io,
10    net::{SocketAddr, ToSocketAddrs},
11    pin::Pin,
12    sync::{Arc, Mutex},
13    task::{self, Poll},
14    time::Instant,
15    vec,
16};
17
18use eyre::{Context, Result, eyre};
19use hyper_util::client::legacy::connect::{HttpConnector, dns::Name};
20use once_cell::sync::OnceCell;
21use tokio::task::JoinHandle;
22use tokio_rustls::rustls::ClientConfig;
23use tonic::transport::{Channel, Endpoint, Uri};
24use tower::Service;
25use tracing::{info, trace};
26
27use crate::{
28    config::Config,
29    multiaddr::{Multiaddr, Protocol, parse_dns, parse_ip4, parse_ip6},
30};
31
32pub async fn connect(address: &Multiaddr, tls_config: ClientConfig) -> Result<Channel> {
33    let channel = endpoint_from_multiaddr(address, tls_config)?
34        .connect()
35        .await?;
36    Ok(channel)
37}
38
39pub fn connect_lazy(address: &Multiaddr, tls_config: ClientConfig) -> Result<Channel> {
40    let channel = endpoint_from_multiaddr(address, tls_config)?.connect_lazy();
41    Ok(channel)
42}
43
44pub(crate) async fn connect_with_config(
45    address: &Multiaddr,
46    tls_config: ClientConfig,
47    config: &Config,
48) -> Result<Channel> {
49    let channel = endpoint_from_multiaddr(address, tls_config)?
50        .apply_config(config)
51        .connect()
52        .await?;
53    Ok(channel)
54}
55
56pub(crate) fn connect_lazy_with_config(
57    address: &Multiaddr,
58    tls_config: ClientConfig,
59    config: &Config,
60) -> Result<Channel> {
61    let channel = endpoint_from_multiaddr(address, tls_config)?
62        .apply_config(config)
63        .connect_lazy();
64    Ok(channel)
65}
66
67fn endpoint_from_multiaddr(addr: &Multiaddr, tls_config: ClientConfig) -> Result<MyEndpoint> {
68    let mut iter = addr.iter();
69
70    let channel = match iter.next().ok_or_else(|| eyre!("address is empty"))? {
71        Protocol::Dns(_) => {
72            let (dns_name, tcp_port, http_or_https) = parse_dns(addr)?;
73            let uri = format!("{http_or_https}://{dns_name}:{tcp_port}");
74            MyEndpoint::try_from_uri(uri, tls_config)?
75        }
76        Protocol::Ip4(_) => {
77            let (socket_addr, http_or_https) = parse_ip4(addr)?;
78            let uri = format!("{http_or_https}://{socket_addr}");
79            MyEndpoint::try_from_uri(uri, tls_config)?
80        }
81        Protocol::Ip6(_) => {
82            let (socket_addr, http_or_https) = parse_ip6(addr)?;
83            let uri = format!("{http_or_https}://{socket_addr}");
84            MyEndpoint::try_from_uri(uri, tls_config)?
85        }
86        unsupported => return Err(eyre!("unsupported protocol {unsupported}")),
87    };
88
89    Ok(channel)
90}
91
92struct MyEndpoint {
93    endpoint: Endpoint,
94    tls_config: ClientConfig,
95}
96
97static DISABLE_CACHING_RESOLVER: OnceCell<bool> = OnceCell::new();
98
99impl MyEndpoint {
100    fn new(endpoint: Endpoint, tls_config: ClientConfig) -> Self {
101        Self {
102            endpoint,
103            tls_config,
104        }
105    }
106
107    fn try_from_uri(uri: String, tls_config: ClientConfig) -> Result<Self> {
108        let uri: Uri = uri
109            .parse()
110            .with_context(|| format!("unable to create Uri from '{uri}'"))?;
111        let endpoint = Endpoint::from(uri);
112        Ok(Self::new(endpoint, tls_config))
113    }
114
115    fn apply_config(mut self, config: &Config) -> Self {
116        self.endpoint = apply_config_to_endpoint(config, self.endpoint);
117        self
118    }
119
120    fn connect_lazy(self) -> Channel {
121        let disable_caching_resolver = *DISABLE_CACHING_RESOLVER.get_or_init(|| {
122            let disable_caching_resolver = std::env::var("DISABLE_CACHING_RESOLVER").is_ok();
123            info!("DISABLE_CACHING_RESOLVER: {disable_caching_resolver}");
124            disable_caching_resolver
125        });
126
127        // Propagate the endpoint's configured connect timeout (Config::connect_timeout
128        // applied in apply_config_to_endpoint) down to the HTTP connector. tonic's
129        // Channel::new path does not forward Endpoint::connect_timeout to a provided
130        // connector, so without this the lazy channel can sit in State::Connecting
131        // forever when a peer is unreachable (msim with dropped SYNs, real network
132        // partitions, etc), blocking every gRPC call on that channel.
133        let connect_timeout = self.endpoint.get_connect_timeout();
134
135        if disable_caching_resolver {
136            let mut http = HttpConnector::new();
137            http.enforce_http(false);
138            http.set_nodelay(true);
139            http.set_keepalive(None);
140            http.set_connect_timeout(connect_timeout);
141
142            Channel::new(
143                hyper_rustls::HttpsConnectorBuilder::new()
144                    .with_tls_config(self.tls_config)
145                    .https_only()
146                    .enable_http2()
147                    .wrap_connector(http),
148                self.endpoint,
149            )
150        } else {
151            let mut http = HttpConnector::new_with_resolver(CachingResolver::new());
152            http.enforce_http(false);
153            http.set_nodelay(true);
154            http.set_keepalive(None);
155            http.set_connect_timeout(connect_timeout);
156
157            let https = hyper_rustls::HttpsConnectorBuilder::new()
158                .with_tls_config(self.tls_config)
159                .https_only()
160                .enable_http2()
161                .wrap_connector(http);
162            Channel::new(https, self.endpoint)
163        }
164    }
165
166    async fn connect(self) -> Result<Channel> {
167        let https_connector = hyper_rustls::HttpsConnectorBuilder::new()
168            .with_tls_config(self.tls_config)
169            .https_only()
170            .enable_http2()
171            .build();
172        Channel::connect(https_connector, self.endpoint)
173            .await
174            .map_err(Into::into)
175    }
176}
177
178fn apply_config_to_endpoint(config: &Config, mut endpoint: Endpoint) -> Endpoint {
179    if let Some(limit) = config.concurrency_limit_per_connection {
180        endpoint = endpoint.concurrency_limit(limit);
181    }
182
183    if let Some(timeout) = config.request_timeout {
184        endpoint = endpoint.timeout(timeout);
185    }
186
187    if let Some(timeout) = config.connect_timeout {
188        endpoint = endpoint.connect_timeout(timeout);
189    }
190
191    if let Some(tcp_nodelay) = config.tcp_nodelay {
192        endpoint = endpoint.tcp_nodelay(tcp_nodelay);
193    }
194
195    if let Some(http2_keepalive_interval) = config.http2_keepalive_interval {
196        endpoint = endpoint.http2_keep_alive_interval(http2_keepalive_interval);
197    }
198
199    if let Some(http2_keepalive_timeout) = config.http2_keepalive_timeout {
200        endpoint = endpoint.keep_alive_timeout(http2_keepalive_timeout);
201    }
202
203    if let Some((limit, duration)) = config.rate_limit {
204        endpoint = endpoint.rate_limit(limit, duration);
205    }
206
207    endpoint
208        .initial_stream_window_size(config.http2_initial_stream_window_size)
209        .initial_connection_window_size(config.http2_initial_connection_window_size)
210        .tcp_keepalive(config.tcp_keepalive)
211}
212
213type CacheEntry = (Instant, Vec<SocketAddr>);
214
215/// A caching resolver based on hyper_util GaiResolver
216#[derive(Clone)]
217pub struct CachingResolver {
218    cache: Arc<Mutex<HashMap<Name, CacheEntry>>>,
219}
220
221type SocketAddrs = vec::IntoIter<SocketAddr>;
222
223pub struct CachingFuture {
224    inner: JoinHandle<Result<SocketAddrs, io::Error>>,
225}
226
227impl CachingResolver {
228    pub fn new() -> Self {
229        CachingResolver {
230            cache: Arc::new(Mutex::new(HashMap::new())),
231        }
232    }
233}
234
235impl Default for CachingResolver {
236    fn default() -> Self {
237        Self::new()
238    }
239}
240
241impl Service<Name> for CachingResolver {
242    type Response = SocketAddrs;
243    type Error = io::Error;
244    type Future = CachingFuture;
245
246    fn poll_ready(&mut self, _cx: &mut task::Context<'_>) -> Poll<Result<(), io::Error>> {
247        Poll::Ready(Ok(()))
248    }
249
250    fn call(&mut self, name: Name) -> Self::Future {
251        let blocking = {
252            let cache = self.cache.clone();
253            tokio::task::spawn_blocking(move || {
254                let entry = cache.lock().unwrap().get(&name).cloned();
255
256                if let Some((when, addrs)) = entry {
257                    trace!("cached host={:?}", name.as_str());
258
259                    if when.elapsed().as_secs() > 60 {
260                        trace!("refreshing cache for host={:?}", name.as_str());
261                        // Start a new task to update the cache later.
262                        tokio::task::spawn_blocking(move || {
263                            if let Ok(addrs) = (name.as_str(), 0).to_socket_addrs() {
264                                let addrs: Vec<_> = addrs.collect();
265                                trace!("updating cached host={:?}", name.as_str());
266                                cache.lock().unwrap().insert(name, (Instant::now(), addrs));
267                            }
268                        });
269                    }
270
271                    Ok(addrs.into_iter())
272                } else {
273                    trace!("resolving host={:?}", name.as_str());
274                    match (name.as_str(), 0).to_socket_addrs() {
275                        Ok(addrs) => {
276                            let addrs: Vec<_> = addrs.collect();
277                            cache
278                                .lock()
279                                .unwrap()
280                                .insert(name, (Instant::now(), addrs.clone()));
281                            Ok(addrs.into_iter())
282                        }
283                        res => res,
284                    }
285                }
286            })
287        };
288
289        CachingFuture { inner: blocking }
290    }
291}
292
293impl fmt::Debug for CachingResolver {
294    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
295        f.pad("CachingResolver")
296    }
297}
298
299impl Future for CachingFuture {
300    type Output = Result<SocketAddrs, io::Error>;
301
302    fn poll(mut self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> Poll<Self::Output> {
303        Pin::new(&mut self.inner).poll(cx).map(|res| match res {
304            Ok(Ok(addrs)) => Ok(addrs),
305            Ok(Err(err)) => Err(err),
306            Err(join_err) => {
307                if join_err.is_cancelled() {
308                    Err(io::Error::new(io::ErrorKind::Interrupted, join_err))
309                } else {
310                    panic!("background task failed: {join_err:?}")
311                }
312            }
313        })
314    }
315}
316
317impl fmt::Debug for CachingFuture {
318    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
319        f.pad("CachingFuture")
320    }
321}
322
323impl Drop for CachingFuture {
324    fn drop(&mut self) {
325        self.inner.abort();
326    }
327}
iota_network_stack/client.rs

iota_network_stack/
client.rs
