use dashmap::DashMap;
use libp2p::PeerId;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SessionState {
Creating,
Active,
Idle,
Closing,
Closed,
}
impl SessionState {
pub fn is_active(&self) -> bool {
matches!(self, Self::Creating | Self::Active)
}
pub fn is_terminated(&self) -> bool {
matches!(self, Self::Closing | Self::Closed)
}
}
#[derive(Debug, Clone)]
pub struct Session {
pub peer_id: PeerId,
pub state: SessionState,
pub created_at: Instant,
pub last_activity: Instant,
pub closed_at: Option<Instant>,
pub bytes_sent: u64,
pub bytes_received: u64,
pub messages_sent: u64,
pub messages_received: u64,
pub metadata: SessionMetadata,
}
impl Session {
pub fn new(peer_id: PeerId) -> Self {
let now = Instant::now();
Self {
peer_id,
state: SessionState::Creating,
created_at: now,
last_activity: now,
closed_at: None,
bytes_sent: 0,
bytes_received: 0,
messages_sent: 0,
messages_received: 0,
metadata: SessionMetadata::default(),
}
}
pub fn duration(&self) -> Duration {
if let Some(closed_at) = self.closed_at {
closed_at.duration_since(self.created_at)
} else {
Instant::now().duration_since(self.created_at)
}
}
pub fn idle_duration(&self) -> Duration {
Instant::now().duration_since(self.last_activity)
}
pub fn is_idle(&self, timeout: Duration) -> bool {
self.idle_duration() > timeout
}
pub fn mark_activity(&mut self) {
self.last_activity = Instant::now();
if self.state == SessionState::Idle {
self.state = SessionState::Active;
}
}
pub fn add_bytes_sent(&mut self, bytes: u64) {
self.bytes_sent += bytes;
self.mark_activity();
}
pub fn add_bytes_received(&mut self, bytes: u64) {
self.bytes_received += bytes;
self.mark_activity();
}
pub fn record_message_sent(&mut self) {
self.messages_sent += 1;
self.mark_activity();
}
pub fn record_message_received(&mut self) {
self.messages_received += 1;
self.mark_activity();
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct SessionMetadata {
pub endpoint: Option<String>,
pub protocol: Option<String>,
pub tags: Vec<String>,
pub quality_score: Option<f64>,
}
#[derive(Debug, Clone)]
pub struct SessionConfig {
pub idle_timeout: Duration,
pub max_sessions: usize,
pub auto_cleanup: bool,
pub cleanup_interval: Duration,
}
impl Default for SessionConfig {
fn default() -> Self {
Self {
idle_timeout: Duration::from_secs(300), max_sessions: 1000,
auto_cleanup: true,
cleanup_interval: Duration::from_secs(60), }
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct SessionStats {
pub total_created: u64,
pub active_sessions: usize,
pub idle_sessions: usize,
pub total_closed: u64,
pub total_bytes_sent: u64,
pub total_bytes_received: u64,
pub total_messages_sent: u64,
pub total_messages_received: u64,
pub avg_session_duration: Duration,
}
pub struct SessionManager {
config: SessionConfig,
sessions: Arc<DashMap<PeerId, Session>>,
stats: Arc<parking_lot::RwLock<SessionStats>>,
}
impl SessionManager {
pub fn new(config: SessionConfig) -> Self {
Self {
config,
sessions: Arc::new(DashMap::new()),
stats: Arc::new(parking_lot::RwLock::new(SessionStats::default())),
}
}
pub fn create_session(&self, peer_id: PeerId) -> bool {
if self.sessions.len() >= self.config.max_sessions {
return false;
}
let session = Session::new(peer_id);
let inserted = self.sessions.insert(peer_id, session).is_none();
if inserted {
let mut stats = self.stats.write();
stats.total_created += 1;
}
inserted
}
pub fn get_session(&self, peer_id: &PeerId) -> Option<Session> {
self.sessions.get(peer_id).map(|s| s.clone())
}
pub fn activate_session(&self, peer_id: &PeerId) {
if let Some(mut session) = self.sessions.get_mut(peer_id) {
session.state = SessionState::Active;
session.mark_activity();
}
}
pub fn mark_activity(&self, peer_id: &PeerId) {
if let Some(mut session) = self.sessions.get_mut(peer_id) {
session.mark_activity();
}
}
pub fn update_bandwidth(&self, peer_id: &PeerId, sent: u64, received: u64) {
if let Some(mut session) = self.sessions.get_mut(peer_id) {
session.add_bytes_sent(sent);
session.add_bytes_received(received);
let mut stats = self.stats.write();
stats.total_bytes_sent += sent;
stats.total_bytes_received += received;
}
}
pub fn record_message(&self, peer_id: &PeerId, sent: bool) {
if let Some(mut session) = self.sessions.get_mut(peer_id) {
if sent {
session.record_message_sent();
let mut stats = self.stats.write();
stats.total_messages_sent += 1;
} else {
session.record_message_received();
let mut stats = self.stats.write();
stats.total_messages_received += 1;
}
}
}
pub fn update_metadata(&self, peer_id: &PeerId, metadata: SessionMetadata) {
if let Some(mut session) = self.sessions.get_mut(peer_id) {
session.metadata = metadata;
}
}
pub fn close_session(&self, peer_id: &PeerId) {
if let Some(mut session) = self.sessions.get_mut(peer_id) {
session.state = SessionState::Closing;
session.closed_at = Some(Instant::now());
let mut stats = self.stats.write();
stats.total_closed += 1;
let total_duration = stats.avg_session_duration.as_secs_f64()
* (stats.total_closed - 1) as f64
+ session.duration().as_secs_f64();
stats.avg_session_duration =
Duration::from_secs_f64(total_duration / stats.total_closed as f64);
}
}
pub fn remove_session(&self, peer_id: &PeerId) -> Option<Session> {
self.sessions.remove(peer_id).map(|(_, s)| s)
}
pub fn get_all_sessions(&self) -> Vec<Session> {
self.sessions.iter().map(|entry| entry.clone()).collect()
}
pub fn get_sessions_by_state(&self, state: SessionState) -> Vec<Session> {
self.sessions
.iter()
.filter(|entry| entry.state == state)
.map(|entry| entry.clone())
.collect()
}
pub fn check_idle_sessions(&self) -> Vec<PeerId> {
let mut idle_peers = Vec::new();
for mut entry in self.sessions.iter_mut() {
if entry.state == SessionState::Active && entry.is_idle(self.config.idle_timeout) {
entry.state = SessionState::Idle;
idle_peers.push(entry.peer_id);
}
}
idle_peers
}
pub fn cleanup_closed_sessions(&self) -> usize {
let closed_sessions: Vec<PeerId> = self
.sessions
.iter()
.filter(|entry| entry.state == SessionState::Closed)
.map(|entry| entry.peer_id)
.collect();
let count = closed_sessions.len();
for peer_id in closed_sessions {
self.sessions.remove(&peer_id);
}
count
}
pub fn stats(&self) -> SessionStats {
let mut stats = self.stats.read().clone();
stats.active_sessions = self
.sessions
.iter()
.filter(|e| e.state == SessionState::Active)
.count();
stats.idle_sessions = self
.sessions
.iter()
.filter(|e| e.state == SessionState::Idle)
.count();
stats
}
pub fn session_count(&self) -> usize {
self.sessions.len()
}
pub fn active_session_count(&self) -> usize {
self.sessions.iter().filter(|e| e.state.is_active()).count()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_session_creation() {
let peer_id = PeerId::random();
let session = Session::new(peer_id);
assert_eq!(session.peer_id, peer_id);
assert_eq!(session.state, SessionState::Creating);
assert_eq!(session.bytes_sent, 0);
assert_eq!(session.bytes_received, 0);
}
#[test]
fn test_session_state() {
assert!(SessionState::Creating.is_active());
assert!(SessionState::Active.is_active());
assert!(!SessionState::Idle.is_active());
assert!(SessionState::Closing.is_terminated());
assert!(SessionState::Closed.is_terminated());
}
#[test]
fn test_session_activity() {
let peer_id = PeerId::random();
let mut session = Session::new(peer_id);
let initial_time = session.last_activity;
std::thread::sleep(std::time::Duration::from_millis(10));
session.mark_activity();
assert!(session.last_activity > initial_time);
}
#[test]
fn test_session_bandwidth() {
let peer_id = PeerId::random();
let mut session = Session::new(peer_id);
session.add_bytes_sent(1024);
session.add_bytes_received(2048);
assert_eq!(session.bytes_sent, 1024);
assert_eq!(session.bytes_received, 2048);
}
#[test]
fn test_session_messages() {
let peer_id = PeerId::random();
let mut session = Session::new(peer_id);
session.record_message_sent();
session.record_message_sent();
session.record_message_received();
assert_eq!(session.messages_sent, 2);
assert_eq!(session.messages_received, 1);
}
#[test]
fn test_session_manager_creation() {
let config = SessionConfig::default();
let manager = SessionManager::new(config);
let peer_id = PeerId::random();
assert!(manager.create_session(peer_id));
assert_eq!(manager.session_count(), 1);
let stats = manager.stats();
assert_eq!(stats.total_created, 1);
}
#[test]
fn test_session_manager_max_sessions() {
let config = SessionConfig {
max_sessions: 2,
..Default::default()
};
let manager = SessionManager::new(config);
let peer1 = PeerId::random();
let peer2 = PeerId::random();
let peer3 = PeerId::random();
assert!(manager.create_session(peer1));
assert!(manager.create_session(peer2));
assert!(!manager.create_session(peer3));
assert_eq!(manager.session_count(), 2);
}
#[test]
fn test_session_manager_activity() {
let manager = SessionManager::new(SessionConfig::default());
let peer_id = PeerId::random();
manager.create_session(peer_id);
manager.activate_session(&peer_id);
let session = manager
.get_session(&peer_id)
.expect("test: session should exist after activate_session");
assert_eq!(session.state, SessionState::Active);
}
#[test]
fn test_session_manager_bandwidth() {
let manager = SessionManager::new(SessionConfig::default());
let peer_id = PeerId::random();
manager.create_session(peer_id);
manager.update_bandwidth(&peer_id, 1024, 2048);
let session = manager
.get_session(&peer_id)
.expect("test: session should exist after update_bandwidth");
assert_eq!(session.bytes_sent, 1024);
assert_eq!(session.bytes_received, 2048);
let stats = manager.stats();
assert_eq!(stats.total_bytes_sent, 1024);
assert_eq!(stats.total_bytes_received, 2048);
}
#[test]
fn test_session_manager_close() {
let manager = SessionManager::new(SessionConfig::default());
let peer_id = PeerId::random();
manager.create_session(peer_id);
manager.close_session(&peer_id);
let session = manager
.get_session(&peer_id)
.expect("test: session should exist after close_session");
assert_eq!(session.state, SessionState::Closing);
assert!(session.closed_at.is_some());
let stats = manager.stats();
assert_eq!(stats.total_closed, 1);
}
#[test]
fn test_session_manager_remove() {
let manager = SessionManager::new(SessionConfig::default());
let peer_id = PeerId::random();
manager.create_session(peer_id);
assert_eq!(manager.session_count(), 1);
manager.remove_session(&peer_id);
assert_eq!(manager.session_count(), 0);
}
#[test]
fn test_session_manager_filter_by_state() {
let manager = SessionManager::new(SessionConfig::default());
let peer1 = PeerId::random();
let peer2 = PeerId::random();
manager.create_session(peer1);
manager.create_session(peer2);
manager.activate_session(&peer1);
let active = manager.get_sessions_by_state(SessionState::Active);
assert_eq!(active.len(), 1);
let creating = manager.get_sessions_by_state(SessionState::Creating);
assert_eq!(creating.len(), 1);
}
#[test]
fn test_session_manager_cleanup() {
let manager = SessionManager::new(SessionConfig::default());
let peer1 = PeerId::random();
let peer2 = PeerId::random();
manager.create_session(peer1);
manager.create_session(peer2);
manager.close_session(&peer1);
if let Some(mut session) = manager.sessions.get_mut(&peer1) {
session.state = SessionState::Closed;
}
let cleaned = manager.cleanup_closed_sessions();
assert_eq!(cleaned, 1);
assert_eq!(manager.session_count(), 1);
}
#[test]
fn test_session_idle_detection() {
let peer_id = PeerId::random();
let session = Session::new(peer_id);
assert!(!session.is_idle(Duration::from_secs(1)));
let mut old_session = Session::new(peer_id);
old_session.last_activity = Instant::now() - Duration::from_secs(10);
assert!(old_session.is_idle(Duration::from_secs(5)));
}
}