add heartbeat-ttl presence to mqp2p peer discovery

crates/mqp2p/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mqp2p"
-version = "0.1.0"
+version = "0.2.0"
 edition = "2024"
 rust-version = "1.88"
 license = "Apache-2.0"
@@ -28,3 +28,4 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 clap = { version = "4", features = ["derive"] }
 tempfile = "3"
 rand = "0.10.0"
+mqtt5 = "0.38"

crates/mqp2p/src/peer.rs

@@ -1,7 +1,8 @@
 use crate::error::{Error, Result};
 use crate::quic::{self, CertIdentity, QuicEndpoint};
 use crate::signaling::{
-    Candidate, CandidateKind, ConnectionOffer, ConnectionSync, PeerInfo, SignalingClient,
+    Candidate, CandidateKind, ConnectionOffer, ConnectionSync, MIN_LEASE_SECS, PEER_LEASE_SECS,
+    PeerInfo, SignalingClient, heartbeat_interval_secs,
 };
 use crate::stun;
 use crate::transfer::{self, FileOffer, TransferProgress, TransferResult};
@@ -18,6 +19,7 @@ pub struct PeerConfig {
     pub bind_addr: SocketAddr,
     pub stun_server: Option<String>,
     pub credentials: Option<(String, String)>,
+    pub lease_secs: u64,
 }
 
 impl PeerConfig {
@@ -28,9 +30,15 @@ impl PeerConfig {
             bind_addr: "0.0.0.0:0".parse().unwrap_or_else(|_| unreachable!()),
             stun_server: Some(stun::DEFAULT_STUN_SERVER.into()),
             credentials: None,
+            lease_secs: PEER_LEASE_SECS,
         }
     }
 
+    pub fn with_lease_secs(mut self, lease_secs: u64) -> Self {
+        self.lease_secs = lease_secs;
+        self
+    }
+
     pub fn with_bind_addr(mut self, addr: SocketAddr) -> Self {
         self.bind_addr = addr;
         self
@@ -63,10 +71,13 @@ pub struct Peer {
     identity: CertIdentity,
     signaling: SignalingClient,
     quic: QuicEndpoint,
+    record_id: String,
     peer_id: Option<String>,
     std_socket: std::net::UdpSocket,
     host_addr: SocketAddr,
     srflx_addr: Option<SocketAddr>,
+    lease_secs: u64,
+    heartbeat: Option<tokio::task::JoinHandle<()>>,
 }
 
 impl Peer {
@@ -131,26 +142,67 @@ impl Peer {
             identity,
             signaling,
             quic,
+            record_id: uuid::Uuid::new_v4().to_string(),
             peer_id: None,
             std_socket,
             host_addr,
             srflx_addr,
+            lease_secs: config.lease_secs.max(MIN_LEASE_SECS),
+            heartbeat: None,
         })
     }
 
     pub async fn register(&mut self) -> Result<PeerId> {
         let quic_port = self.quic.local_addr()?.port();
-
-        let peer_id = self
-            .signaling
-            .register_peer(&self.name, quic_port, &self.identity.fingerprint)
+        let peer_id = self.record_id.clone();
+        let public_addr = self.srflx_addr.map(|a| a.to_string());
+
+        self.signaling
+            .register_peer(
+                &peer_id,
+                &self.name,
+                quic_port,
+                &self.identity.fingerprint,
+                public_addr.as_deref(),
+                self.lease_secs,
+            )
             .await?;
 
-        if let Some(addr) = self.srflx_addr {
-            self.signaling.update_peer_addr(&peer_id, addr).await?;
+        self.peer_id = Some(peer_id.clone());
+
+        if let Some(handle) = self.heartbeat.take() {
+            handle.abort();
         }
 
-        self.peer_id = Some(peer_id.clone());
+        let signaling = self.signaling.clone();
+        let renew_id = peer_id.clone();
+        let renew_name = self.name.clone();
+        let renew_fingerprint = self.identity.fingerprint.clone();
+        let renew_public_addr = public_addr;
+        let lease_secs = self.lease_secs;
+        self.heartbeat = Some(tokio::spawn(async move {
+            let mut ticker = tokio::time::interval(std::time::Duration::from_secs(
+                heartbeat_interval_secs(lease_secs),
+            ));
+            ticker.tick().await;
+            loop {
+                ticker.tick().await;
+                if let Err(e) = signaling
+                    .ensure_registered(
+                        &renew_id,
+                        &renew_name,
+                        quic_port,
+                        &renew_fingerprint,
+                        renew_public_addr.as_deref(),
+                        lease_secs,
+                    )
+                    .await
+                {
+                    warn!(peer_id = renew_id, error = %e, "peer presence renewal failed");
+                }
+            }
+        }));
+
         info!(peer_id, name = self.name, "peer registered");
         Ok(peer_id)
     }
@@ -371,6 +423,7 @@ impl Peer {
             quic_port: remote_addr.port(),
             cert_fingerprint: offer.cert_fingerprint,
             public_addr: Some(remote_addr.to_string()),
+            expires_at: None,
         };
 
         info!(
@@ -385,7 +438,10 @@ impl Peer {
         })
     }
 
