Fix native ping connection lifecycle

client/src/connection.rs

@@ -159,7 +159,7 @@ impl MTPConnection {
     }
 }
 
-pub(crate) fn connection_from_parts(
+pub(crate) async fn connection_from_parts(
     config: ClientConfig,
     sender: mtp_transport::Sender,
     receiver: mtp_transport::Receiver,
@@ -168,7 +168,7 @@ pub(crate) fn connection_from_parts(
     #[cfg(feature = "crypto")] client_id: u64,
 ) -> MTPConnection {
     let remote_addr = sender.handle().remote_addr();
-    let ping = start_ping_session(&config, sender.clone(), &receiver);
+    let ping = start_ping_session(&config, sender.clone(), &receiver).await;
 
     #[cfg(feature = "pipes")]
     {

client/src/lib.rs

@@ -142,9 +142,10 @@ impl MTPClient {
             negotiated,
             error::AuthState::Unauthenticated,
             client_id,
-        ));
+        )
+        .await);
         #[cfg(not(feature = "crypto"))]
-        Ok(connection_from_parts(config, sender, receiver, negotiated))
+        Ok(connection_from_parts(config, sender, receiver, negotiated).await)
     }
 }
 
@@ -283,7 +284,8 @@ impl MTPClient {
             crypto::negotiated_version(&response)?,
             error::AuthState::Authenticated,
             client_id,
-        ))
+        )
+        .await)
     }
 
     pub async fn auth_register(
@@ -440,7 +442,8 @@ impl MTPClient {
             crypto::negotiated_version(&response)?,
             error::AuthState::Authenticated,
             assigned_id,
-        ))
+        )
+        .await)
     }
 }
 

client/src/ping.rs

@@ -1,5 +1,4 @@
 use rand::RngExt;
-use std::collections::HashMap;
 use std::sync::Arc;
 use tokio::sync::{Mutex, mpsc};
 use tokio::time::{Duration, Instant};
@@ -12,6 +11,38 @@ pub(crate) struct PingSession {
     pub(crate) task: tokio::task::JoinHandle<()>,
 }
 
+#[derive(Default)]
+struct PingTracker {
+    pending: Option<(u32, Instant)>,
+    missed_pings: usize,
+}
+
+impl PingTracker {
+    fn begin_round(&mut self) -> usize {
+        if self.pending.take().is_some() {
+            self.missed_pings += 1;
+        }
+        self.missed_pings
+    }
+
+    fn sent(&mut self, id: u32) {
+        self.pending = Some((id, Instant::now()));
+    }
+
+    fn received(&mut self, id: u32) -> Option<Duration> {
+        if !self
+            .pending
+            .as_ref()
+            .is_some_and(|(pending, _)| *pending == id)
+        {
+            return None;
+        }
+        let (_, sent_at) = self.pending.take()?;
+        self.missed_pings = 0;
+        Some(sent_at.elapsed())
+    }
+}
+
 impl PingSession {
     pub(crate) fn get_ping(&self) -> Option<Duration> {
         self.last_ping.try_lock().ok().and_then(|ping| *ping)
@@ -24,7 +55,7 @@ impl Drop for PingSession {
     }
 }
 
-pub(crate) fn start_ping_session(
+pub(crate) async fn start_ping_session(
     config: &crate::config::ClientConfig,
     sender: Sender,
     receiver: &Receiver,
@@ -34,7 +65,7 @@ pub(crate) fn start_ping_session(
     }
 
     let (pong_tx, mut pong_rx) = mpsc::unbounded_channel();
-    receiver.observe_pongs(pong_tx);
+    receiver.observe_pongs(pong_tx).await;
     let last_ping = Arc::new(Mutex::new(None));
     let ping_state = last_ping.clone();
     let interval = config.ping_interval;
@@ -46,7 +77,7 @@ pub(crate) fn start_ping_session(
     let task = tokio::spawn(async move {
         let mut ticker = tokio::time::interval(interval);
         ticker.tick().await;
-        let mut pending = HashMap::new();
+        let mut tracker = PingTracker::default();
 
         loop {
             tokio::select! {
@@ -56,7 +87,8 @@ pub(crate) fn start_ping_session(
                     }
                 }
                 _ = ticker.tick() => {
-                    if max_missed_pings > 0 && !pending.is_empty() && pending.len() >= max_missed_pings {
+                    let missed_pings = tracker.begin_round();
+                    if max_missed_pings > 0 && missed_pings >= max_missed_pings {
                         sender.close().await;
                         break;
                     }
@@ -83,13 +115,13 @@ pub(crate) fn start_ping_session(
                         sender.close().await;
                         break;
                     }
-                    pending.insert(id, Instant::now());
+                    tracker.sent(id);
                 }
                 pong = pong_rx.recv() => match pong {
                     Some(pong) => {
-                        if let Some(sent_at) = pending.remove(&pong.get_id()) {
+                        if let Some(ping) = tracker.received(pong.get_id()) {
                             let mut last_ping = ping_state.lock().await;
-                            *last_ping = Some(sent_at.elapsed());
+                            *last_ping = Some(ping);
                         }
                     }
                     None => break,
@@ -100,3 +132,32 @@ pub(crate) fn start_ping_session(
 
     Some(PingSession { last_ping, task })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::PingTracker;
+
+    #[test]
+    fn successful_pong_resets_consecutive_misses() {
+        let mut tracker = PingTracker::default();
+        tracker.sent(1);
+        assert_eq!(tracker.begin_round(), 1);
+
+        tracker.sent(2);
+        assert!(tracker.received(2).is_some());
+
+        tracker.sent(3);
+        assert_eq!(tracker.begin_round(), 1);
+    }
+
+    #[test]
+    fn stale_pong_does_not_acknowledge_current_round() {
+        let mut tracker = PingTracker::default();
+        tracker.sent(1);
+        assert_eq!(tracker.begin_round(), 1);
+        tracker.sent(2);
+
+        assert!(tracker.received(1).is_none());
+        assert_eq!(tracker.begin_round(), 2);
+    }
+}