2 files changed, 54 insertions, 46 deletions

diff --git a/crates/windlass/src/transport/mod.rs b/crates/windlass/src/transport/mod.rs
index 6f2cc7c..818a594 100644
--- a/crates/windlass/src/transport/mod.rs
+++ b/crates/windlass/src/transport/mod.rs
@@ -1,5 +1,6 @@
 use read::{FrameReader, ReceivedFrame};
-use std::{collections::VecDeque, sync::Arc, time::Duration};
+use rtt::RttState;
+use std::{collections::VecDeque, sync::Arc};
 use tokio::io::AsyncReadExt;
 use tokio::{
     io::{AsyncBufRead, AsyncWrite, AsyncWriteExt},
@@ -12,6 +13,9 @@ use tokio_util::{bytes::BytesMut, sync::CancellationToken};
 
 use crate::encoding::crc16;
 
+mod read;
+mod rtt;
+
 pub(crate) const MESSAGE_HEADER_SIZE: usize = 2;
 pub(crate) const MESSAGE_TRAILER_SIZE: usize = 3;
 pub(crate) const MESSAGE_LENGTH_MIN: usize = MESSAGE_HEADER_SIZE + MESSAGE_TRAILER_SIZE;
@@ -23,8 +27,6 @@ pub(crate) const MESSAGE_VALUE_SYNC: u8 = 0x7E;
 pub(crate) const MESSAGE_DEST: u8 = 0x10;
 pub(crate) const MESSAGE_SEQ_MASK: u8 = 0x0F;
 
-mod read;
-
 /// Wrapper around a connection to a klipper firmware MCU, which deals with
 /// retransmission, flow control, etc.
 ///
@@ -97,48 +99,6 @@ pub enum TransportError {
     IOError(#[from] std::io::Error),
 }
 
-const MIN_RTO: f32 = 0.025;
-const MAX_RTO: f32 = 5.000;
-
-/// State for estimating the round trip time of the connection
-#[derive(Debug)]
-struct RttState {
-    srtt: f32,
-    rttvar: f32,
-    rto: f32,
-}
-
-impl Default for RttState {
-    fn default() -> Self {
-        Self {
-            srtt: 0.0,
-            rttvar: 0.0,
-            rto: MIN_RTO,
-        }
-    }
-}
-
-impl RttState {
-    /// Get the current recommended retransmission timeout
-    fn rto(&self) -> Duration {
-        Duration::from_secs_f32(self.rto)
-    }
-
-    /// Update the RTT estimation given a new observation
-    fn update(&mut self, rtt: Duration) {
-        let r = rtt.as_secs_f32();
-        if self.srtt == 0.0 {
-            self.rttvar = r / 2.0;
-            self.srtt = r * 10.0; // Klipper uses this, we'll copy it
-        } else {
-            self.rttvar = (3.0 * self.rttvar + (self.srtt - r).abs()) / 4.0;
-            self.srtt = (7.0 * self.srtt + r) / 8.0;
-        }
-        let rttvar4 = (self.rttvar * 4.0).max(0.001);
-        self.rto = (self.srtt + rttvar4).clamp(MIN_RTO, MAX_RTO);
-    }
-}
-
 /// State for the task which deals with transport state
 #[derive(Debug)]
 struct TransportState<R, W> {
@@ -383,7 +343,7 @@ impl<R: AsyncBufRead + Unpin, W: AsyncWrite + Unpin> TransportState<R, W> {
             }
             self.retransmit_now = false;
         } else {
-            self.rtt_state.rto = (self.rtt_state.rto * 2.0).clamp(MIN_RTO, MAX_RTO);
+            self.rtt_state.double();
             self.ignore_nak_seq = self.send_sequence;
         }
         self.retransmit_seq = self.send_sequence;
diff --git a/crates/windlass/src/transport/rtt.rs b/crates/windlass/src/transport/rtt.rs
new file mode 100644
index 0000000..69841f6
--- /dev/null
+++ b/crates/windlass/src/transport/rtt.rs
@@ -0,0 +1,48 @@
+use std::time::Duration;
+
+const MIN_RTO: f32 = 0.025;
+const MAX_RTO: f32 = 5.000;
+
+/// State for estimating the round trip time of the connection
+#[derive(Debug)]
+pub struct RttState {
+    srtt: f32,
+    rttvar: f32,
+    rto: f32,
+}
+
+impl Default for RttState {
+    fn default() -> Self {
+        Self {
+            srtt: 0.0,
+            rttvar: 0.0,
+            rto: MIN_RTO,
+        }
+    }
+}
+
+impl RttState {
+    /// Get the current recommended retransmission timeout
+    pub fn rto(&self) -> Duration {
+        Duration::from_secs_f32(self.rto)
+    }
+
+    /// Update the RTT estimation given a new observation
+    pub fn update(&mut self, rtt: Duration) {
+        let r = rtt.as_secs_f32();
+        if self.srtt == 0.0 {
+            self.rttvar = r / 2.0;
+            self.srtt = r * 10.0; // Klipper uses this, we'll copy it
+        } else {
+            self.rttvar = (3.0 * self.rttvar + (self.srtt - r).abs()) / 4.0;
+            self.srtt = (7.0 * self.srtt + r) / 8.0;
+        }
+        let rttvar4 = (self.rttvar * 4.0).max(0.001);
+        self.rto = (self.srtt + rttvar4).clamp(MIN_RTO, MAX_RTO);
+    }
+
+    /// Double the RTO estimation
+    pub fn double(&mut self) {
+        self.rto = (self.rto * 2.0).clamp(MIN_RTO, MAX_RTO);
+    }
+}