-    pub async fn shutdown(self) -> Result<()> {
+    pub async fn shutdown(mut self) -> Result<()> {
+        if let Some(handle) = self.heartbeat.take() {
+            handle.abort();
+        }
         if let Some(peer_id) = &self.peer_id {
             let _ = self.signaling.deregister_peer(peer_id).await;
         }
@@ -395,6 +451,14 @@ impl Peer {
     }
 }
 
+impl Drop for Peer {
+    fn drop(&mut self) {
+        if let Some(handle) = self.heartbeat.take() {
+            handle.abort();
+        }
+    }
+}
+
 pub struct P2pConnection {
     connection: Connection,
     remote_peer: PeerInfo,

crates/mqp2p/src/signaling.rs

@@ -4,11 +4,30 @@ use mqtt5::types::{PublishOptions, PublishProperties};
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 use std::net::SocketAddr;
-use std::time::Duration;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use tracing::{debug, error, trace};
 
 const DEFAULT_TIMEOUT_MS: u64 = 5000;
 
+pub const PEER_LEASE_SECS: u64 = 15;
+
+pub const MIN_LEASE_SECS: u64 = 2;
+
+pub fn heartbeat_interval_secs(lease_secs: u64) -> u64 {
+    (lease_secs / 3).max(1)
+}
+
+fn unix_now_secs() -> u64 {
+    SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .map(|d| d.as_secs())
+        .unwrap_or(0)
+}
+
+fn peer_is_live(expires_at: Option<u64>, now: u64) -> bool {
+    expires_at.map(|e| e > now).unwrap_or(true)
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct PeerInfo {
     pub id: String,
@@ -18,6 +37,8 @@ pub struct PeerInfo {
     pub cert_fingerprint: String,
     #[serde(default)]
     pub public_addr: Option<String>,
+    #[serde(rename = "_expires_at", default)]
+    pub expires_at: Option<u64>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -49,6 +70,7 @@ pub enum CandidateKind {
     Srflx,
 }
 
+#[derive(Clone)]
 pub struct SignalingClient {
     client: MqttClient,
 }
@@ -73,31 +95,73 @@ impl SignalingClient {
 
     pub async fn register_peer(
         &self,
+        id: &str,
         name: &str,
         quic_port: u16,
         cert_fingerprint: &str,
-    ) -> Result<String> {
-        let payload = serde_json::json!({
+        public_addr: Option<&str>,
+        lease_secs: u64,
+    ) -> Result<()> {
+        let mut payload = serde_json::json!({
+            "id": id,
             "name": name,
             "status": "online",
             "quic_port": quic_port,
             "cert_fingerprint": cert_fingerprint,
+            "_expires_at": unix_now_secs() + lease_secs,
         });
+        if let Some(addr) = public_addr {
+            payload["public_addr"] = Value::String(addr.to_string());
+        }
 
         let response = self
             .publish_and_wait("$DB/peers/create", &serde_json::to_vec(&payload)?)
             .await?;
 
         check_response(&response)?;
+        debug!(peer_id = id, name, "peer registered");
+        Ok(())
+    }
 
-        let id = response
-            .pointer("/data/id")
-            .and_then(|v| v.as_str())
-            .ok_or_else(|| Error::Signaling("no peer ID in create response".into()))?
-            .to_string();
+    pub async fn renew_lease(&self, peer_id: &str, lease_secs: u64) -> Result<()> {
+        let topic = format!("$DB/peers/{peer_id}/update");
+        let payload = serde_json::json!({
+            "_expires_at": unix_now_secs() + lease_secs,
+        });
 
-        debug!(peer_id = id, name, "peer registered");
-        Ok(id)
+        let response = self
+            .publish_and_wait(&topic, &serde_json::to_vec(&payload)?)
+            .await?;
+
+        check_response(&response)?;
+        trace!(peer_id, "peer lease renewed");
+        Ok(())
+    }
+
+    pub async fn ensure_registered(
+        &self,
+        id: &str,
+        name: &str,
+        quic_port: u16,
+        cert_fingerprint: &str,
+        public_addr: Option<&str>,
+        lease_secs: u64,
+    ) -> Result<()> {
+        match self.renew_lease(id, lease_secs).await {
+            Ok(()) => Ok(()),
+            Err(e) => {
+                debug!(peer_id = id, error = %e, "lease renewal failed, re-creating record");
+                self.register_peer(
+                    id,
+                    name,
+                    quic_port,
+                    cert_fingerprint,
+                    public_addr,
+                    lease_secs,
+                )
+                .await
+            }
+        }
     }
 
     pub async fn update_peer_addr(&self, peer_id: &str, public_addr: SocketAddr) -> Result<()> {
@@ -131,9 +195,11 @@ impl SignalingClient {
             other => vec![other],
         };
 
+        let now = unix_now_secs();
         let peers: Vec<PeerInfo> = items
             .into_iter()
-            .filter_map(|v| serde_json::from_value(v).ok())
+            .filter_map(|v| serde_json::from_value::<PeerInfo>(v).ok())
+            .filter(|p| peer_is_live(p.expires_at, now))
             .collect();
 
         debug!(count = peers.len(), "peers discovered");
@@ -306,3 +372,66 @@ fn check_response(response: &Value) -> Result<()> {
         Err(Error::Signaling(format!("DB error: {msg}")))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn leaseless_peer_is_live() {
+        assert!(peer_is_live(None, 1000));
+    }
+
+    #[test]
+    fn future_lease_is_live() {
+        assert!(peer_is_live(Some(1001), 1000));
+    }
+
+    #[test]
+    fn lease_at_current_second_is_stale() {
+        assert!(!peer_is_live(Some(1000), 1000));
+    }
+
+    #[test]
+    fn past_lease_is_stale() {
+        assert!(!peer_is_live(Some(999), 1000));
+    }
+
+    #[test]
+    fn heartbeat_interval_is_a_third_of_lease_min_one() {
+        assert_eq!(heartbeat_interval_secs(15), 5);
+        assert_eq!(heartbeat_interval_secs(3), 1);
+        assert_eq!(heartbeat_interval_secs(2), 1);
+        assert_eq!(heartbeat_interval_secs(1), 1);
+    }
+
+    #[test]
+    fn expires_at_round_trips_as_mqdb_field() {
+        let info = PeerInfo {
+            id: "p1".into(),
+            name: "edge".into(),
+            status: "online".into(),
+            quic_port: 4433,
+            cert_fingerprint: "ab".into(),
+            public_addr: None,
+            expires_at: Some(1234),
+        };
+        let json = serde_json::to_value(&info).expect("serialize");
+        assert_eq!(json.get("_expires_at").and_then(Value::as_u64), Some(1234));
+        let back: PeerInfo = serde_json::from_value(json).expect("deserialize");
+        assert_eq!(back.expires_at, Some(1234));
+    }
+
+    #[test]
+    fn missing_expires_at_deserializes_to_none() {
+        let json = serde_json::json!({
+            "id": "p1",
+            "name": "edge",
+            "status": "online",
+            "quic_port": 4433,
+            "cert_fingerprint": "ab",
+        });
+        let info: PeerInfo = serde_json::from_value(json).expect("deserialize");
+        assert_eq!(info.expires_at, None);
+    }
+}