fix: exclude cold connection setup from upload probes

Author: bellabaelfire

apps/desktop/src-tauri/Cargo.toml

@@ -180,3 +180,4 @@ nix = { version = "0.29.0", features = ["fs"] }
 
 [dev-dependencies]
 tauri = { workspace = true, features = ["test"] }
+tokio = { workspace = true, features = ["test-util"] }

apps/desktop/src-tauri/src/upload_health.rs

@@ -11,6 +11,10 @@ use crate::{
     web_api::{AuthedApiError, ManagerExt},
 };
 
+mod timing;
+
+use timing::{measure_warm_probe_rtt, upload_elapsed_after_rtt, upload_mbps_for_bytes};
+
 const PROBE_BYTES: usize = 256 * 1024;
 const HEALTH_FRESH_FOR: Duration = Duration::from_secs(10 * 60);
 const HEALTH_REQUEST_TIMEOUT: Duration = Duration::from_secs(8);
@@ -168,25 +172,8 @@ async fn measure_probe_rtt(app: &AppHandle) -> Option<Duration> {
     }
 }
 
-fn upload_elapsed_after_rtt(total_elapsed: Duration, rtt_elapsed: Option<Duration>) -> Duration {
-    let total_elapsed = total_elapsed.max(Duration::from_millis(1));
-
-    let Some(rtt_elapsed) = rtt_elapsed else {
-        return total_elapsed;
-    };
-
-    match total_elapsed.checked_sub(rtt_elapsed) {
-        Some(adjusted_elapsed) if adjusted_elapsed >= Duration::from_millis(50) => adjusted_elapsed,
-        _ => total_elapsed,
-    }
-}
-
-fn upload_mbps_for_bytes(byte_count: usize, elapsed: Duration) -> f64 {
-    (byte_count as f64 * 8.0) / elapsed.max(Duration::from_millis(1)).as_secs_f64() / 1_000_000.0
-}
-
 async fn run_probe(app: &AppHandle) -> UploadHealthSnapshot {
-    let rtt_elapsed = measure_probe_rtt(app).await;
+    let rtt_elapsed = measure_warm_probe_rtt(HEALTH_RTT_TIMEOUT, || measure_probe_rtt(app)).await;
     let payload = probe_payload();
     let payload_len = payload.len();
 
@@ -373,20 +360,4 @@ mod tests {
 
         assert_eq!(cache.fresh_instant_resolution_cap().await, Some(1280));
     }
-
-    #[test]
-    fn subtracts_rtt_from_probe_elapsed_when_safe() {
-        assert_eq!(
-            upload_elapsed_after_rtt(Duration::from_millis(700), Some(Duration::from_millis(500))),
-            Duration::from_millis(200)
-        );
-    }
-
-    #[test]
-    fn keeps_total_elapsed_when_rtt_would_overcorrect() {
-        assert_eq!(
-            upload_elapsed_after_rtt(Duration::from_millis(520), Some(Duration::from_millis(500))),
-            Duration::from_millis(520)
-        );
-    }
 }

apps/desktop/src-tauri/src/upload_health/timing.rs

