use dashmap::DashMap;
use libp2p::{Multiaddr, PeerId};
use parking_lot::RwLock;
use std::sync::Arc;
use std::time::{Duration, Instant};
use thiserror::Error;
use tracing::{debug, info, warn};
#[derive(Debug, Error)]
pub enum MigrationError {
#[error("No active connection to migrate")]
NoActiveConnection,
#[error("Migration already in progress")]
MigrationInProgress,
#[error("Migration failed: {0}")]
MigrationFailed(String),
#[error("Timeout during migration")]
MigrationTimeout,
#[error("Invalid migration state")]
InvalidState,
#[error("No suitable migration path available")]
NoMigrationPath,
}
#[derive(Debug, Clone)]
pub struct MigrationConfig {
pub auto_migrate: bool,
pub migration_timeout: Duration,
pub max_retry_attempts: usize,
pub retry_backoff: Duration,
pub migration_cooldown: Duration,
pub keep_old_path: bool,
pub validate_new_path: bool,
}
impl Default for MigrationConfig {
fn default() -> Self {
Self {
auto_migrate: true,
migration_timeout: Duration::from_secs(30),
max_retry_attempts: 3,
retry_backoff: Duration::from_secs(2),
migration_cooldown: Duration::from_secs(10),
keep_old_path: true,
validate_new_path: true,
}
}
}
impl MigrationConfig {
pub fn mobile() -> Self {
Self {
auto_migrate: true,
migration_timeout: Duration::from_secs(15),
max_retry_attempts: 5,
retry_backoff: Duration::from_millis(500),
migration_cooldown: Duration::from_secs(5),
keep_old_path: true,
validate_new_path: true,
}
}
pub fn conservative() -> Self {
Self {
auto_migrate: false,
migration_timeout: Duration::from_secs(60),
max_retry_attempts: 2,
retry_backoff: Duration::from_secs(5),
migration_cooldown: Duration::from_secs(30),
keep_old_path: true,
validate_new_path: true,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MigrationState {
Idle,
Initiated,
Validating,
Migrating,
Completed,
Failed,
}
#[derive(Debug, Clone)]
pub struct MigrationAttempt {
pub peer_id: PeerId,
pub old_address: Multiaddr,
pub new_address: Multiaddr,
pub state: MigrationState,
pub started_at: Instant,
pub retry_count: usize,
pub error: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct MigrationStats {
pub total_attempts: usize,
pub successful_migrations: usize,
pub failed_migrations: usize,
pub in_progress: usize,
pub avg_duration_ms: u64,
pub total_retries: usize,
}
pub struct ConnectionMigrationManager {
config: MigrationConfig,
active_migrations: Arc<DashMap<PeerId, MigrationAttempt>>,
last_migration: Arc<DashMap<PeerId, Instant>>,
stats: Arc<RwLock<MigrationStats>>,
migration_durations: Arc<RwLock<Vec<u64>>>,
}
impl ConnectionMigrationManager {
pub fn new(config: MigrationConfig) -> Self {
Self {
config,
active_migrations: Arc::new(DashMap::new()),
last_migration: Arc::new(DashMap::new()),
stats: Arc::new(RwLock::new(MigrationStats::default())),
migration_durations: Arc::new(RwLock::new(Vec::new())),
}
}
pub fn mobile() -> Self {
Self::new(MigrationConfig::mobile())
}
pub fn conservative() -> Self {
Self::new(MigrationConfig::conservative())
}
pub fn initiate_migration(
&self,
peer_id: PeerId,
old_address: Multiaddr,
new_address: Multiaddr,
) -> Result<(), MigrationError> {
if self.active_migrations.contains_key(&peer_id) {
return Err(MigrationError::MigrationInProgress);
}
if let Some(last) = self.last_migration.get(&peer_id) {
if last.elapsed() < self.config.migration_cooldown {
debug!(
"Migration cooldown active for peer {} ({:?} remaining)",
peer_id,
self.config.migration_cooldown - last.elapsed()
);
return Err(MigrationError::InvalidState);
}
}
info!(
"Initiating connection migration for peer {} from {} to {}",
peer_id, old_address, new_address
);
let attempt = MigrationAttempt {
peer_id,
old_address,
new_address,
state: MigrationState::Initiated,
started_at: Instant::now(),
retry_count: 0,
error: None,
};
self.active_migrations.insert(peer_id, attempt);
let mut stats = self.stats.write();
stats.total_attempts += 1;
stats.in_progress += 1;
Ok(())
}
pub fn update_migration_state(
&self,
peer_id: &PeerId,
new_state: MigrationState,
) -> Result<(), MigrationError> {
if let Some(mut attempt) = self.active_migrations.get_mut(peer_id) {
debug!(
"Migration state change for peer {}: {:?} -> {:?}",
peer_id, attempt.state, new_state
);
attempt.state = new_state;
Ok(())
} else {
Err(MigrationError::NoActiveConnection)
}
}
pub fn complete_migration(&self, peer_id: &PeerId) -> Result<(), MigrationError> {
if let Some((_, attempt)) = self.active_migrations.remove(peer_id) {
let duration_ms = attempt.started_at.elapsed().as_millis() as u64;
info!(
"Migration completed for peer {} in {} ms",
peer_id, duration_ms
);
self.last_migration.insert(*peer_id, Instant::now());
{
let mut stats = self.stats.write();
stats.successful_migrations += 1;
stats.in_progress = stats.in_progress.saturating_sub(1);
stats.total_retries += attempt.retry_count;
}
{
let mut durations = self.migration_durations.write();
durations.push(duration_ms);
if durations.len() > 100 {
durations.remove(0);
}
let avg = durations.iter().sum::<u64>() / durations.len() as u64;
self.stats.write().avg_duration_ms = avg;
}
Ok(())
} else {
Err(MigrationError::NoActiveConnection)
}
}
pub fn fail_migration(&self, peer_id: &PeerId, error: String) -> Result<(), MigrationError> {
if let Some((_, mut attempt)) = self.active_migrations.remove(peer_id) {
warn!("Migration failed for peer {}: {}", peer_id, error);
attempt.error = Some(error);
attempt.state = MigrationState::Failed;
let mut stats = self.stats.write();
stats.failed_migrations += 1;
stats.in_progress = stats.in_progress.saturating_sub(1);
stats.total_retries += attempt.retry_count;
Ok(())
} else {
Err(MigrationError::NoActiveConnection)
}
}
pub fn retry_migration(&self, peer_id: &PeerId) -> Result<(), MigrationError> {
if let Some(mut attempt) = self.active_migrations.get_mut(peer_id) {
if attempt.retry_count >= self.config.max_retry_attempts {
return Err(MigrationError::MigrationFailed(
"Max retry attempts reached".to_string(),
));
}
attempt.retry_count += 1;
attempt.state = MigrationState::Initiated;
attempt.error = None;
info!(
"Retrying migration for peer {} (attempt {})",
peer_id,
attempt.retry_count + 1
);
Ok(())
} else {
Err(MigrationError::NoActiveConnection)
}
}
pub fn is_migrating(&self, peer_id: &PeerId) -> bool {
self.active_migrations.contains_key(peer_id)
}
pub fn get_migration_state(&self, peer_id: &PeerId) -> Option<MigrationState> {
self.active_migrations
.get(peer_id)
.map(|attempt| attempt.state)
}
pub fn get_active_migrations(&self) -> Vec<MigrationAttempt> {
self.active_migrations
.iter()
.map(|entry| entry.value().clone())
.collect()
}
pub fn stats(&self) -> MigrationStats {
self.stats.read().clone()
}
pub fn can_migrate(&self, peer_id: &PeerId) -> bool {
if self.active_migrations.contains_key(peer_id) {
return false;
}
if let Some(last) = self.last_migration.get(peer_id) {
last.elapsed() >= self.config.migration_cooldown
} else {
true
}
}
pub fn config(&self) -> &MigrationConfig {
&self.config
}
pub fn cleanup_timeouts(&self) {
let timeout = self.config.migration_timeout;
let mut timed_out = Vec::new();
for entry in self.active_migrations.iter() {
if entry.value().started_at.elapsed() > timeout {
timed_out.push(*entry.key());
}
}
for peer_id in timed_out {
self.fail_migration(&peer_id, "Migration timeout".to_string())
.ok();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::str::FromStr;
fn test_peer_id() -> PeerId {
PeerId::random()
}
fn test_addr() -> Multiaddr {
Multiaddr::from_str("/ip4/127.0.0.1/tcp/4001").expect("test: valid multiaddr literal")
}
fn test_addr2() -> Multiaddr {
Multiaddr::from_str("/ip4/192.168.1.1/tcp/4001").expect("test: valid multiaddr literal")
}
#[test]
fn test_migration_config_default() {
let config = MigrationConfig::default();
assert!(config.auto_migrate);
assert!(config.keep_old_path);
assert!(config.validate_new_path);
}
#[test]
fn test_migration_config_mobile() {
let config = MigrationConfig::mobile();
assert!(config.auto_migrate);
assert_eq!(config.max_retry_attempts, 5);
assert_eq!(config.migration_cooldown, Duration::from_secs(5));
}
#[test]
fn test_migration_config_conservative() {
let config = MigrationConfig::conservative();
assert!(!config.auto_migrate);
assert_eq!(config.max_retry_attempts, 2);
assert_eq!(config.migration_cooldown, Duration::from_secs(30));
}
#[test]
fn test_initiate_migration() {
let manager = ConnectionMigrationManager::new(MigrationConfig::default());
let peer = test_peer_id();
let old_addr = test_addr();
let new_addr = test_addr2();
let result = manager.initiate_migration(peer, old_addr, new_addr);
assert!(result.is_ok());
let stats = manager.stats();
assert_eq!(stats.total_attempts, 1);
assert_eq!(stats.in_progress, 1);
}
#[test]
fn test_migration_in_progress_error() {
let manager = ConnectionMigrationManager::new(MigrationConfig::default());
let peer = test_peer_id();
let old_addr = test_addr();
let new_addr = test_addr2();
manager
.initiate_migration(peer, old_addr.clone(), new_addr.clone())
.expect("test: first migration initiation should succeed");
let result = manager.initiate_migration(peer, old_addr, new_addr);
assert!(matches!(result, Err(MigrationError::MigrationInProgress)));
}
#[test]
fn test_complete_migration() {
let manager = ConnectionMigrationManager::new(MigrationConfig::default());
let peer = test_peer_id();
manager
.initiate_migration(peer, test_addr(), test_addr2())
.expect("test: migration initiation should succeed");
let result = manager.complete_migration(&peer);
assert!(result.is_ok());
let stats = manager.stats();
assert_eq!(stats.successful_migrations, 1);
assert_eq!(stats.in_progress, 0);
}
#[test]
fn test_fail_migration() {
let manager = ConnectionMigrationManager::new(MigrationConfig::default());
let peer = test_peer_id();
manager
.initiate_migration(peer, test_addr(), test_addr2())
.expect("test: migration initiation should succeed");
let result = manager.fail_migration(&peer, "Test error".to_string());
assert!(result.is_ok());
let stats = manager.stats();
assert_eq!(stats.failed_migrations, 1);
assert_eq!(stats.in_progress, 0);
}
#[test]
fn test_retry_migration() {
let manager = ConnectionMigrationManager::new(MigrationConfig::default());
let peer = test_peer_id();
manager
.initiate_migration(peer, test_addr(), test_addr2())
.expect("test: migration initiation should succeed");
let result = manager.retry_migration(&peer);
assert!(result.is_ok());
let attempt = manager
.active_migrations
.get(&peer)
.expect("test: active migration entry should exist");
assert_eq!(attempt.retry_count, 1);
}
#[test]
fn test_retry_limit() {
let config = MigrationConfig {
max_retry_attempts: 2,
..Default::default()
};
let manager = ConnectionMigrationManager::new(config);
let peer = test_peer_id();
manager
.initiate_migration(peer, test_addr(), test_addr2())
.expect("test: migration initiation should succeed");
assert!(manager.retry_migration(&peer).is_ok());
assert!(manager.retry_migration(&peer).is_ok());
assert!(matches!(
manager.retry_migration(&peer),
Err(MigrationError::MigrationFailed(_))
));
}
#[test]
fn test_is_migrating() {
let manager = ConnectionMigrationManager::new(MigrationConfig::default());
let peer = test_peer_id();
assert!(!manager.is_migrating(&peer));
manager
.initiate_migration(peer, test_addr(), test_addr2())
.expect("test: migration initiation should succeed");
assert!(manager.is_migrating(&peer));
manager
.complete_migration(&peer)
.expect("test: complete_migration should succeed");
assert!(!manager.is_migrating(&peer));
}
#[test]
fn test_migration_state_updates() {
let manager = ConnectionMigrationManager::new(MigrationConfig::default());
let peer = test_peer_id();
manager
.initiate_migration(peer, test_addr(), test_addr2())
.expect("test: migration initiation should succeed");
assert_eq!(
manager.get_migration_state(&peer),
Some(MigrationState::Initiated)
);
manager
.update_migration_state(&peer, MigrationState::Validating)
.expect("test: update_migration_state to Validating should succeed");
assert_eq!(
manager.get_migration_state(&peer),
Some(MigrationState::Validating)
);
manager
.update_migration_state(&peer, MigrationState::Migrating)
.expect("test: update_migration_state to Migrating should succeed");
assert_eq!(
manager.get_migration_state(&peer),
Some(MigrationState::Migrating)
);
}
#[test]
fn test_can_migrate() {
let config = MigrationConfig {
migration_cooldown: Duration::from_millis(100),
..Default::default()
};
let manager = ConnectionMigrationManager::new(config);
let peer = test_peer_id();
assert!(manager.can_migrate(&peer));
manager
.initiate_migration(peer, test_addr(), test_addr2())
.expect("test: migration initiation should succeed");
assert!(!manager.can_migrate(&peer));
manager
.complete_migration(&peer)
.expect("test: complete_migration should succeed");
assert!(!manager.can_migrate(&peer));
std::thread::sleep(Duration::from_millis(150));
assert!(manager.can_migrate(&peer)); }
#[test]
fn test_get_active_migrations() {
let manager = ConnectionMigrationManager::new(MigrationConfig::default());
let peer1 = test_peer_id();
let peer2 = test_peer_id();
manager
.initiate_migration(peer1, test_addr(), test_addr2())
.expect("test: peer1 migration initiation should succeed");
manager
.initiate_migration(peer2, test_addr(), test_addr2())
.expect("test: peer2 migration initiation should succeed");
let active = manager.get_active_migrations();
assert_eq!(active.len(), 2);
}
#[test]
fn test_average_duration_calculation() {
let manager = ConnectionMigrationManager::new(MigrationConfig::default());
let peer1 = test_peer_id();
let peer2 = test_peer_id();
manager
.initiate_migration(peer1, test_addr(), test_addr2())
.expect("test: peer1 migration initiation should succeed");
std::thread::sleep(Duration::from_millis(10));
manager
.complete_migration(&peer1)
.expect("test: peer1 complete_migration should succeed");
manager
.initiate_migration(peer2, test_addr(), test_addr2())
.expect("test: peer2 migration initiation should succeed");
std::thread::sleep(Duration::from_millis(10));
manager
.complete_migration(&peer2)
.expect("test: peer2 complete_migration should succeed");
let stats = manager.stats();
assert!(stats.avg_duration_ms > 0);
}
}