use std::collections::BTreeSet;
use std::fmt;
use std::future::Future;
use std::net::SocketAddr;
use std::time::Duration;
use crate::candidate::PeerCandidates;
use crate::dial::{dial_order, NoCommonFamily};
use crate::family::Family;
use crate::local::LocalStack;
#[derive(Debug, Clone, Copy)]
pub struct DialConfig {
pub per_attempt_timeout: Duration,
pub attempt_delay: Duration,
}
impl Default for DialConfig {
fn default() -> Self {
DialConfig {
per_attempt_timeout: Duration::from_secs(10),
attempt_delay: Duration::from_millis(250),
}
}
}
#[derive(Debug)]
pub struct DialWinner<C> {
pub conn: C,
pub addr: SocketAddr,
pub family: Family,
}
#[derive(Debug)]
pub enum ConnectError<E> {
NoCommonFamily(NoCommonFamily),
AllFailed(Vec<(SocketAddr, E)>),
}
impl<E: fmt::Display> fmt::Display for ConnectError<E> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
ConnectError::NoCommonFamily(e) => write!(f, "{e}"),
ConnectError::AllFailed(errs) => {
let joined = errs
.iter()
.map(|(a, e)| format!("{a}: {e}"))
.collect::<Vec<_>>()
.join("; ");
write!(f, "all candidates failed: [{joined}]")
}
}
}
}
impl<E: fmt::Display + fmt::Debug> std::error::Error for ConnectError<E> {}
pub async fn connect<C, E, F, Fut>(
local: &LocalStack,
peer: &PeerCandidates,
config: DialConfig,
dial_fn: F,
) -> Result<DialWinner<C>, ConnectError<E>>
where
E: fmt::Display + FromTimeout,
F: Fn(SocketAddr) -> Fut + Sync,
Fut: Future<Output = Result<C, E>> + Send,
C: Send,
{
let ordered = dial_order(local, peer).map_err(ConnectError::NoCommonFamily)?;
let total = ordered.len();
type Attempt<'f, U, Er> =
std::pin::Pin<Box<dyn Future<Output = (usize, SocketAddr, Result<U, Er>)> + Send + 'f>>;
let mut attempts: futures::stream::FuturesUnordered<Attempt<'_, C, E>> =
futures::stream::FuturesUnordered::new();
let mut live: BTreeSet<usize> = BTreeSet::new();
let mut errors: Vec<(SocketAddr, E)> = Vec::with_capacity(total);
let mut next_prio = 0usize;
let mut best_success: Option<(usize, SocketAddr, C)> = None;
macro_rules! launch {
() => {
if next_prio < total {
let prio = next_prio;
let addr = ordered[prio];
next_prio += 1;
live.insert(prio);
let fut = &dial_fn;
attempts.push(Box::pin(async move {
let res =
match tokio::time::timeout(config.per_attempt_timeout, fut(addr)).await {
Ok(Ok(conn)) => Ok(conn),
Ok(Err(e)) => Err(e),
Err(_) => Err(TimedOut.into_e()),
};
(prio, addr, res)
}));
}
};
}
launch!();
loop {
if let Some((p, _, _)) = &best_success {
let more_preferred_live = live.iter().next().map(|lo| *lo < *p).unwrap_or(false);
let more_preferred_unlaunched = next_prio <= *p;
if !more_preferred_live && !more_preferred_unlaunched {
let (_, addr, conn) = best_success.take().unwrap();
return Ok(DialWinner {
conn,
addr,
family: Family::of(&addr),
});
}
}
if live.is_empty() && next_prio >= total {
break;
}
let stagger = tokio::time::sleep(config.attempt_delay);
tokio::select! {
biased;
finished = futures::StreamExt::next(&mut attempts), if !live.is_empty() => {
match finished {
Some((prio, addr, Ok(conn))) => {
live.remove(&prio);
if prio == 0 {
return Ok(DialWinner { conn, addr, family: Family::of(&addr) });
}
let keep = best_success.as_ref().map(|(bp, _, _)| prio < *bp).unwrap_or(true);
if keep {
best_success = Some((prio, addr, conn));
}
launch!();
}
Some((prio, addr, Err(e))) => {
live.remove(&prio);
errors.push((addr, e));
launch!();
}
None => break,
}
}
_ = stagger, if !live.is_empty() && next_prio < total => {
launch!();
}
}
}
if let Some((_, addr, conn)) = best_success {
return Ok(DialWinner {
conn,
addr,
family: Family::of(&addr),
});
}
Err(ConnectError::AllFailed(errors))
}
struct TimedOut;
impl TimedOut {
fn into_e<E: FromTimeout>(self) -> E {
E::timed_out()
}
}
pub trait FromTimeout {
fn timed_out() -> Self;
}
impl FromTimeout for String {
fn timed_out() -> Self {
"connect attempt timed out".to_string()
}
}
impl FromTimeout for std::io::Error {
fn timed_out() -> Self {
std::io::Error::new(std::io::ErrorKind::TimedOut, "connect attempt timed out")
}
}