diff --git a/relay/src/pool.rs b/relay/src/pool.rs index 8ca074f..b7a8774 100644 --- a/relay/src/pool.rs +++ b/relay/src/pool.rs @@ -109,6 +109,32 @@ struct RelayEntry { reconnect_attempts: u32, last_connection_error: Option, last_publish_error: Option, + /// Earliest time a reconnect attempt is allowed (backoff gate). + next_reconnect_at: Option, + /// When the current connection was established (for stability reset). + connected_since: Option, + /// When the last ping was sent (health monitoring). + last_ping_sent: Option, + /// Whether a ping is awaiting its pong. + ping_pending: bool, +} + +impl RelayEntry { + fn new(url: &str) -> Self { + RelayEntry { + client: NostrWsClient::new(url), + url: url.to_string(), + status: RelayStatus::Disconnected, + stats: RelayStats::default(), + reconnect_attempts: 0, + last_connection_error: None, + last_publish_error: None, + next_reconnect_at: None, + connected_since: None, + last_ping_sent: None, + ping_pending: false, + } + } } // ── Subscription ──────────────────────────────────────────────────────────── @@ -257,15 +283,7 @@ impl RelayPool { if relays.iter().any(|r| r.url == url) { return Ok(()); // Already added } - relays.push(RelayEntry { - client: NostrWsClient::new(url), - url: url.to_string(), - status: RelayStatus::Disconnected, - stats: RelayStats::default(), - reconnect_attempts: 0, - last_connection_error: None, - last_publish_error: None, - }); + relays.push(RelayEntry::new(url)); Ok(()) } @@ -283,12 +301,34 @@ impl RelayPool { } /// Connect to all relays in the pool. + /// + /// Each connection attempt is bounded by a timeout so a single + /// unresponsive relay cannot block the whole pool. pub async fn connect_all(&self) { + self.connect_all_with_timeout(10_000).await; + } + + /// Connect to all relays, bounding each attempt by `timeout_ms`. + pub async fn connect_all_with_timeout(&self, timeout_ms: u64) { let relays = self.relays.read().await; for entry in relays.iter() { if entry.status == RelayStatus::Disconnected { - if let Err(e) = entry.client.connect().await { - warn!("Failed to connect to {}: {:?}", entry.url, e); + let url = entry.url.clone(); + match tokio::time::timeout( + Duration::from_millis(timeout_ms), + entry.client.connect(), + ) + .await + { + Ok(Ok(())) => {} + Ok(Err(e)) => { + warn!("Failed to connect to {}: {:?}", url, e); + eprintln!("[relay] connect failed: {}", url); + } + Err(_) => { + warn!("Connect to {} timed out", url); + eprintln!("[relay] connect timed out: {}", url); + } } } } @@ -330,6 +370,23 @@ impl RelayPool { relays.iter().map(|r| (r.url.clone(), r.status.clone())).collect() } + /// Get the URLs of relays whose WebSocket is currently connected. + /// + /// Useful for querying only reachable relays — a subscription in + /// `EoseResultMode::FullSet` waits for EOSE from every relay it was + /// given, so including a relay that failed to connect would stall the + /// query until its timeout. + pub async fn connected_relay_urls(&self) -> Vec { + let relays = self.relays.read().await; + let mut urls = Vec::new(); + for entry in relays.iter() { + if entry.client.state().await == WsState::Connected { + urls.push(entry.url.clone()); + } + } + urls + } + /// Set NIP-42 authentication credentials. pub async fn set_auth(&self, private_key: &SecretKey, enable: bool) { let mut key = self.auth_key.lock().await; @@ -378,8 +435,11 @@ impl RelayPool { let relays = self.relays.read().await; for url in relay_urls { if let Some(entry) = relays.iter().find(|r| r.url == *url) { - // Connect if not connected - if entry.client.state().await == WsState::Disconnected { + // Connect if not connected (covers Disconnected AND Error — + // a relay that died mid-session must be reconnectable here, + // not just one that was never connected) + let state = entry.client.state().await; + if state != WsState::Connected { if let Err(e) = entry.client.connect().await { warn!("Failed to connect to {} for subscription: {:?}", url, e); continue; @@ -421,27 +481,173 @@ impl RelayPool { // ── Event Loop ────────────────────────────────────────────────────── /// Run the event loop to process incoming messages from all relays. - /// This should be called in a background task. + /// + /// This should be called in a background task. Each iteration: + /// + /// 1. Attempts reconnection of any non-connected relay whose backoff + /// window has elapsed (auto-reconnect with exponential backoff). + /// 2. Sends health-check pings on the configured interval and detects + /// pong timeouts / dead transports. + /// 3. Syncs the pool-level `RelayStatus` from the live ws state so + /// `list_relays()` / `get_relay_status()` report reality. + /// 4. Receives and dispatches incoming messages from connected relays. pub async fn run(&self, timeout_ms: u64) { loop { - let relays = self.relays.read().await; let mut has_activity = false; - for entry in relays.iter() { - if entry.client.state().await != WsState::Connected { - continue; - } + // Phase 1: reconnect + health check + status sync (needs write access) + { + let mut relays = self.relays.write().await; + for entry in relays.iter_mut() { + let ws_state = entry.client.state().await; - match entry.client.receive_text(timeout_ms).await { - Ok(Some(text)) => { - has_activity = true; - self.process_message(&entry.url, &text).await; + // Sync pool status from live ws state + let new_status = match &ws_state { + WsState::Connected => RelayStatus::Connected, + WsState::Connecting => RelayStatus::Connecting, + WsState::Disconnected => RelayStatus::Disconnected, + WsState::Closing => RelayStatus::Disconnected, + WsState::Error(e) => RelayStatus::Error(e.clone()), + }; + if entry.status != new_status { + debug!("relay {} status: {:?} -> {:?}", entry.url, entry.status, new_status); + entry.status = new_status; } - Ok(None) => { - // Timeout or connection closed + + if ws_state == WsState::Connected { + // Track connection stability: reset backoff after the + // connection has been up long enough. + if let Some(since) = entry.connected_since { + if since.elapsed().as_secs() + >= self.reconnect_config.reconnect_reset_stability_secs + && entry.reconnect_attempts > 0 + { + entry.reconnect_attempts = 0; + entry.next_reconnect_at = None; + } + } else { + entry.connected_since = Some(Instant::now()); + } + + // Health monitoring: send ping on interval + if self.reconnect_config.ping_interval_seconds > 0 { + let should_ping = match entry.last_ping_sent { + None => true, + Some(t) => { + t.elapsed().as_secs() + >= self.reconnect_config.ping_interval_seconds + } + }; + if should_ping && !entry.ping_pending { + match entry.client.ping().await { + Ok(()) => { + entry.last_ping_sent = Some(Instant::now()); + entry.ping_pending = true; + } + Err(_) => { + // Ping send failed: transport is dead. + // ws.rs already set the client state to + // Error; force-close so the reconnect + // path picks it up next iteration. + entry.client.close().await.ok(); + entry.connected_since = None; + entry.ping_pending = false; + entry.stats.connection_failures += 1; + warn!("relay {} ping failed; marked for reconnect", entry.url); + } + } + } + + // Pong timeout: connection is dead + if entry.ping_pending { + if let Some(t) = entry.last_ping_sent { + if t.elapsed().as_secs() + > self.reconnect_config.pong_timeout_seconds + { + warn!("relay {} pong timeout; reconnecting", entry.url); + entry.client.close().await.ok(); + entry.connected_since = None; + entry.ping_pending = false; + entry.stats.connection_failures += 1; + } + } + } + } + } else { + // Not connected: reset per-connection bookkeeping + entry.connected_since = None; + entry.ping_pending = false; + + // Auto-reconnect with backoff + if self.reconnect_config.enable_auto_reconnect + && entry.reconnect_attempts + < self.reconnect_config.max_reconnect_attempts + { + let due = match entry.next_reconnect_at { + None => true, + Some(t) => Instant::now() >= t, + }; + if due { + entry.reconnect_attempts += 1; + entry.stats.connection_attempts += 1; + debug!( + "relay {} reconnect attempt {}/{}", + entry.url, + entry.reconnect_attempts, + self.reconnect_config.max_reconnect_attempts + ); + match entry.client.connect().await { + Ok(()) => { + entry.connected_since = Some(Instant::now()); + entry.last_ping_sent = None; + entry.ping_pending = false; + entry.last_connection_error = None; + info!("relay {} reconnected", entry.url); + // Keep reconnect_attempts until the + // connection proves stable (reset above). + let delay_ms = self.reconnect_delay_ms(entry.reconnect_attempts); + entry.next_reconnect_at = + Some(Instant::now() + Duration::from_millis(delay_ms)); + } + Err(e) => { + entry.last_connection_error = + Some(format!("{:?}", e)); + let delay_ms = self.reconnect_delay_ms(entry.reconnect_attempts); + entry.next_reconnect_at = + Some(Instant::now() + Duration::from_millis(delay_ms)); + } + } + } + } } - Err(_) => { - // Error receiving + } + } + + // Phase 2: receive and dispatch messages (needs write access for + // ping_pending bookkeeping on inbound traffic) + { + let mut relays = self.relays.write().await; + for entry in relays.iter_mut() { + if entry.client.state().await != WsState::Connected { + continue; + } + + match entry.client.receive_text(timeout_ms).await { + Ok(Some(text)) => { + has_activity = true; + // Any inbound traffic proves liveness; the + // ping_pending flag is reset by the Phase-1 + // writer pass on the next iteration. + self.process_message(&entry.url, &text).await; + } + Ok(None) => { + // Timeout or connection closed — ws.rs updates + // state on clean close; nothing to do here. + } + Err(_) => { + // Receive error — ws.rs set state to Error; + // the reconnect path picks it up next iteration. + } } } } @@ -452,6 +658,18 @@ impl RelayPool { } } + /// Compute the reconnect backoff delay (ms) for the given attempt number. + fn reconnect_delay_ms(&self, attempt: u32) -> u64 { + let mut delay = self.reconnect_config.initial_reconnect_delay_ms; + for _ in 1..attempt { + delay = (delay as f64 * self.reconnect_config.reconnect_backoff_multiplier) as u64; + if delay >= self.reconnect_config.max_reconnect_delay_ms { + return self.reconnect_config.max_reconnect_delay_ms; + } + } + delay.min(self.reconnect_config.max_reconnect_delay_ms) + } + /// Process a received relay message. async fn process_message(&self, relay_url: &str, text: &str) { // Parse the JSON array @@ -737,8 +955,10 @@ impl RelayPool { let relays = self.relays.read().await; for url in &relay_urls_copy { if let Some(entry) = relays.iter().find(|r| r.url == *url) { - // Connect if not connected - if entry.client.state().await == WsState::Disconnected { + // Connect if not connected (covers Disconnected AND Error — + // a relay whose transport died must be reconnectable here) + let state = entry.client.state().await; + if state != WsState::Connected { if let Err(e) = entry.client.connect().await { warn!("Failed to connect to {} for publish: {:?}", url, e); continue; @@ -840,4 +1060,90 @@ mod tests { assert!(sub.is_ok()); } + + #[test] + fn test_reconnect_backoff_calculation() { + let pool = RelayPool::new(None); + + // Default config: initial 1000ms, multiplier 2.0, max 60_000ms + assert_eq!(pool.reconnect_delay_ms(1), 1000); + assert_eq!(pool.reconnect_delay_ms(2), 2000); + assert_eq!(pool.reconnect_delay_ms(3), 4000); + assert_eq!(pool.reconnect_delay_ms(4), 8000); + assert_eq!(pool.reconnect_delay_ms(5), 16000); + assert_eq!(pool.reconnect_delay_ms(6), 32000); + // Clamped at max + assert_eq!(pool.reconnect_delay_ms(7), 60000); + assert_eq!(pool.reconnect_delay_ms(20), 60000); + } + + #[test] + fn test_reconnect_backoff_custom_config() { + let pool = RelayPool::new(Some(ReconnectConfig { + enable_auto_reconnect: true, + max_reconnect_attempts: 5, + initial_reconnect_delay_ms: 500, + max_reconnect_delay_ms: 5_000, + reconnect_backoff_multiplier: 3.0, + reconnect_reset_stability_secs: 30, + ping_interval_seconds: 59, + pong_timeout_seconds: 10, + })); + + assert_eq!(pool.reconnect_delay_ms(1), 500); + assert_eq!(pool.reconnect_delay_ms(2), 1500); + assert_eq!(pool.reconnect_delay_ms(3), 4500); + // 13500 would exceed max 5000 → clamped + assert_eq!(pool.reconnect_delay_ms(4), 5000); + } + + #[tokio::test] + async fn test_status_initially_disconnected() { + // The pool-level RelayStatus must start Disconnected and reflect + // the live ws state after run()'s sync phase — not a stale + // snapshot from add_relay time. + let pool = RelayPool::new(None); + pool.add_relay("wss://relay.damus.io").await.unwrap(); + + let status = pool.get_relay_status("wss://relay.damus.io").await; + assert_eq!(status, Some(RelayStatus::Disconnected)); + + let relays = pool.relays.read().await; + let entry = relays.iter().next().unwrap(); + let ws_state = entry.client.state().await; + drop(relays); + + assert_eq!(ws_state, WsState::Disconnected); + } + + #[tokio::test] + async fn test_connect_failure_sets_error_state() { + // Connecting to an unreachable relay must leave the ws client in + // Error state (not Connected), so the reconnect path can retry. + let pool = RelayPool::new(None); + pool.add_relay("wss://invalid.relay.example.invalid").await.unwrap(); + + let result = pool.connect_relay("wss://invalid.relay.example.invalid").await; + assert!(result.is_err()); + + let relays = pool.relays.read().await; + let entry = relays.iter().next().unwrap(); + let ws_state = entry.client.state().await; + drop(relays); + + assert!(matches!(ws_state, WsState::Error(_))); + } + + #[tokio::test] + async fn test_relay_entry_defaults() { + // New entries must start with clean reconnect bookkeeping so the + // first reconnect attempt is immediate (no backoff gate). + let entry = RelayEntry::new("wss://relay.example.com"); + assert_eq!(entry.reconnect_attempts, 0); + assert!(entry.next_reconnect_at.is_none()); + assert!(entry.connected_since.is_none()); + assert!(!entry.ping_pending); + assert!(entry.last_ping_sent.is_none()); + assert_eq!(entry.status, RelayStatus::Disconnected); + } } diff --git a/relay/src/ws.rs b/relay/src/ws.rs index 174f2a9..a49c9c8 100644 --- a/relay/src/ws.rs +++ b/relay/src/ws.rs @@ -100,16 +100,27 @@ impl NostrWsClient { } /// Send a text message to the relay. + /// + /// On send failure the connection state is set to `Error` so callers + /// (and the pool's reconnect logic) can detect a dead transport. + /// Without this, a relay that dropped the TCP connection would keep + /// reporting `Connected` and every publish would silently fail. pub async fn send_text(&self, message: &str) -> NostrResult<()> { let mut stream_guard = self.stream.lock().await; - let stream = stream_guard - .as_mut() - .ok_or(NostrError::NetworkFailed)?; + let stream = match stream_guard.as_mut() { + Some(s) => s, + None => { + let mut state = self.state.lock().await; + *state = WsState::Error("send on disconnected stream".to_string()); + return Err(NostrError::NetworkFailed); + } + }; - stream - .send(Message::Text(message.to_string())) - .await - .map_err(|_| NostrError::NetworkFailed)?; + if let Err(e) = stream.send(Message::Text(message.to_string())).await { + let mut state = self.state.lock().await; + *state = WsState::Error(format!("send error: {}", e)); + return Err(NostrError::NetworkFailed); + } Ok(()) } @@ -225,16 +236,25 @@ impl NostrWsClient { } /// Send a ping frame to keep the connection alive. + /// + /// On failure the connection state is set to `Error` so the pool's + /// health check can detect a dead transport and trigger reconnect. pub async fn ping(&self) -> NostrResult<()> { let mut stream_guard = self.stream.lock().await; - let stream = stream_guard - .as_mut() - .ok_or(NostrError::NetworkFailed)?; + let stream = match stream_guard.as_mut() { + Some(s) => s, + None => { + let mut state = self.state.lock().await; + *state = WsState::Error("ping on disconnected stream".to_string()); + return Err(NostrError::NetworkFailed); + } + }; - stream - .send(Message::Ping(vec![])) - .await - .map_err(|_| NostrError::NetworkFailed)?; + if let Err(e) = stream.send(Message::Ping(vec![])).await { + let mut state = self.state.lock().await; + *state = WsState::Error(format!("ping error: {}", e)); + return Err(NostrError::NetworkFailed); + } Ok(()) }