1use std::sync::Arc;
7
8use crate::auth::AuthHandler;
9use crate::config::RoutingMode;
10use crate::metrics::MetricsCollector;
11use crate::pool::BufferPool;
12use crate::router::BackendSelector;
13use crate::types::{ClientAddress, ClientId};
14
15use super::{AuthState, ModeState, SessionMode};
16
17pub struct ClientSession {
19 pub(super) client_addr: ClientAddress,
20 pub(super) buffer_pool: BufferPool,
21 pub(super) client_id: ClientId,
23 pub(super) router: Option<Arc<BackendSelector>>,
25 pub(super) mode_state: ModeState,
27 pub(super) auth_handler: Arc<AuthHandler>,
29 pub(super) auth_state: AuthState,
31 pub(super) metrics: crate::metrics::MetricsCollector,
33
34 pub(super) connection_stats: Option<crate::metrics::ConnectionStatsAggregator>,
36
37 pub(super) cache: Arc<crate::cache::UnifiedCache>,
39
40 pub(super) cache_articles: bool,
42
43 pub(super) adaptive_precheck: bool,
45}
46
47pub struct ClientSessionBuilder {
81 client_addr: ClientAddress,
82 buffer_pool: BufferPool,
83 router: Option<Arc<BackendSelector>>,
84 routing_mode: RoutingMode,
85 auth_handler: Arc<AuthHandler>,
86 metrics: MetricsCollector,
87 connection_stats: Option<crate::metrics::ConnectionStatsAggregator>,
88 cache: BuilderCache,
89 cache_articles: bool,
90 adaptive_precheck: bool,
91}
92
93enum BuilderCache {
94 DefaultAvailability,
95 Shared(Arc<crate::cache::UnifiedCache>),
96}
97
98impl BuilderCache {
99 fn into_cache(self) -> Arc<crate::cache::UnifiedCache> {
100 match self {
101 Self::DefaultAvailability => ClientSession::default_cache(),
102 Self::Shared(cache) => cache,
103 }
104 }
105}
106
107impl ClientSessionBuilder {
108 #[must_use]
113 pub fn with_router(mut self, router: Arc<BackendSelector>) -> Self {
114 self.router = Some(router);
115 self
116 }
117
118 #[must_use]
125 pub const fn with_routing_mode(mut self, mode: RoutingMode) -> Self {
126 self.routing_mode = mode;
127 self
128 }
129
130 #[must_use]
132 pub fn with_auth_handler(mut self, auth_handler: Arc<AuthHandler>) -> Self {
133 self.auth_handler = auth_handler;
134 self
135 }
136
137 #[must_use]
139 pub fn with_connection_stats(
140 mut self,
141 connection_stats: crate::metrics::ConnectionStatsAggregator,
142 ) -> Self {
143 self.connection_stats = Some(connection_stats);
144 self
145 }
146
147 #[must_use]
149 pub fn with_cache(mut self, cache: Arc<crate::cache::UnifiedCache>) -> Self {
150 self.cache = BuilderCache::Shared(cache);
151 self
152 }
153
154 #[must_use]
159 pub const fn with_cache_articles(mut self, cache: bool) -> Self {
160 self.cache_articles = cache;
161 self
162 }
163
164 #[must_use]
166 pub const fn with_adaptive_precheck(mut self, enable: bool) -> Self {
167 self.adaptive_precheck = enable;
168 self
169 }
170
171 #[must_use]
176 pub fn build(self) -> ClientSession {
177 let (mode, routing_mode) = match (&self.router, self.routing_mode) {
178 (Some(_), RoutingMode::PerCommand | RoutingMode::Hybrid) => {
180 (SessionMode::PerCommand, self.routing_mode)
181 }
182 (Some(_), RoutingMode::Stateful) | (None, _) => {
184 (SessionMode::Stateful, RoutingMode::Stateful)
185 }
186 };
187
188 ClientSession {
189 client_addr: self.client_addr,
190 buffer_pool: self.buffer_pool,
191 client_id: ClientId::new(),
192 router: self.router,
193 mode_state: ModeState::new(mode, routing_mode),
194 auth_handler: self.auth_handler,
195 auth_state: AuthState::new(),
196 metrics: self.metrics,
197 connection_stats: self.connection_stats,
198 cache: self.cache.into_cache(),
199 cache_articles: self.cache_articles,
200 adaptive_precheck: self.adaptive_precheck,
201 }
202 }
203}
204
205impl ClientSession {
206 fn default_cache() -> Arc<crate::cache::UnifiedCache> {
208 Arc::new(crate::cache::UnifiedCache::availability(
209 std::time::Duration::MAX,
210 ))
211 }
212
213 #[must_use]
215 pub fn new(
216 client_addr: ClientAddress,
217 buffer_pool: BufferPool,
218 auth_handler: Arc<AuthHandler>,
219 metrics: MetricsCollector,
220 ) -> Self {
221 Self {
222 client_addr,
223 buffer_pool,
224 client_id: ClientId::new(),
225 router: None,
226 mode_state: ModeState::new(SessionMode::Stateful, RoutingMode::Stateful),
227 auth_handler,
228 auth_state: AuthState::new(),
229 metrics,
230 connection_stats: None,
231 cache: Self::default_cache(),
232 cache_articles: false,
233 adaptive_precheck: false,
234 }
235 }
236
237 #[must_use]
242 pub fn new_with_router(
243 client_addr: ClientAddress,
244 buffer_pool: BufferPool,
245 router: Arc<BackendSelector>,
246 routing_mode: RoutingMode,
247 auth_handler: Arc<AuthHandler>,
248 metrics: MetricsCollector,
249 ) -> Self {
250 Self {
251 client_addr,
252 buffer_pool,
253 client_id: ClientId::new(),
254 router: Some(router),
255 mode_state: ModeState::new(SessionMode::PerCommand, routing_mode),
256 auth_handler,
257 auth_state: AuthState::new(),
258 metrics,
259 connection_stats: None,
260 cache: Self::default_cache(),
261 cache_articles: false,
262 adaptive_precheck: false,
263 }
264 }
265
266 #[must_use]
288 pub fn builder(
289 client_addr: ClientAddress,
290 buffer_pool: BufferPool,
291 auth_handler: Arc<AuthHandler>,
292 metrics: MetricsCollector,
293 ) -> ClientSessionBuilder {
294 ClientSessionBuilder {
295 client_addr,
296 buffer_pool,
297 router: None,
298 routing_mode: RoutingMode::Stateful,
299 auth_handler,
300 metrics,
301 connection_stats: None,
302 cache: BuilderCache::DefaultAvailability,
303 cache_articles: false,
304 adaptive_precheck: false,
305 }
306 }
307
308 #[must_use]
312 #[inline]
313 pub const fn client_id(&self) -> ClientId {
314 self.client_id
315 }
316
317 #[must_use]
323 #[inline]
324 pub const fn is_per_command_routing(&self) -> bool {
325 self.router.is_some()
326 }
327
328 #[must_use]
330 #[inline]
331 pub fn mode(&self) -> SessionMode {
332 self.mode_state.mode()
333 }
334
335 #[inline]
340 #[must_use]
341 pub fn username(&self) -> Option<&str> {
342 self.auth_state.username()
343 }
344
345 pub(crate) fn set_username(&self, username: Option<String>) {
350 if let Some(name) = username {
351 self.auth_state.mark_authenticated(name);
352 }
353 }
354
355 #[must_use]
357 #[inline]
358 pub(crate) const fn connection_stats(
359 &self,
360 ) -> Option<&crate::metrics::ConnectionStatsAggregator> {
361 self.connection_stats.as_ref()
362 }
363
364 #[inline]
372 pub(crate) fn is_authenticated_cached(&self, skip_auth_check: bool) -> bool {
373 self.auth_state.is_authenticated_or_skipped(skip_auth_check)
374 }
375}
376
377#[cfg(test)]
378mod tests {
379 use super::*;
380 use crate::auth::AuthHandler;
381 use crate::metrics::MetricsCollector;
382 use crate::types::BufferSize;
383 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
384 use std::sync::Arc;
385
386 fn test_auth_handler() -> Arc<AuthHandler> {
388 Arc::new(AuthHandler::new(None, None).unwrap())
389 }
390
391 fn test_metrics() -> MetricsCollector {
393 MetricsCollector::new(1)
394 }
395
396 #[test]
397 fn test_client_session_creation() {
398 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
399 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
400 let session = ClientSession::new(
401 addr.into(),
402 buffer_pool,
403 test_auth_handler(),
404 test_metrics(),
405 );
406
407 assert_eq!(session.client_addr.port(), 8080);
408 assert_eq!(session.client_addr.ip(), IpAddr::V4(Ipv4Addr::LOCALHOST));
409 }
410
411 #[test]
412 fn test_client_session_with_different_ports() {
413 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
414
415 let addr1 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
416 let session1 = ClientSession::new(
417 addr1.into(),
418 buffer_pool.clone(),
419 test_auth_handler(),
420 test_metrics(),
421 );
422
423 let addr2 = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9090);
424 let session2 = ClientSession::new(
425 addr2.into(),
426 buffer_pool,
427 test_auth_handler(),
428 test_metrics(),
429 );
430
431 assert_ne!(session1.client_addr.port(), session2.client_addr.port());
432 assert_eq!(session1.client_addr.port(), 8080);
433 assert_eq!(session2.client_addr.port(), 9090);
434 }
435
436 #[test]
437 fn test_client_session_with_ipv6() {
438 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
439 let addr = SocketAddr::new(IpAddr::V6("::1".parse().unwrap()), 8119);
440 let session = ClientSession::new(
441 addr.into(),
442 buffer_pool,
443 test_auth_handler(),
444 test_metrics(),
445 );
446
447 assert_eq!(session.client_addr.port(), 8119);
448 assert!(session.client_addr.is_ipv6());
449 }
450
451 #[test]
452 fn test_buffer_pool_cloning() {
453 let buffer_pool = BufferPool::new(BufferSize::try_new(8192).unwrap(), 10);
454 let buffer_pool_clone = buffer_pool.clone();
455
456 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)), 1234);
457 let _session1 = ClientSession::new(
458 addr.into(),
459 buffer_pool,
460 test_auth_handler(),
461 test_metrics(),
462 );
463 let _session2 = ClientSession::new(
464 addr.into(),
465 buffer_pool_clone,
466 test_auth_handler(),
467 test_metrics(),
468 );
469 }
470
471 #[test]
472 fn test_session_addr_formatting() {
473 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), 5555);
474 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
475 let session = ClientSession::new(
476 addr.into(),
477 buffer_pool,
478 test_auth_handler(),
479 test_metrics(),
480 );
481
482 let addr_str = format!("{}", session.client_addr);
483 assert!(addr_str.contains("10.0.0.1"));
484 assert!(addr_str.contains("5555"));
485 }
486
487 #[test]
488 fn test_multiple_sessions_same_buffer_pool() {
489 let buffer_pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 8);
490 let sessions: Vec<_> = (0..5)
491 .map(|i| {
492 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8000 + i);
493 ClientSession::new(
494 addr.into(),
495 buffer_pool.clone(),
496 test_auth_handler(),
497 test_metrics(),
498 )
499 })
500 .collect();
501
502 assert_eq!(sessions.len(), 5);
503 for (i, session) in sessions.iter().enumerate() {
504 assert_eq!(session.client_addr.port(), 8000 + i as u16);
505 }
506 }
507
508 #[test]
509 fn test_loopback_address() {
510 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
511 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8119);
512 let session = ClientSession::new(
513 addr.into(),
514 buffer_pool,
515 test_auth_handler(),
516 test_metrics(),
517 );
518
519 assert!(session.client_addr.ip().is_loopback());
520 }
521
522 #[test]
523 fn test_unspecified_address() {
524 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
525 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), 0);
526 let session = ClientSession::new(
527 addr.into(),
528 buffer_pool,
529 test_auth_handler(),
530 test_metrics(),
531 );
532
533 assert!(session.client_addr.ip().is_unspecified());
534 assert_eq!(session.client_addr.port(), 0);
535 }
536
537 #[test]
538 fn test_session_without_router() {
539 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
540 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
541 let session = ClientSession::new(
542 addr.into(),
543 buffer_pool,
544 test_auth_handler(),
545 test_metrics(),
546 );
547
548 assert!(!session.is_per_command_routing());
549 assert_eq!(session.client_addr.port(), 8080);
550 }
551
552 #[test]
553 fn test_session_with_router() {
554 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
555 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
556 let router = Arc::new(BackendSelector::new());
557 let session = ClientSession::new_with_router(
558 addr.into(),
559 buffer_pool,
560 router,
561 RoutingMode::PerCommand,
562 test_auth_handler(),
563 test_metrics(),
564 );
565
566 assert!(session.is_per_command_routing());
567 assert_eq!(session.client_addr.port(), 8080);
568 }
569
570 #[test]
571 fn test_client_id_uniqueness() {
572 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
573 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
574
575 let session1 = ClientSession::new(
576 addr.into(),
577 buffer_pool.clone(),
578 test_auth_handler(),
579 test_metrics(),
580 );
581 let session2 = ClientSession::new(
582 addr.into(),
583 buffer_pool,
584 test_auth_handler(),
585 test_metrics(),
586 );
587
588 assert_ne!(session1.client_id(), session2.client_id());
589 }
590
591 #[test]
592 fn test_session_mode_enum() {
593 assert_eq!(SessionMode::PerCommand, SessionMode::PerCommand);
594 assert_eq!(SessionMode::Stateful, SessionMode::Stateful);
595 assert_ne!(SessionMode::PerCommand, SessionMode::Stateful);
596
597 let per_command = format!("{:?}", SessionMode::PerCommand);
598 let stateful = format!("{:?}", SessionMode::Stateful);
599 assert!(per_command.contains("PerCommand"));
600 assert!(stateful.contains("Stateful"));
601 }
602
603 #[test]
604 fn test_hybrid_session_creation() {
605 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
606 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
607 let router = Arc::new(BackendSelector::new());
608
609 let session = ClientSession::new_with_router(
610 addr.into(),
611 buffer_pool,
612 router,
613 RoutingMode::Hybrid,
614 test_auth_handler(),
615 test_metrics(),
616 );
617
618 assert!(session.is_per_command_routing());
619 assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
620 assert_eq!(session.mode(), SessionMode::PerCommand);
621 }
622
623 #[test]
624 fn test_routing_mode_configurations() {
625 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
626 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
627 let router = Arc::new(BackendSelector::new());
628
629 let session = ClientSession::new_with_router(
631 addr.into(),
632 buffer_pool.clone(),
633 router.clone(),
634 RoutingMode::Stateful,
635 test_auth_handler(),
636 test_metrics(),
637 );
638 assert!(session.is_per_command_routing());
639 assert_eq!(session.mode_state.routing_mode(), RoutingMode::Stateful);
640
641 let session = ClientSession::new_with_router(
643 addr.into(),
644 buffer_pool.clone(),
645 router.clone(),
646 RoutingMode::PerCommand,
647 test_auth_handler(),
648 test_metrics(),
649 );
650 assert!(session.is_per_command_routing());
651 assert_eq!(session.mode_state.routing_mode(), RoutingMode::PerCommand);
652 assert_eq!(session.mode(), SessionMode::PerCommand);
653
654 let session = ClientSession::new_with_router(
656 addr.into(),
657 buffer_pool,
658 router,
659 RoutingMode::Hybrid,
660 test_auth_handler(),
661 test_metrics(),
662 );
663 assert!(session.is_per_command_routing());
664 assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
665 assert_eq!(session.mode(), SessionMode::PerCommand);
666 }
667
668 #[test]
669 fn test_hybrid_mode_initial_state() {
670 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
671 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
672 let router = Arc::new(BackendSelector::new());
673
674 let session = ClientSession::new_with_router(
675 addr.into(),
676 buffer_pool,
677 router,
678 RoutingMode::Hybrid,
679 test_auth_handler(),
680 test_metrics(),
681 );
682
683 assert_eq!(session.mode(), SessionMode::PerCommand);
684 assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
685 assert!(session.is_per_command_routing());
686 }
687
688 #[test]
689 fn test_is_per_command_routing_logic() {
690 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
691 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
692 let router = Arc::new(BackendSelector::new());
693
694 let session = ClientSession::new_with_router(
696 addr.into(),
697 buffer_pool.clone(),
698 router.clone(),
699 RoutingMode::Stateful,
700 test_auth_handler(),
701 test_metrics(),
702 );
703 assert!(session.is_per_command_routing());
704
705 let session = ClientSession::new_with_router(
707 addr.into(),
708 buffer_pool.clone(),
709 router.clone(),
710 RoutingMode::PerCommand,
711 test_auth_handler(),
712 test_metrics(),
713 );
714 assert!(session.is_per_command_routing());
715
716 let session = ClientSession::new_with_router(
718 addr.into(),
719 buffer_pool.clone(),
720 router,
721 RoutingMode::Hybrid,
722 test_auth_handler(),
723 test_metrics(),
724 );
725 assert!(session.is_per_command_routing());
726
727 let session = ClientSession::new(
729 addr.into(),
730 buffer_pool,
731 test_auth_handler(),
732 test_metrics(),
733 );
734 assert!(!session.is_per_command_routing());
735 }
736
737 #[test]
740 fn test_builder_basic_construction() {
741 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
742 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
743 let auth_handler = test_auth_handler();
744
745 let session =
746 ClientSession::builder(addr.into(), buffer_pool, auth_handler, test_metrics()).build();
747
748 assert_eq!(*session.client_addr, addr);
749 assert!(!session.is_per_command_routing());
750 assert_eq!(session.mode(), SessionMode::Stateful);
751 assert_eq!(session.mode_state.routing_mode(), RoutingMode::Stateful);
752 }
753
754 #[test]
755 fn test_builder_with_router() {
756 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
757 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
758 let router = Arc::new(BackendSelector::new());
759
760 let session = ClientSession::builder(
761 addr.into(),
762 buffer_pool,
763 test_auth_handler(),
764 test_metrics(),
765 )
766 .with_router(router)
767 .with_routing_mode(RoutingMode::PerCommand)
768 .build();
769
770 assert!(session.is_per_command_routing());
771 assert_eq!(session.mode(), SessionMode::PerCommand);
772 assert_eq!(session.mode_state.routing_mode(), RoutingMode::PerCommand);
773 }
774
775 #[test]
776 fn test_builder_with_hybrid_mode() {
777 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
778 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
779 let router = Arc::new(BackendSelector::new());
780
781 let session = ClientSession::builder(
782 addr.into(),
783 buffer_pool,
784 test_auth_handler(),
785 test_metrics(),
786 )
787 .with_router(router)
788 .with_routing_mode(RoutingMode::Hybrid)
789 .build();
790
791 assert!(session.is_per_command_routing());
792 assert_eq!(session.mode(), SessionMode::PerCommand);
793 assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
794 }
795
796 #[test]
797 fn test_builder_with_stateful_mode() {
798 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
799 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
800 let router = Arc::new(BackendSelector::new());
801
802 let session = ClientSession::builder(
804 addr.into(),
805 buffer_pool,
806 test_auth_handler(),
807 test_metrics(),
808 )
809 .with_router(router)
810 .with_routing_mode(RoutingMode::Stateful)
811 .build();
812
813 assert!(session.is_per_command_routing());
814 assert_eq!(session.mode(), SessionMode::Stateful);
815 assert_eq!(session.mode_state.routing_mode(), RoutingMode::Stateful);
816 }
817
818 #[test]
819 fn test_builder_without_router_ignores_mode() {
820 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
821 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
822
823 let session = ClientSession::builder(
825 addr.into(),
826 buffer_pool,
827 test_auth_handler(),
828 test_metrics(),
829 )
830 .with_routing_mode(RoutingMode::PerCommand)
831 .build();
832
833 assert!(!session.is_per_command_routing());
834 assert_eq!(session.mode(), SessionMode::Stateful);
835 assert_eq!(session.mode_state.routing_mode(), RoutingMode::Stateful);
836 }
837
838 #[test]
839 fn test_builder_with_auth_handler() {
840 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
841 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
842 let auth_handler =
843 Arc::new(AuthHandler::new(Some("user".to_string()), Some("pass".to_string())).unwrap());
844
845 let session = ClientSession::builder(
846 addr.into(),
847 buffer_pool,
848 test_auth_handler(),
849 test_metrics(),
850 )
851 .with_auth_handler(auth_handler)
852 .build();
853
854 assert_eq!(*session.client_addr, addr);
856 }
857
858 #[test]
859 fn test_builder_with_metrics() {
860 use crate::metrics::MetricsCollector;
861
862 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
863 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
864 let metrics = MetricsCollector::new(1); let session =
867 ClientSession::builder(addr.into(), buffer_pool, test_auth_handler(), metrics).build();
868
869 let _ = session.metrics; }
872
873 #[test]
874 fn test_builder_with_connection_stats() {
875 use crate::metrics::ConnectionStatsAggregator;
876
877 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
878 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
879 let stats = ConnectionStatsAggregator::default();
880
881 let session = ClientSession::builder(
882 addr.into(),
883 buffer_pool,
884 test_auth_handler(),
885 test_metrics(),
886 )
887 .with_connection_stats(stats)
888 .build();
889
890 assert!(session.connection_stats().is_some());
891 }
892
893 #[test]
894 fn test_builder_with_cache() {
895 use crate::cache::UnifiedCache;
896 use std::time::Duration;
897
898 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
899 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
900 let cache = Arc::new(UnifiedCache::memory(100, Duration::from_secs(3600)));
901
902 let session = ClientSession::builder(
903 addr.into(),
904 buffer_pool,
905 test_auth_handler(),
906 test_metrics(),
907 )
908 .with_cache(cache)
909 .build();
910
911 assert_eq!(session.cache.capacity(), 100);
913 }
914
915 #[test]
916 fn test_builder_method_chaining() {
917 use crate::metrics::MetricsCollector;
918
919 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
920 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
921 let router = Arc::new(BackendSelector::new());
922 let metrics = MetricsCollector::new(1); let session =
926 ClientSession::builder(addr.into(), buffer_pool, test_auth_handler(), metrics)
927 .with_router(router)
928 .with_routing_mode(RoutingMode::Hybrid)
929 .build();
930
931 assert!(session.is_per_command_routing());
932 assert_eq!(session.mode(), SessionMode::PerCommand);
933 assert_eq!(session.mode_state.routing_mode(), RoutingMode::Hybrid);
934 let _ = session.metrics; }
937
938 #[test]
941 fn test_mode_getter() {
942 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
943 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
944 let router = Arc::new(BackendSelector::new());
945
946 let session = ClientSession::new_with_router(
948 addr.into(),
949 buffer_pool.clone(),
950 router,
951 RoutingMode::PerCommand,
952 test_auth_handler(),
953 test_metrics(),
954 );
955 assert_eq!(session.mode(), SessionMode::PerCommand);
956
957 let session = ClientSession::new(
959 addr.into(),
960 buffer_pool,
961 test_auth_handler(),
962 test_metrics(),
963 );
964 assert_eq!(session.mode(), SessionMode::Stateful);
965 }
966
967 #[test]
968 fn test_username_initially_none() {
969 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
970 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
971 let session = ClientSession::new(
972 addr.into(),
973 buffer_pool,
974 test_auth_handler(),
975 test_metrics(),
976 );
977
978 assert!(session.username().is_none());
979 }
980
981 #[test]
982 fn test_set_username_and_get() {
983 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
984 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
985 let session = ClientSession::new(
986 addr.into(),
987 buffer_pool,
988 test_auth_handler(),
989 test_metrics(),
990 );
991
992 session.set_username(Some("testuser".to_string()));
993
994 let username = session.username();
995 assert!(username.is_some());
996 assert_eq!(username, Some("testuser"));
997 }
998
999 #[test]
1000 fn test_set_username_none() {
1001 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1002 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1003 let session = ClientSession::new(
1004 addr.into(),
1005 buffer_pool,
1006 test_auth_handler(),
1007 test_metrics(),
1008 );
1009
1010 session.set_username(None);
1011
1012 assert!(session.username().is_none());
1013 }
1014
1015 #[test]
1016 fn test_username_cheap_clone() {
1017 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1018 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1019 let session = ClientSession::new(
1020 addr.into(),
1021 buffer_pool,
1022 test_auth_handler(),
1023 test_metrics(),
1024 );
1025
1026 session.set_username(Some("testuser".to_string()));
1027
1028 let username1 = session.username();
1029 let username2 = session.username();
1030
1031 assert!(username1.is_some());
1032 assert!(username2.is_some());
1033 assert_eq!(username1, username2);
1034 assert_eq!(username1, Some("testuser"));
1035 }
1036
1037 #[test]
1038 fn test_connection_stats_none_by_default() {
1039 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1040 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1041 let session = ClientSession::new(
1042 addr.into(),
1043 buffer_pool,
1044 test_auth_handler(),
1045 test_metrics(),
1046 );
1047
1048 assert!(session.connection_stats().is_none());
1049 }
1050
1051 #[test]
1052 fn test_connection_stats_with_aggregator() {
1053 use crate::metrics::ConnectionStatsAggregator;
1054
1055 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1056 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1057 let stats = ConnectionStatsAggregator::default();
1058
1059 let session = ClientSession::builder(
1060 addr.into(),
1061 buffer_pool,
1062 test_auth_handler(),
1063 test_metrics(),
1064 )
1065 .with_connection_stats(stats)
1066 .build();
1067
1068 assert!(session.connection_stats().is_some());
1069 }
1070
1071 #[test]
1074 fn test_metrics_direct_calls_no_metrics() {
1075 use crate::types::BackendId;
1076
1077 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1078 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1079 let session = ClientSession::new(
1080 addr.into(),
1081 buffer_pool,
1082 test_auth_handler(),
1083 test_metrics(),
1084 );
1085
1086 session.metrics.record_command(BackendId::from_index(0));
1088 session.metrics.user_command(session.username());
1089 session.metrics.stateful_session_started();
1090 session.metrics.stateful_session_ended();
1091 session.metrics.user_bytes_sent(session.username(), 1024);
1092 session
1093 .metrics
1094 .user_bytes_received(session.username(), 2048);
1095 }
1096
1097 #[test]
1098 fn test_metrics_direct_calls_with_metrics() {
1099 use crate::metrics::MetricsCollector;
1100 use crate::types::BackendId;
1101
1102 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1103 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1104 let metrics = MetricsCollector::new(1); let session = ClientSession::builder(
1107 addr.into(),
1108 buffer_pool,
1109 test_auth_handler(),
1110 metrics.clone(),
1111 )
1112 .build();
1113
1114 session.set_username(Some("testuser".to_string()));
1115
1116 session.metrics.record_command(BackendId::from_index(0));
1118 session.metrics.user_command(session.username());
1119 session.metrics.stateful_session_started();
1120 session.metrics.stateful_session_ended();
1121 session.metrics.user_bytes_sent(session.username(), 1024);
1122 session
1123 .metrics
1124 .user_bytes_received(session.username(), 2048);
1125
1126 let snapshot = metrics.snapshot(None);
1128 assert!(!snapshot.backend_stats.is_empty());
1129 assert!(snapshot.backend_stats[0].total_commands.get() > 0);
1130 }
1131
1132 #[test]
1133 fn test_record_command_with_metrics() {
1134 use crate::metrics::MetricsCollector;
1135 use crate::types::BackendId;
1136
1137 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1138 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1139 let metrics = MetricsCollector::new(1); let session = ClientSession::builder(
1142 addr.into(),
1143 buffer_pool,
1144 test_auth_handler(),
1145 metrics.clone(),
1146 )
1147 .build();
1148
1149 let backend_id = BackendId::from_index(0);
1150 session.metrics.record_command(backend_id);
1151
1152 let snapshot = metrics.snapshot(None);
1153 assert_eq!(snapshot.backend_stats[0].total_commands.get(), 1);
1154 }
1155
1156 #[test]
1157 fn test_user_bytes_tracking() {
1158 use crate::metrics::MetricsCollector;
1159
1160 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1161 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1162 let metrics = MetricsCollector::new(1); let session = ClientSession::builder(
1165 addr.into(),
1166 buffer_pool,
1167 test_auth_handler(),
1168 metrics.clone(),
1169 )
1170 .build();
1171
1172 session.set_username(Some("testuser".to_string()));
1173 session.metrics.user_bytes_sent(session.username(), 1024);
1174 session
1175 .metrics
1176 .user_bytes_received(session.username(), 2048);
1177
1178 let snapshot = metrics.snapshot(None);
1179 let user_stats = snapshot
1180 .user_stats
1181 .iter()
1182 .find(|s| s.username == "testuser");
1183 assert!(user_stats.is_some());
1184
1185 let stats = user_stats.unwrap();
1186 assert_eq!(stats.bytes_sent.as_u64(), 1024);
1187 assert_eq!(stats.bytes_received.as_u64(), 2048);
1188 }
1189
1190 #[test]
1191 fn test_stateful_session_tracking() {
1192 use crate::metrics::MetricsCollector;
1193
1194 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 8080);
1195 let buffer_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 4);
1196 let metrics = MetricsCollector::new(1); let session = ClientSession::builder(
1199 addr.into(),
1200 buffer_pool,
1201 test_auth_handler(),
1202 metrics.clone(),
1203 )
1204 .build();
1205
1206 session.metrics.stateful_session_started();
1207
1208 let snapshot = metrics.snapshot(None);
1209 assert_eq!(snapshot.stateful_sessions, 1);
1210
1211 session.metrics.stateful_session_ended();
1212
1213 let snapshot = metrics.snapshot(None);
1214 assert_eq!(snapshot.stateful_sessions, 0);
1215 }
1216}