Skip to main content

msg_transport/tcp/
stats.rs

1use std::{os::fd::AsRawFd, time::Duration};
2
3use tokio::net::TcpStream;
4
5#[derive(Debug, Default)]
6pub struct TcpStats {
7    /// The congestion window in bytes.
8    pub congestion_window: u32,
9    /// Our receive window in bytes.
10    pub receive_window: u32,
11    /// Our send window (= the peer's advertised receive window) in bytes.
12    #[cfg(target_os = "macos")]
13    pub send_window: u32,
14    /// The most recent RTT sample.
15    #[cfg(target_os = "macos")]
16    pub last_rtt: Duration,
17    /// The smoothed round-trip time.
18    pub smoothed_rtt: Duration,
19    /// The round-trip time variance.
20    pub rtt_variance: Duration,
21    /// Total bytes sent on the socket.
22    #[cfg(target_os = "macos")]
23    pub tx_bytes: u64,
24    /// Total bytes received on the socket.
25    #[cfg(target_os = "macos")]
26    pub rx_bytes: u64,
27    /// Total sender retransmitted bytes on the socket.
28    pub retransmitted_bytes: u64,
29    /// Total sender retransmitted packets on the socket.
30    pub retransmitted_packets: u64,
31    /// The current retransmission timeout.
32    pub retransmission_timeout: Duration,
33}
34
35#[cfg(target_os = "macos")]
36impl TryFrom<&TcpStream> for TcpStats {
37    type Error = std::io::Error;
38
39    /// Gathers stats from the given TCP socket file descriptor, sourced from the OS with
40    /// [`libc::getsockopt`].
41    fn try_from(stream: &TcpStream) -> Result<Self, Self::Error> {
42        let info = getsockopt::<libc::tcp_connection_info>(stream, libc::TCP_CONNECTION_INFO)?;
43
44        Ok(info.into())
45    }
46}
47
48#[cfg(target_os = "macos")]
49impl From<libc::tcp_connection_info> for TcpStats {
50    /// Converts a [`libc::tcp_connection_info`] into [`TcpStats`].
51    fn from(info: libc::tcp_connection_info) -> Self {
52        // Window sizes
53        let congestion_window = info.tcpi_snd_cwnd;
54        let receive_window = info.tcpi_rcv_wnd;
55        let send_window = info.tcpi_snd_wnd;
56
57        // RTT
58        let last_rtt = Duration::from_millis(info.tcpi_rttcur as u64);
59        let smoothed_rtt = Duration::from_millis(info.tcpi_srtt as u64);
60        let rtt_variance = Duration::from_millis(info.tcpi_rttvar as u64);
61
62        // Volumes
63        let tx_bytes = info.tcpi_txbytes;
64        let rx_bytes = info.tcpi_rxbytes;
65
66        // Retransmissions
67        let retransmitted_bytes = info.tcpi_txretransmitbytes;
68        let retransmitted_packets = info.tcpi_rxretransmitpackets;
69        let retransmission_timeout = Duration::from_millis(info.tcpi_rto as u64);
70
71        Self {
72            congestion_window,
73            receive_window,
74            send_window,
75            last_rtt,
76            smoothed_rtt,
77            rtt_variance,
78            tx_bytes,
79            rx_bytes,
80            retransmitted_bytes,
81            retransmitted_packets,
82            retransmission_timeout,
83        }
84    }
85}
86
87#[cfg(target_os = "linux")]
88impl TryFrom<&TcpStream> for TcpStats {
89    type Error = std::io::Error;
90
91    /// Gathers stats from the given TCP socket file descriptor, sourced from the OS with
92    /// [`libc::getsockopt`].
93    fn try_from(stream: &TcpStream) -> Result<Self, Self::Error> {
94        let info = getsockopt::<libc::tcp_info>(stream, libc::TCP_INFO)?;
95
96        Ok(info.into())
97    }
98}
99
100#[cfg(target_os = "linux")]
101impl From<libc::tcp_info> for TcpStats {
102    /// Converts a [`libc::tcp_info`] into [`TcpStats`].
103    fn from(info: libc::tcp_info) -> Self {
104        // On Linux, tcpi_snd_cwnd is in segments; convert to bytes using snd_mss.
105        let congestion_window = info.tcpi_snd_cwnd.saturating_mul(info.tcpi_snd_mss);
106        // Local advertised receive window (bytes).
107        let receive_window = info.tcpi_rcv_space;
108
109        // RTT fields are reported in microseconds.
110        let smoothed_rtt = Duration::from_micros(info.tcpi_rtt as u64);
111        let rtt_variance = Duration::from_micros(info.tcpi_rttvar as u64);
112
113        // Retransmissions
114        let retransmitted_packets = info.tcpi_total_retrans as u64;
115        let retransmitted_bytes = retransmitted_packets.saturating_mul(info.tcpi_snd_mss as u64);
116        // RTO is in microseconds.
117        let retransmission_timeout = Duration::from_micros(info.tcpi_rto as u64);
118
119        Self {
120            congestion_window,
121            receive_window,
122            smoothed_rtt,
123            rtt_variance,
124            retransmitted_bytes,
125            retransmitted_packets,
126            retransmission_timeout,
127        }
128    }
129}
130
131/// Helper function to get a socket option from a TCP stream.
132fn getsockopt<T>(stream: &TcpStream, option: libc::c_int) -> std::io::Result<T> {
133    let mut info = unsafe { std::mem::zeroed::<T>() };
134    let mut len = std::mem::size_of::<T>() as libc::socklen_t;
135    let dst = &mut info as *mut _ as *mut _;
136
137    let result =
138        unsafe { libc::getsockopt(stream.as_raw_fd(), libc::IPPROTO_TCP, option, dst, &mut len) };
139
140    if result != 0 {
141        return Err(std::io::Error::last_os_error());
142    }
143
144    Ok(info)
145}