1 files changed, 168 insertions, 0 deletions
diff --git a/src/swarm/throughput.rs b/src/swarm/throughput.rs
new file mode 100644
index 0000000..f1b9b14
--- /dev/null
+++ b/src/swarm/throughput.rs
@@ -0,0 +1,168 @@
+//! Sync throughput with peers, measured from sync exchanges we have anyway: from sending
+//! a sync message until the peer's answer arrives, the sync bytes that went either way
+//! with that peer. No transfers of our own.
+//!
+//! Only exchanges of at least `MIN_BYTES` count; smaller ones measure the round trip,
+//! not the throughput. The round trip (`swarm::rtt`) is taken off the elapsed time, but
+//! never more than half of it, so a stale round trip can't inflate a sample. What
+//! remains also includes the time the peer took to apply and answer: this is the rate a
+//! peer delivers at in practice, not the raw capacity of the path.
+//!
+//! Samples are smoothed (exponential moving average, weight 1/4). Large exchanges are
+//! rare (catching up, big writes), so estimates are kept however old; they come with
+//! their age.
+
+use std::collections::HashMap;
+use std::time::Duration;
+
+use tokio::time::Instant;
+
+use crate::transport::traffic::Traffic;
+
+/// Smallest exchange that counts: some dozens of Mycelium chunks.
+const MIN_BYTES: u64 = 32 * 1024;
+const WEIGHT: f64 = 0.25;
+
+/// Sync bytes exchanged with a peer so far (from the traffic counters).
+#[derive(Debug, Default, Clone, Copy)]
+pub(crate) struct SyncBytes {
+    sent: u64,
+    received: u64,
+}
+
+impl SyncBytes {
+    pub(crate) fn of(traffic: &Traffic, peer: &str) -> Self {
+        traffic
+            .peer(peer)
+            .map(|t| Self {
+                sent: t.sent.sync.bytes,
+                received: t.received.sync.bytes,
+            })
+            .unwrap_or_default()
+    }
+
+    /// Bytes either way since `earlier` (0 if the counters were dropped meanwhile).
+    fn since(&self, earlier: &SyncBytes) -> u64 {
+        self.sent.saturating_sub(earlier.sent) + self.received.saturating_sub(earlier.received)
+    }
+}
+
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct Estimate {
+    pub(crate) bytes_per_sec: f64,
+    pub(crate) samples: u32,
+    /// When the last sample was taken.
+    pub(crate) at: Instant,
+}
+
+#[derive(Debug, Default)]
+pub(crate) struct ThroughputTable {
+    peers: HashMap<String, Estimate>,
+}
+
+impl ThroughputTable {
+    /// A sync exchange with `peer` finished: it started at `started` with the counters
+    /// at `before`, and its answer arrived at `arrived` with the counters at `after`.
+    pub(crate) fn exchange(
+        &mut self,
+        peer: &str,
+        (started, before): (Instant, SyncBytes),
+        (arrived, after): (Instant, SyncBytes),
+        rtt: Option<Duration>,
+    ) {
+        let bytes = after.since(&before);
+        if bytes < MIN_BYTES {
+            return;
+        }
+        let elapsed = arrived.saturating_duration_since(started);
+        let transfer = elapsed
+            .saturating_sub(rtt.unwrap_or_default())
+            .max(elapsed / 2);
+        if transfer.is_zero() {
+            return;
+        }
+        let sample = bytes as f64 / transfer.as_secs_f64();
+        let estimate = self.peers.entry(peer.to_owned()).or_insert(Estimate {
+            bytes_per_sec: sample,
+            samples: 0,
+            at: arrived,
+        });
+        if estimate.samples > 0 {
+            estimate.bytes_per_sec += WEIGHT * (sample - estimate.bytes_per_sec);
+        }
+        estimate.samples += 1;
+        estimate.at = arrived;
+    }
+
+    pub(crate) fn get(&self, peer: &str) -> Option<Estimate> {
+        self.peers.get(peer).copied()
+    }
+
+    pub(crate) fn retain(&mut self, mut keep: impl FnMut(&str) -> bool) {
+        self.peers.retain(|peer, _| keep(peer));
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn bytes(sent: u64, received: u64) -> SyncBytes {
+        SyncBytes { sent, received }
+    }
+
+    #[test]
+    fn small_exchanges_dont_count() {
+        let mut table = ThroughputTable::default();
+        let t0 = Instant::now();
+        table.exchange(
+            "a",
+            (t0, bytes(0, 0)),
+            (t0 + Duration::from_millis(200), bytes(500, 10_000)),
+            None,
+        );
+        assert!(table.get("a").is_none());
+    }
+
+    #[test]
+    fn round_trip_is_taken_off_but_at_most_half() {
+        let mut table = ThroughputTable::default();
+        let t0 = Instant::now();
+        // 1 MB in 2 s, of which 1 s round trip: 1 MB/s.
+        table.exchange(
+            "a",
+            (t0, bytes(100, 0)),
+            (t0 + Duration::from_secs(2), bytes(200, 999_900)),
+            Some(Duration::from_secs(1)),
+        );
+        let a = table.get("a").unwrap();
+        assert!((a.bytes_per_sec - 1_000_000.0).abs() < 1.0);
+
+        // A round trip longer than the exchange takes off only half.
+        table.exchange(
+            "b",
+            (t0, bytes(0, 0)),
+            (t0 + Duration::from_secs(2), bytes(0, 1_000_000)),
+            Some(Duration::from_secs(5)),
+        );
+        assert!((table.get("b").unwrap().bytes_per_sec - 1_000_000.0).abs() < 1.0);
+    }
+
+    #[test]
+    fn smoothed() {
+        let mut table = ThroughputTable::default();
+        let t0 = Instant::now();
+        for received in [500_000u64, 1_000_000] {
+            table.exchange(
+                "a",
+                (t0, bytes(0, 0)),
+                (t0 + Duration::from_secs(1), bytes(0, received)),
+                None,
+            );
+        }
+        // Samples of 500 kB/s, then 1 MB/s: 500 kB/s + (1 MB/s - 500 kB/s) / 4.
+        let a = table.get("a").unwrap();
+        assert_eq!(a.samples, 2);
+        assert!((a.bytes_per_sec - 625_000.0).abs() < 1.0);
+    }
+}