Implement proper ping/pong response tracking

Previously ping() just sent a ping and checked if the connection was
still alive after a small delay. Now it properly:
- Registers a pending ping with sequence number before sending
- Awaits the pong response with the configured timeout
- Returns Pong only when actual pong is received

Added pending_pings DashMap to DistributionManager with register_ping()
and complete_ping() methods.

Author: scrogson

crates/ambitious/src/distribution/manager.rs

@@ -16,7 +16,7 @@ use std::path::Path;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::Duration;
-use tokio::sync::mpsc;
+use tokio::sync::{mpsc, oneshot};
 
 /// The type of a connected node.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -135,6 +135,9 @@ pub struct DistributionManager {
     erlang_nodes: DashMap<Atom, ErlangNodeHandle>,
     /// Node type lookup (for all connected nodes).
     node_types: DashMap<Atom, NodeType>,
+    /// Pending ping requests awaiting pong responses.
+    /// Key is (node_atom, sequence_number), value is the oneshot sender.
+    pending_pings: DashMap<(Atom, u64), oneshot::Sender<()>>,
 }
 
 impl DistributionManager {
@@ -153,6 +156,7 @@ impl DistributionManager {
             #[cfg(feature = "erlang-dist")]
             erlang_nodes: DashMap::new(),
             node_types: DashMap::new(),
+            pending_pings: DashMap::new(),
         }
     }
 
@@ -457,6 +461,21 @@ impl DistributionManager {
         self.nodes.get(&node_atom).map(|n| n.tx.clone())
     }
 
+    /// Register a pending ping and return a receiver to await the pong.
+    pub(crate) fn register_ping(&self, node: Atom, seq: u64) -> oneshot::Receiver<()> {
+        let (tx, rx) = oneshot::channel();
+        self.pending_pings.insert((node, seq), tx);
+        rx
+    }
+
+    /// Complete a pending ping (called when pong is received).
+    pub(crate) fn complete_ping(&self, node: Atom, seq: u64) {
+        if let Some((_, tx)) = self.pending_pings.remove(&(node, seq)) {
+            // Send () to signal pong received. Ignore error if receiver dropped.
+            let _ = tx.send(());
+        }
+    }
+
     /// Send a message to an Erlang/BEAM node.
     ///
     /// The payload should be ETF-encoded bytes. This method decodes them to
@@ -915,11 +934,13 @@ async fn handle_incoming_message(from_node: Atom, msg: DistMessage) {
             }
         }
         DistMessage::Pong { seq } => {
-            // Update last seen timestamp - this is the critical heartbeat response
-            if let Some(manager) = DIST_MANAGER.get()
-                && let Some(node) = manager.nodes.get(&from_node)
-            {
-                node.touch();
+            // Update last seen timestamp and complete any pending ping request
+            if let Some(manager) = DIST_MANAGER.get() {
+                if let Some(node) = manager.nodes.get(&from_node) {
+                    node.touch();
+                }
+                // Notify any waiting ping caller that pong was received
+                manager.complete_ping(from_node, seq);
             }
             tracing::trace!(seq, from_node = %from_node, "Received pong");
         }

crates/ambitious/src/node.rs

@@ -221,7 +221,7 @@ pub async fn ping(node: Atom, timeout_ms: u64) -> PingResult {
 }
 
 /// Implementation of ping with Duration timeout.
-async fn ping_impl(node: Atom, _timeout_duration: Duration) -> PingResult {
+async fn ping_impl(node: Atom, timeout_duration: Duration) -> PingResult {
     let manager = match distribution::manager() {
         Some(m) => m,
         None => return PingResult::Pang,
@@ -234,24 +234,20 @@ async fn ping_impl(node: Atom, _timeout_duration: Duration) -> PingResult {
 
     let seq = PING_SEQ.fetch_add(1, Ordering::Relaxed);
 
+    // Register the ping before sending so we can await the pong
+    let pong_rx = manager.register_ping(node, seq);
+
     // Send ping
     let msg = DistMessage::Ping { seq };
     if tx.try_send(msg).is_err() {
         return PingResult::Pang;
     }
 
-    // For now, we just check if the send succeeded.
-    // A proper implementation would wait for the Pong response.
-    // TODO: Implement proper ping/pong with response tracking
-
-    // Small delay to let the ping go through
-    tokio::time::sleep(Duration::from_millis(10)).await;
-
-    // Check if still connected (basic liveness check)
-    if manager.get_node_tx(node).is_some() {
-        PingResult::Pong
-    } else {
-        PingResult::Pang
+    // Wait for pong response with timeout
+    match tokio::time::timeout(timeout_duration, pong_rx).await {
+        Ok(Ok(())) => PingResult::Pong,
+        Ok(Err(_)) => PingResult::Pang, // Sender dropped (node disconnected)
+        Err(_) => PingResult::Pang,     // Timeout
     }
 }