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 s = format!("http://{s}");
53        let addr = url::Url::parse(&s).with_context(|| format!("failed to parse address: {s}"))?;
54        Ok(HostAddr {
55            host: addr.host().context("invalid host")?.to_string(),
56            port: addr.port().context("invalid port")?,
57        })
58    }
59}
60
61impl TryFrom<&String> for HostAddr {
62    type Error = anyhow::Error;
63
64    fn try_from(s: &String) -> Result<Self, Self::Error> {
65        Self::try_from(s.as_str())
66    }
67}
68
69impl FromStr for HostAddr {
70    type Err = anyhow::Error;
71
72    fn from_str(s: &str) -> Result<Self, Self::Err> {
73        Self::try_from(s)
74    }
75}
76
77impl From<&PbHostAddress> for HostAddr {
78    fn from(addr: &PbHostAddress) -> Self {
79        HostAddr {
80            host: addr.get_host().clone(),
81            port: addr.get_port() as u16,
82        }
83    }
84}
85
86impl HostAddr {
87    pub fn to_protobuf(&self) -> PbHostAddress {
88        PbHostAddress {
89            host: self.host.clone(),
90            port: self.port as i32,
91        }
92    }
93}
94
95pub fn is_local_address(server_addr: &HostAddr, peer_addr: &HostAddr) -> bool {
96    server_addr == peer_addr
97}
98
99pub async fn try_resolve_dns(host: &str, port: i32) -> Result<SocketAddr, String> {
100    let addr = format!("{}:{}", host, port);
101    let mut backoff = exponential_backoff(Duration::from_millis(100), 5, Duration::from_secs(3));
102    const MAX_RETRY: usize = 20;
103    for i in 1..=MAX_RETRY {
104        let err = match addr.to_socket_addrs() {
105            Ok(mut addr_iter) => {
106                if let Some(addr) = addr_iter.next() {
107                    return Ok(addr);
108                } else {
109                    format!("{} resolved to no addr", addr)
110                }
111            }
112            Err(e) => e.to_report_string(),
113        };
114        // It may happen that the dns information of newly registered worker node
115        // has not been propagated to the meta node and cause error. Wait for a while and retry
116        let delay = backoff.next().unwrap();
117        error!(
118            attempt = i,
119            backoff_delay = ?delay,
120            err,
121            addr,
122            "failed to resolve the worker node address",
123        );
124        sleep(delay).await;
125    }
126    Err(format!("failed to resolve dns: {}", addr))
127}
128
129#[cfg(test)]
130mod tests {
131    use crate::util::addr::{HostAddr, is_local_address};
132
133    #[test]
134    fn test_is_local_address() {
135        let check_local = |a: &str, b: &str, result: bool| {
136            assert_eq!(
137                is_local_address(&a.parse().unwrap(), &b.parse().unwrap()),
138                result
139            );
140        };
141        check_local("localhost:3456", "localhost:3456", true);
142        check_local("10.11.12.13:3456", "10.11.12.13:3456", true);
143        check_local("some.host.in.k8s:3456", "some.host.in.k8s:3456", true);
144        check_local("some.host.in.k8s:3456", "other.host.in.k8s:3456", false);
145        check_local("some.host.in.k8s:3456", "some.host.in.k8s:4567", false);
146    }
147
148    #[test]
149    fn test_host_addr_convert() {
150        let addr = "1.2.3.4:567";
151        assert_eq!(
152            addr.parse::<HostAddr>().unwrap(),
153            HostAddr {
154                host: String::from("1.2.3.4"),
155                port: 567
156            }
157        );
158        let addr = "test.test:12345";
159        assert_eq!(
160            addr.parse::<HostAddr>().unwrap(),
161            HostAddr {
162                host: String::from("test.test"),
163                port: 12345
164            }
165        );
166        let addr = "test.test";
167        assert!(addr.parse::<HostAddr>().is_err());
168        let addr = "test.test:65537";
169        assert!(addr.parse::<HostAddr>().is_err());
170        let addr = "test.test:";
171        assert!(addr.parse::<HostAddr>().is_err());
172        let addr = "test.test:12345:12345";
173        assert!(addr.parse::<HostAddr>().is_err());
174    }
175}
risingwave_common/util/addr.rs

risingwave_common/util/
addr.rs
