use std::collections::HashMap;
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct PeerLifetimeConfig {
pub disconnected_timeout: Duration,
pub connected_timeout: Duration,
pub cleanup_interval: Duration,
}
impl Default for PeerLifetimeConfig {
fn default() -> Self {
Self {
disconnected_timeout: Duration::from_secs(30),
connected_timeout: Duration::from_secs(60),
cleanup_interval: Duration::from_secs(10),
}
}
}
impl PeerLifetimeConfig {
pub fn new(
disconnected_timeout: Duration,
connected_timeout: Duration,
cleanup_interval: Duration,
) -> Self {
Self {
disconnected_timeout,
connected_timeout,
cleanup_interval,
}
}
pub fn fast() -> Self {
Self {
disconnected_timeout: Duration::from_secs(5),
connected_timeout: Duration::from_secs(10),
cleanup_interval: Duration::from_secs(2),
}
}
pub fn relaxed() -> Self {
Self {
disconnected_timeout: Duration::from_secs(60),
connected_timeout: Duration::from_secs(120),
cleanup_interval: Duration::from_secs(30),
}
}
}
#[derive(Debug, Clone)]
struct PeerState {
connected: bool,
last_seen: Instant,
first_seen: Instant,
disconnected_at: Option<Instant>,
}
impl PeerState {
fn new(connected: bool) -> Self {
let now = Instant::now();
Self {
connected,
last_seen: now,
first_seen: now,
disconnected_at: if connected { None } else { Some(now) },
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StaleReason {
DisconnectedTimeout,
ConnectedTimeout,
}
#[derive(Debug, Clone)]
pub struct StalePeerInfo {
pub address: String,
pub reason: StaleReason,
pub time_since_last_seen: Duration,
pub was_connected: bool,
}
#[derive(Debug)]
pub struct PeerLifetimeManager {
config: PeerLifetimeConfig,
peers: HashMap<String, PeerState>,
}
impl PeerLifetimeManager {
pub fn new(config: PeerLifetimeConfig) -> Self {
Self {
config,
peers: HashMap::new(),
}
}
pub fn with_defaults() -> Self {
Self::new(PeerLifetimeConfig::default())
}
pub fn on_peer_activity(&mut self, address: &str, connected: bool) {
let now = Instant::now();
if let Some(state) = self.peers.get_mut(address) {
state.last_seen = now;
if connected && !state.connected {
state.connected = true;
state.disconnected_at = None;
log::debug!("Peer {} connected", address);
} else if !connected && state.connected {
state.connected = false;
state.disconnected_at = Some(now);
log::debug!("Peer {} disconnected", address);
}
} else {
log::debug!("New peer {} (connected: {})", address, connected);
self.peers
.insert(address.to_string(), PeerState::new(connected));
}
}
pub fn on_peer_disconnected(&mut self, address: &str) {
if let Some(state) = self.peers.get_mut(address) {
if state.connected {
state.connected = false;
state.disconnected_at = Some(Instant::now());
log::debug!("Peer {} marked as disconnected", address);
}
}
}
pub fn is_tracked(&self, address: &str) -> bool {
self.peers.contains_key(address)
}
pub fn is_connected(&self, address: &str) -> bool {
self.peers
.get(address)
.map(|s| s.connected)
.unwrap_or(false)
}
pub fn get_stale_peers(&self) -> Vec<StalePeerInfo> {
self.peers
.iter()
.filter_map(|(address, state)| {
let time_since_last_seen = state.last_seen.elapsed();
let (is_stale, reason) = if state.connected {
let is_stale = time_since_last_seen > self.config.connected_timeout;
(is_stale, StaleReason::ConnectedTimeout)
} else {
let is_stale = time_since_last_seen > self.config.disconnected_timeout;
(is_stale, StaleReason::DisconnectedTimeout)
};
if is_stale {
Some(StalePeerInfo {
address: address.clone(),
reason,
time_since_last_seen,
was_connected: state.connected,
})
} else {
None
}
})
.collect()
}
pub fn get_stale_peer_addresses(&self) -> Vec<String> {
self.get_stale_peers()
.into_iter()
.map(|info| info.address)
.collect()
}
pub fn remove_peer(&mut self, address: &str) -> bool {
if self.peers.remove(address).is_some() {
log::debug!("Removed peer {} from lifetime tracking", address);
true
} else {
false
}
}
pub fn cleanup_stale_peers(&mut self) -> Vec<StalePeerInfo> {
let stale = self.get_stale_peers();
for info in &stale {
self.peers.remove(&info.address);
}
if !stale.is_empty() {
log::debug!("Cleaned up {} stale peers", stale.len());
}
stale
}
pub fn stats(&self) -> PeerLifetimeStats {
let mut connected = 0;
let mut disconnected = 0;
for state in self.peers.values() {
if state.connected {
connected += 1;
} else {
disconnected += 1;
}
}
PeerLifetimeStats {
total_tracked: self.peers.len(),
connected,
disconnected,
}
}
pub fn get_peer_info(&self, address: &str) -> Option<PeerInfo> {
self.peers.get(address).map(|state| PeerInfo {
connected: state.connected,
time_since_last_seen: state.last_seen.elapsed(),
time_since_first_seen: state.first_seen.elapsed(),
time_since_disconnect: state.disconnected_at.map(|t| t.elapsed()),
})
}
pub fn clear(&mut self) {
let count = self.peers.len();
self.peers.clear();
if count > 0 {
log::debug!("Cleared {} peers from lifetime tracking", count);
}
}
pub fn tracked_count(&self) -> usize {
self.peers.len()
}
pub fn cleanup_interval(&self) -> Duration {
self.config.cleanup_interval
}
}
#[derive(Debug, Clone, Copy)]
pub struct PeerLifetimeStats {
pub total_tracked: usize,
pub connected: usize,
pub disconnected: usize,
}
#[derive(Debug, Clone)]
pub struct PeerInfo {
pub connected: bool,
pub time_since_last_seen: Duration,
pub time_since_first_seen: Duration,
pub time_since_disconnect: Option<Duration>,
}
#[cfg(test)]
mod tests {
use super::*;
use std::thread::sleep;
#[test]
fn test_new_peer_tracking() {
let mut manager = PeerLifetimeManager::with_defaults();
assert!(!manager.is_tracked("test"));
manager.on_peer_activity("test", true);
assert!(manager.is_tracked("test"));
assert!(manager.is_connected("test"));
}
#[test]
fn test_peer_disconnect() {
let mut manager = PeerLifetimeManager::with_defaults();
manager.on_peer_activity("test", true);
assert!(manager.is_connected("test"));
manager.on_peer_disconnected("test");
assert!(!manager.is_connected("test"));
}
#[test]
fn test_stale_peer_detection() {
let config = PeerLifetimeConfig {
disconnected_timeout: Duration::from_millis(50),
connected_timeout: Duration::from_millis(100),
cleanup_interval: Duration::from_millis(10),
};
let mut manager = PeerLifetimeManager::new(config);
manager.on_peer_activity("test", false);
assert!(manager.get_stale_peers().is_empty());
sleep(Duration::from_millis(60));
let stale = manager.get_stale_peers();
assert_eq!(stale.len(), 1);
assert_eq!(stale[0].address, "test");
assert_eq!(stale[0].reason, StaleReason::DisconnectedTimeout);
}
#[test]
fn test_cleanup_stale_peers() {
let config = PeerLifetimeConfig {
disconnected_timeout: Duration::from_millis(10),
connected_timeout: Duration::from_millis(100),
cleanup_interval: Duration::from_millis(5),
};
let mut manager = PeerLifetimeManager::new(config);
manager.on_peer_activity("peer1", false);
manager.on_peer_activity("peer2", true);
sleep(Duration::from_millis(20));
let cleaned = manager.cleanup_stale_peers();
assert_eq!(cleaned.len(), 1);
assert_eq!(cleaned[0].address, "peer1");
assert!(!manager.is_tracked("peer1"));
assert!(manager.is_tracked("peer2"));
}
#[test]
fn test_stats() {
let mut manager = PeerLifetimeManager::with_defaults();
manager.on_peer_activity("connected1", true);
manager.on_peer_activity("connected2", true);
manager.on_peer_activity("disconnected1", false);
let stats = manager.stats();
assert_eq!(stats.total_tracked, 3);
assert_eq!(stats.connected, 2);
assert_eq!(stats.disconnected, 1);
}
#[test]
fn test_kotlin_timeout_values() {
let config = PeerLifetimeConfig::new(
Duration::from_secs(120),
Duration::from_secs(300),
Duration::from_secs(30),
);
assert_eq!(config.disconnected_timeout, Duration::from_secs(120));
assert_eq!(config.connected_timeout, Duration::from_secs(300));
let mut manager = PeerLifetimeManager::new(config);
manager.on_peer_activity("peer1", false);
assert!(manager.get_stale_peers().is_empty());
manager.on_peer_activity("peer2", true);
assert!(manager.get_stale_peers().is_empty());
}
#[test]
fn test_disconnect_does_not_update_last_seen() {
let config = PeerLifetimeConfig {
disconnected_timeout: Duration::from_millis(50),
connected_timeout: Duration::from_millis(200),
cleanup_interval: Duration::from_millis(10),
};
let mut manager = PeerLifetimeManager::new(config);
manager.on_peer_activity("test", true);
sleep(Duration::from_millis(30));
manager.on_peer_disconnected("test");
sleep(Duration::from_millis(30));
let stale = manager.get_stale_peers();
assert_eq!(stale.len(), 1);
assert_eq!(stale[0].address, "test");
}
#[test]
fn test_activity_resets_stale_timer() {
let config = PeerLifetimeConfig {
disconnected_timeout: Duration::from_millis(50),
connected_timeout: Duration::from_millis(100),
cleanup_interval: Duration::from_millis(10),
};
let mut manager = PeerLifetimeManager::new(config);
manager.on_peer_activity("test", false);
sleep(Duration::from_millis(40));
manager.on_peer_activity("test", false);
assert!(manager.get_stale_peers().is_empty());
sleep(Duration::from_millis(20));
assert!(manager.get_stale_peers().is_empty());
sleep(Duration::from_millis(40));
assert_eq!(manager.get_stale_peers().len(), 1);
}
#[test]
fn test_connected_peer_longer_timeout() {
let config = PeerLifetimeConfig {
disconnected_timeout: Duration::from_millis(30),
connected_timeout: Duration::from_millis(80),
cleanup_interval: Duration::from_millis(10),
};
let mut manager = PeerLifetimeManager::new(config);
manager.on_peer_activity("connected", true);
manager.on_peer_activity("disconnected", false);
sleep(Duration::from_millis(40));
let stale = manager.get_stale_peers();
assert_eq!(stale.len(), 1);
assert_eq!(stale[0].address, "disconnected");
assert_eq!(stale[0].reason, StaleReason::DisconnectedTimeout);
sleep(Duration::from_millis(50));
let stale = manager.get_stale_peers();
assert_eq!(stale.len(), 2);
}
}