1#![allow(deprecated)] use std::collections::HashMap;
4use std::sync::Arc;
5use std::time::{Duration, Instant};
6use tokio::sync::RwLock;
7use tokio::sync::mpsc;
8
9#[derive(Debug, Clone)]
35pub struct PingPongConfig {
36 ping_interval: Duration,
38 pong_timeout: Duration,
40}
41
42impl Default for PingPongConfig {
43 fn default() -> Self {
44 Self {
45 ping_interval: Duration::from_secs(30),
46 pong_timeout: Duration::from_secs(10),
47 }
48 }
49}
50
51impl PingPongConfig {
52 pub fn new(ping_interval: Duration, pong_timeout: Duration) -> Self {
73 Self {
74 ping_interval,
75 pong_timeout,
76 }
77 }
78
79 pub fn ping_interval(&self) -> Duration {
81 self.ping_interval
82 }
83
84 pub fn pong_timeout(&self) -> Duration {
86 self.pong_timeout
87 }
88
89 pub fn with_ping_interval(mut self, interval: Duration) -> Self {
102 self.ping_interval = interval;
103 self
104 }
105
106 pub fn with_pong_timeout(mut self, timeout: Duration) -> Self {
119 self.pong_timeout = timeout;
120 self
121 }
122}
123
124#[deprecated(
151 since = "0.2.0",
152 note = "Use `ConnectionSettings` with the `#[settings]` macro instead."
153)]
154#[derive(Debug, Clone)]
155pub struct ConnectionConfig {
156 idle_timeout: Duration,
157 handshake_timeout: Duration,
158 cleanup_interval: Duration,
159 max_connections: Option<usize>,
161 ping_config: PingPongConfig,
163}
164
165impl Default for ConnectionConfig {
166 fn default() -> Self {
167 Self {
168 idle_timeout: Duration::from_secs(300), handshake_timeout: Duration::from_secs(10), cleanup_interval: Duration::from_secs(30), max_connections: None, ping_config: PingPongConfig::default(),
173 }
174 }
175}
176
177impl ConnectionConfig {
178 pub fn new() -> Self {
192 Self::default()
193 }
194
195 pub fn with_idle_timeout(mut self, timeout: Duration) -> Self {
213 self.idle_timeout = timeout;
214 self
215 }
216
217 pub fn with_handshake_timeout(mut self, timeout: Duration) -> Self {
235 self.handshake_timeout = timeout;
236 self
237 }
238
239 pub fn with_cleanup_interval(mut self, interval: Duration) -> Self {
257 self.cleanup_interval = interval;
258 self
259 }
260
261 pub fn idle_timeout(&self) -> Duration {
263 self.idle_timeout
264 }
265
266 pub fn handshake_timeout(&self) -> Duration {
268 self.handshake_timeout
269 }
270
271 pub fn cleanup_interval(&self) -> Duration {
273 self.cleanup_interval
274 }
275
276 pub fn with_max_connections(mut self, max: Option<usize>) -> Self {
293 self.max_connections = max;
294 self
295 }
296
297 pub fn max_connections(&self) -> Option<usize> {
299 self.max_connections
300 }
301
302 pub fn with_ping_config(mut self, config: PingPongConfig) -> Self {
325 self.ping_config = config;
326 self
327 }
328
329 pub fn ping_config(&self) -> &PingPongConfig {
331 &self.ping_config
332 }
333
334 pub fn no_timeout() -> Self {
346 Self {
347 idle_timeout: Duration::MAX,
348 handshake_timeout: Duration::MAX,
349 cleanup_interval: Duration::from_secs(30),
350 max_connections: None,
351 ping_config: PingPongConfig::default(),
352 }
353 }
354
355 pub fn strict() -> Self {
373 Self {
374 idle_timeout: Duration::from_secs(30),
375 handshake_timeout: Duration::from_secs(5),
376 cleanup_interval: Duration::from_secs(10),
377 max_connections: None,
378 ping_config: PingPongConfig::new(Duration::from_secs(10), Duration::from_secs(5)),
379 }
380 }
381
382 pub fn permissive() -> Self {
400 Self {
401 idle_timeout: Duration::from_secs(3600),
402 handshake_timeout: Duration::from_secs(30),
403 cleanup_interval: Duration::from_secs(60),
404 max_connections: None,
405 ping_config: PingPongConfig::new(Duration::from_secs(60), Duration::from_secs(30)),
406 }
407 }
408}
409
410#[derive(Debug, thiserror::Error)]
412pub enum WebSocketError {
413 #[error("Connection error")]
415 Connection(String),
416 #[error("Send failed")]
418 Send(String),
419 #[error("Receive failed")]
421 Receive(String),
422 #[error("Protocol error")]
424 Protocol(String),
425 #[error("Internal error")]
427 Internal(String),
428 #[error("Connection timed out")]
430 Timeout(Duration),
431 #[error("Reconnection failed")]
433 ReconnectFailed(u32),
434 #[error("Invalid binary payload: {0}")]
436 BinaryPayload(String),
437 #[error("Heartbeat timeout: no pong received within {0:?}")]
439 HeartbeatTimeout(Duration),
440 #[error("Slow consumer: send timed out after {0:?}")]
442 SlowConsumer(Duration),
443}
444
445impl WebSocketError {
446 pub fn client_message(&self) -> &'static str {
451 match self {
452 Self::Connection(_) => "Connection error",
453 Self::Send(_) => "Failed to send message",
454 Self::Receive(_) => "Failed to receive message",
455 Self::Protocol(_) => "Protocol error",
456 Self::Internal(_) => "Internal server error",
457 Self::Timeout(_) => "Connection timed out",
458 Self::ReconnectFailed(_) => "Reconnection failed",
459 Self::BinaryPayload(_) => "Invalid message format",
460 Self::HeartbeatTimeout(_) => "Connection timed out",
461 Self::SlowConsumer(_) => "Server overloaded",
462 }
463 }
464
465 pub fn internal_detail(&self) -> String {
470 match self {
471 Self::Connection(msg) => format!("Connection error: {}", msg),
472 Self::Send(msg) => format!("Send error: {}", msg),
473 Self::Receive(msg) => format!("Receive error: {}", msg),
474 Self::Protocol(msg) => format!("Protocol error: {}", msg),
475 Self::Internal(msg) => format!("Internal error: {}", msg),
476 Self::Timeout(d) => format!("Connection timeout: idle for {:?}", d),
477 Self::ReconnectFailed(n) => format!("Reconnection failed after {} attempts", n),
478 Self::BinaryPayload(msg) => format!("Invalid binary payload: {}", msg),
479 Self::HeartbeatTimeout(d) => format!("Heartbeat timeout: no pong within {:?}", d),
480 Self::SlowConsumer(d) => format!("Slow consumer: send timed out after {:?}", d),
481 }
482 }
483}
484
485pub type WebSocketResult<T> = Result<T, WebSocketError>;
487
488#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
490#[serde(tag = "type")]
491pub enum Message {
492 Text {
494 data: String,
496 },
497 Binary {
499 data: Vec<u8>,
501 },
502 Ping,
504 Pong,
506 Close {
508 code: u16,
510 reason: String,
512 },
513}
514
515impl Message {
516 pub fn text(data: String) -> Self {
530 Self::Text { data }
531 }
532 pub fn binary(data: Vec<u8>) -> Self {
547 Self::Binary { data }
548 }
549 pub fn json<T: serde::Serialize>(data: &T) -> WebSocketResult<Self> {
578 let json =
579 serde_json::to_string(data).map_err(|e| WebSocketError::Protocol(e.to_string()))?;
580 Ok(Self::text(json))
581 }
582 pub fn parse_json<T: serde::de::DeserializeOwned>(&self) -> WebSocketResult<T> {
602 match self {
603 Message::Text { data } => {
604 serde_json::from_str(data).map_err(|e| WebSocketError::Protocol(e.to_string()))
605 }
606 _ => Err(WebSocketError::Protocol("Not a text message".to_string())),
607 }
608 }
609}
610
611pub struct WebSocketConnection {
613 id: String,
614 tx: mpsc::UnboundedSender<Message>,
615 closed: Arc<RwLock<bool>>,
616 subprotocol: Option<String>,
618 last_activity: Arc<RwLock<Instant>>,
620 config: ConnectionConfig,
622}
623
624impl WebSocketConnection {
625 pub fn new(id: String, tx: mpsc::UnboundedSender<Message>) -> Self {
640 Self {
641 id,
642 tx,
643 closed: Arc::new(RwLock::new(false)),
644 subprotocol: None,
645 last_activity: Arc::new(RwLock::new(Instant::now())),
646 config: ConnectionConfig::default(),
647 }
648 }
649
650 pub fn with_config(
668 id: String,
669 tx: mpsc::UnboundedSender<Message>,
670 config: ConnectionConfig,
671 ) -> Self {
672 Self {
673 id,
674 tx,
675 closed: Arc::new(RwLock::new(false)),
676 subprotocol: None,
677 last_activity: Arc::new(RwLock::new(Instant::now())),
678 config,
679 }
680 }
681
682 pub fn with_subprotocol(
700 id: String,
701 tx: mpsc::UnboundedSender<Message>,
702 subprotocol: Option<String>,
703 ) -> Self {
704 Self {
705 id,
706 tx,
707 closed: Arc::new(RwLock::new(false)),
708 subprotocol,
709 last_activity: Arc::new(RwLock::new(Instant::now())),
710 config: ConnectionConfig::default(),
711 }
712 }
713
714 pub fn subprotocol(&self) -> Option<&str> {
731 self.subprotocol.as_deref()
732 }
733
734 pub fn id(&self) -> &str {
747 &self.id
748 }
749
750 pub fn config(&self) -> &ConnectionConfig {
752 &self.config
753 }
754
755 pub async fn record_activity(&self) {
776 *self.last_activity.write().await = Instant::now();
777 }
778
779 pub async fn idle_duration(&self) -> Duration {
797 self.last_activity.read().await.elapsed()
798 }
799
800 pub async fn is_idle(&self) -> bool {
817 self.idle_duration().await > self.config.idle_timeout
818 }
819
820 pub async fn send(&self, message: Message) -> WebSocketResult<()> {
842 if *self.closed.read().await {
843 return Err(WebSocketError::Send("Connection closed".to_string()));
844 }
845
846 let result = self
847 .tx
848 .send(message)
849 .map_err(|e| WebSocketError::Send(e.to_string()));
850
851 if result.is_ok() {
852 self.record_activity().await;
853 }
854
855 result
856 }
857 pub async fn send_text(&self, text: String) -> WebSocketResult<()> {
879 self.send(Message::text(text)).await
880 }
881 pub async fn send_binary(&self, data: Vec<u8>) -> WebSocketResult<()> {
904 self.send(Message::binary(data)).await
905 }
906 pub async fn send_json<T: serde::Serialize>(&self, data: &T) -> WebSocketResult<()> {
936 let message = Message::json(data)?;
937 self.send(message).await
938 }
939 pub async fn close(&self) -> WebSocketResult<()> {
960 *self.closed.write().await = true;
962
963 self.tx
965 .send(Message::Close {
966 code: 1000,
967 reason: "Normal closure".to_string(),
968 })
969 .map_err(|e| WebSocketError::Send(e.to_string()))
970 }
971 pub async fn close_with_reason(&self, code: u16, reason: String) -> WebSocketResult<()> {
1001 *self.closed.write().await = true;
1003
1004 self.tx
1006 .send(Message::Close { code, reason })
1007 .map_err(|e| WebSocketError::Send(e.to_string()))
1008 }
1009
1010 pub async fn force_close(&self) {
1030 *self.closed.write().await = true;
1031 }
1032
1033 pub async fn is_closed(&self) -> bool {
1049 *self.closed.read().await
1050 }
1051}
1052
1053pub struct ConnectionTimeoutMonitor {
1083 connections: Arc<RwLock<HashMap<String, Arc<WebSocketConnection>>>>,
1084 config: ConnectionConfig,
1085}
1086
1087impl ConnectionTimeoutMonitor {
1088 pub fn new(config: ConnectionConfig) -> Self {
1090 Self {
1091 connections: Arc::new(RwLock::new(HashMap::new())),
1092 config,
1093 }
1094 }
1095
1096 pub async fn register(
1101 &self,
1102 connection: Arc<WebSocketConnection>,
1103 ) -> Result<(), WebSocketError> {
1104 let mut connections = self.connections.write().await;
1105
1106 if let Some(max) = self.config.max_connections
1107 && connections.len() >= max
1108 {
1109 return Err(WebSocketError::Connection(format!(
1110 "maximum connection limit reached ({})",
1111 max
1112 )));
1113 }
1114
1115 connections.insert(connection.id().to_string(), connection);
1116 Ok(())
1117 }
1118
1119 pub async fn unregister(&self, connection_id: &str) {
1121 self.connections.write().await.remove(connection_id);
1122 }
1123
1124 pub async fn connection_count(&self) -> usize {
1126 self.connections.read().await.len()
1127 }
1128
1129 pub async fn check_idle_connections(&self) -> Vec<String> {
1133 let connections = self.connections.read().await;
1134 let mut timed_out = Vec::new();
1135
1136 for (id, conn) in connections.iter() {
1137 if conn.is_closed().await {
1138 timed_out.push(id.clone());
1139 continue;
1140 }
1141
1142 let idle_duration = conn.idle_duration().await;
1143 if idle_duration > self.config.idle_timeout {
1144 let reason = format!(
1145 "Idle timeout: connection idle for {}s (limit: {}s)",
1146 idle_duration.as_secs(),
1147 self.config.idle_timeout.as_secs()
1148 );
1149 let _ = conn.close_with_reason(1001, reason).await;
1151 timed_out.push(id.clone());
1152 }
1153 }
1154
1155 drop(connections);
1156
1157 if !timed_out.is_empty() {
1159 let mut connections = self.connections.write().await;
1160 for id in &timed_out {
1161 connections.remove(id);
1162 }
1163 }
1164
1165 timed_out
1166 }
1167
1168 pub async fn shutdown_all(&self) -> Vec<String> {
1176 let mut connections = self.connections.write().await;
1177 let mut shut_down = Vec::with_capacity(connections.len());
1178
1179 for (id, conn) in connections.drain() {
1180 if !conn.is_closed().await {
1181 let _ = conn
1182 .close_with_reason(1001, "Server shutting down".to_string())
1183 .await;
1184 }
1185 shut_down.push(id);
1186 }
1187
1188 shut_down
1189 }
1190
1191 pub fn start(self: &Arc<Self>) -> tokio::task::JoinHandle<()> {
1217 let monitor = Arc::clone(self);
1218 tokio::spawn(async move {
1219 let mut interval = tokio::time::interval(monitor.config.cleanup_interval);
1220 loop {
1221 interval.tick().await;
1222 monitor.check_idle_connections().await;
1223 }
1224 })
1225 }
1226}
1227
1228#[derive(Debug, Clone)]
1248pub struct HeartbeatConfig {
1249 ping_interval: Duration,
1251 pong_timeout: Duration,
1253}
1254
1255impl HeartbeatConfig {
1256 pub fn new(ping_interval: Duration, pong_timeout: Duration) -> Self {
1258 Self {
1259 ping_interval,
1260 pong_timeout,
1261 }
1262 }
1263
1264 pub fn ping_interval(&self) -> Duration {
1266 self.ping_interval
1267 }
1268
1269 pub fn pong_timeout(&self) -> Duration {
1271 self.pong_timeout
1272 }
1273}
1274
1275impl Default for HeartbeatConfig {
1276 fn default() -> Self {
1277 Self {
1278 ping_interval: Duration::from_secs(30),
1279 pong_timeout: Duration::from_secs(10),
1280 }
1281 }
1282}
1283
1284pub struct HeartbeatMonitor {
1309 connection: Arc<WebSocketConnection>,
1310 config: HeartbeatConfig,
1311 last_pong: Arc<RwLock<Instant>>,
1312 timed_out: Arc<RwLock<bool>>,
1313 pong_notify: Arc<tokio::sync::Notify>,
1314}
1315
1316impl HeartbeatMonitor {
1317 pub fn new(connection: Arc<WebSocketConnection>, config: HeartbeatConfig) -> Self {
1319 Self {
1320 connection,
1321 config,
1322 last_pong: Arc::new(RwLock::new(Instant::now())),
1323 timed_out: Arc::new(RwLock::new(false)),
1324 pong_notify: Arc::new(tokio::sync::Notify::new()),
1325 }
1326 }
1327
1328 pub async fn record_pong(&self) {
1333 *self.last_pong.write().await = Instant::now();
1334 self.pong_notify.notify_one();
1335 }
1336
1337 pub async fn time_since_last_pong(&self) -> Duration {
1339 self.last_pong.read().await.elapsed()
1340 }
1341
1342 pub async fn is_timed_out(&self) -> bool {
1344 *self.timed_out.read().await
1345 }
1346
1347 pub async fn check_heartbeat(&self) -> bool {
1352 let since_pong = self.time_since_last_pong().await;
1353
1354 if since_pong > self.config.pong_timeout {
1355 self.connection.force_close().await;
1356 *self.timed_out.write().await = true;
1357 return true;
1358 }
1359
1360 false
1361 }
1362
1363 pub async fn send_ping(&self) -> WebSocketResult<()> {
1368 self.connection.send(Message::Ping).await
1369 }
1370
1371 pub fn config(&self) -> &HeartbeatConfig {
1373 &self.config
1374 }
1375
1376 pub fn connection(&self) -> &Arc<WebSocketConnection> {
1378 &self.connection
1379 }
1380
1381 pub fn start(self: &Arc<Self>) -> tokio::task::JoinHandle<()> {
1388 let monitor = Arc::clone(self);
1389 tokio::spawn(async move {
1390 let mut interval = tokio::time::interval(monitor.config.ping_interval);
1391 loop {
1392 interval.tick().await;
1393
1394 if monitor.connection.is_closed().await {
1395 break;
1396 }
1397
1398 let _ = monitor.send_ping().await;
1400
1401 tokio::select! {
1405 () = tokio::time::sleep(monitor.config.pong_timeout) => {
1406 if monitor.check_heartbeat().await {
1408 break;
1409 }
1410 }
1411 () = monitor.pong_notify.notified() => {
1412 }
1414 }
1415 }
1416 })
1417 }
1418}
1419
1420#[cfg(test)]
1421mod tests {
1422 use super::*;
1423 use rstest::rstest;
1424
1425 #[rstest]
1426 fn test_message_text() {
1427 let text = "Hello".to_string();
1429
1430 let msg = Message::text(text);
1432
1433 match msg {
1435 Message::Text { data } => assert_eq!(data, "Hello"),
1436 _ => panic!("Expected text message"),
1437 }
1438 }
1439
1440 #[rstest]
1441 fn test_message_json() {
1442 #[derive(serde::Serialize)]
1444 struct TestData {
1445 value: i32,
1446 }
1447 let data = TestData { value: 42 };
1448
1449 let msg = Message::json(&data).unwrap();
1451
1452 match msg {
1454 Message::Text { data } => assert!(data.contains("42")),
1455 _ => panic!("Expected text message"),
1456 }
1457 }
1458
1459 #[rstest]
1460 #[tokio::test]
1461 async fn test_connection_send() {
1462 let (tx, mut rx) = mpsc::unbounded_channel();
1464 let conn = WebSocketConnection::new("test".to_string(), tx);
1465
1466 conn.send_text("Hello".to_string()).await.unwrap();
1468
1469 let received = rx.recv().await.unwrap();
1471 match received {
1472 Message::Text { data } => assert_eq!(data, "Hello"),
1473 _ => panic!("Expected text message"),
1474 }
1475 }
1476
1477 #[rstest]
1478 fn test_connection_config_default() {
1479 let config = ConnectionConfig::new();
1481
1482 assert_eq!(config.idle_timeout(), Duration::from_secs(300));
1484 assert_eq!(config.handshake_timeout(), Duration::from_secs(10));
1485 assert_eq!(config.cleanup_interval(), Duration::from_secs(30));
1486 }
1487
1488 #[rstest]
1489 fn test_connection_config_strict() {
1490 let config = ConnectionConfig::strict();
1492
1493 assert_eq!(config.idle_timeout(), Duration::from_secs(30));
1495 assert_eq!(config.handshake_timeout(), Duration::from_secs(5));
1496 assert_eq!(config.cleanup_interval(), Duration::from_secs(10));
1497 }
1498
1499 #[rstest]
1500 fn test_connection_config_permissive() {
1501 let config = ConnectionConfig::permissive();
1503
1504 assert_eq!(config.idle_timeout(), Duration::from_secs(3600));
1506 assert_eq!(config.handshake_timeout(), Duration::from_secs(30));
1507 assert_eq!(config.cleanup_interval(), Duration::from_secs(60));
1508 }
1509
1510 #[rstest]
1511 fn test_connection_config_no_timeout() {
1512 let config = ConnectionConfig::no_timeout();
1514
1515 assert_eq!(config.idle_timeout(), Duration::MAX);
1517 assert_eq!(config.handshake_timeout(), Duration::MAX);
1518 }
1519
1520 #[rstest]
1521 fn test_connection_config_builder() {
1522 let config = ConnectionConfig::new()
1524 .with_idle_timeout(Duration::from_secs(120))
1525 .with_handshake_timeout(Duration::from_secs(15))
1526 .with_cleanup_interval(Duration::from_secs(20));
1527
1528 assert_eq!(config.idle_timeout(), Duration::from_secs(120));
1530 assert_eq!(config.handshake_timeout(), Duration::from_secs(15));
1531 assert_eq!(config.cleanup_interval(), Duration::from_secs(20));
1532 }
1533
1534 #[rstest]
1535 #[tokio::test]
1536 async fn test_connection_with_config() {
1537 let config = ConnectionConfig::new().with_idle_timeout(Duration::from_secs(60));
1539 let (tx, _rx) = mpsc::unbounded_channel();
1540
1541 let conn = WebSocketConnection::with_config("test".to_string(), tx, config);
1543
1544 assert_eq!(conn.config().idle_timeout(), Duration::from_secs(60));
1546 assert!(!conn.is_idle().await);
1547 }
1548
1549 #[rstest]
1550 #[tokio::test]
1551 async fn test_connection_record_activity_resets_idle() {
1552 let config = ConnectionConfig::new().with_idle_timeout(Duration::from_millis(50));
1554 let (tx, _rx) = mpsc::unbounded_channel();
1555 let conn = WebSocketConnection::with_config("test".to_string(), tx, config);
1556
1557 tokio::time::sleep(Duration::from_millis(60)).await;
1559 assert!(conn.is_idle().await);
1560
1561 conn.record_activity().await;
1563
1564 assert!(!conn.is_idle().await);
1566 }
1567
1568 #[rstest]
1569 #[tokio::test]
1570 async fn test_connection_becomes_idle_after_timeout() {
1571 let config = ConnectionConfig::new().with_idle_timeout(Duration::from_millis(50));
1573 let (tx, _rx) = mpsc::unbounded_channel();
1574 let conn = WebSocketConnection::with_config("test".to_string(), tx, config);
1575
1576 tokio::time::sleep(Duration::from_millis(60)).await;
1578
1579 assert!(conn.is_idle().await);
1581 assert!(conn.idle_duration().await >= Duration::from_millis(50));
1582 }
1583
1584 #[rstest]
1585 #[tokio::test]
1586 async fn test_send_resets_activity() {
1587 let config = ConnectionConfig::new().with_idle_timeout(Duration::from_millis(100));
1589 let (tx, mut _rx) = mpsc::unbounded_channel();
1590 let conn = WebSocketConnection::with_config("test".to_string(), tx, config);
1591
1592 tokio::time::sleep(Duration::from_millis(50)).await;
1594 conn.send_text("ping".to_string()).await.unwrap();
1595
1596 assert!(conn.idle_duration().await < Duration::from_millis(30));
1598 assert!(!conn.is_idle().await);
1599 }
1600
1601 #[rstest]
1602 #[tokio::test]
1603 async fn test_close_with_reason() {
1604 let (tx, mut rx) = mpsc::unbounded_channel();
1606 let conn = WebSocketConnection::new("test".to_string(), tx);
1607
1608 conn.close_with_reason(1001, "Idle timeout".to_string())
1610 .await
1611 .unwrap();
1612
1613 assert!(conn.is_closed().await);
1615 let msg = rx.recv().await.unwrap();
1616 match msg {
1617 Message::Close { code, reason } => {
1618 assert_eq!(code, 1001);
1619 assert_eq!(reason, "Idle timeout");
1620 }
1621 _ => panic!("Expected close message"),
1622 }
1623 }
1624
1625 #[rstest]
1626 #[tokio::test]
1627 async fn test_timeout_monitor_register_and_count() {
1628 let config = ConnectionConfig::new();
1630 let monitor = ConnectionTimeoutMonitor::new(config);
1631 let (tx, _rx) = mpsc::unbounded_channel();
1632 let conn = Arc::new(WebSocketConnection::new("conn_1".to_string(), tx));
1633
1634 monitor.register(conn).await.unwrap();
1636
1637 assert_eq!(monitor.connection_count().await, 1);
1639 }
1640
1641 #[rstest]
1642 #[tokio::test]
1643 async fn test_timeout_monitor_unregister() {
1644 let config = ConnectionConfig::new();
1646 let monitor = ConnectionTimeoutMonitor::new(config);
1647 let (tx, _rx) = mpsc::unbounded_channel();
1648 let conn = Arc::new(WebSocketConnection::new("conn_1".to_string(), tx));
1649 monitor.register(conn).await.unwrap();
1650
1651 monitor.unregister("conn_1").await;
1653
1654 assert_eq!(monitor.connection_count().await, 0);
1656 }
1657
1658 #[rstest]
1659 #[tokio::test]
1660 async fn test_timeout_monitor_closes_idle_connections() {
1661 let config = ConnectionConfig::new().with_idle_timeout(Duration::from_millis(50));
1663 let monitor = ConnectionTimeoutMonitor::new(config);
1664
1665 let (tx1, mut rx1) = mpsc::unbounded_channel();
1666 let conn1 = Arc::new(WebSocketConnection::with_config(
1667 "idle_conn".to_string(),
1668 tx1,
1669 ConnectionConfig::new().with_idle_timeout(Duration::from_millis(50)),
1670 ));
1671
1672 let (tx2, _rx2) = mpsc::unbounded_channel();
1673 let conn2 = Arc::new(WebSocketConnection::with_config(
1674 "active_conn".to_string(),
1675 tx2,
1676 ConnectionConfig::new().with_idle_timeout(Duration::from_secs(300)),
1677 ));
1678
1679 monitor.register(conn1).await.unwrap();
1680 monitor.register(conn2.clone()).await.unwrap();
1681
1682 tokio::time::sleep(Duration::from_millis(60)).await;
1684 conn2.record_activity().await;
1686
1687 let timed_out = monitor.check_idle_connections().await;
1688
1689 assert_eq!(timed_out.len(), 1);
1691 assert_eq!(timed_out[0], "idle_conn");
1692 assert_eq!(monitor.connection_count().await, 1);
1693
1694 let msg = rx1.recv().await.unwrap();
1696 match msg {
1697 Message::Close { code, reason } => {
1698 assert_eq!(code, 1001);
1699 assert!(reason.contains("Idle timeout"));
1700 }
1701 _ => panic!("Expected close message for idle connection"),
1702 }
1703 }
1704
1705 #[rstest]
1706 #[tokio::test]
1707 async fn test_timeout_monitor_removes_already_closed_connections() {
1708 let config = ConnectionConfig::new();
1710 let monitor = ConnectionTimeoutMonitor::new(config);
1711 let (tx, _rx) = mpsc::unbounded_channel();
1712 let conn = Arc::new(WebSocketConnection::new("conn_1".to_string(), tx));
1713 conn.close().await.unwrap();
1714 monitor.register(conn).await.unwrap();
1715
1716 let timed_out = monitor.check_idle_connections().await;
1718
1719 assert_eq!(timed_out.len(), 1);
1721 assert_eq!(timed_out[0], "conn_1");
1722 assert_eq!(monitor.connection_count().await, 0);
1723 }
1724
1725 #[rstest]
1726 #[tokio::test]
1727 async fn test_timeout_monitor_background_task() {
1728 let config = ConnectionConfig::new()
1730 .with_idle_timeout(Duration::from_millis(30))
1731 .with_cleanup_interval(Duration::from_millis(20));
1732 let monitor = Arc::new(ConnectionTimeoutMonitor::new(config));
1733
1734 let (tx, mut rx) = mpsc::unbounded_channel();
1735 let conn = Arc::new(WebSocketConnection::with_config(
1736 "bg_conn".to_string(),
1737 tx,
1738 ConnectionConfig::new().with_idle_timeout(Duration::from_millis(30)),
1739 ));
1740 monitor.register(conn).await.unwrap();
1741
1742 let handle = monitor.start();
1744
1745 tokio::time::sleep(Duration::from_millis(120)).await;
1747
1748 assert_eq!(monitor.connection_count().await, 0);
1750
1751 let msg = rx.recv().await.unwrap();
1753 assert!(matches!(msg, Message::Close { .. }));
1754
1755 handle.abort();
1757 }
1758
1759 #[rstest]
1760 fn test_ping_pong_config_default() {
1761 let config = PingPongConfig::default();
1763
1764 assert_eq!(config.ping_interval(), Duration::from_secs(30));
1766 assert_eq!(config.pong_timeout(), Duration::from_secs(10));
1767 }
1768
1769 #[rstest]
1770 fn test_ping_pong_config_custom() {
1771 let config = PingPongConfig::new(Duration::from_secs(15), Duration::from_secs(5));
1773
1774 assert_eq!(config.ping_interval(), Duration::from_secs(15));
1776 assert_eq!(config.pong_timeout(), Duration::from_secs(5));
1777 }
1778
1779 #[rstest]
1780 fn test_ping_pong_config_builder() {
1781 let config = PingPongConfig::default()
1783 .with_ping_interval(Duration::from_secs(60))
1784 .with_pong_timeout(Duration::from_secs(20));
1785
1786 assert_eq!(config.ping_interval(), Duration::from_secs(60));
1788 assert_eq!(config.pong_timeout(), Duration::from_secs(20));
1789 }
1790
1791 #[rstest]
1792 fn test_connection_config_has_default_ping_config() {
1793 let config = ConnectionConfig::new();
1795
1796 assert_eq!(
1798 config.ping_config().ping_interval(),
1799 Duration::from_secs(30)
1800 );
1801 assert_eq!(config.ping_config().pong_timeout(), Duration::from_secs(10));
1802 }
1803
1804 #[rstest]
1805 fn test_connection_config_with_custom_ping_config() {
1806 let ping_config = PingPongConfig::new(Duration::from_secs(15), Duration::from_secs(5));
1808
1809 let config = ConnectionConfig::new().with_ping_config(ping_config);
1811
1812 assert_eq!(
1814 config.ping_config().ping_interval(),
1815 Duration::from_secs(15)
1816 );
1817 assert_eq!(config.ping_config().pong_timeout(), Duration::from_secs(5));
1818 }
1819
1820 #[rstest]
1821 fn test_strict_config_has_aggressive_ping() {
1822 let config = ConnectionConfig::strict();
1824
1825 assert_eq!(
1827 config.ping_config().ping_interval(),
1828 Duration::from_secs(10)
1829 );
1830 assert_eq!(config.ping_config().pong_timeout(), Duration::from_secs(5));
1831 }
1832
1833 #[rstest]
1834 fn test_permissive_config_has_relaxed_ping() {
1835 let config = ConnectionConfig::permissive();
1837
1838 assert_eq!(
1840 config.ping_config().ping_interval(),
1841 Duration::from_secs(60)
1842 );
1843 assert_eq!(config.ping_config().pong_timeout(), Duration::from_secs(30));
1844 }
1845
1846 #[rstest]
1847 #[tokio::test]
1848 async fn test_timeout_monitor_rejects_when_max_connections_reached() {
1849 let config = ConnectionConfig::new().with_max_connections(Some(1));
1851 let monitor = ConnectionTimeoutMonitor::new(config);
1852
1853 let (tx1, _rx1) = mpsc::unbounded_channel();
1854 let conn1 = Arc::new(WebSocketConnection::new("conn_1".to_string(), tx1));
1855
1856 let (tx2, _rx2) = mpsc::unbounded_channel();
1857 let conn2 = Arc::new(WebSocketConnection::new("conn_2".to_string(), tx2));
1858
1859 monitor.register(conn1).await.unwrap();
1861 let result = monitor.register(conn2).await;
1862
1863 assert!(result.is_err());
1865 assert_eq!(monitor.connection_count().await, 1);
1866 }
1867
1868 #[rstest]
1869 #[tokio::test]
1870 async fn test_force_close_marks_connection_closed() {
1871 let (tx, _rx) = mpsc::unbounded_channel();
1873 let conn = WebSocketConnection::new("test".to_string(), tx);
1874
1875 conn.force_close().await;
1877
1878 assert!(conn.is_closed().await);
1880 }
1881
1882 #[rstest]
1883 #[tokio::test]
1884 async fn test_close_marks_closed_even_when_channel_dropped() {
1885 let (tx, rx) = mpsc::unbounded_channel();
1887 let conn = WebSocketConnection::new("test".to_string(), tx);
1888
1889 drop(rx);
1891
1892 let result = conn.close().await;
1894
1895 assert!(result.is_err()); assert!(conn.is_closed().await); }
1899
1900 #[rstest]
1901 #[tokio::test]
1902 async fn test_close_with_reason_marks_closed_even_when_channel_dropped() {
1903 let (tx, rx) = mpsc::unbounded_channel();
1905 let conn = WebSocketConnection::new("test".to_string(), tx);
1906
1907 drop(rx);
1909
1910 let result = conn
1912 .close_with_reason(1006, "Abnormal close".to_string())
1913 .await;
1914
1915 assert!(result.is_err());
1917 assert!(conn.is_closed().await);
1918 }
1919
1920 #[rstest]
1921 #[tokio::test]
1922 async fn test_send_after_force_close_returns_error() {
1923 let (tx, _rx) = mpsc::unbounded_channel();
1925 let conn = WebSocketConnection::new("test".to_string(), tx);
1926 conn.force_close().await;
1927
1928 let result = conn.send_text("should fail".to_string()).await;
1930
1931 assert!(result.is_err());
1933 assert!(matches!(result.unwrap_err(), WebSocketError::Send(_)));
1934 }
1935
1936 #[rstest]
1937 fn test_heartbeat_config_default() {
1938 let config = HeartbeatConfig::default();
1940
1941 assert_eq!(config.ping_interval(), Duration::from_secs(30));
1943 assert_eq!(config.pong_timeout(), Duration::from_secs(10));
1944 }
1945
1946 #[rstest]
1947 fn test_heartbeat_config_custom() {
1948 let config = HeartbeatConfig::new(Duration::from_secs(15), Duration::from_secs(5));
1950
1951 assert_eq!(config.ping_interval(), Duration::from_secs(15));
1953 assert_eq!(config.pong_timeout(), Duration::from_secs(5));
1954 }
1955
1956 #[rstest]
1957 #[tokio::test]
1958 async fn test_heartbeat_monitor_initial_state() {
1959 let (tx, _rx) = mpsc::unbounded_channel();
1961 let conn = Arc::new(WebSocketConnection::new("hb_test".to_string(), tx));
1962 let config = HeartbeatConfig::default();
1963
1964 let monitor = HeartbeatMonitor::new(conn, config);
1966
1967 assert!(!monitor.is_timed_out().await);
1969 assert!(monitor.time_since_last_pong().await < Duration::from_secs(1));
1970 }
1971
1972 #[rstest]
1973 #[tokio::test]
1974 async fn test_heartbeat_monitor_record_pong_resets_timer() {
1975 let (tx, _rx) = mpsc::unbounded_channel();
1977 let conn = Arc::new(WebSocketConnection::new("hb_pong".to_string(), tx));
1978 let config = HeartbeatConfig::new(Duration::from_millis(50), Duration::from_millis(30));
1979 let monitor = HeartbeatMonitor::new(conn, config);
1980
1981 tokio::time::sleep(Duration::from_millis(20)).await;
1983 monitor.record_pong().await;
1984
1985 assert!(monitor.time_since_last_pong().await < Duration::from_millis(10));
1987 }
1988
1989 #[rstest]
1990 #[tokio::test]
1991 async fn test_heartbeat_monitor_timeout_closes_connection() {
1992 let (tx, _rx) = mpsc::unbounded_channel();
1994 let conn = Arc::new(WebSocketConnection::new("hb_timeout".to_string(), tx));
1995 let config = HeartbeatConfig::new(Duration::from_millis(50), Duration::from_millis(30));
1996 let monitor = HeartbeatMonitor::new(conn.clone(), config);
1997
1998 tokio::time::sleep(Duration::from_millis(40)).await;
2000 let timed_out = monitor.check_heartbeat().await;
2001
2002 assert!(timed_out);
2004 assert!(monitor.is_timed_out().await);
2005 assert!(conn.is_closed().await);
2006 }
2007
2008 #[rstest]
2009 #[tokio::test]
2010 async fn test_heartbeat_monitor_no_timeout_when_pong_received() {
2011 let (tx, _rx) = mpsc::unbounded_channel();
2013 let conn = Arc::new(WebSocketConnection::new("hb_ok".to_string(), tx));
2014 let config = HeartbeatConfig::new(Duration::from_millis(100), Duration::from_millis(50));
2015 let monitor = HeartbeatMonitor::new(conn.clone(), config);
2016
2017 tokio::time::sleep(Duration::from_millis(20)).await;
2019 monitor.record_pong().await;
2020 let timed_out = monitor.check_heartbeat().await;
2021
2022 assert!(!timed_out);
2024 assert!(!monitor.is_timed_out().await);
2025 assert!(!conn.is_closed().await);
2026 }
2027
2028 #[rstest]
2029 #[tokio::test]
2030 async fn test_heartbeat_monitor_send_ping() {
2031 let (tx, mut rx) = mpsc::unbounded_channel();
2033 let conn = Arc::new(WebSocketConnection::new("hb_ping".to_string(), tx));
2034 let config = HeartbeatConfig::default();
2035 let monitor = HeartbeatMonitor::new(conn, config);
2036
2037 monitor.send_ping().await.unwrap();
2039
2040 let msg = rx.recv().await.unwrap();
2042 assert!(matches!(msg, Message::Ping));
2043 }
2044
2045 #[rstest]
2046 #[tokio::test]
2047 async fn test_heartbeat_monitor_early_pong_skips_full_sleep() {
2048 let (tx, _rx) = mpsc::unbounded_channel();
2050 let conn = Arc::new(WebSocketConnection::new("hb_early".to_string(), tx));
2051 let config = HeartbeatConfig {
2053 ping_interval: Duration::from_secs(60),
2054 pong_timeout: Duration::from_secs(10),
2055 };
2056 let monitor = Arc::new(HeartbeatMonitor::new(conn, config));
2057
2058 let monitor_clone = Arc::clone(&monitor);
2060 tokio::spawn(async move {
2061 tokio::time::sleep(Duration::from_millis(50)).await;
2062 monitor_clone.record_pong().await;
2063 });
2064
2065 let _ = monitor.send_ping().await;
2067 let start = Instant::now();
2068
2069 tokio::select! {
2070 () = tokio::time::sleep(monitor.config.pong_timeout) => {
2071 panic!("Should not reach full timeout");
2072 }
2073 () = monitor.pong_notify.notified() => {
2074 }
2076 }
2077
2078 let elapsed = start.elapsed();
2080 assert!(
2081 elapsed < Duration::from_secs(2),
2082 "Expected early wakeup but elapsed {:?}",
2083 elapsed
2084 );
2085 }
2086
2087 #[rstest]
2088 fn test_websocket_error_binary_payload_variant() {
2089 let err = WebSocketError::BinaryPayload("invalid data".to_string());
2091
2092 assert_eq!(err.to_string(), "Invalid binary payload: invalid data");
2094 }
2095
2096 #[rstest]
2097 fn test_websocket_error_heartbeat_timeout_variant() {
2098 let err = WebSocketError::HeartbeatTimeout(Duration::from_secs(10));
2100
2101 assert_eq!(
2103 err.to_string(),
2104 "Heartbeat timeout: no pong received within 10s"
2105 );
2106 }
2107
2108 #[rstest]
2109 fn test_websocket_error_slow_consumer_variant() {
2110 let err = WebSocketError::SlowConsumer(Duration::from_secs(5));
2112
2113 assert_eq!(err.to_string(), "Slow consumer: send timed out after 5s");
2115 }
2116}