forge_orchestration/sdk/
session.rs

1//! Session and connection tracking for game servers
2//!
3//! Provides primitives for tracking player sessions, connections, and graceful migration.
4
5use std::collections::HashMap;
6use std::sync::Arc;
7use std::time::Duration;
8use parking_lot::RwLock;
9use serde::{Deserialize, Serialize};
10use tracing::{debug, info, warn};
11
12/// Unique session identifier
13#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
14pub struct SessionId(String);
15
16impl SessionId {
17    /// Create a new session ID
18    pub fn new(id: impl Into<String>) -> Self {
19        Self(id.into())
20    }
21
22    /// Generate a random session ID
23    pub fn generate() -> Self {
24        Self(uuid::Uuid::new_v4().to_string())
25    }
26
27    /// Get the inner string
28    pub fn as_str(&self) -> &str {
29        &self.0
30    }
31}
32
33impl std::fmt::Display for SessionId {
34    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
35        write!(f, "{}", self.0)
36    }
37}
38
39impl From<String> for SessionId {
40    fn from(s: String) -> Self {
41        Self::new(s)
42    }
43}
44
45impl From<&str> for SessionId {
46    fn from(s: &str) -> Self {
47        Self::new(s)
48    }
49}
50
51/// Session state
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
53pub enum SessionState {
54    /// Session is connecting
55    Connecting,
56    /// Session is active
57    Active,
58    /// Session is idle (no recent activity)
59    Idle,
60    /// Session is being migrated to another server
61    Migrating,
62    /// Session is disconnecting
63    Disconnecting,
64    /// Session has ended
65    Ended,
66}
67
68/// Session metadata
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct Session {
71    /// Session ID
72    pub id: SessionId,
73    /// Current state
74    pub state: SessionState,
75    /// User/player identifier (optional)
76    pub user_id: Option<String>,
77    /// Client address
78    pub client_addr: Option<String>,
79    /// Session creation time
80    pub created_at: chrono::DateTime<chrono::Utc>,
81    /// Last activity time
82    pub last_activity: chrono::DateTime<chrono::Utc>,
83    /// Custom metadata
84    pub metadata: HashMap<String, String>,
85}
86
87impl Session {
88    /// Create a new session
89    pub fn new(id: SessionId) -> Self {
90        let now = chrono::Utc::now();
91        Self {
92            id,
93            state: SessionState::Connecting,
94            user_id: None,
95            client_addr: None,
96            created_at: now,
97            last_activity: now,
98            metadata: HashMap::new(),
99        }
100    }
101
102    /// Set user ID
103    pub fn with_user(mut self, user_id: impl Into<String>) -> Self {
104        self.user_id = Some(user_id.into());
105        self
106    }
107
108    /// Set client address
109    pub fn with_addr(mut self, addr: impl Into<String>) -> Self {
110        self.client_addr = Some(addr.into());
111        self
112    }
113
114    /// Add metadata
115    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
116        self.metadata.insert(key.into(), value.into());
117        self
118    }
119
120    /// Update last activity time
121    pub fn touch(&mut self) {
122        self.last_activity = chrono::Utc::now();
123    }
124
125    /// Get session duration
126    pub fn duration(&self) -> chrono::Duration {
127        chrono::Utc::now() - self.created_at
128    }
129
130    /// Get idle duration
131    pub fn idle_duration(&self) -> chrono::Duration {
132        chrono::Utc::now() - self.last_activity
133    }
134}
135
136/// Session tracker for managing active sessions
137#[derive(Clone)]
138pub struct SessionTracker {
139    sessions: Arc<RwLock<HashMap<SessionId, Session>>>,
140    config: SessionConfig,
141}
142
143/// Session tracker configuration
144#[derive(Debug, Clone)]
145pub struct SessionConfig {
146    /// Maximum sessions allowed
147    pub max_sessions: usize,
148    /// Idle timeout before marking session as idle
149    pub idle_timeout: Duration,
150    /// Disconnect timeout after idle
151    pub disconnect_timeout: Duration,
152    /// Enable automatic cleanup of ended sessions
153    pub auto_cleanup: bool,
154}
155
156impl Default for SessionConfig {
157    fn default() -> Self {
158        Self {
159            max_sessions: 1000,
160            idle_timeout: Duration::from_secs(60),
161            disconnect_timeout: Duration::from_secs(300),
162            auto_cleanup: true,
163        }
164    }
165}
166
167impl SessionTracker {
168    /// Create a new session tracker
169    pub fn new(config: SessionConfig) -> Self {
170        Self {
171            sessions: Arc::new(RwLock::new(HashMap::new())),
172            config,
173        }
174    }
175
176    /// Create with default config
177    pub fn default_config() -> Self {
178        Self::new(SessionConfig::default())
179    }
180
181    /// Create a new session
182    pub fn create_session(&self) -> Option<Session> {
183        let mut sessions = self.sessions.write();
184        
185        if sessions.len() >= self.config.max_sessions {
186            warn!(max = self.config.max_sessions, "Session limit reached");
187            return None;
188        }
189
190        let id = SessionId::generate();
191        let session = Session::new(id.clone());
192        sessions.insert(id, session.clone());
193        
194        info!(session_id = %session.id, "Session created");
195        Some(session)
196    }
197
198    /// Create a session with a specific ID
199    pub fn create_session_with_id(&self, id: SessionId) -> Option<Session> {
200        let mut sessions = self.sessions.write();
201        
202        if sessions.len() >= self.config.max_sessions {
203            warn!(max = self.config.max_sessions, "Session limit reached");
204            return None;
205        }
206
207        if sessions.contains_key(&id) {
208            warn!(session_id = %id, "Session ID already exists");
209            return None;
210        }
211
212        let session = Session::new(id.clone());
213        sessions.insert(id, session.clone());
214        
215        info!(session_id = %session.id, "Session created");
216        Some(session)
217    }
218
219    /// Get a session by ID
220    pub fn get(&self, id: &SessionId) -> Option<Session> {
221        self.sessions.read().get(id).cloned()
222    }
223
224    /// Update session state
225    pub fn set_state(&self, id: &SessionId, state: SessionState) -> bool {
226        let mut sessions = self.sessions.write();
227        if let Some(session) = sessions.get_mut(id) {
228            session.state = state;
229            session.touch();
230            debug!(session_id = %id, state = ?state, "Session state updated");
231            true
232        } else {
233            false
234        }
235    }
236
237    /// Mark session as active
238    pub fn activate(&self, id: &SessionId) -> bool {
239        self.set_state(id, SessionState::Active)
240    }
241
242    /// Touch session (update last activity)
243    pub fn touch(&self, id: &SessionId) -> bool {
244        let mut sessions = self.sessions.write();
245        if let Some(session) = sessions.get_mut(id) {
246            session.touch();
247            true
248        } else {
249            false
250        }
251    }
252
253    /// End a session
254    pub fn end_session(&self, id: &SessionId) -> Option<Session> {
255        let mut sessions = self.sessions.write();
256        if let Some(mut session) = sessions.remove(id) {
257            session.state = SessionState::Ended;
258            info!(session_id = %id, duration = ?session.duration(), "Session ended");
259            Some(session)
260        } else {
261            None
262        }
263    }
264
265    /// Get all active sessions
266    pub fn active_sessions(&self) -> Vec<Session> {
267        self.sessions
268            .read()
269            .values()
270            .filter(|s| s.state == SessionState::Active)
271            .cloned()
272            .collect()
273    }
274
275    /// Get session count
276    pub fn count(&self) -> usize {
277        self.sessions.read().len()
278    }
279
280    /// Get active session count
281    pub fn active_count(&self) -> usize {
282        self.sessions
283            .read()
284            .values()
285            .filter(|s| s.state == SessionState::Active)
286            .count()
287    }
288
289    /// Check for idle sessions and update their state
290    pub fn check_idle(&self) -> Vec<SessionId> {
291        let mut idle_sessions = Vec::new();
292        let idle_threshold = chrono::Duration::from_std(self.config.idle_timeout)
293            .unwrap_or(chrono::Duration::seconds(60));
294
295        let mut sessions = self.sessions.write();
296        for (id, session) in sessions.iter_mut() {
297            if session.state == SessionState::Active && session.idle_duration() > idle_threshold {
298                session.state = SessionState::Idle;
299                idle_sessions.push(id.clone());
300                debug!(session_id = %id, "Session marked idle");
301            }
302        }
303
304        idle_sessions
305    }
306
307    /// Cleanup ended/disconnected sessions
308    pub fn cleanup(&self) -> usize {
309        let mut sessions = self.sessions.write();
310        let before = sessions.len();
311        sessions.retain(|_, s| s.state != SessionState::Ended);
312        let removed = before - sessions.len();
313        if removed > 0 {
314            debug!(removed = removed, "Cleaned up ended sessions");
315        }
316        removed
317    }
318
319    /// Prepare all sessions for migration (e.g., before server shutdown)
320    pub fn prepare_migration(&self) -> Vec<Session> {
321        let mut sessions = self.sessions.write();
322        let mut migrating = Vec::new();
323
324        for session in sessions.values_mut() {
325            if session.state == SessionState::Active || session.state == SessionState::Idle {
326                session.state = SessionState::Migrating;
327                migrating.push(session.clone());
328            }
329        }
330
331        info!(count = migrating.len(), "Sessions prepared for migration");
332        migrating
333    }
334
335    /// Export all sessions for migration
336    pub fn export_sessions(&self) -> Vec<Session> {
337        self.sessions.read().values().cloned().collect()
338    }
339
340    /// Import sessions (e.g., after migration)
341    pub fn import_sessions(&self, sessions: Vec<Session>) -> usize {
342        let mut store = self.sessions.write();
343        let mut imported = 0;
344
345        for mut session in sessions {
346            if !store.contains_key(&session.id) && store.len() < self.config.max_sessions {
347                session.state = SessionState::Active;
348                session.touch();
349                store.insert(session.id.clone(), session);
350                imported += 1;
351            }
352        }
353
354        info!(imported = imported, "Sessions imported");
355        imported
356    }
357}
358
359/// Connection info for network tracking
360#[derive(Debug, Clone, Serialize, Deserialize)]
361pub struct ConnectionInfo {
362    /// Session this connection belongs to
363    pub session_id: SessionId,
364    /// Remote address
365    pub remote_addr: String,
366    /// Local port
367    pub local_port: u16,
368    /// Protocol (TCP, UDP, QUIC)
369    pub protocol: String,
370    /// Connection established time
371    pub connected_at: chrono::DateTime<chrono::Utc>,
372    /// Bytes sent
373    pub bytes_sent: u64,
374    /// Bytes received
375    pub bytes_received: u64,
376    /// Packets sent (for UDP)
377    pub packets_sent: u64,
378    /// Packets received (for UDP)
379    pub packets_received: u64,
380    /// Round-trip time estimate
381    pub rtt_ms: Option<u32>,
382}
383
384impl ConnectionInfo {
385    /// Create new connection info
386    pub fn new(session_id: SessionId, remote_addr: impl Into<String>, protocol: impl Into<String>) -> Self {
387        Self {
388            session_id,
389            remote_addr: remote_addr.into(),
390            local_port: 0,
391            protocol: protocol.into(),
392            connected_at: chrono::Utc::now(),
393            bytes_sent: 0,
394            bytes_received: 0,
395            packets_sent: 0,
396            packets_received: 0,
397            rtt_ms: None,
398        }
399    }
400
401    /// Update traffic stats
402    pub fn update_stats(&mut self, sent: u64, received: u64) {
403        self.bytes_sent += sent;
404        self.bytes_received += received;
405    }
406
407    /// Update packet stats
408    pub fn update_packets(&mut self, sent: u64, received: u64) {
409        self.packets_sent += sent;
410        self.packets_received += received;
411    }
412
413    /// Set RTT
414    pub fn set_rtt(&mut self, rtt_ms: u32) {
415        self.rtt_ms = Some(rtt_ms);
416    }
417
418    /// Get connection duration
419    pub fn duration(&self) -> chrono::Duration {
420        chrono::Utc::now() - self.connected_at
421    }
422}
423
424#[cfg(test)]
425mod tests {
426    use super::*;
427
428    #[test]
429    fn test_session_creation() {
430        let tracker = SessionTracker::default_config();
431        let session = tracker.create_session().unwrap();
432        
433        assert_eq!(session.state, SessionState::Connecting);
434        assert_eq!(tracker.count(), 1);
435    }
436
437    #[test]
438    fn test_session_lifecycle() {
439        let tracker = SessionTracker::default_config();
440        let session = tracker.create_session().unwrap();
441        let id = session.id.clone();
442
443        // Activate
444        assert!(tracker.activate(&id));
445        let session = tracker.get(&id).unwrap();
446        assert_eq!(session.state, SessionState::Active);
447
448        // End
449        let ended = tracker.end_session(&id).unwrap();
450        assert_eq!(ended.state, SessionState::Ended);
451        assert_eq!(tracker.count(), 0);
452    }
453
454    #[test]
455    fn test_session_limit() {
456        let config = SessionConfig {
457            max_sessions: 2,
458            ..Default::default()
459        };
460        let tracker = SessionTracker::new(config);
461
462        assert!(tracker.create_session().is_some());
463        assert!(tracker.create_session().is_some());
464        assert!(tracker.create_session().is_none()); // Should fail
465    }
466
467    #[test]
468    fn test_migration() {
469        let tracker = SessionTracker::default_config();
470        
471        let s1 = tracker.create_session().unwrap();
472        tracker.activate(&s1.id);
473        
474        let s2 = tracker.create_session().unwrap();
475        tracker.activate(&s2.id);
476
477        let migrating = tracker.prepare_migration();
478        assert_eq!(migrating.len(), 2);
479
480        for s in &migrating {
481            assert_eq!(s.state, SessionState::Migrating);
482        }
483    }
484}