@@ -0,0 +1,158 @@
+use std::{future::Future, time::Duration};
+
+pub(super) async fn measure_warm_probe_rtt<Probe, ProbeFuture>(
+    budget: Duration,
+    mut probe: Probe,
+) -> Option<Duration>
+where
+    Probe: FnMut() -> ProbeFuture,
+    ProbeFuture: Future<Output = Option<Duration>>,
+{
+    tokio::time::timeout(budget, async {
+        // The first HEAD can include DNS/TCP/TLS setup that the pooled POST will not pay.
+        probe().await?;
+        probe().await
+    })
+    .await
+    .ok()
+    .flatten()
+}
+
+pub(super) fn upload_elapsed_after_rtt(
+    total_elapsed: Duration,
+    rtt_elapsed: Option<Duration>,
+) -> Duration {
+    let total_elapsed = total_elapsed.max(Duration::from_millis(1));
+
+    let Some(rtt_elapsed) = rtt_elapsed else {
+        return total_elapsed;
+    };
+
+    match total_elapsed.checked_sub(rtt_elapsed) {
+        Some(adjusted_elapsed) if adjusted_elapsed >= Duration::from_millis(50) => adjusted_elapsed,
+        _ => total_elapsed,
+    }
+}
+
+pub(super) fn upload_mbps_for_bytes(byte_count: usize, elapsed: Duration) -> f64 {
+    (byte_count as f64 * 8.0) / elapsed.max(Duration::from_millis(1)).as_secs_f64() / 1_000_000.0
+}
+
+#[cfg(test)]
+mod tests {
+    use std::future::{pending, ready};
+
+    use super::*;
+
+    #[tokio::test]
+    async fn cold_connection_time_does_not_inflate_upload_speed() {
+        let mut samples = [
+            Some(Duration::from_millis(500)),
+            Some(Duration::from_millis(100)),
+        ]
+        .into_iter();
+        let rtt = measure_warm_probe_rtt(Duration::from_secs(2), || {
+            ready(samples.next().expect("unexpected extra HEAD request"))
+        })
+        .await;
+
+        assert_eq!(rtt, Some(Duration::from_millis(100)));
+        assert_eq!(samples.next(), None);
+        let elapsed = upload_elapsed_after_rtt(Duration::from_millis(700), rtt);
+        assert_eq!(elapsed, Duration::from_millis(600));
+        let mbps = upload_mbps_for_bytes(256 * 1024, elapsed);
+        assert!(mbps > 3.0 && mbps < 4.0, "unexpected upload rate: {mbps}");
+    }
+
+    #[tokio::test]
+    async fn failed_warmup_does_not_start_another_head() {
+        let mut calls = 0;
+        let rtt = measure_warm_probe_rtt(Duration::from_secs(2), || {
+            calls += 1;
+            ready(None)
+        })
+        .await;
+
+        assert_eq!(rtt, None);
+        assert_eq!(calls, 1);
+    }
+
+    #[tokio::test]
+    async fn failed_measurement_does_not_reuse_the_cold_sample() {
+        let mut samples = [Some(Duration::from_millis(500)), None].into_iter();
+        let rtt = measure_warm_probe_rtt(Duration::from_secs(2), || {
+            ready(samples.next().expect("unexpected extra HEAD request"))
+        })
+        .await;
+
+        assert_eq!(rtt, None);
+        assert_eq!(
+            upload_elapsed_after_rtt(Duration::from_millis(700), rtt),
+            Duration::from_millis(700)
+        );
+    }
+
+    #[tokio::test]
+    async fn budget_expires_while_warming_up() {
+        let mut calls = 0;
+        let rtt = measure_warm_probe_rtt(Duration::ZERO, || {
+            calls += 1;
+            pending::<Option<Duration>>()
+        })
+        .await;
+
+        assert_eq!(rtt, None);
+        assert_eq!(calls, 1);
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn warmup_and_measurement_share_one_deadline() {
+        let mut calls = 0;
+        let started = tokio::time::Instant::now();
+        let rtt = measure_warm_probe_rtt(Duration::from_secs(2), || {
+            calls += 1;
+            async {
+                tokio::time::sleep(Duration::from_millis(1500)).await;
+                Some(Duration::from_millis(1500))
+            }
+        })
+        .await;
+
+        assert_eq!(rtt, None);
+        assert_eq!(calls, 2);
+        assert_eq!(started.elapsed(), Duration::from_secs(2));
+    }
+
+    #[test]
+    fn subtracts_warm_rtt_from_upload_elapsed() {
+        assert_eq!(
+            upload_elapsed_after_rtt(Duration::from_millis(700), Some(Duration::from_millis(100))),
+            Duration::from_millis(600)
+        );
+    }
+
+    #[test]
+    fn keeps_total_elapsed_when_rtt_would_overcorrect() {
+        for rtt in [500, 520, 600] {
+            assert_eq!(
+                upload_elapsed_after_rtt(
+                    Duration::from_millis(520),
+                    Some(Duration::from_millis(rtt))
+                ),
+                Duration::from_millis(520)
+            );
+        }
+    }
+
+    #[test]
+    fn elapsed_is_positive_without_an_rtt_sample() {
+        assert_eq!(
+            upload_elapsed_after_rtt(Duration::ZERO, None),
+            Duration::from_millis(1)
+        );
+        assert_eq!(
+            upload_elapsed_after_rtt(Duration::from_millis(700), None),
+            Duration::from_millis(700)
+        );
+    }
+}