use std::fmt;
use std::future::Future;
use std::net::SocketAddr;
use std::time::Duration;
use futures::StreamExt;
use futures::stream::FuturesUnordered;
use crate::resolve::Candidates;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Failure<E> {
pub addr: SocketAddr,
pub error: E,
}
impl<E: fmt::Display> fmt::Display for Failure<E> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}: {}", self.addr, self.error)
}
}
impl<E: std::error::Error + 'static> std::error::Error for Failure<E> {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
Some(&self.error)
}
}
pub(crate) trait Aggregate: Sized {
fn aggregate(failures: Vec<Failure<Self>>) -> Self;
fn resolve(error: Option<std::io::Error>) -> Self;
}
pub(crate) fn describe<E: fmt::Display>(failures: &[Failure<E>]) -> String {
failures.iter().map(|f| f.to_string()).collect::<Vec<_>>().join("; ")
}
pub(crate) async fn race<C, E, F, Fut>(mut candidates: Candidates, delay: Duration, mut dial: F) -> Result<C, E>
where
F: FnMut(SocketAddr) -> Fut,
Fut: Future<Output = Result<C, E>>,
E: Aggregate + fmt::Display,
{
let mut attempts = FuturesUnordered::new();
let mut failures: Vec<(usize, Failure<E>)> = Vec::new();
let mut exhausted = false;
let mut ready = tokio::time::Instant::now();
let mut next_index = 0;
let mut start = |addr: SocketAddr, attempts: &mut FuturesUnordered<_>| {
let index = next_index;
next_index += 1;
tracing::debug!(%addr, index, "dialing");
let attempt = dial(addr);
attempts.push(async move { (index, addr, attempt.await) });
};
loop {
if exhausted && attempts.is_empty() {
if failures.is_empty() {
return Err(E::resolve(candidates.failure()));
}
failures.sort_by_key(|(index, _)| *index);
return Err(collapse(failures.into_iter().map(|(_, failure)| failure).collect()));
}
tokio::select! {
biased;
Some((index, addr, res)) = attempts.next(), if !attempts.is_empty() => {
match res {
Ok(conn) => {
tracing::debug!(%addr, index, "connected");
return Ok(conn);
}
Err(err) => {
tracing::debug!(%addr, index, %err, "connection attempt failed");
failures.push((index, Failure { addr, error: err }));
ready = tokio::time::Instant::now();
}
}
}
addr = pull(&mut candidates, ready), if !exhausted => {
match addr {
Some(addr) => {
start(addr, &mut attempts);
ready = tokio::time::Instant::now() + delay;
}
None => exhausted = true,
}
}
}
}
}
async fn pull(candidates: &mut Candidates, ready: tokio::time::Instant) -> Option<SocketAddr> {
tokio::time::sleep_until(ready).await;
candidates.next().await
}
fn collapse<E: Aggregate>(mut failures: Vec<Failure<E>>) -> E {
match failures.len() {
1 => failures.pop().expect("checked len").error,
_ => E::aggregate(failures),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::client::DEFAULT_FAILOVER_DELAY;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
fn addr(s: &str) -> SocketAddr {
s.parse().unwrap()
}
#[derive(Debug, PartialEq, Eq)]
enum TestError {
Dial(&'static str),
All(Vec<Failure<TestError>>),
}
impl fmt::Display for TestError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Dial(err) => write!(f, "{err}"),
Self::All(failures) => write!(f, "all {} attempts failed: {}", failures.len(), describe(failures)),
}
}
}
impl Aggregate for TestError {
fn aggregate(failures: Vec<Failure<Self>>) -> Self {
Self::All(failures)
}
fn resolve(error: Option<std::io::Error>) -> Self {
match error {
Some(_) => Self::Dial("lookup failed"),
None => Self::Dial("no addresses"),
}
}
}
fn failed(dest: &str, err: &'static str) -> Failure<TestError> {
Failure {
addr: addr(dest),
error: TestError::Dial(err),
}
}
#[tokio::test(start_paused = true)]
async fn first_success_returns_immediately() {
let dials = Arc::new(AtomicUsize::new(0));
let counter = dials.clone();
let res: Result<&str, TestError> = race(
Candidates::fixed([addr("1.1.1.1:1"), addr("2.2.2.2:2")]),
DEFAULT_FAILOVER_DELAY,
move |_| {
counter.fetch_add(1, Ordering::SeqCst);
async { Ok("winner") }
},
)
.await;
assert_eq!(res, Ok("winner"));
assert_eq!(dials.load(Ordering::SeqCst), 1, "no second dial after a fast success");
}
#[tokio::test(start_paused = true)]
async fn second_wins_when_first_hangs() {
let start = tokio::time::Instant::now();
let res: Result<&str, TestError> = race(
Candidates::fixed([addr("1.1.1.1:1"), addr("2.2.2.2:2")]),
DEFAULT_FAILOVER_DELAY,
|dest| async move {
if dest == addr("1.1.1.1:1") {
std::future::pending().await
} else {
Ok("second")
}
},
)
.await;
assert_eq!(res, Ok("second"));
assert_eq!(
start.elapsed(),
DEFAULT_FAILOVER_DELAY,
"second dial waits out the stagger"
);
}
#[tokio::test(start_paused = true)]
async fn failure_starts_the_next_attempt_immediately() {
let start = tokio::time::Instant::now();
let res: Result<&str, TestError> = race(
Candidates::fixed([addr("1.1.1.1:1"), addr("2.2.2.2:2")]),
DEFAULT_FAILOVER_DELAY,
|dest| async move {
if dest == addr("1.1.1.1:1") {
Err(TestError::Dial("boom"))
} else {
Ok("second")
}
},
)
.await;
assert_eq!(res, Ok("second"));
assert_eq!(start.elapsed(), Duration::ZERO, "failure must not wait for the timer");
}
#[tokio::test(start_paused = true)]
async fn all_failures_are_reported_when_the_preferred_fails_first() {
let res: Result<&str, TestError> = race(
Candidates::fixed([addr("1.1.1.1:1"), addr("2.2.2.2:2")]),
Duration::from_millis(10),
|dest| async move {
if dest == addr("1.1.1.1:1") {
Err(TestError::Dial("network unreachable"))
} else {
tokio::time::sleep(Duration::from_secs(1)).await;
Err(TestError::Dial("invalid peer certificate"))
}
},
)
.await;
assert_eq!(
res,
Err(TestError::All(vec![
failed("1.1.1.1:1", "network unreachable"),
failed("2.2.2.2:2", "invalid peer certificate"),
]))
);
}
#[tokio::test(start_paused = true)]
async fn all_failures_are_reported_when_the_preferred_times_out_last() {
let res: Result<&str, TestError> = race(
Candidates::fixed([addr("1.1.1.1:1"), addr("2.2.2.2:2")]),
Duration::from_millis(10),
|dest| async move {
if dest == addr("1.1.1.1:1") {
tokio::time::sleep(Duration::from_secs(30)).await;
Err(TestError::Dial("timed out"))
} else {
Err(TestError::Dial("invalid peer certificate"))
}
},
)
.await;
assert_eq!(
res,
Err(TestError::All(vec![
failed("1.1.1.1:1", "timed out"),
failed("2.2.2.2:2", "invalid peer certificate"),
]))
);
}
#[tokio::test(start_paused = true)]
async fn a_lone_failure_is_returned_unwrapped() {
let res: Result<&str, TestError> = race(
Candidates::fixed([addr("1.1.1.1:1")]),
DEFAULT_FAILOVER_DELAY,
|_| async { Err(TestError::Dial("invalid peer certificate")) },
)
.await;
assert_eq!(res, Err(TestError::Dial("invalid peer certificate")));
}
#[test]
fn describe_lists_every_attempt() {
let failures = [failed("1.1.1.1:1", "timed out"), failed("2.2.2.2:2", "bad cert")];
assert_eq!(describe(&failures), "1.1.1.1:1: timed out; 2.2.2.2:2: bad cert");
}
#[tokio::test(start_paused = true)]
async fn losers_are_dropped_on_success() {
struct Guard(Arc<AtomicUsize>);
impl Drop for Guard {
fn drop(&mut self) {
self.0.fetch_add(1, Ordering::SeqCst);
}
}
let dropped = Arc::new(AtomicUsize::new(0));
let count = dropped.clone();
let res: Result<&str, TestError> = race(
Candidates::fixed([addr("1.1.1.1:1"), addr("2.2.2.2:2")]),
Duration::ZERO,
move |dest| {
let guard = Guard(count.clone());
async move {
if dest == addr("1.1.1.1:1") {
let _guard = guard;
std::future::pending().await
} else {
drop(guard);
tokio::time::sleep(Duration::from_millis(1)).await;
Ok("second")
}
}
},
)
.await;
assert_eq!(res, Ok("second"));
assert_eq!(dropped.load(Ordering::SeqCst), 2, "the hung attempt was not aborted");
}
#[tokio::test(start_paused = true)]
async fn zero_delay_dials_all_at_once() {
let start = tokio::time::Instant::now();
let res: Result<&str, TestError> = race(
Candidates::fixed([addr("1.1.1.1:1"), addr("2.2.2.2:2")]),
Duration::ZERO,
|dest| async move {
if dest == addr("1.1.1.1:1") {
std::future::pending().await
} else {
Ok("second")
}
},
)
.await;
assert_eq!(res, Ok("second"));
assert_eq!(start.elapsed(), Duration::ZERO);
}
#[tokio::test(start_paused = true)]
async fn an_empty_resolution_reports_why() {
let res: Result<&str, TestError> = race(Candidates::fixed([]), DEFAULT_FAILOVER_DELAY, |_| async {
unreachable!("dialed without an address")
})
.await;
assert_eq!(res, Err(TestError::Dial("no addresses")));
}
#[tokio::test(start_paused = true)]
async fn dials_the_first_address_to_resolve() {
let start = tokio::time::Instant::now();
let res: Result<&str, TestError> = race(
Candidates::slow(
(&[], Duration::from_secs(30)),
(&[addr("1.1.1.1:1")], Duration::from_millis(100)),
),
DEFAULT_FAILOVER_DELAY,
|_| async { Ok("winner") },
)
.await;
assert_eq!(res, Ok("winner"));
assert_eq!(
start.elapsed(),
Duration::from_millis(100),
"waited for the other query"
);
}
#[tokio::test(start_paused = true)]
async fn a_late_candidate_starts_as_soon_as_it_resolves() {
let start = tokio::time::Instant::now();
let res: Result<&str, TestError> = race(
Candidates::slow(
(&[addr("[2001:db8::1]:1"), addr("1.1.1.1:1")], Duration::from_secs(1)),
(&[addr("1.1.1.1:1")], Duration::ZERO),
),
DEFAULT_FAILOVER_DELAY,
|dest| async move {
match dest.is_ipv6() {
true => Ok("second"),
false => std::future::pending().await,
}
},
)
.await;
assert_eq!(res, Ok("second"));
assert_eq!(start.elapsed(), Duration::from_secs(1));
}
}