1use 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#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
14pub struct SessionId(String);
15
16impl SessionId {
17 pub fn new(id: impl Into<String>) -> Self {
19 Self(id.into())
20 }
21
22 pub fn generate() -> Self {
24 Self(uuid::Uuid::new_v4().to_string())
25 }
26
27 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
53pub enum SessionState {
54 Connecting,
56 Active,
58 Idle,
60 Migrating,
62 Disconnecting,
64 Ended,
66}
67
68#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct Session {
71 pub id: SessionId,
73 pub state: SessionState,
75 pub user_id: Option<String>,
77 pub client_addr: Option<String>,
79 pub created_at: chrono::DateTime<chrono::Utc>,
81 pub last_activity: chrono::DateTime<chrono::Utc>,
83 pub metadata: HashMap<String, String>,
85}
86
87impl Session {
88 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 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 pub fn with_addr(mut self, addr: impl Into<String>) -> Self {
110 self.client_addr = Some(addr.into());
111 self
112 }
113
114 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 pub fn touch(&mut self) {
122 self.last_activity = chrono::Utc::now();
123 }
124
125 pub fn duration(&self) -> chrono::Duration {
127 chrono::Utc::now() - self.created_at
128 }
129
130 pub fn idle_duration(&self) -> chrono::Duration {
132 chrono::Utc::now() - self.last_activity
133 }
134}
135
136#[derive(Clone)]
138pub struct SessionTracker {
139 sessions: Arc<RwLock<HashMap<SessionId, Session>>>,
140 config: SessionConfig,
141}
142
143#[derive(Debug, Clone)]
145pub struct SessionConfig {
146 pub max_sessions: usize,
148 pub idle_timeout: Duration,
150 pub disconnect_timeout: Duration,
152 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 pub fn new(config: SessionConfig) -> Self {
170 Self {
171 sessions: Arc::new(RwLock::new(HashMap::new())),
172 config,
173 }
174 }
175
176 pub fn default_config() -> Self {
178 Self::new(SessionConfig::default())
179 }
180
181 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 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 pub fn get(&self, id: &SessionId) -> Option<Session> {
221 self.sessions.read().get(id).cloned()
222 }
223
224 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 pub fn activate(&self, id: &SessionId) -> bool {
239 self.set_state(id, SessionState::Active)
240 }
241
242 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 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 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 pub fn count(&self) -> usize {
277 self.sessions.read().len()
278 }
279
280 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 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 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 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 pub fn export_sessions(&self) -> Vec<Session> {
337 self.sessions.read().values().cloned().collect()
338 }
339
340 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#[derive(Debug, Clone, Serialize, Deserialize)]
361pub struct ConnectionInfo {
362 pub session_id: SessionId,
364 pub remote_addr: String,
366 pub local_port: u16,
368 pub protocol: String,
370 pub connected_at: chrono::DateTime<chrono::Utc>,
372 pub bytes_sent: u64,
374 pub bytes_received: u64,
376 pub packets_sent: u64,
378 pub packets_received: u64,
380 pub rtt_ms: Option<u32>,
382}
383
384impl ConnectionInfo {
385 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 pub fn update_stats(&mut self, sent: u64, received: u64) {
403 self.bytes_sent += sent;
404 self.bytes_received += received;
405 }
406
407 pub fn update_packets(&mut self, sent: u64, received: u64) {
409 self.packets_sent += sent;
410 self.packets_received += received;
411 }
412
413 pub fn set_rtt(&mut self, rtt_ms: u32) {
415 self.rtt_ms = Some(rtt_ms);
416 }
417
418 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 assert!(tracker.activate(&id));
445 let session = tracker.get(&id).unwrap();
446 assert_eq!(session.state, SessionState::Active);
447
448 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()); }
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}