risingwave_common/util/
addr.rs

1// Copyright 2022 RisingWave Labs
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::net::{SocketAddr, ToSocketAddrs};
16use std::str::FromStr;
17use std::time::Duration;
18
19use anyhow::Context;
20use risingwave_pb::common::PbHostAddress;
21use thiserror_ext::AsReport;
22use tokio::time::sleep;
23use tracing::error;
24
25use crate::util::retry::exponential_backoff;
26
27/// General host address and port.
28#[derive(Debug, Clone, PartialEq, Eq, Hash)]
29pub struct HostAddr {
30    pub host: String,
31    pub port: u16,
32}
33
34impl std::fmt::Display for HostAddr {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        write!(f, "{}:{}", self.host, self.port)
37    }
38}
39impl From<SocketAddr> for HostAddr {
40    fn from(addr: SocketAddr) -> Self {
41        HostAddr {
42            host: addr.ip().to_string(),
43            port: addr.port(),
44        }
45    }
46}
47
48impl TryFrom<&str> for HostAddr {
49    type Error = anyhow::Error;
50
51    fn try_from(s: &str) -> Result<Self, Self::Error> {
52        let url = url::Url::parse(&format!("http://{s}"))
53            .with_context(|| format!("failed to parse address: {s}"))?;
54        let port = match url.port() {
55            Some(port) => port,
56            // The parser normalizes a port equal to the scheme's default (`:80` for `http`)
57            // to `None`. Re-parse with `https` to recover it; a bare host is `None` under both.
58            None => url::Url::parse(&format!("https://{s}"))
59                .with_context(|| format!("failed to parse address: {s}"))?
60                .port()
61                .context("invalid port")?,
62        };
63        Ok(HostAddr {
64            host: url.host().context("invalid host")?.to_string(),
65            port,
66        })
67    }
68}
69
70impl TryFrom<&String> for HostAddr {
71    type Error = anyhow::Error;
72
73    fn try_from(s: &String) -> Result<Self, Self::Error> {
74        Self::try_from(s.as_str())
75    }
76}
77
78impl FromStr for HostAddr {
79    type Err = anyhow::Error;
80
81    fn from_str(s: &str) -> Result<Self, Self::Err> {
82        Self::try_from(s)
83    }
84}
85
86impl From<&PbHostAddress> for HostAddr {
87    fn from(addr: &PbHostAddress) -> Self {
88        HostAddr {
89            host: addr.get_host().clone(),
90            port: addr.get_port() as u16,
91        }
92    }
93}
94
95impl HostAddr {
96    pub fn to_protobuf(&self) -> PbHostAddress {
97        PbHostAddress {
98            host: self.host.clone(),
99            port: self.port as i32,
100        }
101    }
102}
103
104pub fn is_local_address(server_addr: &HostAddr, peer_addr: &HostAddr) -> bool {
105    server_addr == peer_addr
106}
107
108pub async fn try_resolve_dns(host: &str, port: i32) -> Result<SocketAddr, String> {
109    let addr = format!("{}:{}", host, port);
110    let mut backoff = exponential_backoff(Duration::from_millis(100), 5, Duration::from_secs(3));
111    const MAX_RETRY: usize = 20;
112    for i in 1..=MAX_RETRY {
113        let err = match addr.to_socket_addrs() {
114            Ok(mut addr_iter) => {
115                if let Some(addr) = addr_iter.next() {
116                    return Ok(addr);
117                } else {
118                    format!("{} resolved to no addr", addr)
119                }
120            }
121            Err(e) => e.to_report_string(),
122        };
123        // It may happen that the dns information of newly registered worker node
124        // has not been propagated to the meta node and cause error. Wait for a while and retry
125        let delay = backoff.next().unwrap();
126        error!(
127            attempt = i,
128            backoff_delay = ?delay,
129            err,
130            addr,
131            "failed to resolve the worker node address",
132        );
133        sleep(delay).await;
134    }
135    Err(format!("failed to resolve dns: {}", addr))
136}
137
138#[cfg(test)]
139mod tests {
140    use crate::util::addr::{HostAddr, is_local_address};
141
142    #[test]
143    fn test_is_local_address() {
144        let check_local = |a: &str, b: &str, result: bool| {
145            assert_eq!(
146                is_local_address(&a.parse().unwrap(), &b.parse().unwrap()),
147                result
148            );
149        };
150        check_local("localhost:3456", "localhost:3456", true);
151        check_local("10.11.12.13:3456", "10.11.12.13:3456", true);
152        check_local("some.host.in.k8s:3456", "some.host.in.k8s:3456", true);
153        check_local("some.host.in.k8s:3456", "other.host.in.k8s:3456", false);
154        check_local("some.host.in.k8s:3456", "some.host.in.k8s:4567", false);
155    }
156
157    #[test]
158    fn test_host_addr_convert() {
159        let addr = "1.2.3.4:567";
160        assert_eq!(
161            addr.parse::<HostAddr>().unwrap(),
162            HostAddr {
163                host: String::from("1.2.3.4"),
164                port: 567
165            }
166        );
167        let addr = "test.test:12345";
168        assert_eq!(
169            addr.parse::<HostAddr>().unwrap(),
170            HostAddr {
171                host: String::from("test.test"),
172                port: 12345
173            }
174        );
175        // An explicit port equal to a scheme's default must not be normalized away.
176        let addr = "test.test:80";
177        assert_eq!(
178            addr.parse::<HostAddr>().unwrap(),
179            HostAddr {
180                host: String::from("test.test"),
181                port: 80
182            }
183        );
184        let addr = "test.test:443";
185        assert_eq!(
186            addr.parse::<HostAddr>().unwrap(),
187            HostAddr {
188                host: String::from("test.test"),
189                port: 443
190            }
191        );
192        let addr = "[::1]:80";
193        assert_eq!(
194            addr.parse::<HostAddr>().unwrap(),
195            HostAddr {
196                host: String::from("[::1]"),
197                port: 80
198            }
199        );
200        let addr = "test.test";
201        assert!(addr.parse::<HostAddr>().is_err());
202        let addr = "test.test:65537";
203        assert!(addr.parse::<HostAddr>().is_err());
204        let addr = "test.test:";
205        assert!(addr.parse::<HostAddr>().is_err());
206        let addr = "test.test:12345:12345";
207        assert!(addr.parse::<HostAddr>().is_err());
208    }
209}
risingwave_common/util/addr.rs

risingwave_common/util/
addr.rs
