use crate::transport::{
Connection, Transport, TransportError, TransportSelectionStrategy, TransportSelector,
TransportType,
};
use dashmap::DashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use tracing::{debug, info, warn};
pub struct ConnectionAttempt {
pub connection: Option<Box<dyn Connection>>,
pub transport_type: TransportType,
pub fallback: bool,
pub error: Option<TransportError>,
}
#[derive(Debug, Clone)]
pub struct MultiTransportConfig {
pub strategy: TransportSelectionStrategy,
pub enable_tcp_fallback: bool,
pub max_connection_attempts: usize,
pub remember_transport_per_peer: bool,
}
impl Default for MultiTransportConfig {
fn default() -> Self {
Self {
strategy: TransportSelectionStrategy::PreferType(TransportType::Quic),
enable_tcp_fallback: true,
max_connection_attempts: 3,
remember_transport_per_peer: true,
}
}
}
pub struct MultiTransportManager {
selector: Arc<TransportSelector>,
config: MultiTransportConfig,
peer_transports: Arc<DashMap<SocketAddr, TransportType>>,
}
impl MultiTransportManager {
pub fn new(config: MultiTransportConfig) -> Self {
let selector = TransportSelector::new(config.strategy);
Self {
selector: Arc::new(selector),
config,
peer_transports: Arc::new(DashMap::new()),
}
}
pub fn register_transport(&mut self, transport: Box<dyn Transport>) {
let transport_type = transport.transport_type();
info!("Registering transport: {}", transport_type);
let selector = Arc::get_mut(&mut self.selector)
.expect("Cannot register transport after manager is shared");
selector.register(transport);
}
pub fn available_transports(&self) -> Vec<TransportType> {
self.selector.available_transports()
}
pub async fn connect(&self, addr: SocketAddr) -> Result<Box<dyn Connection>, TransportError> {
if self.config.remember_transport_per_peer {
if let Some(entry) = self.peer_transports.get(&addr) {
let preferred_type = *entry;
debug!("Using remembered transport {} for {}", preferred_type, addr);
if let Some(transport) = self.find_transport(preferred_type) {
match transport.connect(addr).await {
Ok(conn) => {
info!(
"Connected to {} using remembered transport {}",
addr, preferred_type
);
return Ok(conn);
}
Err(e) => {
warn!(
"Remembered transport {} failed for {}: {}",
preferred_type, addr, e
);
self.peer_transports.remove(&addr);
}
}
}
}
}
let mut attempt_count = 0;
let mut last_error = None;
let transports_to_try = self.get_ordered_transports();
for transport_type in transports_to_try {
if attempt_count >= self.config.max_connection_attempts {
break;
}
attempt_count += 1;
if let Some(transport) = self.find_transport(transport_type) {
debug!(
"Attempting connection to {} via {} (attempt {}/{})",
addr, transport_type, attempt_count, self.config.max_connection_attempts
);
match transport.connect(addr).await {
Ok(conn) => {
info!("Successfully connected to {} via {}", addr, transport_type);
if self.config.remember_transport_per_peer {
self.peer_transports.insert(addr, transport_type);
}
return Ok(conn);
}
Err(e) => {
warn!(
"Connection to {} via {} failed: {}",
addr, transport_type, e
);
last_error = Some(e);
}
}
}
}
Err(last_error
.unwrap_or_else(|| TransportError::NotAvailable("No transports available".to_string())))
}
pub async fn connect_detailed(&self, addr: SocketAddr) -> ConnectionAttempt {
match self.connect(addr).await {
Ok(connection) => {
let transport_type = connection.transport_type();
ConnectionAttempt {
connection: Some(connection),
transport_type,
fallback: transport_type == TransportType::Tcp,
error: None,
}
}
Err(error) => ConnectionAttempt {
connection: None,
transport_type: TransportType::Tcp, fallback: false,
error: Some(error),
},
}
}
fn get_ordered_transports(&self) -> Vec<TransportType> {
let mut transports = Vec::new();
if let Some(transport) = self.selector.select() {
transports.push(transport.transport_type());
}
if self.config.enable_tcp_fallback {
let available = self.available_transports();
for &transport_type in &[
TransportType::Quic,
TransportType::WebTransport,
TransportType::Tcp,
TransportType::WebSocket,
] {
if !transports.contains(&transport_type) && available.contains(&transport_type) {
transports.push(transport_type);
}
}
}
transports
}
fn find_transport(&self, transport_type: TransportType) -> Option<&dyn Transport> {
self.selector
.select()
.filter(|t| t.transport_type() == transport_type)
}
pub fn clear_peer_memory(&self) {
self.peer_transports.clear();
info!("Cleared peer transport memory");
}
pub fn forget_peer(&self, addr: &SocketAddr) {
self.peer_transports.remove(addr);
debug!("Forgot transport for peer {}", addr);
}
pub fn transport_usage_stats(&self) -> Vec<(TransportType, usize)> {
let mut stats = std::collections::HashMap::new();
for entry in self.peer_transports.iter() {
*stats.entry(*entry.value()).or_insert(0) += 1;
}
let mut result: Vec<_> = stats.into_iter().collect();
result.sort_by_key(|(_, count)| std::cmp::Reverse(*count));
result
}
}
pub struct MultiTransportManagerBuilder {
config: MultiTransportConfig,
transports: Vec<Box<dyn Transport>>,
}
impl MultiTransportManagerBuilder {
pub fn new() -> Self {
Self {
config: MultiTransportConfig::default(),
transports: Vec::new(),
}
}
pub fn strategy(mut self, strategy: TransportSelectionStrategy) -> Self {
self.config.strategy = strategy;
self
}
pub fn enable_tcp_fallback(mut self, enable: bool) -> Self {
self.config.enable_tcp_fallback = enable;
self
}
pub fn max_attempts(mut self, max: usize) -> Self {
self.config.max_connection_attempts = max;
self
}
pub fn remember_transports(mut self, remember: bool) -> Self {
self.config.remember_transport_per_peer = remember;
self
}
pub fn add_transport(mut self, transport: Box<dyn Transport>) -> Self {
self.transports.push(transport);
self
}
pub fn build(self) -> MultiTransportManager {
let mut manager = MultiTransportManager::new(self.config);
for transport in self.transports {
manager.register_transport(transport);
}
manager
}
}
impl Default for MultiTransportManagerBuilder {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_config_default() {
let config = MultiTransportConfig::default();
assert!(config.enable_tcp_fallback);
assert_eq!(config.max_connection_attempts, 3);
assert!(config.remember_transport_per_peer);
}
#[test]
fn test_builder() {
let builder = MultiTransportManagerBuilder::new()
.strategy(TransportSelectionStrategy::FirstAvailable)
.enable_tcp_fallback(false)
.max_attempts(5)
.remember_transports(false);
let manager = builder.build();
assert!(!manager.config.enable_tcp_fallback);
assert_eq!(manager.config.max_connection_attempts, 5);
assert!(!manager.config.remember_transport_per_peer);
}
#[test]
fn test_peer_memory() {
let manager = MultiTransportManager::new(MultiTransportConfig::default());
let addr: SocketAddr = "127.0.0.1:8080".parse().expect("test: valid socket addr");
manager.peer_transports.insert(addr, TransportType::Quic);
let stats = manager.transport_usage_stats();
assert_eq!(stats.len(), 1);
assert_eq!(stats[0].0, TransportType::Quic);
assert_eq!(stats[0].1, 1);
manager.forget_peer(&addr);
let stats = manager.transport_usage_stats();
assert_eq!(stats.len(), 0);
}
#[test]
fn test_transport_ordering() {
let config = MultiTransportConfig {
strategy: TransportSelectionStrategy::PreferType(TransportType::Quic),
enable_tcp_fallback: true,
max_connection_attempts: 3,
remember_transport_per_peer: false,
};
let manager = MultiTransportManager::new(config);
let ordered = manager.get_ordered_transports();
assert!(!ordered.is_empty() || ordered.is_empty()); }
}