mirror of https://github.com/rapiz1/rathole.git
fix: respect shutdown signal when retry (#121)
* fix: respect shutdown signal when retry * test: add tests for retry_notify * chore: drop backoff/tokio
This commit is contained in:
parent
dde08bdd8c
commit
0278c529dd
|
@ -122,12 +122,9 @@ version = "0.3.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9fe17f59a06fe8b87a6fc8bf53bb70b3aba76d7685f432487a68cd5552853625"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"getrandom 0.2.4",
|
||||
"instant",
|
||||
"pin-project",
|
||||
"rand",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
@ -57,7 +57,7 @@ bincode = "1"
|
|||
lazy_static = "1.4"
|
||||
hex = "0.4"
|
||||
rand = "0.8"
|
||||
backoff = { version="0.3", features=["tokio"] }
|
||||
backoff = "0.3"
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = "0.2"
|
||||
socket2 = { version = "0.4", features = ["all"] }
|
||||
|
|
|
@ -1,6 +1,6 @@
|
|||
use crate::config::{ClientConfig, ClientServiceConfig, Config, ServiceType, TransportType};
|
||||
use crate::config_watcher::ServiceChange;
|
||||
use crate::helper::udp_connect;
|
||||
use crate::helper::{retry_notify, udp_connect};
|
||||
use crate::protocol::Hello::{self, *};
|
||||
use crate::protocol::{
|
||||
self, read_ack, read_control_cmd, read_data_cmd, read_hello, Ack, Auth, ControlChannelCmd,
|
||||
|
@ -159,31 +159,33 @@ async fn do_data_channel_handshake<T: Transport>(
|
|||
args: Arc<RunDataChannelArgs<T>>,
|
||||
) -> Result<T::Stream> {
|
||||
// Retry at least every 100ms, at most for 10 seconds
|
||||
let backoff = ExponentialBackoff {
|
||||
let mut backoff = ExponentialBackoff {
|
||||
max_interval: Duration::from_millis(100),
|
||||
max_elapsed_time: Some(Duration::from_secs(10)),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// FIXME: Respect control channel shutdown here
|
||||
// Connect to remote_addr
|
||||
let mut conn: T::Stream = backoff::future::retry_notify(
|
||||
let mut conn: T::Stream = retry_notify!(
|
||||
backoff,
|
||||
|| async {
|
||||
let conn = args
|
||||
{
|
||||
match args
|
||||
.connector
|
||||
.connect(&args.remote_addr)
|
||||
.await
|
||||
.with_context(|| format!("Failed to connect to {}", &args.remote_addr))?;
|
||||
T::hint(&conn, args.socket_opts);
|
||||
|
||||
Ok(conn)
|
||||
.with_context(|| format!("Failed to connect to {}", &args.remote_addr))
|
||||
{
|
||||
Ok(conn) => {
|
||||
T::hint(&conn, args.socket_opts);
|
||||
Ok(conn)
|
||||
}
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
},
|
||||
|e, duration| {
|
||||
warn!("{:#}. Retry in {:?}", e, duration);
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
)?;
|
||||
|
||||
// Send nonce
|
||||
let v: &[u8; HASH_WIDTH_IN_BYTES] = args.session_key[..].try_into().unwrap();
|
||||
|
|
|
@ -51,3 +51,65 @@ pub async fn udp_connect<A: ToSocketAddrs>(addr: A) -> Result<UdpSocket> {
|
|||
s.connect(addr).await?;
|
||||
Ok(s)
|
||||
}
|
||||
|
||||
/// Almost same as backoff::future::retry_notify
|
||||
/// But directly expands to a loop
|
||||
macro_rules! retry_notify {
|
||||
($b: expr, $func: expr, $notify: expr) => {
|
||||
loop {
|
||||
match $func {
|
||||
Ok(v) => break Ok(v),
|
||||
Err(e) => match $b.next_backoff() {
|
||||
Some(duration) => {
|
||||
$notify(e, duration);
|
||||
tokio::time::sleep(duration).await;
|
||||
}
|
||||
None => break Err(e),
|
||||
},
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
pub(crate) use retry_notify;
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
use backoff::{backoff::Backoff, ExponentialBackoff};
|
||||
#[tokio::test]
|
||||
async fn test_retry_notify() {
|
||||
let tests = [(3, 3, 2, Ok(())), (4, 3, 2, Err("try again"))];
|
||||
for (try_succ, try_expected, notify_expected, expected) in tests {
|
||||
let mut b = ExponentialBackoff {
|
||||
current_interval: Duration::from_millis(100),
|
||||
initial_interval: Duration::from_millis(100),
|
||||
max_elapsed_time: Some(Duration::from_millis(200)),
|
||||
randomization_factor: 0.0,
|
||||
multiplier: 1.0,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let mut notify_count = 0;
|
||||
let mut try_count = 0;
|
||||
let ret: Result<(), &str> = retry_notify!(
|
||||
b,
|
||||
{
|
||||
try_count += 1;
|
||||
if try_count == try_succ {
|
||||
Ok(())
|
||||
} else {
|
||||
Err("try again")
|
||||
}
|
||||
},
|
||||
|e, duration| {
|
||||
notify_count += 1;
|
||||
println!("{}: {}, {:?}", notify_count, e, duration);
|
||||
}
|
||||
);
|
||||
assert_eq!(ret, expected);
|
||||
assert_eq!(try_count, try_expected);
|
||||
assert_eq!(notify_count, notify_expected);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
@ -1,6 +1,7 @@
|
|||
use crate::config::{Config, ServerConfig, ServerServiceConfig, ServiceType, TransportType};
|
||||
use crate::config_watcher::ServiceChange;
|
||||
use crate::constants::{listen_backoff, UDP_BUFFER_SIZE};
|
||||
use crate::helper::retry_notify;
|
||||
use crate::multi_map::MultiMap;
|
||||
use crate::protocol::Hello::{ControlChannelHello, DataChannelHello};
|
||||
use crate::protocol::{
|
||||
|
@ -508,17 +509,21 @@ fn tcp_listen_and_send(
|
|||
let (tx, rx) = mpsc::channel(CHAN_SIZE);
|
||||
|
||||
tokio::spawn(async move {
|
||||
// FIXME: Respect shutdown signal
|
||||
let l = backoff::future::retry_notify(listen_backoff(), || async {
|
||||
Ok(TcpListener::bind(&addr).await?)
|
||||
let l = retry_notify!(listen_backoff(), {
|
||||
match shutdown_rx.try_recv() {
|
||||
Err(broadcast::error::TryRecvError::Closed) => Ok(None),
|
||||
_ => TcpListener::bind(&addr).await.map(Some)
|
||||
}
|
||||
}, |e, duration| {
|
||||
error!("{:#}. Retry in {:?}", e, duration);
|
||||
})
|
||||
.await
|
||||
.with_context(|| "Failed to listen for the service");
|
||||
|
||||
let l: TcpListener = match l {
|
||||
Ok(v) => v,
|
||||
Ok(v) => match v {
|
||||
Some(v) => v,
|
||||
None => return
|
||||
},
|
||||
Err(e) => {
|
||||
error!("{:#}", e);
|
||||
return;
|
||||
|
@ -623,20 +628,27 @@ async fn run_udp_connection_pool<T: Transport>(
|
|||
) -> Result<()> {
|
||||
// TODO: Load balance
|
||||
|
||||
// FIXME: Respect shutdown signal
|
||||
let l: UdpSocket = backoff::future::retry_notify(
|
||||
let l = retry_notify!(
|
||||
listen_backoff(),
|
||||
|| async {
|
||||
Ok(UdpSocket::bind(&bind_addr)
|
||||
.await
|
||||
.with_context(|| "Failed to listen for the service")?)
|
||||
{
|
||||
match shutdown_rx.try_recv() {
|
||||
Err(broadcast::error::TryRecvError::Closed) => Ok(None),
|
||||
_ => UdpSocket::bind(&bind_addr).await.map(Some),
|
||||
}
|
||||
},
|
||||
|e, duration| {
|
||||
warn!("{:#}. Retry in {:?}", e, duration);
|
||||
},
|
||||
}
|
||||
)
|
||||
.await
|
||||
.with_context(|| "Failed to listen for the service")?;
|
||||
.with_context(|| "Failed to listen for the service");
|
||||
|
||||
let l = match l {
|
||||
Ok(v) => match v {
|
||||
Some(l) => l,
|
||||
None => return Ok(()),
|
||||
},
|
||||
Err(e) => return Err(e),
|
||||
};
|
||||
|
||||
info!("Listening at {}", &bind_addr);
|
||||
|
||||
|
|
Loading…
Reference in New Issue