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

risingwave_common/util/
addr.rs
