1use std::collections::{HashMap, HashSet, VecDeque};
9use std::future::Future;
10use std::net::SocketAddr;
11use std::pin::Pin;
12use std::time::{Duration, Instant};
13
14use bytes::Bytes;
15use sipx_sip::{Limits, Message, StreamParser};
16use tokio::io::{AsyncReadExt, AsyncWriteExt};
17use tokio::net::{TcpListener, TcpStream};
18use tokio::sync::mpsc;
19use tokio::task::JoinSet;
20use tokio_util::sync::CancellationToken;
21
22use crate::error::Result;
23use crate::policy::{ConnectionState, ObservationHub, connection_event};
24use crate::target::{ConnectionKey, TransportKind};
25
26#[derive(Debug)]
28pub enum Event {
29 Message {
31 message: Box<Message>,
33 source: SocketAddr,
35 transport: TransportKind,
39 id: u64,
41 #[cfg(feature = "quic")]
43 quic_reply: Option<crate::quic::Reply>,
44 },
45 FramingFailed {
56 key: ConnectionKey,
58 },
59 Pong {
65 key: ConnectionKey,
67 id: u64,
69 },
70 ConnectFailed {
76 key: ConnectionKey,
78 id: u64,
80 kind: std::io::ErrorKind,
82 detail: String,
84 },
85 #[cfg(feature = "tls")]
91 HandshakeFailed {
92 key: ConnectionKey,
94 id: u64,
96 detail: String,
98 },
99 #[cfg(feature = "quic")]
101 QuicClosed {
102 key: ConnectionKey,
104 id: u64,
106 detail: String,
108 },
109 Closed {
115 key: ConnectionKey,
118 id: u64,
120 },
121}
122
123#[derive(Debug, Clone, Copy)]
125pub struct PoolConfig {
126 pub max_connections: usize,
128 pub idle_timeout: Duration,
130 pub reuse_inbound_for_outbound: bool,
136}
137
138impl Default for PoolConfig {
139 fn default() -> Self {
140 Self {
141 max_connections: 1024,
142 idle_timeout: Duration::from_secs(300),
143 reuse_inbound_for_outbound: false,
144 }
145 }
146}
147
148#[derive(Debug, Clone, Copy, PartialEq, Eq)]
151pub enum Origin {
152 Outbound,
154 Inbound,
156}
157
158#[derive(Debug)]
159struct Pooled {
160 writer: mpsc::Sender<Bytes>,
161 cancel: CancellationToken,
162 id: u64,
163 origin: Origin,
164 last_used: Instant,
165 active: bool,
166 admission_generation: Option<u64>,
167 has_sent: bool,
168}
169
170type ConnectionTask = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
171
172struct Pending {
173 key: ConnectionKey,
174 id: u64,
175 cancel: CancellationToken,
176 task: ConnectionTask,
177 admission_generation: Option<u64>,
178}
179
180struct Registration {
181 writer: mpsc::Receiver<Bytes>,
182 cancel: CancellationToken,
183 id: u64,
184 start_now: bool,
185 admission_generation: Option<u64>,
186}
187
188pub(crate) fn observe_state(
189 observations: Option<&std::sync::Arc<ObservationHub>>,
190 key: &ConnectionKey,
191 id: u64,
192 admission_generation: Option<u64>,
193 state: ConnectionState,
194) {
195 if let Some(observations) = observations {
196 observations.emit(connection_event(
197 key.clone(),
198 id,
199 admission_generation,
200 state,
201 ));
202 }
203}
204
205pub(crate) fn observe_ready(
206 observations: Option<&std::sync::Arc<ObservationHub>>,
207 key: &ConnectionKey,
208 id: u64,
209 admission_generation: Option<u64>,
210 authenticated: bool,
211) {
212 if authenticated {
213 observe_state(
214 observations,
215 key,
216 id,
217 admission_generation,
218 ConnectionState::Authenticated,
219 );
220 }
221 observe_state(
222 observations,
223 key,
224 id,
225 admission_generation,
226 ConnectionState::Opened,
227 );
228 observe_state(
229 observations,
230 key,
231 id,
232 admission_generation,
233 ConnectionState::Pooled,
234 );
235}
236
237pub struct Pool {
239 connections: HashMap<ConnectionKey, Pooled>,
240 config: PoolConfig,
241 events: mpsc::Sender<Event>,
242 limits: Limits,
243 tasks: JoinSet<()>,
244 next_id: u64,
245 shutdown: CancellationToken,
246 retiring: HashSet<(ConnectionKey, u64)>,
248 pending: VecDeque<Pending>,
250 observations: Option<std::sync::Arc<ObservationHub>>,
251}
252
253impl std::fmt::Debug for Pool {
254 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
255 formatter
256 .debug_struct("Pool")
257 .field("connections", &self.connections)
258 .field("config", &self.config)
259 .field("tasks", &self.tasks.len())
260 .field("retiring", &self.retiring.len())
261 .field("pending", &self.pending.len())
262 .finish_non_exhaustive()
263 }
264}
265
266impl Pool {
267 #[must_use]
269 pub fn new(config: PoolConfig, limits: Limits, events: mpsc::Sender<Event>) -> Self {
270 Self {
271 connections: HashMap::new(),
272 config,
273 events,
274 limits,
275 tasks: JoinSet::new(),
276 next_id: 1,
277 shutdown: CancellationToken::new(),
278 retiring: HashSet::new(),
279 pending: VecDeque::new(),
280 observations: None,
281 }
282 }
283
284 pub(crate) fn new_observed(
285 config: PoolConfig,
286 limits: Limits,
287 events: mpsc::Sender<Event>,
288 observations: std::sync::Arc<ObservationHub>,
289 ) -> Self {
290 let mut pool = Self::new(config, limits, events);
291 pool.observations = Some(observations);
292 pool
293 }
294
295 #[must_use]
297 pub fn len(&self) -> usize {
298 self.tasks.len()
299 }
300
301 #[must_use]
303 pub fn is_empty(&self) -> bool {
304 self.tasks.is_empty()
305 }
306
307 pub fn accept(&mut self, stream: TcpStream, peer: SocketAddr) {
309 self.accept_admitted(stream, peer, 0);
310 }
311
312 pub(crate) fn accept_admitted(
313 &mut self,
314 stream: TcpStream,
315 peer: SocketAddr,
316 admission_generation: u64,
317 ) {
318 let key = ConnectionKey::new(peer, TransportKind::Tcp);
319 if self
320 .insert_stream_admitted(
321 stream,
322 key,
323 Origin::Inbound,
324 Some(admission_generation),
325 false,
326 )
327 .is_err()
328 {
329 tracing::debug!(%peer, "refused inbound TCP connection at capacity");
332 }
333 }
334
335 pub async fn send(&mut self, key: &ConnectionKey, bytes: Bytes) -> Result<()> {
343 self.send_generation(key, bytes).await.map(|_| ())
344 }
345
346 pub(crate) async fn send_generation(
347 &mut self,
348 key: &ConnectionKey,
349 bytes: Bytes,
350 ) -> Result<u64> {
351 if !self.reusable(key) {
352 self.dial(key.clone())?;
355 }
356 self.queue(key, bytes).await?;
357 self.generation(key)
358 .ok_or(crate::error::Error::EndpointClosed)
359 }
360
361 #[cfg(feature = "tls")]
363 pub fn accept_tls(
364 &mut self,
365 stream: tokio_rustls::server::TlsStream<TcpStream>,
366 peer: SocketAddr,
367 ) {
368 self.accept_tls_admitted(stream, peer, 0);
369 }
370
371 #[cfg(feature = "tls")]
372 pub(crate) fn accept_tls_admitted(
373 &mut self,
374 stream: tokio_rustls::server::TlsStream<TcpStream>,
375 peer: SocketAddr,
376 admission_generation: u64,
377 ) {
378 let key = ConnectionKey::new(peer, TransportKind::Tls);
379 if self
380 .insert_stream_admitted(
381 stream,
382 key.clone(),
383 Origin::Inbound,
384 Some(admission_generation),
385 true,
386 )
387 .is_err()
388 {
389 tracing::debug!(peer = %key.peer, "refused inbound TLS connection at capacity");
392 }
393 }
394
395 #[cfg(feature = "quic")]
397 pub fn accept_quic(&mut self, connection: quinn::Connection, peer: SocketAddr) {
398 self.accept_quic_admitted(connection, peer, 0);
399 }
400
401 #[cfg(feature = "quic")]
402 pub(crate) fn accept_quic_admitted(
403 &mut self,
404 connection: quinn::Connection,
405 peer: SocketAddr,
406 admission_generation: u64,
407 ) {
408 let key = ConnectionKey::new(peer, TransportKind::Quic);
409 if self
410 .insert_quic(
411 connection,
412 key.clone(),
413 Origin::Inbound,
414 Some(admission_generation),
415 )
416 .is_err()
417 {
418 tracing::debug!(peer = %key.peer, "refused inbound QUIC connection at capacity");
420 }
421 }
422
423 #[cfg(feature = "quic")]
425 pub(crate) async fn send_quic_generation(
426 &mut self,
427 key: &ConnectionKey,
428 verification_name: &str,
429 client: &crate::tls::ClientTls,
430 endpoint: &quinn::Endpoint,
431 bytes: Bytes,
432 ) -> Result<u64> {
433 if !self.reusable(key) {
434 self.dial_quic(key.clone(), verification_name, client, endpoint)?;
435 }
436 self.queue(key, bytes).await?;
437 self.generation(key)
438 .ok_or(crate::error::Error::EndpointClosed)
439 }
440
441 #[cfg(feature = "tls")]
447 pub async fn send_tls(
448 &mut self,
449 key: &ConnectionKey,
450 verification_name: &str,
451 client: &crate::tls::ClientTls,
452 bytes: Bytes,
453 ) -> Result<()> {
454 self.send_tls_generation(key, verification_name, client, bytes)
455 .await
456 .map(|_| ())
457 }
458
459 #[cfg(feature = "tls")]
460 pub(crate) async fn send_tls_generation(
461 &mut self,
462 key: &ConnectionKey,
463 verification_name: &str,
464 client: &crate::tls::ClientTls,
465 bytes: Bytes,
466 ) -> Result<u64> {
467 if !self.reusable(key) {
468 self.dial_tls(key.clone(), verification_name, client)?;
469 }
470 self.queue(key, bytes).await?;
471 self.generation(key)
472 .ok_or(crate::error::Error::EndpointClosed)
473 }
474
475 #[cfg(feature = "ws")]
477 pub fn accept_ws<S>(
478 &mut self,
479 ws: crate::ws::Socket<S>,
480 key: ConnectionKey,
481 keepalive: Duration,
482 ) where
483 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
484 {
485 self.accept_ws_admitted(ws, key, keepalive, 0);
486 }
487
488 #[cfg(feature = "ws")]
489 pub(crate) fn accept_ws_admitted<S>(
490 &mut self,
491 ws: crate::ws::Socket<S>,
492 key: ConnectionKey,
493 keepalive: Duration,
494 admission_generation: u64,
495 ) where
496 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
497 {
498 let peer = key.peer;
499 let authenticated = key.transport == TransportKind::Wss;
500 if self
501 .spawn_ws(
502 ws,
503 key,
504 Origin::Inbound,
505 keepalive,
506 Some(admission_generation),
507 authenticated,
508 )
509 .is_err()
510 {
511 tracing::debug!(%peer, "refused inbound WebSocket connection at capacity");
514 }
515 }
516
517 #[cfg(feature = "ws")]
522 pub async fn send_ws(
523 &mut self,
524 key: &ConnectionKey,
525 authority: &str,
526 keepalive: Duration,
527 #[cfg(feature = "wss")] client: Option<&crate::tls::ClientTls>,
528 bytes: Bytes,
529 ) -> Result<()> {
530 self.send_ws_generation(
531 key,
532 authority,
533 keepalive,
534 #[cfg(feature = "wss")]
535 client,
536 bytes,
537 )
538 .await
539 .map(|_| ())
540 }
541
542 #[cfg(feature = "ws")]
543 pub(crate) async fn send_ws_generation(
544 &mut self,
545 key: &ConnectionKey,
546 authority: &str,
547 keepalive: Duration,
548 #[cfg(feature = "wss")] client: Option<&crate::tls::ClientTls>,
549 bytes: Bytes,
550 ) -> Result<u64> {
551 if !self.reusable(key) {
552 self.dial_ws(
553 key.clone(),
554 authority,
555 keepalive,
556 #[cfg(feature = "wss")]
557 client,
558 )?;
559 }
560 self.queue(key, bytes).await?;
561 self.generation(key)
562 .ok_or(crate::error::Error::EndpointClosed)
563 }
564
565 pub fn remove(&mut self, key: &ConnectionKey, id: u64) -> bool {
567 let removed = if self
568 .connections
569 .get(key)
570 .is_some_and(|pooled| pooled.id == id)
571 {
572 self.connections.remove(key);
573 true
574 } else {
575 self.retiring.remove(&(key.clone(), id))
576 };
577 self.reap_finished();
578 self.activate_pending();
579 removed
580 }
581
582 #[must_use]
584 pub(crate) fn generation(&self, key: &ConnectionKey) -> Option<u64> {
585 self.connections.get(key).map(|pooled| pooled.id)
586 }
587
588 pub(crate) fn observation_generation(&self, key: &ConnectionKey) -> Option<(u64, Option<u64>)> {
590 self.connections
591 .get(key)
592 .map(|pooled| (pooled.id, pooled.admission_generation))
593 }
594
595 #[must_use]
597 pub fn holds(&self, key: &ConnectionKey) -> bool {
598 self.connections.contains_key(key)
599 }
600
601 pub async fn send_on_existing(&mut self, key: &ConnectionKey, bytes: Bytes) -> bool {
608 self.send_on_existing_generation(key, bytes).await.is_some()
609 }
610
611 pub(crate) async fn send_on_existing_generation(
612 &mut self,
613 key: &ConnectionKey,
614 bytes: Bytes,
615 ) -> Option<u64> {
616 let pooled = self.connections.get_mut(key)?;
617 if pooled.writer.send(bytes).await.is_err() {
618 self.retire(key);
619 return None;
620 }
621 pooled.last_used = Instant::now();
622 Some(pooled.id)
623 }
624
625 pub fn evict_idle(&mut self) -> Vec<ConnectionKey> {
627 let deadline = Instant::now();
628 let idle_timeout = self.config.idle_timeout;
629 let evicted: Vec<ConnectionKey> = self
630 .connections
631 .iter()
632 .filter(|(_, c)| deadline.duration_since(c.last_used) > idle_timeout)
633 .map(|(key, _)| key.clone())
634 .collect();
635 for key in &evicted {
636 self.retire(key);
637 }
638 evicted
639 }
640
641 fn reusable(&self, key: &ConnectionKey) -> bool {
643 self.connections.get(key).is_some_and(|pooled| {
644 !pooled.writer.is_closed()
645 && (pooled.origin == Origin::Outbound || self.config.reuse_inbound_for_outbound)
646 })
647 }
648
649 async fn queue(&mut self, key: &ConnectionKey, bytes: Bytes) -> Result<()> {
651 let Some(pooled) = self.connections.get_mut(key) else {
652 return Err(crate::error::Error::EndpointClosed);
653 };
654 pooled.last_used = Instant::now();
655 let reused = pooled.has_sent;
656 pooled.has_sent = true;
657 let id = pooled.id;
658 let admission_generation = pooled.admission_generation;
659 let writer = pooled.writer.clone();
660 if reused && let Some(observations) = &self.observations {
661 observations.emit(connection_event(
662 key.clone(),
663 id,
664 admission_generation,
665 ConnectionState::Reused,
666 ));
667 }
668 writer
669 .send(bytes)
670 .await
671 .map_err(|_| crate::error::Error::EndpointClosed)
672 }
673
674 fn register(
679 &mut self,
680 key: &ConnectionKey,
681 origin: Origin,
682 admission_generation: Option<u64>,
683 ) -> Result<Registration> {
684 self.reap_finished();
685 self.activate_pending();
686 if self.connections.contains_key(key) {
690 self.retire(key);
691 } else if self.connections.len() >= self.config.max_connections {
692 self.evict_least_recently_used();
693 }
694 self.reap_finished();
695 self.activate_pending();
696 if self.connections.len() >= self.config.max_connections {
697 return Err(crate::error::Error::ConnectionCapacity {
698 max: self.config.max_connections,
699 });
700 }
701 let (writer_tx, writer_rx) = mpsc::channel::<Bytes>(64);
702 let cancel = CancellationToken::new();
703 let id = self.next_id;
704 self.next_id = self.next_id.wrapping_add(1).max(1);
705 self.connections.insert(
706 key.clone(),
707 Pooled {
708 writer: writer_tx,
709 cancel: cancel.clone(),
710 id,
711 origin,
712 last_used: Instant::now(),
713 active: self.tasks.len() < self.config.max_connections,
714 admission_generation,
715 has_sent: false,
716 },
717 );
718 if let Some(observations) = &self.observations
719 && origin == Origin::Inbound
720 {
721 observations.emit(connection_event(
722 key.clone(),
723 id,
724 admission_generation,
725 ConnectionState::Accepted,
726 ));
727 }
728 Ok(Registration {
729 writer: writer_rx,
730 cancel,
731 id,
732 start_now: self.tasks.len() < self.config.max_connections,
733 admission_generation,
734 })
735 }
736
737 fn launch<F>(
738 &mut self,
739 key: ConnectionKey,
740 id: u64,
741 cancel: CancellationToken,
742 start_now: bool,
743 task: F,
744 admission_generation: Option<u64>,
745 ) where
746 F: Future<Output = ()> + Send + 'static,
747 {
748 if start_now {
749 self.track(key, id, cancel, task, admission_generation);
750 } else {
751 self.pending.push_back(Pending {
752 key,
753 id,
754 cancel,
755 task: Box::pin(task),
756 admission_generation,
757 });
758 }
759 }
760
761 fn activate_pending(&mut self) {
762 while self.tasks.len() < self.config.max_connections {
763 let Some(pending) = self.pending.pop_front() else {
764 break;
765 };
766 let is_current = self
767 .connections
768 .get_mut(&pending.key)
769 .filter(|pooled| pooled.id == pending.id);
770 let Some(pooled) = is_current else {
771 continue;
772 };
773 pooled.active = true;
774 self.track(
775 pending.key,
776 pending.id,
777 pending.cancel,
778 pending.task,
779 pending.admission_generation,
780 );
781 }
782 }
783
784 fn track<F>(
785 &mut self,
786 key: ConnectionKey,
787 id: u64,
788 cancel: CancellationToken,
789 task: F,
790 admission_generation: Option<u64>,
791 ) where
792 F: std::future::Future<Output = ()> + Send + 'static,
793 {
794 let events = self.events.clone();
795 let observations = self.observations.clone();
796 let shutdown = self.shutdown.clone();
797 self.tasks.spawn(async move {
798 let report = tokio::select! {
799 biased;
800 () = shutdown.cancelled() => false,
801 () = cancel.cancelled() => true,
802 () = task => true,
803 };
804 if report {
805 if let Some(observations) = observations {
806 observations.emit(connection_event(
807 key.clone(),
808 id,
809 admission_generation,
810 ConnectionState::Closed,
811 ));
812 }
813 tokio::select! {
816 biased;
817 () = shutdown.cancelled() => {}
818 result = events.send(Event::Closed { key, id }) => {
819 let _ = result;
820 }
821 }
822 }
823 });
824 }
825
826 fn dial(&mut self, key: ConnectionKey) -> Result<()> {
827 let Registration {
828 writer: writer_rx,
829 cancel,
830 id,
831 start_now,
832 admission_generation,
833 } = self.register(&key, Origin::Outbound, None)?;
834 let events = self.events.clone();
835 let observations = self.observations.clone();
836 let limits = self.limits;
837 let task_key = key.clone();
838 self.launch(
839 key,
840 id,
841 cancel,
842 start_now,
843 async move {
844 match TcpStream::connect(task_key.peer).await {
845 Ok(stream) => {
846 observe_ready(
847 observations.as_ref(),
848 &task_key,
849 id,
850 admission_generation,
851 false,
852 );
853 pump(stream, task_key, id, writer_rx, events, limits).await;
854 }
855 Err(error) => {
856 tracing::debug!(%error, peer = %task_key.peer, "connect failed");
857 observe_state(
858 observations.as_ref(),
859 &task_key,
860 id,
861 admission_generation,
862 ConnectionState::Failed,
863 );
864 let _ = events
867 .send(Event::ConnectFailed {
868 key: task_key,
869 id,
870 kind: error.kind(),
871 detail: error.to_string(),
872 })
873 .await;
874 }
875 }
876 },
877 admission_generation,
878 );
879 Ok(())
880 }
881
882 #[cfg(feature = "tls")]
883 fn dial_tls(
884 &mut self,
885 key: ConnectionKey,
886 verification_name: &str,
887 client: &crate::tls::ClientTls,
888 ) -> Result<()> {
889 let name = crate::tls::verification_name(verification_name)?;
890 let connector = client.connector();
891 let Registration {
892 writer: writer_rx,
893 cancel,
894 id,
895 start_now,
896 admission_generation,
897 } = self.register(&key, Origin::Outbound, None)?;
898 let events = self.events.clone();
899 let observations = self.observations.clone();
900 let limits = self.limits;
901 let task_key = key.clone();
902
903 self.launch(
904 key,
905 id,
906 cancel,
907 start_now,
908 async move {
909 let stream = match TcpStream::connect(task_key.peer).await {
910 Ok(stream) => stream,
911 Err(error) => {
912 tracing::debug!(%error, peer = %task_key.peer, "connect failed");
913 let _ = events
916 .send(Event::ConnectFailed {
917 key: task_key,
918 id,
919 kind: error.kind(),
920 detail: error.to_string(),
921 })
922 .await;
923 return;
924 }
925 };
926 match connector.connect(name, stream).await {
927 Ok(tls) => {
928 observe_ready(
929 observations.as_ref(),
930 &task_key,
931 id,
932 admission_generation,
933 true,
934 );
935 pump(tls, task_key, id, writer_rx, events, limits).await;
936 }
937 Err(error) => {
938 tracing::warn!(%error, peer = %task_key.peer, "TLS handshake failed");
942 let _ = events
945 .send(Event::HandshakeFailed {
946 key: task_key,
947 id,
948 detail: error.to_string(),
949 })
950 .await;
951 }
952 }
953 },
954 admission_generation,
955 );
956 Ok(())
957 }
958
959 #[cfg(feature = "quic")]
960 fn dial_quic(
961 &mut self,
962 key: ConnectionKey,
963 verification_name: &str,
964 client: &crate::tls::ClientTls,
965 endpoint: &quinn::Endpoint,
966 ) -> Result<()> {
967 crate::tls::verification_name(verification_name)?;
968 let mut client_config = client.quic_config()?;
969 client_config.transport_config(crate::quic::transport_config());
970 let connecting = endpoint
971 .connect_with(client_config, key.peer, verification_name)
972 .map_err(|error| crate::tls::TlsError::Handshake {
973 peer: key.peer.to_string(),
974 detail: error.to_string(),
975 })?;
976 let Registration {
977 writer: writer_rx,
978 cancel,
979 id,
980 start_now,
981 admission_generation,
982 } = self.register(&key, Origin::Outbound, None)?;
983 let events = self.events.clone();
984 let observations = self.observations.clone();
985 let limits = self.limits;
986 let task_key = key.clone();
987 self.launch(
988 key,
989 id,
990 cancel,
991 start_now,
992 async move {
993 match connecting.await {
994 Ok(connection) => {
995 observe_ready(
996 observations.as_ref(),
997 &task_key,
998 id,
999 admission_generation,
1000 true,
1001 );
1002 crate::quic::pump(connection, task_key, id, writer_rx, events, limits)
1003 .await;
1004 }
1005 Err(error) => {
1006 tracing::warn!(%error, peer = %task_key.peer, "QUIC handshake failed");
1007 let _ = events
1009 .send(Event::HandshakeFailed {
1010 key: task_key,
1011 id,
1012 detail: error.to_string(),
1013 })
1014 .await;
1015 }
1016 }
1017 },
1018 admission_generation,
1019 );
1020 Ok(())
1021 }
1022
1023 #[cfg(feature = "ws")]
1024 fn dial_ws(
1025 &mut self,
1026 key: ConnectionKey,
1027 authority: &str,
1028 keepalive: Duration,
1029 #[cfg(feature = "wss")] client: Option<&crate::tls::ClientTls>,
1030 ) -> Result<()> {
1031 #[cfg(feature = "wss")]
1032 let secure = if key.transport == TransportKind::Wss {
1033 let client = client.ok_or(crate::error::Error::UnsupportedTransport(
1034 "WSS (no client configuration, so no outbound connection can be verified)",
1035 ))?;
1036 let verify = key
1042 .identity
1043 .as_deref()
1044 .map_or_else(|| key.peer.ip().to_string(), str::to_owned);
1045 Some((client.connector(), crate::tls::verification_name(&verify)?))
1046 } else {
1047 None
1048 };
1049
1050 let authority = authority.to_owned();
1051 let Registration {
1052 writer: writer_rx,
1053 cancel,
1054 id,
1055 start_now,
1056 admission_generation,
1057 } = self.register(&key, Origin::Outbound, None)?;
1058 let events = self.events.clone();
1059 let observations = self.observations.clone();
1060 let limits = self.limits;
1061 let task_key = key.clone();
1062
1063 self.launch(
1064 key,
1065 id,
1066 cancel,
1067 start_now,
1068 async move {
1069 let stream = match TcpStream::connect(task_key.peer).await {
1070 Ok(stream) => stream,
1071 Err(error) => {
1072 tracing::debug!(%error, peer = %task_key.peer, "connect failed");
1073 let _ = events
1076 .send(Event::ConnectFailed {
1077 key: task_key,
1078 id,
1079 kind: error.kind(),
1080 detail: error.to_string(),
1081 })
1082 .await;
1083 return;
1084 }
1085 };
1086
1087 #[cfg(feature = "wss")]
1088 if let Some((connector, name)) = secure {
1089 match connector.connect(name, stream).await {
1090 Ok(tls) => {
1091 crate::ws::dial(
1092 tls,
1093 &authority,
1094 task_key,
1095 id,
1096 writer_rx,
1097 events,
1098 limits,
1099 keepalive,
1100 observations,
1101 admission_generation,
1102 true,
1103 )
1104 .await;
1105 }
1106 Err(error) => {
1107 tracing::warn!(%error, peer = %task_key.peer, "TLS handshake failed");
1108 let _ = events
1111 .send(Event::HandshakeFailed {
1112 key: task_key,
1113 id,
1114 detail: error.to_string(),
1115 })
1116 .await;
1117 }
1118 }
1119 return;
1120 }
1121
1122 crate::ws::dial(
1123 stream,
1124 &authority,
1125 task_key,
1126 id,
1127 writer_rx,
1128 events,
1129 limits,
1130 keepalive,
1131 observations,
1132 admission_generation,
1133 false,
1134 )
1135 .await;
1136 },
1137 admission_generation,
1138 );
1139 Ok(())
1140 }
1141
1142 #[cfg(test)]
1143 fn insert_stream<S>(&mut self, stream: S, key: ConnectionKey, origin: Origin) -> Result<()>
1144 where
1145 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static,
1146 {
1147 self.insert_stream_admitted(stream, key, origin, None, false)
1148 }
1149
1150 fn insert_stream_admitted<S>(
1151 &mut self,
1152 stream: S,
1153 key: ConnectionKey,
1154 origin: Origin,
1155 admission_generation: Option<u64>,
1156 authenticated: bool,
1157 ) -> Result<()>
1158 where
1159 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static,
1160 {
1161 let Registration {
1162 writer: writer_rx,
1163 cancel,
1164 id,
1165 start_now,
1166 admission_generation,
1167 } = self.register(&key, origin, admission_generation)?;
1168 observe_ready(
1169 self.observations.as_ref(),
1170 &key,
1171 id,
1172 admission_generation,
1173 authenticated,
1174 );
1175 let events = self.events.clone();
1176 let limits = self.limits;
1177 let task_key = key.clone();
1178 self.launch(
1179 key,
1180 id,
1181 cancel,
1182 start_now,
1183 pump(stream, task_key, id, writer_rx, events, limits),
1184 admission_generation,
1185 );
1186 Ok(())
1187 }
1188
1189 #[cfg(feature = "quic")]
1190 fn insert_quic(
1191 &mut self,
1192 connection: quinn::Connection,
1193 key: ConnectionKey,
1194 origin: Origin,
1195 admission_generation: Option<u64>,
1196 ) -> Result<()> {
1197 let Registration {
1198 writer: writer_rx,
1199 cancel,
1200 id,
1201 start_now,
1202 admission_generation,
1203 } = self.register(&key, origin, admission_generation)?;
1204 observe_ready(
1205 self.observations.as_ref(),
1206 &key,
1207 id,
1208 admission_generation,
1209 true,
1210 );
1211 let events = self.events.clone();
1212 let limits = self.limits;
1213 let task_key = key.clone();
1214 self.launch(
1215 key,
1216 id,
1217 cancel,
1218 start_now,
1219 crate::quic::pump(connection, task_key, id, writer_rx, events, limits),
1220 admission_generation,
1221 );
1222 Ok(())
1223 }
1224
1225 #[cfg(feature = "ws")]
1226 fn spawn_ws<S>(
1227 &mut self,
1228 ws: crate::ws::Socket<S>,
1229 key: ConnectionKey,
1230 origin: Origin,
1231 keepalive: Duration,
1232 admission_generation: Option<u64>,
1233 authenticated: bool,
1234 ) -> Result<()>
1235 where
1236 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
1237 {
1238 let Registration {
1239 writer: writer_rx,
1240 cancel,
1241 id,
1242 start_now,
1243 admission_generation,
1244 } = self.register(&key, origin, admission_generation)?;
1245 observe_ready(
1246 self.observations.as_ref(),
1247 &key,
1248 id,
1249 admission_generation,
1250 authenticated,
1251 );
1252 let events = self.events.clone();
1253 let limits = self.limits;
1254 let task_key = key.clone();
1255 self.launch(
1256 key,
1257 id,
1258 cancel,
1259 start_now,
1260 crate::ws::pump(ws, task_key, id, writer_rx, events, limits, keepalive),
1261 admission_generation,
1262 );
1263 Ok(())
1264 }
1265
1266 fn retire(&mut self, key: &ConnectionKey) {
1267 if let Some(pooled) = self.connections.remove(key) {
1268 pooled.cancel.cancel();
1269 if pooled.active {
1270 self.retiring.insert((key.clone(), pooled.id));
1271 } else if let Some(index) = self
1272 .pending
1273 .iter()
1274 .position(|pending| pending.key == *key && pending.id == pooled.id)
1275 {
1276 self.pending.remove(index);
1277 }
1278 }
1279 }
1280
1281 fn reap_finished(&mut self) {
1282 while self.tasks.try_join_next().is_some() {}
1283 }
1284
1285 pub async fn shutdown(&mut self) {
1287 self.shutdown.cancel();
1288 for pooled in self.connections.values() {
1289 pooled.cancel.cancel();
1290 }
1291 self.connections.clear();
1292 self.pending.clear();
1293 self.retiring.clear();
1294 while self.tasks.join_next().await.is_some() {}
1295 }
1296
1297 fn evict_least_recently_used(&mut self) {
1298 let Some(victim) = self
1299 .connections
1300 .iter()
1301 .min_by_key(|(_, c)| c.last_used)
1302 .map(|(key, _)| key.clone())
1303 else {
1304 return;
1305 };
1306 self.retire(&victim);
1307 }
1308}
1309
1310impl Drop for Pool {
1311 fn drop(&mut self) {
1312 self.shutdown.cancel();
1313 for pooled in self.connections.values() {
1314 pooled.cancel.cancel();
1315 }
1316 self.pending.clear();
1317 }
1320}
1321
1322async fn pump<S>(
1328 stream: S,
1329 key: ConnectionKey,
1330 id: u64,
1331 mut outgoing: mpsc::Receiver<Bytes>,
1332 events: mpsc::Sender<Event>,
1333 limits: Limits,
1334) where
1335 S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Send + 'static,
1336{
1337 let (peer, transport) = (key.peer, key.transport);
1338 let (mut read_half, mut write_half) = tokio::io::split(stream);
1339 let mut parser = StreamParser::new(limits);
1340 let mut buf = vec![0u8; 8192];
1341
1342 loop {
1343 tokio::select! {
1344 read = read_half.read(&mut buf) => match read {
1345 Ok(0) | Err(_) => break,
1346 Ok(n) => {
1347 let chunk = buf.get(..n).unwrap_or(&[]);
1348 match parser.push(chunk) {
1349 Ok(messages) => {
1350 for _ in 0..parser.take_keepalives() {
1353 if events
1354 .send(Event::Pong {
1355 key: key.clone(),
1356 id,
1357 })
1358 .await
1359 .is_err()
1360 {
1361 return;
1362 }
1363 }
1364 for message in messages {
1365 if events
1366 .send(Event::Message {
1367 message: Box::new(message),
1368 source: peer,
1369 transport,
1370 id,
1371 #[cfg(feature = "quic")]
1372 quic_reply: None,
1373 })
1374 .await
1375 .is_err()
1376 {
1377 return;
1378 }
1379 }
1380 }
1381 Err(error) => {
1382 tracing::debug!(%error, %peer, "closing connection on framing error");
1391 let _ = events.send(Event::FramingFailed { key: key.clone() }).await;
1395 break;
1396 }
1397 }
1398 }
1399 },
1400 Some(bytes) = outgoing.recv() => {
1401 if write_half.write_all(&bytes).await.is_err() {
1402 break;
1403 }
1404 }
1405 else => break,
1406 }
1407 }
1408}
1409
1410pub async fn accept_loop(listener: TcpListener, incoming: mpsc::Sender<(TcpStream, SocketAddr)>) {
1414 loop {
1415 match listener.accept().await {
1416 Ok((stream, peer)) => {
1417 if incoming.send((stream, peer)).await.is_err() {
1418 return;
1419 }
1420 }
1421 Err(error) => {
1422 tracing::warn!(%error, "accept failed");
1423 return;
1424 }
1425 }
1426 }
1427}
1428
1429#[cfg(test)]
1430#[allow(
1431 clippy::unwrap_used,
1432 clippy::expect_used,
1433 clippy::panic,
1434 clippy::indexing_slicing
1435)]
1436mod tests {
1437 use super::*;
1438
1439 const MESSAGE: &str = "OPTIONS sip:a@b.com SIP/2.0\r\n\
1440 Via: SIP/2.0/TCP h.example.com;branch=z9hG4bKx\r\n\
1441 To: <sip:a@b.com>\r\n\
1442 From: <sip:c@d.net>;tag=1\r\n\
1443 Call-ID: x@y\r\n\
1444 CSeq: 1 OPTIONS\r\n\
1445 Content-Length: 0\r\n\r\n";
1446
1447 async fn pool_with_listener() -> (Pool, SocketAddr, mpsc::Receiver<Event>, TcpListener) {
1448 let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
1449 let addr = listener.local_addr().expect("has an address");
1450 let (tx, rx) = mpsc::channel(64);
1451 let pool = Pool::new(PoolConfig::default(), Limits::stream(), tx);
1452 (pool, addr, rx, listener)
1453 }
1454
1455 fn tcp(peer: SocketAddr) -> ConnectionKey {
1456 ConnectionKey::new(peer, TransportKind::Tcp)
1457 }
1458
1459 async fn wait_until_no_live_tasks(pool: &mut Pool) {
1460 tokio::time::timeout(Duration::from_secs(2), async {
1461 loop {
1462 pool.reap_finished();
1463 if pool.is_empty() {
1464 return;
1465 }
1466 tokio::task::yield_now().await;
1467 }
1468 })
1469 .await
1470 .expect("connection tasks terminate");
1471 }
1472
1473 #[tokio::test]
1475 async fn tcp_message_split_across_segments_is_assembled() {
1476 let (mut pool, addr, mut events, listener) = pool_with_listener().await;
1477
1478 let accepted = tokio::spawn(async move {
1479 let (stream, peer) = listener.accept().await.expect("accepts");
1480 (stream, peer)
1481 });
1482
1483 let mut client = TcpStream::connect(addr).await.expect("connects");
1484 let (stream, peer) = accepted.await.expect("accepted");
1485 pool.accept(stream, peer);
1486
1487 for byte in MESSAGE.as_bytes() {
1490 client.write_all(&[*byte]).await.expect("writes");
1491 client.flush().await.expect("flushes");
1492 }
1493
1494 let event = tokio::time::timeout(Duration::from_secs(2), events.recv())
1495 .await
1496 .expect("no timeout")
1497 .expect("an event");
1498 match event {
1499 Event::Message { message, .. } => {
1500 assert_eq!(message.to_bytes().as_ref(), MESSAGE.as_bytes());
1501 }
1502 Event::Pong { .. }
1503 | Event::ConnectFailed { .. }
1504 | Event::Closed { .. }
1505 | Event::FramingFailed { .. }
1506 | Event::HandshakeFailed { .. } => {
1507 panic!("expected a message")
1508 }
1509 #[cfg(feature = "quic")]
1510 Event::QuicClosed { .. } => {
1511 panic!("expected a message")
1512 }
1513 }
1514 }
1515
1516 #[tokio::test]
1518 async fn two_messages_in_one_segment_are_both_delivered() {
1519 let (mut pool, addr, mut events, listener) = pool_with_listener().await;
1520 let accepted = tokio::spawn(async move { listener.accept().await.expect("accepts") });
1521 let mut client = TcpStream::connect(addr).await.expect("connects");
1522 let (stream, peer) = accepted.await.expect("accepted");
1523 pool.accept(stream, peer);
1524
1525 let both = format!("{MESSAGE}{MESSAGE}");
1526 client.write_all(both.as_bytes()).await.expect("writes");
1527
1528 for _ in 0..2 {
1529 let event = tokio::time::timeout(Duration::from_secs(2), events.recv())
1530 .await
1531 .expect("no timeout")
1532 .expect("an event");
1533 assert!(matches!(event, Event::Message { .. }));
1534 }
1535 }
1536
1537 #[tokio::test]
1540 async fn a_closed_connection_is_reported_immediately() {
1541 let (mut pool, addr, mut events, listener) = pool_with_listener().await;
1542 let accepted = tokio::spawn(async move { listener.accept().await.expect("accepts") });
1543 let client = TcpStream::connect(addr).await.expect("connects");
1544 let (stream, peer) = accepted.await.expect("accepted");
1545 pool.accept(stream, peer);
1546
1547 drop(client);
1548
1549 let event = tokio::time::timeout(Duration::from_secs(2), events.recv())
1550 .await
1551 .expect("no timeout")
1552 .expect("an event");
1553 assert!(matches!(event, Event::Closed { .. }));
1554 }
1555
1556 #[tokio::test]
1559 async fn a_framing_error_closes_the_connection() {
1560 let (mut pool, addr, mut events, listener) = pool_with_listener().await;
1561 let accepted = tokio::spawn(async move { listener.accept().await.expect("accepts") });
1562 let mut client = TcpStream::connect(addr).await.expect("connects");
1563 let (stream, peer) = accepted.await.expect("accepted");
1564 pool.accept(stream, peer);
1565
1566 client
1567 .write_all(b"OPTIONS sip:a@b SIP/2.0\r\nContent-Length: -1\r\n\r\n")
1568 .await
1569 .expect("writes");
1570
1571 let first = tokio::time::timeout(Duration::from_secs(2), events.recv())
1577 .await
1578 .expect("no timeout")
1579 .expect("an event");
1580 assert!(
1581 matches!(first, Event::FramingFailed { .. }),
1582 "a framing error must report the loss, not only the close: {first:?}"
1583 );
1584
1585 let second = tokio::time::timeout(Duration::from_secs(2), events.recv())
1586 .await
1587 .expect("no timeout")
1588 .expect("an event");
1589 assert!(matches!(second, Event::Closed { .. }), "{second:?}");
1590 }
1591
1592 #[tokio::test]
1595 async fn an_inbound_connection_is_not_reused_for_outbound_requests_by_default() {
1596 let (mut pool, addr, _events, listener) = pool_with_listener().await;
1597 let accepted = tokio::spawn(async move { listener.accept().await.expect("accepts") });
1598 let _client = TcpStream::connect(addr).await.expect("connects");
1599 let (stream, peer) = accepted.await.expect("accepted");
1600 pool.accept(stream, peer);
1601 assert_eq!(pool.len(), 1);
1602
1603 assert!(
1605 pool.send_on_existing(&tcp(peer), Bytes::from_static(b"x"))
1606 .await,
1607 "a response must go back over the connection it came in on"
1608 );
1609
1610 let before = pool.len();
1615 let _ = pool.send(&tcp(peer), Bytes::from_static(b"y")).await;
1616 assert!(
1617 pool.connections.len() <= before,
1618 "the inbound connection must not have been adopted for outbound use"
1619 );
1620 }
1621
1622 #[tokio::test]
1623 async fn idle_connections_are_evicted() {
1624 let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
1625 let addr = listener.local_addr().expect("has an address");
1626 let (tx, mut rx) = mpsc::channel(64);
1627 let mut pool = Pool::new(
1628 PoolConfig {
1629 idle_timeout: Duration::ZERO,
1630 ..PoolConfig::default()
1631 },
1632 Limits::stream(),
1633 tx,
1634 );
1635
1636 let accepted = tokio::spawn(async move { listener.accept().await.expect("accepts") });
1637 let mut client = TcpStream::connect(addr).await.expect("connects");
1638 let (stream, peer) = accepted.await.expect("accepted");
1639 pool.accept(stream, peer);
1640 assert_eq!(pool.len(), 1);
1641
1642 tokio::time::sleep(Duration::from_millis(5)).await;
1646 assert_eq!(pool.evict_idle().len(), 1);
1647 let mut byte = [0u8; 1];
1648 let read = tokio::time::timeout(Duration::from_secs(2), client.read(&mut byte))
1649 .await
1650 .expect("eviction closes promptly")
1651 .expect("read reports EOF");
1652 assert_eq!(read, 0, "the evicted peer must observe EOF");
1653 let closed = rx.recv().await.expect("connection reports completion");
1654 let Event::Closed { key, id } = closed else {
1655 panic!("expected a close event, got {closed:?}");
1656 };
1657 pool.remove(&key, id);
1658 wait_until_no_live_tasks(&mut pool).await;
1659 assert!(pool.is_empty(), "the live task releases its pool slot");
1660 }
1661
1662 #[tokio::test]
1663 async fn capacity_eviction_closes_the_socket_before_reusing_its_slot() {
1664 let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
1665 let addr = listener.local_addr().expect("has an address");
1666 let (tx, mut events) = mpsc::channel(64);
1667 let mut pool = Pool::new(
1668 PoolConfig {
1669 max_connections: 1,
1670 ..PoolConfig::default()
1671 },
1672 Limits::stream(),
1673 tx,
1674 );
1675
1676 let mut first = TcpStream::connect(addr).await.expect("first connects");
1677 let (first_server, first_peer) = listener.accept().await.expect("first accepts");
1678 pool.accept(first_server, first_peer);
1679
1680 let mut replacement = TcpStream::connect(addr).await.expect("second connects");
1681 let (second_server, second_peer) = listener.accept().await.expect("second accepts");
1682 pool.accept(second_server, second_peer);
1683 assert_eq!(pool.len(), 1, "a retiring task still owns the sole slot");
1684
1685 let mut byte = [0u8; 1];
1686 assert_eq!(
1687 tokio::time::timeout(Duration::from_secs(2), first.read(&mut byte))
1688 .await
1689 .expect("eviction closes promptly")
1690 .expect("first read completes"),
1691 0
1692 );
1693 let Event::Closed { key, id } = events.recv().await.expect("close reported") else {
1694 panic!("expected close");
1695 };
1696 assert!(pool.remove(&key, id));
1697 assert_eq!(
1698 pool.len(),
1699 1,
1700 "the reserved replacement takes the released slot"
1701 );
1702 assert!(
1703 pool.send_on_existing(&tcp(second_peer), Bytes::from_static(b"r"))
1704 .await
1705 );
1706 assert_eq!(
1707 replacement
1708 .read(&mut byte)
1709 .await
1710 .expect("replacement reads"),
1711 1
1712 );
1713 pool.shutdown().await;
1714 }
1715
1716 #[cfg(feature = "ws")]
1717 #[tokio::test]
1718 async fn idle_websocket_eviction_closes_the_peer_and_finishes_the_task() {
1719 use futures_util::StreamExt as _;
1720 use tokio_tungstenite::tungstenite::protocol::Role;
1721
1722 let (server_io, client_io) = tokio::io::duplex(1024);
1723 let server = crate::ws::Socket::from_raw_socket(server_io, Role::Server, None).await;
1724 let mut client = crate::ws::Socket::from_raw_socket(client_io, Role::Client, None).await;
1725 let (tx, mut events) = mpsc::channel(64);
1726 let mut pool = Pool::new(
1727 PoolConfig {
1728 idle_timeout: Duration::ZERO,
1729 ..PoolConfig::default()
1730 },
1731 Limits::stream(),
1732 tx,
1733 );
1734 let key = ConnectionKey::new(
1735 "127.0.0.1:5090".parse().expect("address"),
1736 TransportKind::Ws,
1737 );
1738 pool.accept_ws(server, key, Duration::from_secs(60));
1739 tokio::time::sleep(Duration::from_millis(5)).await;
1742 assert_eq!(pool.evict_idle().len(), 1);
1743
1744 let peer_end = tokio::time::timeout(Duration::from_secs(2), client.next())
1745 .await
1746 .expect("WebSocket eviction closes promptly");
1747 assert!(
1748 peer_end.is_none()
1749 || matches!(
1750 peer_end,
1751 Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_)))
1752 )
1753 || matches!(peer_end, Some(Err(_)))
1754 );
1755 let Event::Closed { key, id } = events.recv().await.expect("close reported") else {
1756 panic!("expected close");
1757 };
1758 pool.remove(&key, id);
1759 wait_until_no_live_tasks(&mut pool).await;
1760 assert!(pool.is_empty());
1761 }
1762
1763 #[cfg(feature = "ws")]
1764 #[tokio::test]
1765 async fn capacity_eviction_closes_websocket_peers_without_exceeding_the_live_limit() {
1766 use futures_util::StreamExt as _;
1767 use tokio_tungstenite::tungstenite::protocol::Role;
1768
1769 let (first_server_io, first_client_io) = tokio::io::duplex(1024);
1770 let first_server =
1771 crate::ws::Socket::from_raw_socket(first_server_io, Role::Server, None).await;
1772 let mut first_client =
1773 crate::ws::Socket::from_raw_socket(first_client_io, Role::Client, None).await;
1774 let (second_server_io, second_client_io) = tokio::io::duplex(1024);
1775 let second_server =
1776 crate::ws::Socket::from_raw_socket(second_server_io, Role::Server, None).await;
1777 let mut second_client =
1778 crate::ws::Socket::from_raw_socket(second_client_io, Role::Client, None).await;
1779 let (tx, mut events) = mpsc::channel(64);
1780 let mut pool = Pool::new(
1781 PoolConfig {
1782 max_connections: 1,
1783 ..PoolConfig::default()
1784 },
1785 Limits::stream(),
1786 tx,
1787 );
1788 let first_key = ConnectionKey::new(
1789 "127.0.0.1:5091".parse().expect("address"),
1790 TransportKind::Ws,
1791 );
1792 let second_key = ConnectionKey::new(
1793 "127.0.0.1:5092".parse().expect("address"),
1794 TransportKind::Ws,
1795 );
1796 pool.accept_ws(first_server, first_key, Duration::from_secs(60));
1797 pool.accept_ws(second_server, second_key, Duration::from_secs(60));
1798 assert_eq!(pool.len(), 1, "the retiring WebSocket still owns the slot");
1799
1800 let first_end = tokio::time::timeout(Duration::from_secs(2), first_client.next())
1801 .await
1802 .expect("eviction closes the first WebSocket");
1803 assert!(
1804 first_end.is_none()
1805 || matches!(&first_end, Some(Err(_)))
1806 || matches!(
1807 &first_end,
1808 Some(Ok(tokio_tungstenite::tungstenite::Message::Close(_)))
1809 )
1810 );
1811 let Event::Closed { key, id } = events.recv().await.expect("close reported") else {
1812 panic!("expected close");
1813 };
1814 assert!(pool.remove(&key, id));
1815 assert_eq!(pool.len(), 1, "the WebSocket reservation activates");
1816 assert!(
1817 tokio::time::timeout(Duration::from_millis(20), second_client.next())
1818 .await
1819 .is_err(),
1820 "the replacement remains live rather than being refused"
1821 );
1822 pool.shutdown().await;
1823 }
1824
1825 #[tokio::test]
1826 async fn quiet_connection_churn_never_exceeds_the_live_task_limit() {
1827 let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds");
1828 let addr = listener.local_addr().expect("address");
1829 let (tx, mut events) = mpsc::channel(64);
1830 let mut pool = Pool::new(
1831 PoolConfig {
1832 max_connections: 2,
1833 ..PoolConfig::default()
1834 },
1835 Limits::stream(),
1836 tx,
1837 );
1838
1839 for _ in 0..24 {
1840 let client = TcpStream::connect(addr).await.expect("connects");
1841 let (server, peer) = listener.accept().await.expect("accepts");
1842 pool.accept(server, peer);
1843 assert!(pool.len() <= 2, "live tasks exceeded the configured limit");
1844 if !pool.holds(&tcp(peer)) {
1845 drop(client);
1846 if let Some(Event::Closed { key, id }) = events.recv().await {
1847 pool.remove(&key, id);
1848 }
1849 if let Some(joined) = pool.tasks.join_next().await {
1850 joined.expect("evicted task exits");
1851 }
1852 }
1853 }
1854 pool.shutdown().await;
1855 }
1856
1857 #[tokio::test]
1858 async fn stale_close_generation_cannot_remove_or_fail_its_live_replacement() {
1859 let (first_socket, mut first_peer) = tokio::io::duplex(1024);
1860 let (replacement_socket, _replacement_peer) = tokio::io::duplex(1024);
1861 let (tx, mut events) = mpsc::channel(8);
1862 let mut pool = Pool::new(PoolConfig::default(), Limits::stream(), tx);
1863 let key = tcp("127.0.0.1:5093".parse().expect("address"));
1864
1865 pool.insert_stream(first_socket, key.clone(), Origin::Inbound)
1866 .expect("first generation starts");
1867 let old_id = pool.connections.get(&key).expect("registered").id;
1868 pool.retire(&key);
1869 let mut byte = [0u8; 1];
1870 assert_eq!(first_peer.read(&mut byte).await.expect("EOF"), 0);
1871 let Event::Closed { id, .. } = events.recv().await.expect("old close arrives") else {
1872 panic!("expected old close");
1873 };
1874 assert_eq!(id, old_id);
1875
1876 pool.insert_stream(replacement_socket, key.clone(), Origin::Inbound)
1877 .expect("replacement starts after old task exits");
1878 let replacement_id = pool.connections.get(&key).expect("replacement held").id;
1879 assert_ne!(replacement_id, old_id);
1880
1881 assert!(pool.remove(&key, old_id), "the retirement is acknowledged");
1884 assert!(
1885 !pool.remove(&key, old_id),
1886 "the acknowledgement is exactly once"
1887 );
1888 assert!(pool.holds(&key), "the replacement remains routable");
1889 pool.shutdown().await;
1890 }
1891
1892 #[tokio::test]
1893 async fn same_key_replacement_at_capacity_does_not_evict_an_unrelated_connection() {
1894 let (a_socket, mut a_peer) = tokio::io::duplex(1024);
1895 let (b_socket, mut b_peer) = tokio::io::duplex(1024);
1896 let (replacement_socket, _replacement_peer) = tokio::io::duplex(1024);
1897 let (tx, mut events) = mpsc::channel(8);
1898 let mut pool = Pool::new(
1899 PoolConfig {
1900 max_connections: 2,
1901 ..PoolConfig::default()
1902 },
1903 Limits::stream(),
1904 tx,
1905 );
1906 let a = tcp("127.0.0.1:5094".parse().expect("address"));
1907 let b = tcp("127.0.0.1:5095".parse().expect("address"));
1908 pool.insert_stream(a_socket, a.clone(), Origin::Inbound)
1909 .expect("A starts");
1910 pool.insert_stream(b_socket, b.clone(), Origin::Inbound)
1911 .expect("B starts");
1912
1913 pool.insert_stream(replacement_socket, a, Origin::Inbound)
1914 .expect("A replacement is reserved while the old task retires");
1915
1916 let mut byte = [0u8; 1];
1917 assert_eq!(a_peer.read(&mut byte).await.expect("A closes"), 0);
1918 let Event::Closed { key, id } = events.recv().await.expect("A close arrives") else {
1919 panic!("expected close");
1920 };
1921 assert!(pool.remove(&key, id));
1922 assert!(pool.holds(&b), "B must not be selected as a second victim");
1923 assert!(
1924 pool.send_on_existing(&b, Bytes::from_static(b"b")).await,
1925 "B remains writable"
1926 );
1927 assert_eq!(b_peer.read(&mut byte).await.expect("B receives"), 1);
1928 assert_eq!(byte[0], b'b');
1929 pool.shutdown().await;
1930 }
1931
1932 #[tokio::test]
1933 async fn public_send_reserves_same_key_replacement_without_evicting_b() {
1934 let (a_socket, mut a_peer) = tokio::io::duplex(1024);
1935 let (b_socket, mut b_peer) = tokio::io::duplex(1024);
1936 let (events, mut closed) = mpsc::channel(8);
1937 let mut pool = Pool::new(
1938 PoolConfig {
1939 max_connections: 2,
1940 ..PoolConfig::default()
1941 },
1942 Limits::stream(),
1943 events,
1944 );
1945 let a = tcp("127.0.0.1:59991".parse().expect("address"));
1946 let b = tcp("127.0.0.1:59992".parse().expect("address"));
1947 pool.insert_stream(a_socket, a.clone(), Origin::Inbound)
1948 .expect("A starts inbound");
1949 pool.insert_stream(b_socket, b.clone(), Origin::Inbound)
1950 .expect("B starts inbound");
1951 let old_a = pool.generation(&a).expect("A generation");
1952
1953 pool.send(&a, Bytes::from_static(b"replacement"))
1954 .await
1955 .expect("public send reserves outbound A");
1956 let replacement = pool.generation(&a).expect("replacement generation");
1957 assert_ne!(replacement, old_a);
1958 assert!(pool.holds(&b), "B is not selected as a second victim");
1959 let mut byte = [0u8; 1];
1960 assert_eq!(a_peer.read(&mut byte).await.expect("old A closes"), 0);
1961 let Event::Closed { key, id } = closed.recv().await.expect("old A reports close") else {
1962 panic!("expected close");
1963 };
1964 assert_eq!((&key, id), (&a, old_a));
1965 assert!(pool.remove(&key, id));
1966 assert!(pool.send_on_existing(&b, Bytes::from_static(b"b")).await);
1967 assert_eq!(b_peer.read(&mut byte).await.expect("B reads"), 1);
1968 pool.shutdown().await;
1969 }
1970
1971 #[cfg(feature = "tls")]
1972 #[tokio::test]
1973 async fn public_send_tls_reserves_same_key_replacement_without_evicting_b() {
1974 use sipx_testkit::certs::Ca;
1975
1976 use crate::tls::{ClientTls, TrustAnchors};
1977
1978 let ca = Ca::new();
1979 let mut anchors = TrustAnchors::only();
1980 anchors.add_pem(ca.pem().as_bytes()).expect("CA loads");
1981 let client = ClientTls::new(&anchors).expect("TLS client");
1982 let (a_socket, mut a_peer) = tokio::io::duplex(1024);
1983 let (b_socket, _b_peer) = tokio::io::duplex(1024);
1984 let (events, _closed) = mpsc::channel(8);
1985 let mut pool = Pool::new(
1986 PoolConfig {
1987 max_connections: 2,
1988 ..PoolConfig::default()
1989 },
1990 Limits::stream(),
1991 events,
1992 );
1993 let a = ConnectionKey::new(
1994 "127.0.0.1:59993".parse().expect("address"),
1995 TransportKind::Tls,
1996 );
1997 let b = tcp("127.0.0.1:59994".parse().expect("address"));
1998 pool.insert_stream(a_socket, a.clone(), Origin::Inbound)
1999 .expect("A starts inbound");
2000 pool.insert_stream(b_socket, b.clone(), Origin::Inbound)
2001 .expect("B starts inbound");
2002
2003 pool.send_tls(&a, "localhost", &client, Bytes::from_static(b"replacement"))
2004 .await
2005 .expect("public TLS send reserves outbound A");
2006 assert!(pool.holds(&b));
2007 let mut byte = [0u8; 1];
2008 assert_eq!(a_peer.read(&mut byte).await.expect("old A closes"), 0);
2009 pool.shutdown().await;
2010 }
2011
2012 #[cfg(feature = "ws")]
2013 #[tokio::test]
2014 async fn public_send_ws_reserves_same_key_replacement_without_evicting_b() {
2015 let (a_socket, mut a_peer) = tokio::io::duplex(1024);
2016 let (b_socket, _b_peer) = tokio::io::duplex(1024);
2017 let (events, _closed) = mpsc::channel(8);
2018 let mut pool = Pool::new(
2019 PoolConfig {
2020 max_connections: 2,
2021 ..PoolConfig::default()
2022 },
2023 Limits::stream(),
2024 events,
2025 );
2026 let a = ConnectionKey::new(
2027 "127.0.0.1:59995".parse().expect("address"),
2028 TransportKind::Ws,
2029 );
2030 let b = tcp("127.0.0.1:59996".parse().expect("address"));
2031 pool.insert_stream(a_socket, a.clone(), Origin::Inbound)
2032 .expect("A starts inbound");
2033 pool.insert_stream(b_socket, b.clone(), Origin::Inbound)
2034 .expect("B starts inbound");
2035
2036 pool.send_ws(
2037 &a,
2038 "localhost",
2039 Duration::from_secs(60),
2040 #[cfg(feature = "wss")]
2041 None,
2042 Bytes::from_static(b"replacement"),
2043 )
2044 .await
2045 .expect("public WebSocket send reserves outbound A");
2046 assert!(pool.holds(&b));
2047 let mut byte = [0u8; 1];
2048 assert_eq!(a_peer.read(&mut byte).await.expect("old A closes"), 0);
2049 pool.shutdown().await;
2050 }
2051
2052 #[tokio::test]
2053 async fn shutdown_cancels_a_final_close_blocked_on_a_full_event_channel() {
2054 let (events, _unread) = mpsc::channel(1);
2055 events
2056 .send(Event::FramingFailed {
2057 key: tcp("127.0.0.1:59997".parse().expect("address")),
2058 })
2059 .await
2060 .expect("fills event channel");
2061 let (socket, mut peer) = tokio::io::duplex(1024);
2062 let mut pool = Pool::new(PoolConfig::default(), Limits::stream(), events);
2063 let key = tcp("127.0.0.1:59998".parse().expect("address"));
2064 pool.insert_stream(socket, key.clone(), Origin::Inbound)
2065 .expect("connection starts");
2066 pool.retire(&key);
2067 let mut byte = [0u8; 1];
2068 assert_eq!(peer.read(&mut byte).await.expect("retirement closes"), 0);
2069 tokio::task::yield_now().await;
2070
2071 tokio::time::timeout(Duration::from_secs(1), pool.shutdown())
2072 .await
2073 .expect("shutdown cancels the blocked final event sender");
2074 }
2075}