librqbit_utp/
traits.rs

1use std::{
2    future::Future,
3    net::SocketAddr,
4    task::{Context, Poll},
5    time::Instant,
6};
7
8use librqbit_dualstack_sockets::{PollSendToVectored, UdpSocket};
9
10use crate::metrics::METRICS;
11
12/// An abstraction for underlying transport. UDP is default, but can be swapped to a custom transport.
13///
14/// Tests use mock transport.
15pub trait Transport: PollSendToVectored + Send + Sync + Unpin + 'static {
16    fn recv_from<'a>(
17        &'a self,
18        buf: &'a mut [u8],
19    ) -> impl Future<Output = std::io::Result<(usize, SocketAddr)>> + Send + Sync + 'a;
20
21    fn send_to<'a>(
22        &'a self,
23        buf: &'a [u8],
24        target: SocketAddr,
25    ) -> impl Future<Output = std::io::Result<usize>> + Send + Sync + 'a;
26
27    fn poll_send_to(
28        &self,
29        cx: &mut Context<'_>,
30        buf: &[u8],
31        target: SocketAddr,
32    ) -> Poll<std::io::Result<usize>>;
33
34    // The local address transport is bound to. Used only for logging.
35    fn bind_addr(&self) -> SocketAddr;
36}
37
38impl Transport for UdpSocket {
39    fn recv_from<'a>(
40        &'a self,
41        buf: &'a mut [u8],
42    ) -> impl Future<Output = std::io::Result<(usize, SocketAddr)>> + Send + Sync + 'a {
43        UdpSocket::recv_from(self, buf)
44    }
45
46    fn send_to<'a>(
47        &'a self,
48        buf: &'a [u8],
49        target: SocketAddr,
50    ) -> impl Future<Output = std::io::Result<usize>> + Send + Sync + 'a {
51        UdpSocket::send_to(self, buf, target)
52    }
53
54    fn poll_send_to(
55        &self,
56        cx: &mut Context<'_>,
57        buf: &[u8],
58        target: SocketAddr,
59    ) -> Poll<std::io::Result<usize>> {
60        let res = UdpSocket::poll_send_to(self, cx, buf, target);
61        if res.is_pending() {
62            METRICS.send_poll_pending.increment(1);
63        }
64        res
65    }
66
67    fn bind_addr(&self) -> SocketAddr {
68        UdpSocket::bind_addr(self)
69    }
70}
71
72// A trait for mocking stuff in tests.
73pub trait UtpEnvironment: Send + Sync + Unpin + 'static {
74    fn now(&self) -> Instant;
75    fn copy(&self) -> Self;
76    fn random_u16(&self) -> u16;
77}
78
79#[derive(Default, Clone, Copy)]
80pub struct DefaultUtpEnvironment {}
81
82impl UtpEnvironment for DefaultUtpEnvironment {
83    fn now(&self) -> Instant {
84        Instant::now()
85    }
86
87    fn copy(&self) -> Self {
88        *self
89    }
90
91    fn random_u16(&self) -> u16 {
92        rand::random()
93    }
94}
librqbit_utp/traits.rs

librqbit_utp/
traits.rs
