use std::collections::HashMap;
use std::sync::{Arc, Mutex, PoisonError};
use std::time::Duration;
use iroh::endpoint::Connection;
use iroh::{Endpoint, EndpointAddr, PublicKey};
use tokio::runtime::Handle;
use tokio_util::task::AbortOnDropHandle;
use tracing::Instrument as _;
use decdn_protocol::ALPN_CLIENT;
use crate::{LocalPullFault, PullTimeout, rate_limited};
pub(crate) async fn dial(
endpoint: &Endpoint,
target: EndpointAddr,
runtime: Option<&Handle>,
stage: &'static str,
) -> anyhow::Result<Connection> {
let Some(runtime) = runtime else {
return endpoint
.connect(target, ALPN_CLIENT)
.await
.map_err(|e| rate_limited::transport_error(stage, e));
};
let endpoint = endpoint.clone();
let connect = async move { endpoint.connect(target, ALPN_CLIENT).await }
.instrument(tracing::Span::current());
AbortOnDropHandle::new(runtime.spawn(connect))
.await
.map_err(|e| {
let outcome = if e.is_panic() {
"panicked"
} else {
"was cancelled by its runtime shutting down"
};
anyhow::anyhow!("{stage}: the dial task {outcome}").context(LocalPullFault)
})?
.map_err(|e| rate_limited::transport_error(stage, e))
}
#[derive(Debug)]
#[doc(hidden)]
pub struct WarmConnection {
conn: Connection,
}
impl WarmConnection {
pub async fn connect(
endpoint: &Endpoint,
target: EndpointAddr,
open: Duration,
dial_runtime: Option<&Handle>,
) -> anyhow::Result<Self> {
let conn = tokio::time::timeout(
open,
dial(endpoint, target, dial_runtime, "warm connect failed"),
)
.await
.map_err(|_| anyhow::Error::new(PullTimeout { after: open }))??;
Ok(Self { conn })
}
pub(crate) const fn connection(&self) -> &Connection {
&self.conn
}
pub(crate) fn is_closed(&self) -> bool {
self.conn.close_reason().is_some()
}
}
#[derive(Debug, Clone)]
pub struct Connections {
inner: Arc<ConnectionsInner>,
}
type Slot = Arc<tokio::sync::Mutex<Option<Arc<WarmConnection>>>>;
#[derive(Debug)]
struct ConnectionsInner {
endpoint: Endpoint,
slots: Mutex<HashMap<PublicKey, Slot>>,
#[cfg(any(test, feature = "test-util"))]
dials: std::sync::atomic::AtomicU64,
}
impl Connections {
#[must_use]
pub fn new(endpoint: Endpoint) -> Self {
Self {
inner: Arc::new(ConnectionsInner {
endpoint,
slots: Mutex::new(HashMap::new()),
#[cfg(any(test, feature = "test-util"))]
dials: std::sync::atomic::AtomicU64::new(0),
}),
}
}
#[doc(hidden)]
pub async fn get(&self, target: &EndpointAddr) -> anyhow::Result<Arc<WarmConnection>> {
let slot = Arc::clone(
self.inner
.slots
.lock()
.unwrap_or_else(PoisonError::into_inner)
.entry(target.id)
.or_default(),
);
let mut held = slot.lock().await;
if let Some(warm) = held.as_ref().filter(|warm| !warm.is_closed()) {
return Ok(Arc::clone(warm));
}
#[cfg(any(test, feature = "test-util"))]
self.inner
.dials
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
let conn = dial(&self.inner.endpoint, target.clone(), None, "connect failed").await?;
let warm = Arc::new(WarmConnection { conn });
*held = Some(Arc::clone(&warm));
Ok(warm)
}
#[doc(hidden)]
pub async fn invalidate(&self, node: PublicKey, conn: &Arc<WarmConnection>) {
let slot = self
.inner
.slots
.lock()
.unwrap_or_else(PoisonError::into_inner)
.get(&node)
.map(Arc::clone);
let Some(slot) = slot else {
return;
};
let mut held = slot.lock().await;
if held
.as_ref()
.is_some_and(|pinned| Arc::ptr_eq(pinned, conn))
{
*held = None;
}
}
#[cfg(any(test, feature = "test-util"))]
#[must_use]
pub fn dials(&self) -> u64 {
self.inner.dials.load(std::sync::atomic::Ordering::SeqCst)
}
}
pub(crate) fn is_transport_fault(err: &anyhow::Error, warm: &WarmConnection) -> bool {
warm.is_closed()
|| err
.chain()
.any(|e| e.is::<PullTimeout>() || e.is::<crate::PullStalled>())
}
#[derive(Debug, Clone)]
pub(crate) enum ConnOwner {
Owned,
#[cfg(any(test, feature = "test-util"))]
Borrowed,
Shared {
warm: Arc<WarmConnection>,
map: Connections,
},
}
impl ConnOwner {
pub(crate) const fn owns(&self) -> bool {
matches!(self, Self::Owned)
}
pub(crate) async fn unpin_on_fault(&self, err: &anyhow::Error) {
if let Self::Shared { warm, map } = self
&& is_transport_fault(err, warm)
{
map.invalidate(warm.connection().remote_id(), warm).await;
}
}
}
impl Drop for WarmConnection {
fn drop(&mut self) {
self.conn.close(0u32.into(), b"warm-connection-dropped");
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)] mod tests;