yamux/connection/
rtt.rs

1// Copyright (c) 2023 Protocol Labs.
2//
3// Licensed under the Apache License, Version 2.0 or MIT license, at your option.
4//
5// A copy of the Apache License, Version 2.0 is included in the software as
6// LICENSE-APACHE and a copy of the MIT license is included in the software
7// as LICENSE-MIT. You may also obtain a copy of the Apache License, Version 2.0
8// at https://www.apache.org/licenses/LICENSE-2.0 and a copy of the MIT license
9// at https://opensource.org/licenses/MIT.
10
11//! Connection round-trip time measurement
12
13use std::sync::Arc;
14
15use parking_lot::Mutex;
16use web_time::{Duration, Instant};
17
18use crate::connection::Action;
19use crate::frame::{header::Ping, Frame};
20
21const PING_INTERVAL: Duration = Duration::from_secs(10);
22
23#[derive(Clone, Debug)]
24pub(crate) struct Rtt(Arc<Mutex<RttInner>>);
25
26impl Rtt {
27    pub(crate) fn new() -> Self {
28        Self(Arc::new(Mutex::new(RttInner {
29            rtt: None,
30            state: RttState::Waiting {
31                next: Instant::now(),
32            },
33        })))
34    }
35
36    pub(crate) fn next_ping(&mut self) -> Option<Frame<Ping>> {
37        let state = &mut self.0.lock().state;
38
39        match state {
40            RttState::AwaitingPong { .. } => return None,
41            RttState::Waiting { next } => {
42                if *next > Instant::now() {
43                    return None;
44                }
45            }
46        }
47
48        let nonce = rand::random();
49        *state = RttState::AwaitingPong {
50            sent_at: Instant::now(),
51            nonce,
52        };
53        log::debug!("sending ping {nonce}");
54        Some(Frame::ping(nonce))
55    }
56
57    pub(crate) fn handle_pong(&mut self, received_nonce: u32) -> Action {
58        let inner = &mut self.0.lock();
59
60        let (sent_at, expected_nonce) = match inner.state {
61            RttState::Waiting { .. } => {
62                log::error!("received unexpected pong {received_nonce}");
63                return Action::Terminate(Frame::protocol_error());
64            }
65            RttState::AwaitingPong { sent_at, nonce } => (sent_at, nonce),
66        };
67
68        if received_nonce != expected_nonce {
69            log::error!("received pong with {received_nonce} but expected {expected_nonce}");
70            return Action::Terminate(Frame::protocol_error());
71        }
72
73        let rtt = sent_at.elapsed();
74        inner.rtt = Some(rtt);
75        log::debug!("received pong {received_nonce}, estimated round-trip-time {rtt:?}");
76
77        inner.state = RttState::Waiting {
78            next: Instant::now() + PING_INTERVAL,
79        };
80
81        Action::None
82    }
83
84    pub(crate) fn get(&self) -> Option<Duration> {
85        self.0.lock().rtt
86    }
87}
88
89#[cfg(test)]
90impl quickcheck::Arbitrary for Rtt {
91    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
92        Self(Arc::new(Mutex::new(RttInner::arbitrary(g))))
93    }
94}
95
96#[derive(Debug)]
97#[cfg_attr(test, derive(Clone))]
98struct RttInner {
99    state: RttState,
100    rtt: Option<Duration>,
101}
102
103#[cfg(test)]
104impl quickcheck::Arbitrary for RttInner {
105    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
106        Self {
107            state: RttState::arbitrary(g),
108            rtt: if bool::arbitrary(g) {
109                Some(Duration::arbitrary(g))
110            } else {
111                None
112            },
113        }
114    }
115}
116
117#[derive(Debug)]
118#[cfg_attr(test, derive(Clone))]
119enum RttState {
120    AwaitingPong { sent_at: Instant, nonce: u32 },
121    Waiting { next: Instant },
122}
123
124#[cfg(test)]
125impl quickcheck::Arbitrary for RttState {
126    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
127        if bool::arbitrary(g) {
128            RttState::AwaitingPong {
129                sent_at: Instant::now(),
130                nonce: u32::arbitrary(g),
131            }
132        } else {
133            RttState::Waiting {
134                next: Instant::now(),
135            }
136        }
137    }
138}
yamux/connection/rtt.rs

yamux/connection/
rtt.rs
