From f54692ae01f3d7719e55177ca92fd3dfc15eefe1 Mon Sep 17 00:00:00 2001 From: Yuki Kishimoto Date: Wed, 7 Oct 2026 06:48:08 +0000 Subject: [PATCH 1/2] sdk: rework relay connection task handling Keep one task per relay and retain connection requests while an old WebSocket closes. Share connection outcomes between `Relay::try_connect` callers and stop the task on ban, shutdown or the last Relay drop. Replaces https://github.com/nostrdevkit/nostr/pull/1479 Pull-Request: https://github.com/nostrdevkit/nostr/pull/1491 Signed-off-by: Yuki Kishimoto --- Cargo.lock | 2 +- Cargo.toml | 2 +- nostr-sdk/CHANGELOG.md | 7 + nostr-sdk/src/client/api/send_event.rs | 87 ++ nostr-sdk/src/client/mod.rs | 36 +- nostr-sdk/src/lib.rs | 1 + nostr-sdk/src/mutex.rs | 21 + nostr-sdk/src/relay/api/count_events.rs | 2 +- nostr-sdk/src/relay/api/send_event.rs | 2 +- nostr-sdk/src/relay/api/sync.rs | 2 +- nostr-sdk/src/relay/api/try_connect.rs | 172 ++- nostr-sdk/src/relay/inner.rs | 1366 ++++++++++++++++++----- nostr-sdk/src/relay/mod.rs | 138 ++- nostr-sdk/src/relay/status.rs | 185 ++- nostr-sdk/src/test_utils.rs | 260 ++++- 15 files changed, 1766 insertions(+), 517 deletions(-) create mode 100644 nostr-sdk/src/mutex.rs diff --git a/Cargo.lock b/Cargo.lock index 1168e64d8..d75403926 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -115,7 +115,7 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "async-utility" version = "0.3.2" -source = "git+https://github.com/yukibtc/async-utility?rev=7d3a79983d0e848aafed9d321637af7334e01349#7d3a79983d0e848aafed9d321637af7334e01349" +source = "git+https://github.com/yukibtc/async-utility?rev=b323ae0c1c1424b78bf8bb8832b95e86771f1dee#b323ae0c1c1424b78bf8bb8832b95e86771f1dee" dependencies = [ "futures-util", "gloo-timers", diff --git a/Cargo.toml b/Cargo.toml index 8b55c9e75..a640ef887 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,7 +40,7 @@ license = "MIT" rust-version = "1.85.0" [workspace.dependencies] -async-utility = { git = "https://github.com/yukibtc/async-utility", rev = "7d3a79983d0e848aafed9d321637af7334e01349" } +async-utility = { git = "https://github.com/yukibtc/async-utility", rev = "b323ae0c1c1424b78bf8bb8832b95e86771f1dee" } atomic-destructor = { version = "0.3", default-features = false } base64 = { version = "0.22", default-features = false } bitcoin_hashes = { version = "1.1", default-features = false } diff --git a/nostr-sdk/CHANGELOG.md b/nostr-sdk/CHANGELOG.md index 89992e829..c7cbae359 100644 --- a/nostr-sdk/CHANGELOG.md +++ b/nostr-sdk/CHANGELOG.md @@ -34,6 +34,11 @@ - Redesign `Relay::count_events` API - Redesign `Client::try_connect_relay` API - Return `NotificationStream` from client and relay `notifications()` methods instead of a boxed stream. +- Remove `RelayStatus::Pending` and rename `RelayStatus::Terminated` to `RelayStatus::Idle` (https://github.com/nostrdevkit/nostr/pull/1491). + +### Changed + +- `Relay::try_connect` waits for an ongoing attempt or scheduled retry (https://github.com/nostrdevkit/nostr/pull/1491). ### Added @@ -50,6 +55,8 @@ - Preserve the broadcast lag or closure cause in event `OK` and authentication waiters instead of returning a generic premature-exit error. - Clean up NIP-77 request-owned subscriptions when reconciliation fails or is canceled. - Fail NIP-77 reconciliation when a download batch is rejected or closes before its requested events arrive. +- Retain reconnect requests while a previous WebSocket closes (https://github.com/nostrdevkit/nostr/pull/1491). +- Allow banning and shutting down idle relays (https://github.com/nostrdevkit/nostr/pull/1491). ### Deprecated diff --git a/nostr-sdk/src/client/api/send_event.rs b/nostr-sdk/src/client/api/send_event.rs index 02556c8e3..b51b4ff5b 100644 --- a/nostr-sdk/src/client/api/send_event.rs +++ b/nostr-sdk/src/client/api/send_event.rs @@ -446,6 +446,8 @@ where #[cfg(test)] mod tests { + use std::time::Duration; + use nostr::nips::nip17::InboxRelayList; use nostr::nips::nip65::RelayList; use nostr::prelude::*; @@ -454,9 +456,12 @@ mod tests { use nostr_gossip_memory::store::NostrGossipMemory; use super::*; + use crate::authenticator::SignerAuthenticator; use crate::client::{GossipConfig, GossipRelayLimits}; use crate::error::ErrorKind; use crate::local_relay::*; + use crate::relay::RelayStatus; + use crate::test_utils::{setup_client, setup_client_with_authenticator}; #[tokio::test] async fn unverified_event_is_rejected_before_sending() { @@ -779,4 +784,86 @@ mod tests { assert_eq!(err.kind(), ErrorKind::State); assert_eq!(err.to_string(), "gossip not configured"); } + + #[tokio::test] + async fn auth_gated_clients_keep_independent_signers_across_removal_and_reconnect() { + let local = LocalRelay::builder() + .nip42(LocalRelayBuilderNip42::write()) + .build(); + local.run().await.unwrap(); + + let url = local.url().await; + + let first_keys = Keys::generate(); + let second_keys = Keys::generate(); + + let first: Client = setup_client_with_authenticator( + url.clone(), + SignerAuthenticator::new(first_keys.clone()), + ) + .await; + let second: Client = setup_client_with_authenticator( + url.clone(), + SignerAuthenticator::new(second_keys.clone()), + ) + .await; + let unauthenticated: Client = setup_client(url.clone()).await; + + for (client, keys, content) in [ + (&first, &first_keys, "first account"), + (&second, &second_keys, "second account"), + ] { + let event = EventBuilder::new(Kind::TextNote, content) + .tag(Tag::protected()) + .finalize(keys) + .unwrap(); + + let output = client + .send_event(&event) + .ok_timeout(Duration::from_secs(2)) + .authentication_timeout(Duration::from_secs(2)) + .await + .unwrap(); + + assert!(output.failed.is_empty()); + assert!(output.success.get(&url).unwrap().is_ack()); + } + + let unauthenticated_event = EventBuilder::new(Kind::TextNote, "no signer") + .tag(Tag::protected()) + .finalize(&Keys::generate()) + .unwrap(); + let output = unauthenticated + .send_event(&unauthenticated_event) + .ok_timeout(Duration::from_secs(2)) + .await + .unwrap(); + + assert!(output.success.is_empty()); + assert!(output.failed.contains_key(&url)); + + first.remove_relay(&url).force().await.unwrap(); + + second.disconnect().await; + second.connect().and_wait(Duration::from_secs(3)).await; + + assert_eq!( + second.relay(&url).await.unwrap().unwrap().status(), + RelayStatus::Connected + ); + + let event = EventBuilder::new(Kind::TextNote, "second account after reconnect") + .tag(Tag::protected()) + .finalize(&second_keys) + .unwrap(); + let output = second + .send_event(&event) + .ok_timeout(Duration::from_secs(2)) + .authentication_timeout(Duration::from_secs(2)) + .await + .unwrap(); + + assert!(output.failed.is_empty(), "{:?}", output.failed); + assert!(output.success.get(&url).unwrap().is_ack()); + } } diff --git a/nostr-sdk/src/client/mod.rs b/nostr-sdk/src/client/mod.rs index 2c39316c7..f75eeb359 100644 --- a/nostr-sdk/src/client/mod.rs +++ b/nostr-sdk/src/client/mod.rs @@ -439,8 +439,8 @@ impl Client { /// If the relay is already connected or currently attempting to connect, /// this method does nothing. /// - /// If the relay is disconnected, sleeping, or otherwise inactive, a - /// background task is spawned to initiate a connection. + /// If the relay is sleeping or otherwise inactive, its persistent + /// connection task is started or woken to initiate a connection. /// /// For further details, see the documentation of [`Relay::connect`]. /// @@ -468,16 +468,14 @@ impl Client { /// Attempts to initiate a connection with a single relay. /// /// At most **one connection per relay** is allowed at any time. - /// If the relay is already connected or currently attempting to connect, - /// this method does nothing. + /// Returns immediately if already connected; otherwise waits for the current + /// attempt or next scheduled retry, or requests a connection if inactive. /// - /// If the initial connection attempt succeeds, a background task is spawned - /// to maintain the connection and handle future reconnections. - /// If the initial attempt fails, no background task is spawned and no - /// automatic retries are scheduled. + /// A newly requested attempt that fails does not schedule automatic retries. + /// The connection task remains available until ban or shutdown. + /// Dropping the future only stops waiting for the result. /// - /// Use [`Client::connect_relay`] if you want to always spawn a background - /// connection task, regardless of whether the initial attempt succeeds. + /// Use [`Client::connect_relay`] to enable retries after an initial failure. /// /// For further details, see the documentation of [`Relay::try_connect`]. /// @@ -522,8 +520,8 @@ impl Client { /// If a relay is already connected or currently attempting to connect, /// this method does nothing for that relay. /// - /// If a relay is disconnected, sleeping, or otherwise inactive, a - /// background task is spawned to initiate a connection. + /// If a relay is sleeping or otherwise inactive, its persistent + /// connection task is started or woken to initiate a connection. /// /// For further details, see the documentation of [`Relay::connect`]. /// @@ -561,16 +559,14 @@ impl Client { /// Attempts to initiate a connection with relays. /// /// At most **one connection per relay** is allowed at any time. - /// If a relay is already connected or currently attempting to connect, - /// this method does nothing for that relay. + /// Returns immediately for connected relays. For other relays, waits for + /// the current attempt or next scheduled retry, or requests a connection if inactive. /// - /// If the initial connection attempt succeeds, a background task is spawned - /// to maintain the connection and handle future reconnections. - /// If the initial attempt fails, no background task is spawned and no - /// automatic retries are scheduled. + /// A newly requested attempt that fails does not schedule automatic retries. + /// Connection tasks remain available until ban or shutdown. + /// Dropping the future only stops waiting for results. /// - /// Use [`Client::connect`] if you want to always spawn a background - /// connection task, regardless of whether the initial attempt succeeds. + /// Use [`Client::connect`] to enable retries after an initial failure. /// /// For further details, see the documentation of [`Relay::try_connect`]. /// diff --git a/nostr-sdk/src/lib.rs b/nostr-sdk/src/lib.rs index 284fb06da..6f0a3ee00 100644 --- a/nostr-sdk/src/lib.rs +++ b/nostr-sdk/src/lib.rs @@ -13,6 +13,7 @@ mod future; #[cfg(any(feature = "local-relay", test))] pub mod local_relay; pub mod monitor; +mod mutex; pub mod policy; mod pool; pub mod prelude; diff --git a/nostr-sdk/src/mutex.rs b/nostr-sdk/src/mutex.rs new file mode 100644 index 000000000..0e35c6d35 --- /dev/null +++ b/nostr-sdk/src/mutex.rs @@ -0,0 +1,21 @@ +pub(crate) use std::sync::MutexGuard as NonPoisoningMutexGuard; +use std::sync::{Mutex, PoisonError}; + +/// Adapter for `std::Mutex` that removes the poisoning aspects +/// from its API. +#[derive(Debug)] +pub(crate) struct NonPoisoningMutex(Mutex) +where + T: ?Sized; + +impl NonPoisoningMutex { + #[inline] + pub(crate) fn new(value: T) -> Self { + Self(Mutex::new(value)) + } + + #[inline] + pub(crate) fn lock(&self) -> NonPoisoningMutexGuard<'_, T> { + self.0.lock().unwrap_or_else(PoisonError::into_inner) + } +} diff --git a/nostr-sdk/src/relay/api/count_events.rs b/nostr-sdk/src/relay/api/count_events.rs index 3febb2fd1..35695763e 100644 --- a/nostr-sdk/src/relay/api/count_events.rs +++ b/nostr-sdk/src/relay/api/count_events.rs @@ -94,7 +94,7 @@ async fn receive_count_reply( } _ => {} }, - RelayNotification::RelayStatus { status } if status.is_disconnected() => { + RelayNotification::RelayStatus { status } if status.is_connection_closed() => { return Err(Error::not_connected()); } _ => {} diff --git a/nostr-sdk/src/relay/api/send_event.rs b/nostr-sdk/src/relay/api/send_event.rs index 948620395..7221293d1 100644 --- a/nostr-sdk/src/relay/api/send_event.rs +++ b/nostr-sdk/src/relay/api/send_event.rs @@ -206,7 +206,7 @@ async fn wait_for_authentication( RelayNotification::AuthenticationFailed => { return Err(Error::authentication_msg("authentication failed")); } - RelayNotification::RelayStatus { status } if status.is_disconnected() => { + RelayNotification::RelayStatus { status } if status.is_connection_closed() => { return Err(Error::not_connected()); } _ => (), diff --git a/nostr-sdk/src/relay/api/sync.rs b/nostr-sdk/src/relay/api/sync.rs index 3fcdf674c..15886dd8a 100644 --- a/nostr-sdk/src/relay/api/sync.rs +++ b/nostr-sdk/src/relay/api/sync.rs @@ -607,7 +607,7 @@ pub(super) async fn sync( last_relevant_msg = Instant::now(); } } - RelayNotification::RelayStatus { status } if status.is_disconnected() => { + RelayNotification::RelayStatus { status } if status.is_connection_closed() => { return Err(Error::not_connected()); } _ => (), diff --git a/nostr-sdk/src/relay/api/try_connect.rs b/nostr-sdk/src/relay/api/try_connect.rs index 99f26bf28..fa629afcf 100644 --- a/nostr-sdk/src/relay/api/try_connect.rs +++ b/nostr-sdk/src/relay/api/try_connect.rs @@ -3,9 +3,7 @@ use std::time::Duration; use crate::error::Error; use crate::future::BoxedFuture; -use crate::policy::AdmitStatus; -use crate::relay::{Relay, RelayStatus}; -use crate::transport::websocket::{WebSocketSink, WebSocketStream}; +use crate::relay::Relay; /// Try to connect relay #[must_use = "Does nothing unless you await!"] @@ -36,57 +34,25 @@ impl<'relay> IntoFuture for TryConnect<'relay> { type IntoFuture = BoxedFuture<'relay, Self::Output>; fn into_future(self) -> Self::IntoFuture { - Box::pin(async move { - let status: RelayStatus = self.relay.status(); - - if status.is_shutdown() { - return Err(Error::shutdown()); - } - - if status.is_banned() { - return Err(Error::banned()); - } - - // Check if relay can't connect - if !status.can_connect() { - return Ok(()); - } - - // Check connection policy - if let AdmitStatus::Rejected { reason } = - self.relay.inner.check_connection_policy().await? - { - // Set status to "terminated" - self.relay.inner.set_status(RelayStatus::Terminated, false); - - // Return error - return Err(Error::connection_rejected(reason)); - } - - // Try to connect - // This will set the status to "terminated" if the connection fails - let stream: (WebSocketSink, WebSocketStream) = self - .relay - .inner - ._try_connect(self.timeout, RelayStatus::Terminated) - .await?; - - // Spawn connection task - self.relay.inner.spawn_connection_task(Some(stream)); - - Ok(()) - }) + Box::pin(self.relay.inner.try_connect(self.timeout)) } } #[cfg(test)] mod tests { - use async_utility::time; use nostr::types::RelayUrl; + use tokio::task::JoinHandle; + use tokio::time; use super::*; use crate::error::ErrorKind; use crate::local_relay::*; + use crate::policy::AdmitStatus; + use crate::relay::RelayStatus; + use crate::test_utils::{ + ControlledConnectionPolicy, ControlledRelay, PolicyReply, TEST_TIMEOUT, fake_connection, + start_try_connect, wait_until, + }; #[tokio::test] async fn test_try_connect() { @@ -122,9 +88,121 @@ mod tests { let res = relay.try_connect().timeout(Duration::from_secs(2)).await; assert_eq!(res.unwrap_err().kind(), ErrorKind::Transport); - assert_eq!(relay.status(), RelayStatus::Terminated); + assert_eq!(relay.status(), RelayStatus::Idle); + + // Connection failed; the persistent task waits for another request. + assert!(relay.inner.is_running()); + } + + async fn terminal_operation_cancels_connection_policy( + stop: fn(&Relay), + expected_status: RelayStatus, + expected_error: &str, + ) { + let (policy, mut checks) = ControlledConnectionPolicy::new(); + + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::with_policy(policy); + + let caller: JoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + let reply: PolicyReply = checks.next_check().await; + + assert_eq!(relay.status(), RelayStatus::Connecting); + + stop(&relay); + + let error: Error = time::timeout(TEST_TIMEOUT, caller) + .await + .expect("timed out waiting for the cancelled connection attempt") + .expect("connection caller panicked") + .unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::State); + assert_eq!(error.to_string(), expected_error); + assert_eq!(relay.status(), expected_status); + assert_eq!(relay.stats().attempts(), 0); + assert_eq!(relay.stats().success(), 0); + + wait_until("connection task to exit", || !relay.inner.is_running()).await; + + assert!(reply.is_closed(), "connection policy was not cancelled"); + transport.assert_no_pending_dials(); + checks.assert_no_pending_checks(); + } + + #[tokio::test] + async fn try_connect_ban_cancels_pending_connection_policy() { + terminal_operation_cancels_connection_policy( + Relay::ban, + RelayStatus::Banned, + "relay banned", + ) + .await; + } + + #[tokio::test] + async fn try_connect_shutdown_cancels_pending_connection_policy() { + terminal_operation_cancels_connection_policy( + Relay::shutdown, + RelayStatus::Shutdown, + "shutdown", + ) + .await; + } + + #[tokio::test] + async fn concurrent_try_connect_shares_connection_policy_and_dial() { + let (policy, mut checks) = ControlledConnectionPolicy::new(); + + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::with_policy(policy); + + let first: JoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + let reply: PolicyReply = checks.next_check().await; + + let mut second: BoxedFuture<'_, Result<(), Error>> = + relay.try_connect().timeout(TEST_TIMEOUT).into_future(); + + assert!(futures::poll!(second.as_mut()).is_pending()); + assert_eq!(relay.status(), RelayStatus::Connecting); + checks.assert_no_pending_checks(); + transport.assert_no_pending_dials(); + + reply + .send(Ok(AdmitStatus::Success)) + .expect("connection policy was cancelled before approval"); + + assert!( + transport + .next_dial() + .await + .send(Ok(fake_connection())) + .is_ok(), + "transport dial was cancelled before completion" + ); + + time::timeout(TEST_TIMEOUT, async { + first + .await + .expect("first connection caller panicked") + .unwrap(); + second.await.unwrap(); + }) + .await + .expect("timed out waiting for concurrent connection callers"); + + assert_eq!(relay.status(), RelayStatus::Connected); + assert_eq!(relay.stats().attempts(), 1); + assert_eq!(relay.stats().success(), 1); + assert!(relay.inner.is_running()); + checks.assert_no_pending_checks(); + transport.assert_no_pending_dials(); - // Connection failed, the connection task is not running - assert!(!relay.inner.is_running()); + relay.shutdown(); + wait_until("connection task to exit", || !relay.inner.is_running()).await; } } diff --git a/nostr-sdk/src/relay/inner.rs b/nostr-sdk/src/relay/inner.rs index 535d34317..85b5fb59d 100644 --- a/nostr-sdk/src/relay/inner.rs +++ b/nostr-sdk/src/relay/inner.rs @@ -1,12 +1,13 @@ use std::borrow::Cow; -use std::cmp; use std::collections::HashMap; use std::net::SocketAddr; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, OnceLock}; use std::time::Duration; +use std::{cmp, ptr}; -use async_utility::{task, time}; +use async_utility::task::{self, JoinHandle}; +use async_utility::time; use async_wsocket::Message; use futures::{self, SinkExt, StreamExt}; use nostr::filter::MatchEventOptions; @@ -18,6 +19,7 @@ use rand::Rng; use rand::RngExt; use rand::rand_core::UnwrapErr; use rand::rngs::SysRng; +use tokio::sync::futures::Notified; use tokio::sync::mpsc::{self, Receiver, Sender}; use tokio::sync::{Mutex, MutexGuard, Notify, RwLock, RwLockWriteGuard, broadcast, oneshot, watch}; use universal_time::Instant; @@ -35,13 +37,14 @@ use super::{ SubscriptionAutoClosedReason, }; use crate::client::ClientNotification; -use crate::error::Error; +use crate::error::{Error, ErrorKind}; +use crate::mutex::{NonPoisoningMutex, NonPoisoningMutexGuard}; use crate::policy::AdmitStatus; -use crate::relay::status::AtomicRelayStatus; use crate::shared::SharedState; use crate::transport::websocket::{WebSocketSink, WebSocketStream}; type ClientMessageJson = String; +type ConnectionResult = Result<(), Arc>; // Skip NIP-50 matches since they may create issues and ban non-malicious relays. const MATCH_EVENT_OPTS: MatchEventOptions = MatchEventOptions::new().nip50(false); @@ -66,11 +69,91 @@ struct JsonMessageItem { confirmation: Option>, } +#[derive(Debug)] +enum RequestConnectionKind { + /// Connection has been requested from a caller not interested in waiting for the outcome. + Background, + /// Connection has been requested from a caller interested in waiting for the outcome. + WaitForResult { deadline: Instant }, +} + +impl RequestConnectionKind { + #[inline] + fn is_background(&self) -> bool { + matches!(self, Self::Background) + } +} + +#[derive(Debug)] +struct ConnectionAttempt { + kind: RequestConnectionKind, + result: OnceLock, + result_ready: Notify, +} + +impl ConnectionAttempt { + #[inline] + fn new(kind: RequestConnectionKind) -> Self { + Self { + kind, + result: OnceLock::new(), + result_ready: Notify::new(), + } + } + + #[inline] + fn is_current(&self, current: &Option>) -> bool { + match current { + Some(current) => ptr::eq(self, current.as_ref()), + None => false, + } + } + + fn finish(&self, result: Result<(), Error>) { + if self.result.get().is_some() { + return; + } + + if self.result.set(result.map_err(Arc::new)).is_ok() { + self.result_ready.notify_waiters(); + } + } + + async fn wait_for_result(&self) -> ConnectionResult { + loop { + // Register before checking the result so completion cannot be missed. + let ready: Notified<'_> = self.result_ready.notified(); + + if let Some(result) = self.result.get() { + return result.clone(); + } + + ready.await; + } + } +} + +// Status and attempt identity change together. Never retain this lock across an await. +#[derive(Debug)] +struct ConnectionState { + status: RelayStatus, + attempt: Option>, +} + +impl Default for ConnectionState { + fn default() -> Self { + Self { + status: RelayStatus::Initialized, + attempt: None, + } + } +} + #[derive(Debug)] struct RelayChannels { nostr: (Sender, Mutex>), + connection_changed: Notify, ping: Notify, - terminate: Notify, } impl RelayChannels { @@ -79,8 +162,8 @@ impl RelayChannels { Self { nostr: (tx_nostr, Mutex::new(rx_nostr)), + connection_changed: Notify::new(), ping: Notify::new(), - terminate: Notify::new(), } } @@ -101,10 +184,6 @@ impl RelayChannels { pub fn ping(&self) { self.ping.notify_one() } - - pub fn terminate(&self) { - self.terminate.notify_one() - } } #[derive(Debug)] @@ -150,10 +229,10 @@ impl SubscriptionData { // put all fields that require an `Arc` here. #[derive(Debug)] pub(super) struct AtomicPrivateData { - status: AtomicRelayStatus, + connection_state: NonPoisoningMutex, + connection_task_handle: OnceLock>, channels: RelayChannels, subscriptions: RwLock>, - running: AtomicBool, } #[derive(Debug, Clone)] @@ -181,10 +260,10 @@ impl InnerRelay { Self { url, atomic: Arc::new(AtomicPrivateData { - status: AtomicRelayStatus::default(), + connection_state: NonPoisoningMutex::new(ConnectionState::default()), + connection_task_handle: OnceLock::new(), channels: RelayChannels::new(), subscriptions: RwLock::new(HashMap::new()), - running: AtomicBool::new(false), }), capabilities: Arc::new(AtomicRelayCapabilities::new(capabilities)), opts, @@ -207,39 +286,44 @@ impl InnerRelay { None } - /// Check if the connection task is running - #[inline] + #[cfg(test)] pub(super) fn is_running(&self) -> bool { - self.atomic.running.load(Ordering::SeqCst) + match self.atomic.connection_task_handle.get() { + Some(handle) => !handle.is_finished(), + None => false, + } } #[inline] - pub fn status(&self) -> RelayStatus { - self.atomic.status.load() - } - - pub(super) fn set_status(&self, status: RelayStatus, log: bool) { - // Change status - self.atomic.status.set(status); - - // Log - if log { - match status { - RelayStatus::Initialized => tracing::trace!(url = %self.url, "Relay initialized."), - RelayStatus::Pending => tracing::trace!(url = %self.url, "Relay is pending."), - RelayStatus::Connecting => tracing::debug!("Connecting to '{}'", self.url), - RelayStatus::Connected => tracing::info!("Connected to '{}'", self.url), - RelayStatus::Disconnected => tracing::info!("Disconnected from '{}'", self.url), - RelayStatus::Terminated => { - tracing::info!("Completely disconnected from '{}'", self.url) - } - RelayStatus::Banned => tracing::info!(url = %self.url, "Relay banned."), - RelayStatus::Sleeping => tracing::info!("Relay '{}' went to sleep.", self.url), - RelayStatus::Shutdown => tracing::info!("Relay '{}' has been shutdown.", self.url), - } + fn connection_state(&self) -> NonPoisoningMutexGuard<'_, ConnectionState> { + self.atomic.connection_state.lock() + } + + #[inline] + pub(super) fn status(&self) -> RelayStatus { + let state = self.connection_state(); + state.status + } + + fn set_status(&self, state: &mut ConnectionState, status: RelayStatus) { + if state.status == status { + return; } - // Send notification + state.status = status; + + match status { + RelayStatus::Initialized => tracing::trace!(url = %self.url, "Relay initialized."), + RelayStatus::Connecting => tracing::debug!("Connecting to '{}'", self.url), + RelayStatus::Connected => tracing::info!("Connected to '{}'", self.url), + RelayStatus::Disconnected => tracing::info!("Disconnected from '{}'", self.url), + RelayStatus::Idle => tracing::info!(url = %self.url, "Relay idle."), + RelayStatus::Banned => tracing::info!(url = %self.url, "Relay banned."), + RelayStatus::Sleeping => tracing::info!("Relay '{}' went to sleep.", self.url), + RelayStatus::Shutdown => tracing::info!("Relay '{}' has been shutdown.", self.url), + } + + // Publish under the same lock so notifications follow the order of transitions. self.send_notification(RelayNotification::RelayStatus { status }, false); // If monitor is enabled, notify status change. @@ -493,7 +577,8 @@ impl InnerRelay { } } - pub(super) async fn check_connection_policy(&self) -> Result { + #[inline] + async fn check_connection_policy(&self) -> Result { match &self.state.admit_policy { Some(policy) => Ok(policy.admit_connection(&self.url).await?), None => Ok(AdmitStatus::Success), @@ -553,126 +638,119 @@ impl InnerRelay { tracing::debug!(url = %self.url, "Waking up sleeping relay."); - // Update status to pending - // TODO: is this needed here? - self.set_status(RelayStatus::Pending, false); + self.request_connect(); - // Spawn a new connection task - self.spawn_connection_task(None); - - // Update the wake-up timestamp self.stats.just_woke_up(); } - pub(super) fn spawn_connection_task(&self, stream: Option<(WebSocketSink, WebSocketStream)>) { - // Check if the connection task is already running - // This is checked also later, but it's checked also here to avoid a full-clone if we know that is already running. - if self.is_running() { - tracing::warn!(url = %self.url, "Connection task is already running."); - return; + #[inline] + pub(super) fn request_connect(&self) { + if let Err(error) = self.request_connection(RequestConnectionKind::Background) { + tracing::debug!(url = %self.url, %error, "Connection request ignored."); } + } + + pub(super) async fn try_connect(&self, timeout: Duration) -> Result<(), Error> { + let kind = RequestConnectionKind::WaitForResult { + deadline: Instant::now() + timeout, + }; - // Full-clone - let relay: InnerRelay = self.clone(); + let Some(attempt) = self.request_connection(kind)? else { + return Ok(()); + }; - // Spawn task - task::spawn(relay.connection_task(stream)); + match time::timeout(Some(timeout), attempt.wait_for_result()).await { + Some(result) => result.map_err(|error| Error::new(error.kind(), error)), + None => Err(Error::timeout()), + } } - /// This **MUST** be called only by the [`InnerRelay::spawn_connection_task`] method! - async fn connection_task(self, mut stream: Option<(WebSocketSink, WebSocketStream)>) { - // Set the connection task as running and get the previous value. - let is_running: bool = self.atomic.running.swap(true, Ordering::SeqCst); + fn request_connection( + &self, + kind: RequestConnectionKind, + ) -> Result>, Error> { + let attempt: Arc = { + let mut state: NonPoisoningMutexGuard<'_, ConnectionState> = self.connection_state(); - // Re-check if the connection task is already running. - // This is required because may happen that two tasks are spawned at the exact same moment. - // Not use the "assert" macro since will cause the task to panic. - if is_running { - tracing::warn!(url = %self.url, "Connection task is already running."); - return; - } + if state.status.is_banned() { + return Err(Error::banned()); + } - // Lock receiver - let mut rx_nostr = self.atomic.channels.rx_nostr().await; + if state.status.is_shutdown() { + return Err(Error::shutdown()); + } - // Last websocket error - // Store it to avoid printing every time the same connection error - let mut last_ws_error = None; + if state.status.is_connected() { + return Ok(None); + } - // Auto-connect loop - loop { - // Check if the connection is allowed - match self.check_connection_policy().await { - Ok(status) => { - // Connection rejected, update status and break the loop. - if let AdmitStatus::Rejected { reason } = status { - if let Some(reason) = reason { - tracing::warn!(reason = %reason, "Connection rejected by admission policy."); - } + match &state.attempt { + // We already have an attempt in progress + Some(attempt) => Arc::clone(attempt), + // No attempt in progress, create a new one + None => { + let attempt: Arc = Arc::new(ConnectionAttempt::new(kind)); - // Set the status to "terminated" and break loop. - self.set_status(RelayStatus::Terminated, false); - break; - } + state.attempt = Some(Arc::clone(&attempt)); + + self.set_status(&mut state, RelayStatus::Connecting); + + self.atomic.channels.connection_changed.notify_one(); + + attempt } - Err(e) => tracing::error!(error = %e, "Impossible to check connection policy."), } + }; - // Connect and run message handler - // The termination requests are handled inside this method! - self.connect_and_run(stream.take(), &mut rx_nostr, &mut last_ws_error) - .await; + let handle: &JoinHandle<()> = self.atomic.connection_task_handle.get_or_init(|| { + let relay: InnerRelay = self.clone(); + task::spawn(relay.connection_task()) + }); - // Get status - let status: RelayStatus = self.status(); + // Ban or shutdown may run before get_or_init publishes the task handle. + let status: RelayStatus = self.status(); + if status.is_terminal() { + handle.abort(); + } - // If the relay is terminated, banned or sleeping, break the loop. - if status.is_terminated() - || status.is_banned() - || status.is_sleeping() - || status.is_shutdown() - { - break; - } + Ok(Some(attempt)) + } - // Check if reconnection is enabled - if self.opts.reconnect { - // Check if the relay is marked as disconnected. If not, update status. - // Check if disconnected to avoid a possible double log - if !status.is_disconnected() { - self.set_status(RelayStatus::Disconnected, true); - } + async fn next_connection(&self) -> Option<(Arc, RelayStatus)> { + loop { + // Only the connection task waits here; the request itself remains in the state. + let changed: Notified<'_> = self.atomic.channels.connection_changed.notified(); - // Sleep before retry to connect - let interval: Duration = self.calculate_retry_interval(); - tracing::debug!( - "Reconnecting to '{}' relay in {} secs", - self.url, - interval.as_secs() - ); + { + let state: NonPoisoningMutexGuard<'_, ConnectionState> = self.connection_state(); - // Sleep before retry to connect - // Handle termination to allow exiting immediately if request is received during the sleep. - tokio::select! { - // Sleep - _ = time::sleep(interval) => {}, - // Handle termination notification - _ = self.handle_terminate() => break, + if state.status.is_terminal() { + return None; } - } else { - // Reconnection disabled, set status to "terminated" - self.set_status(RelayStatus::Terminated, true); - // Break loop and exit - tracing::debug!(url = %self.url, "Reconnection disabled, breaking loop."); - break; + if let Some(attempt) = &state.attempt { + return Some((Arc::clone(attempt), state.status)); + } } + + changed.await; } + } + + async fn connection_changed(&self, attempt: &ConnectionAttempt) { + loop { + let changed: Notified<'_> = self.atomic.channels.connection_changed.notified(); - // Mark the connection task as stopped. - self.atomic.running.store(false, Ordering::SeqCst); + { + let state = self.connection_state(); + if !attempt.is_current(&state.attempt) { + return; + } + } - tracing::debug!(url = %self.url, "Auto connect loop terminated."); + // Notifications only wake the task. Recheck identity after stale or coalesced wakes. + changed.await; + } } /// Depending on attempts and success, use default or incremental retry interval @@ -714,105 +792,186 @@ impl InnerRelay { self.opts.retry_interval } - #[inline] - async fn handle_terminate(&self) { - // Wait to be notified - self.atomic.channels.terminate.notified().await; + async fn connection_task(self) { + let mut rx_nostr: MutexGuard<'_, Receiver> = + self.atomic.channels.rx_nostr().await; + let mut last_ws_error: Option = None; + + while let Some((attempt, status)) = self.next_connection().await { + if status.is_disconnected() { + let interval: Duration = self.calculate_retry_interval(); + + tracing::debug!(url = %self.url, ?interval, "Waiting before reconnect."); + + tokio::select! { + biased; + _ = self.connection_changed(&attempt) => continue, + _ = time::sleep(interval) => {}, + } + } + + self.connect_and_run(&attempt, &mut rx_nostr, &mut last_ws_error) + .await; + } } - pub(super) async fn _try_connect( + fn begin_connection(&self, attempt: &ConnectionAttempt) -> bool { + let mut state: NonPoisoningMutexGuard<'_, ConnectionState> = self.connection_state(); + + if !attempt.is_current(&state.attempt) { + return false; + } + + self.set_status(&mut state, RelayStatus::Connecting); + + true + } + + fn begin_dial(&self, attempt: &ConnectionAttempt) -> bool { + let state: NonPoisoningMutexGuard<'_, ConnectionState> = self.connection_state(); + + if !attempt.is_current(&state.attempt) { + return false; + } + + self.stats.new_attempt(); + + true + } + + async fn dial_connection( &self, - timeout: Duration, - status_on_failure: RelayStatus, + attempt: &ConnectionAttempt, ) -> Result<(WebSocketSink, WebSocketStream), Error> { - // Update status - self.set_status(RelayStatus::Connecting, true); + let timeout: Duration = match attempt.kind { + RequestConnectionKind::Background => self.opts.connect_timeout, + RequestConnectionKind::WaitForResult { deadline } => deadline - Instant::now(), + }; - // Increase the attempts - self.stats.new_attempt(); + if timeout.is_zero() { + return Err(Error::timeout()); + } - // Connect futures - let connect_fut = self - .state - .transport - .connect((&self.url).into(), self.proxy()); - let fut = time::timeout(Some(timeout), connect_fut); + time::timeout(Some(timeout), async { + if let AdmitStatus::Rejected { reason } = self.check_connection_policy().await? { + return Err(Error::connection_rejected(reason)); + } - // Try to connect - // If during connection the termination request is received, abort the connection and return error. - // At this stem is NOT required to close the WebSocket connection. - tokio::select! { - // Connect - res = fut => match res { - Some(Ok((ws_tx, ws_rx))) => { - // Update status - self.set_status(RelayStatus::Connected, true); + if !self.begin_dial(attempt) { + return Err(Error::state_msg("connection attempt superseded")); + } - // Increment success stats - self.stats.new_success(); + self.state + .transport + .connect((&self.url).into(), self.proxy()) + .await + .map_err(Error::transport) + }) + .await + .unwrap_or_else(|| Err(Error::timeout())) + } - Ok((ws_tx, ws_rx)) - } - Some(Err(e)) => { - // Update status - self.set_status(status_on_failure, false); + fn accept_connection(&self, attempt: &ConnectionAttempt) -> bool { + let mut state: NonPoisoningMutexGuard<'_, ConnectionState> = self.connection_state(); - // Return error - Err(Error::transport(e)) - } - None => { - // Update status - self.set_status(status_on_failure, false); + // A superseded dial must not publish Connected or count as a successful connection. + if !attempt.is_current(&state.attempt) { + return false; + } - // Return error - Err(Error::timeout()) - } - }, - // Handle termination notification - _ = self.handle_terminate() => Err(Error::rejected_msg("received termination request")), + self.stats.new_success(); + + self.set_status(&mut state, RelayStatus::Connected); + + attempt.finish(Ok(())); + + true + } + + fn connection_failed( + &self, + attempt: &ConnectionAttempt, + error: Error, + last_ws_error: &mut Option, + ) { + let mut state: NonPoisoningMutexGuard<'_, ConnectionState> = self.connection_state(); + + if !attempt.is_current(&state.attempt) { + return; + } + + let message: String = error.to_string(); + if last_ws_error.as_ref() != Some(&message) { + tracing::error!(url = %self.url, %error, "Connection failed."); + *last_ws_error = Some(message); + } + + let stop: bool = !attempt.kind.is_background() + || error.kind() == ErrorKind::Rejected + || !self.opts.reconnect; + + attempt.finish(Err(error)); + + if stop { + state.attempt = None; + self.set_status(&mut state, RelayStatus::Idle); + return; + } + + state.attempt = Some(Arc::new(ConnectionAttempt::new( + RequestConnectionKind::Background, + ))); + + self.set_status(&mut state, RelayStatus::Disconnected); + } + + fn schedule_reconnect(&self, attempt: &ConnectionAttempt) { + let mut state: NonPoisoningMutexGuard<'_, ConnectionState> = self.connection_state(); + + // A new request may already be waiting while the previous WebSocket closes. + if !attempt.is_current(&state.attempt) { + return; + } + + if !self.opts.reconnect { + state.attempt = None; + self.set_status(&mut state, RelayStatus::Idle); + return; } + + state.attempt = Some(Arc::new(ConnectionAttempt::new( + RequestConnectionKind::Background, + ))); + + self.set_status(&mut state, RelayStatus::Disconnected); } - /// Connect and run message handler - /// - /// If `stream` arg is passed, no connection attempt will be done. async fn connect_and_run( &self, - stream: Option<(WebSocketSink, WebSocketStream)>, + attempt: &ConnectionAttempt, rx_nostr: &mut MutexGuard<'_, Receiver>, last_ws_error: &mut Option, ) { - match stream { - // Already have a stream, go to post-connection stage - Some((ws_tx, ws_rx)) => self.post_connection(ws_tx, ws_rx, rx_nostr).await, - // No stream is passed, try to connect - // Set the status to "disconnected" to allow to automatic retries - None => match self - ._try_connect(self.opts.connect_timeout, RelayStatus::Disconnected) - .await - { - // Connection success, go to post-connection stage - Ok((ws_tx, ws_rx)) => self.post_connection(ws_tx, ws_rx, rx_nostr).await, - // Error during connection - Err(e) => { - // TODO: avoid string allocation. The error is converted to string only to perform the `!=` binary operation. - // Check if error should be logged - let e: String = e.to_string(); - let to_log: bool = match &last_ws_error { - Some(prev_err) => { - // Log only if different from the last one - prev_err != &e - } - None => true, - }; + if !self.begin_connection(attempt) { + return; + } - // Log error and update the last error - if to_log { - tracing::error!(url = %self.url, error= %e, "Connection failed."); - *last_ws_error = Some(e); - } + let connection: Result<(WebSocketSink, WebSocketStream), Error> = tokio::select! { + biased; + _ = self.connection_changed(attempt) => return, + result = self.dial_connection(attempt) => result, + }; + + match connection { + Ok((ws_tx, ws_rx)) => { + if !self.accept_connection(attempt) { + return; } - }, + + self.post_connection(ws_tx, ws_rx, rx_nostr, attempt).await; + self.schedule_reconnect(attempt); + } + Err(error) => self.connection_failed(attempt, error, last_ws_error), } } @@ -823,12 +982,23 @@ impl InnerRelay { mut ws_tx: WebSocketSink, ws_rx: WebSocketStream, rx_nostr: &mut MutexGuard<'_, Receiver>, + attempt: &ConnectionAttempt, ) { - // (Re)subscribe to relay - if self.capabilities.can_read() { - if let Err(e) = self.resubscribe().await { - tracing::error!(url = %self.url, error = %e, "Impossible to subscribe.") + let resubscribe = async { + if self.capabilities.can_read() { + if let Err(error) = self.resubscribe().await { + tracing::error!(url = %self.url, %error, "Impossible to subscribe."); + } } + }; + + tokio::select! { + biased; + _ = self.connection_changed(attempt) => { + let _ = close_ws(&mut ws_tx).await; + return; + }, + _ = resubscribe => {}, } let ping: PingTracker = PingTracker::default(); @@ -855,9 +1025,9 @@ impl InnerRelay { Err(e) => tracing::error!(url = %self.url, error = %e, "Relay ingester exited with error.") }, // Monitor when the relay can go to sleep - _ = self.sleep_when_idle_monitor() => {}, + _ = self.sleep_when_idle_monitor(attempt) => {}, // Termination handler - _ = self.handle_terminate() => {}, + _ = self.connection_changed(attempt) => {}, // Pinger _ = self.pinger() => {} } @@ -1034,15 +1204,20 @@ impl InnerRelay { } /// Monitor if it's time to put the relay in sleep mode. - async fn sleep_when_idle_monitor(&self) { + async fn sleep_when_idle_monitor(&self, attempt: &ConnectionAttempt) { loop { // Sleep time::sleep(SLEEP_INTERVAL).await; // Check if should go to sleep if self.should_sleep().await { - // Update status - self.set_status(RelayStatus::Sleeping, true); + let mut state: NonPoisoningMutexGuard<'_, ConnectionState> = + self.connection_state(); + + if attempt.is_current(&state.attempt) { + state.attempt = None; + self.set_status(&mut state, RelayStatus::Sleeping); + } // Break the loop break; @@ -1335,49 +1510,52 @@ impl InnerRelay { })) } - pub fn disconnect(&self) { - let status = self.status(); - - // Check if it's already terminated, banned or shutdown - if status.is_terminated() || status.is_banned() || status.is_shutdown() { - return; - } + #[inline] + pub(super) fn disconnect(&self) { + self.stop_connection(RelayStatus::Idle); + } - // Notify termination - self.atomic.channels.terminate(); + #[inline] + pub(super) fn ban(&self) { + self.stop_connection(RelayStatus::Banned); + } - // Update status - self.set_status(RelayStatus::Terminated, true); + #[inline] + pub(super) fn shutdown(&self) { + self.stop_connection(RelayStatus::Shutdown); } - pub fn ban(&self) { - let status = self.status(); + fn stop_connection(&self, status: RelayStatus) { + { + let mut state: NonPoisoningMutexGuard<'_, ConnectionState> = self.connection_state(); - // Check if it's already terminated, banned or shutdown - if status.is_terminated() || status.is_banned() || status.is_shutdown() { - return; - } + let previous: RelayStatus = state.status; - // Notify termination - self.atomic.channels.terminate(); + if previous.is_terminal() || previous == status { + return; + } - // Update status - self.set_status(RelayStatus::Banned, true); - } + if let Some(attempt) = state.attempt.take() { + let error: Error = match status { + RelayStatus::Banned => Error::banned(), + RelayStatus::Shutdown => Error::shutdown(), + _ => Error::rejected_msg("received termination request"), + }; + attempt.finish(Err(error)); + } - pub(super) fn shutdown(&self) { - let status = self.status(); + self.set_status(&mut state, status); - // Check if it's already terminated, banned or shutdown - if status.is_terminated() || status.is_banned() || status.is_shutdown() { - return; + // Wake the connection task so it observes the invalidated attempt + // and stops any pending retry, dial, or active session. + self.atomic.channels.connection_changed.notify_one(); } - // Notify termination - self.atomic.channels.terminate(); - - // Update status - self.set_status(RelayStatus::Shutdown, true); + if status.is_terminal() { + if let Some(handle) = self.atomic.connection_task_handle.get() { + handle.abort(); + } + } } #[inline] @@ -1490,7 +1668,7 @@ impl InnerRelay { } } } - RelayNotification::RelayStatus { status } if status.is_disconnected() => { + RelayNotification::RelayStatus { status } if status.is_connection_closed() => { return Err(Error::not_connected()); } _ => (), @@ -1758,7 +1936,7 @@ impl InnerRelay { reason: Some(SubscriptionAutoClosedReason::AuthenticationFailed), }); } - RelayNotification::RelayStatus { status } if status.is_disconnected() => { + RelayNotification::RelayStatus { status } if status.is_connection_closed() => { return Some(HandleAutoClosing { to_close: false, // No need to send CLOSE msg reason: None, @@ -1792,7 +1970,7 @@ impl InnerRelay { } } RelayNotification::RelayStatus { status } - if status.is_disconnected() => + if status.is_connection_closed() => { return Ok(()); } @@ -1879,8 +2057,9 @@ async fn close_ws(tx: &mut WebSocketSink) -> Result<(), Error> { #[cfg(test)] mod tests { use std::borrow::Cow; - use std::error::Error as _; - use std::future::Future; + use std::error::Error as StdError; + use std::future::{Future, IntoFuture}; + use std::io::{Error as IoError, ErrorKind as IoErrorKind}; use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -1890,12 +2069,19 @@ mod tests { use nostr::key::Keys; use nostr::message::SubscriptionId; use nostr::types::RelayUrl; + use tokio::sync::broadcast::Receiver as BroadcastReceiver; + use tokio::task::JoinHandle as TokioJoinHandle; use super::*; use crate::authenticator::SignerAuthenticator; - use crate::error::ErrorKind; + use crate::error::{Error, ErrorKind}; + use crate::future::BoxedFuture; use crate::policy::{AdmitPolicy, AdmitStatus}; use crate::relay::{Relay, RelayOptions}; + use crate::test_utils::{ + ControlledConnectionPolicy, ControlledRelay, DialReply, PolicyReply, TEST_TIMEOUT, + connection_gate_close, fake_connection, start_try_connect, wait_for_status, wait_until, + }; #[derive(Debug)] struct CountingAdmitPolicy(Arc); @@ -2363,6 +2549,664 @@ mod tests { .is_some_and(|source| matches!(source, broadcast::error::RecvError::Closed)) ); } + + async fn wait_for_task_exit(inner: &InnerRelay) { + wait_until("connection task to exit", || !inner.is_running()).await; + } + + enum ReconnectRequest { + Background, + WaitForResult, + } + + async fn reconnect_while_closing(opts: RelayOptions, request: ReconnectRequest) { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::with_opts(opts); + + let initial: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + + let (connection, mut close_gate) = connection_gate_close(fake_connection()); + + assert!(transport.next_dial().await.send(Ok(connection)).is_ok()); + + initial.await.unwrap().unwrap(); + + relay.disconnect(); + + close_gate.wait_until_closing().await; + + assert!(relay.inner.is_running()); + + // Submit the new request while the previous WebSocket is still closing. + let pending: Option>> = match request { + ReconnectRequest::WaitForResult => { + let mut pending: BoxedFuture<'_, Result<(), Error>> = + Box::pin(relay.try_connect().timeout(TEST_TIMEOUT).into_future()); + assert!(futures::poll!(pending.as_mut()).is_pending()); + Some(pending) + } + ReconnectRequest::Background => { + relay.connect(); + relay.connect(); + None + } + }; + + assert_eq!(relay.status(), RelayStatus::Connecting); + + transport.assert_no_pending_dials(); + + close_gate.release(); + + assert!( + transport + .next_dial() + .await + .send(Ok(fake_connection())) + .is_ok() + ); + + if let Some(pending) = pending { + pending.await.unwrap(); + } + wait_for_status(&relay, RelayStatus::Connected).await; + + assert_eq!(relay.stats().attempts(), 2); + assert_eq!(relay.stats().success(), 2); + + transport.assert_no_pending_dials(); + } + + #[tokio::test] + async fn connection_result_wakes_all_callers_and_cannot_be_overwritten() { + let attempt: ConnectionAttempt = ConnectionAttempt::new(RequestConnectionKind::Background); + + let mut first: BoxedFuture<'_, ConnectionResult> = Box::pin(attempt.wait_for_result()); + let mut second: BoxedFuture<'_, ConnectionResult> = Box::pin(attempt.wait_for_result()); + + assert!(futures::poll!(first.as_mut()).is_pending()); + assert!(futures::poll!(second.as_mut()).is_pending()); + + attempt.finish(Err(Error::state_msg("shared failure"))); + attempt.finish(Ok(())); + + let first: Arc = first.await.unwrap_err(); + let second: Arc = second.await.unwrap_err(); + + assert!(Arc::ptr_eq(&first, &second)); + assert_eq!(first.to_string(), "shared failure"); + assert!(Arc::ptr_eq( + &first, + &attempt.wait_for_result().await.unwrap_err() + )); + } + + #[tokio::test] + async fn stale_connection_wakes_do_not_cancel_the_current_dial() { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::new(); + + relay.connect(); + + let reply: DialReply = transport.next_dial().await; + + for _ in 0..3 { + relay.inner.atomic.channels.connection_changed.notify_one(); + } + + assert!(reply.send(Ok(fake_connection())).is_ok()); + + wait_for_status(&relay, RelayStatus::Connected).await; + + assert_eq!(relay.stats().attempts(), 1); + assert_eq!(relay.stats().success(), 1); + + transport.assert_no_pending_dials(); + } + + #[tokio::test] + async fn wait_for_connection_waits_for_retry_but_returns_when_idle() { + let opts: RelayOptions = RelayOptions::default() + .adjust_retry_interval(false) + .retry_interval(Duration::from_secs(1)); + + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::with_opts(opts); + + relay.connect(); + + assert!( + transport + .next_dial() + .await + .send(Err(Error::state_msg("retry"))) + .is_ok() + ); + + wait_for_status(&relay, RelayStatus::Disconnected).await; + + let mut waiting: BoxedFuture<'_, ()> = Box::pin(relay.wait_for_connection(TEST_TIMEOUT)); + + assert!(futures::poll!(waiting.as_mut()).is_pending()); + + relay.disconnect(); + + assert!(futures::poll!(waiting.as_mut()).is_ready()); + assert_eq!(relay.status(), RelayStatus::Idle); + assert!( + futures::poll!(Box::pin(relay.wait_for_connection(TEST_TIMEOUT)).as_mut()).is_ready() + ); + } + + #[tokio::test] + async fn reconnect_request_survives_websocket_close() { + reconnect_while_closing(RelayOptions::default(), ReconnectRequest::Background).await; + } + + #[tokio::test] + async fn explicit_reconnect_survives_close_without_automatic_reconnect() { + reconnect_while_closing( + RelayOptions::default().reconnect(false), + ReconnectRequest::Background, + ) + .await; + } + + #[tokio::test] + async fn try_connect_waits_for_websocket_close() { + reconnect_while_closing( + RelayOptions::default().reconnect(false), + ReconnectRequest::WaitForResult, + ) + .await; + } + + #[tokio::test] + async fn callers_join_connection_policy_and_share_its_error() { + let (policy, mut checks) = ControlledConnectionPolicy::new(); + + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::with_policy(policy); + + let first: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + + let reply: PolicyReply = checks.next_check().await; + + assert_eq!(relay.status(), RelayStatus::Connecting); + + let mut second: BoxedFuture<'_, Result<(), Error>> = + Box::pin(relay.try_connect().timeout(TEST_TIMEOUT).into_future()); + assert!(futures::poll!(second.as_mut()).is_pending()); + + let error: Error = Error::policy(IoError::other("policy failure")); + reply.send(Err(error)).unwrap(); + + for result in [first.await.unwrap(), second.await] { + let error: Error = result.unwrap_err(); + assert_eq!(error.kind(), ErrorKind::Policy); + assert_eq!(error.to_string(), "policy failure"); + } + + assert_eq!(relay.status(), RelayStatus::Idle); + assert_eq!(relay.stats().attempts(), 0); + + transport.assert_no_pending_dials(); + checks.assert_no_pending_checks(); + } + + #[tokio::test] + async fn initial_try_connect_deadline_also_cancels_policy_check() { + let (policy, mut checks) = ControlledConnectionPolicy::new(); + + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::with_policy(policy); + + let caller: TokioJoinHandle> = + start_try_connect(&relay, Duration::from_millis(20)); + + let reply: PolicyReply = checks.next_check().await; + + assert_eq!( + caller.await.unwrap().unwrap_err().kind(), + ErrorKind::Timeout + ); + + wait_for_status(&relay, RelayStatus::Idle).await; + + assert!(reply.is_closed()); + assert_eq!(relay.stats().attempts(), 0); + assert!(relay.inner.is_running()); + + transport.assert_no_pending_dials(); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn concurrent_try_connect_shares_transport_error_and_source() { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::new(); + + let first: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + let reply: DialReply = transport.next_dial().await; + + let mut second: BoxedFuture<'_, Result<(), Error>> = + Box::pin(relay.try_connect().timeout(TEST_TIMEOUT).into_future()); + assert!(futures::poll!(second.as_mut()).is_pending()); + + let cause: IoError = IoError::new(IoErrorKind::ConnectionRefused, "refused"); + let error: Error = Error::other(cause); + assert!(reply.send(Err(error)).is_ok()); + + for result in [first.await.unwrap(), second.await] { + let error: Error = result.unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::Transport); + assert_eq!(error.to_string(), "refused"); + + let mut source: Option<&(dyn StdError + 'static)> = error.source(); + while let Some(error) = source { + if let Some(error) = error.downcast_ref::() { + assert_eq!(error.kind(), IoErrorKind::ConnectionRefused); + break; + } + source = error.source(); + } + + assert!(source.is_some(), "original transport source was lost"); + } + + assert_eq!(relay.status(), RelayStatus::Idle); + assert_eq!(relay.stats().attempts(), 1); + assert!(relay.inner.is_running()); + + transport.assert_no_pending_dials(); + } + + #[tokio::test] + async fn joining_caller_timeout_does_not_cancel_shared_dial() { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::new(); + + let first: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + let reply: DialReply = transport.next_dial().await; + + let error: Error = relay + .try_connect() + .timeout(Duration::from_millis(20)) + .await + .unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::Timeout); + assert_eq!(relay.status(), RelayStatus::Connecting); + assert!(!reply.is_closed()); + + assert!(reply.send(Ok(fake_connection())).is_ok()); + first.await.unwrap().unwrap(); + + assert_eq!(relay.status(), RelayStatus::Connected); + assert_eq!(relay.stats().attempts(), 1); + } + + #[tokio::test] + async fn dropping_try_connect_only_stops_waiting() { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::new(); + + let caller: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + let reply: DialReply = transport.next_dial().await; + + caller.abort(); + + assert!(caller.await.unwrap_err().is_cancelled()); + assert!(!reply.is_closed()); + + assert!(reply.send(Ok(fake_connection())).is_ok()); + + wait_for_status(&relay, RelayStatus::Connected).await; + + assert_eq!(relay.stats().attempts(), 1); + } + + #[tokio::test] + async fn initial_try_connect_timeout_leaves_task_idle_without_retries() { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::new(); + + let caller: TokioJoinHandle> = + start_try_connect(&relay, Duration::from_millis(20)); + let reply: DialReply = transport.next_dial().await; + + assert_eq!( + caller.await.unwrap().unwrap_err().kind(), + ErrorKind::Timeout + ); + + wait_for_status(&relay, RelayStatus::Idle).await; + + assert!(reply.is_closed()); + assert!(relay.inner.is_running()); + + transport.assert_no_pending_dials(); + + let handle: &JoinHandle<()> = relay.inner.atomic.connection_task_handle.get().unwrap(); + + let retry: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + + assert!( + transport + .next_dial() + .await + .send(Ok(fake_connection())) + .is_ok() + ); + + retry.await.unwrap().unwrap(); + + assert_eq!(relay.stats().attempts(), 2); + assert!(ptr::eq( + handle, + relay.inner.atomic.connection_task_handle.get().unwrap() + )); + } + + #[tokio::test] + async fn try_connect_timeout_includes_waiting_for_close() { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::new(); + + let initial: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + + let (connection, mut close_gate) = connection_gate_close(fake_connection()); + + assert!(transport.next_dial().await.send(Ok(connection)).is_ok()); + + initial.await.unwrap().unwrap(); + + relay.disconnect(); + + close_gate.wait_until_closing().await; + + let error: Error = relay + .try_connect() + .timeout(Duration::from_millis(20)) + .await + .unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::Timeout); + + transport.assert_no_pending_dials(); + + close_gate.release(); + + wait_for_status(&relay, RelayStatus::Idle).await; + + assert_eq!(relay.stats().attempts(), 1); + } + + #[tokio::test] + async fn superseded_handshake_cannot_publish_connected() { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::new(); + + let first: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + let reply: DialReply = transport.next_dial().await; + + let mut updates: BroadcastReceiver = + relay.inner.internal_notification_sender.subscribe(); + + // Replace the attempt without yielding after completing its handshake, + // so the connection task cannot accept the superseded connection. + assert!(reply.send(Ok(fake_connection())).is_ok()); + + relay.disconnect(); + relay.connect(); + relay.disconnect(); + relay.connect(); + + let second: DialReply = transport.next_dial().await; + + assert_eq!( + first.await.unwrap().unwrap_err().kind(), + ErrorKind::Rejected + ); + assert_eq!(relay.stats().success(), 0); + + while let Ok(notification) = updates.try_recv() { + assert!(!matches!( + notification, + RelayNotification::RelayStatus { + status: RelayStatus::Connected + } + )); + } + + assert!(second.send(Ok(fake_connection())).is_ok()); + + wait_for_status(&relay, RelayStatus::Connected).await; + + assert_eq!(relay.stats().attempts(), 2); + assert_eq!(relay.stats().success(), 1); + + transport.assert_no_pending_dials(); + } + + #[tokio::test] + async fn try_connect_joins_next_scheduled_retry() { + let opts: RelayOptions = RelayOptions::default() + .adjust_retry_interval(false) + .retry_interval(Duration::from_secs(1)); + + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::with_opts(opts); + + relay.connect(); + + assert!( + transport + .next_dial() + .await + .send(Err(Error::state_msg("first failure"))) + .is_ok() + ); + + wait_for_status(&relay, RelayStatus::Disconnected).await; + + let mut caller: BoxedFuture<'_, Result<(), Error>> = + Box::pin(relay.try_connect().timeout(TEST_TIMEOUT).into_future()); + + assert!(futures::poll!(caller.as_mut()).is_pending()); + transport.assert_no_pending_dials(); + assert!( + transport + .next_dial() + .await + .send(Ok(fake_connection())) + .is_ok() + ); + + caller.await.unwrap(); + + assert_eq!(relay.stats().attempts(), 2); + } + + #[tokio::test] + async fn terminal_operations_cancel_pending_reconnects_and_stop_task() { + for terminal in [RelayStatus::Banned, RelayStatus::Shutdown] { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::new(); + + let first: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + let (connection, mut close_gate) = connection_gate_close(fake_connection()); + + assert!(transport.next_dial().await.send(Ok(connection)).is_ok()); + + first.await.unwrap().unwrap(); + + relay.disconnect(); + close_gate.wait_until_closing().await; + + let mut pending: BoxedFuture<'_, Result<(), Error>> = + Box::pin(relay.try_connect().timeout(TEST_TIMEOUT).into_future()); + + assert!(futures::poll!(pending.as_mut()).is_pending()); + + match terminal { + RelayStatus::Banned => relay.ban(), + _ => relay.shutdown(), + } + let error: Error = pending.await.unwrap_err(); + + assert_eq!(error.kind(), ErrorKind::State); + assert_eq!( + error.to_string(), + match terminal { + RelayStatus::Banned => "relay banned", + _ => "shutdown", + } + ); + + wait_for_task_exit(&relay.inner).await; + + assert!(close_gate.is_cancelled()); + + relay.connect(); + + assert_eq!(relay.status(), terminal); + assert!(relay.inner.connection_state().attempt.is_none()); + assert_eq!(relay.stats().attempts(), 1); + + transport.assert_no_pending_dials(); + } + } + + #[tokio::test] + async fn ban_and_shutdown_stop_idle_tasks() { + for terminal in [RelayStatus::Banned, RelayStatus::Shutdown] { + for idle in [RelayStatus::Idle, RelayStatus::Sleeping] { + let opts: RelayOptions = + RelayOptions::default().sleep_when_idle(SleepWhenIdle::Enabled { + timeout: Duration::ZERO, + }); + + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::with_opts(opts); + + let first: TokioJoinHandle> = + start_try_connect(&relay, TEST_TIMEOUT); + + assert!( + transport + .next_dial() + .await + .send(Ok(fake_connection())) + .is_ok() + ); + + first.await.unwrap().unwrap(); + + match idle { + RelayStatus::Sleeping => wait_for_status(&relay, idle).await, + _ => relay.disconnect(), + } + + match terminal { + RelayStatus::Banned => relay.ban(), + _ => relay.shutdown(), + } + + wait_for_task_exit(&relay.inner).await; + + assert_eq!(relay.status(), terminal); + } + } + } + + #[tokio::test] + async fn last_drop_shuts_down_idle_connection_task() { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::new(); + + let first: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + + assert!( + transport + .next_dial() + .await + .send(Ok(fake_connection())) + .is_ok() + ); + + first.await.unwrap().unwrap(); + + relay.disconnect(); + + let inner: InnerRelay = relay.inner.clone(); + let clone: Relay = relay.clone(); + + drop(relay); + + assert_eq!(inner.status(), RelayStatus::Idle); + assert!(inner.is_running()); + + drop(clone); + wait_for_task_exit(&inner).await; + + assert_eq!(inner.status(), RelayStatus::Shutdown); + assert_eq!(Arc::strong_count(&inner.atomic), 1); + } + + #[tokio::test] + async fn disconnect_before_first_connect_does_not_cancel_new_attempt() { + let ControlledRelay { + relay, + mut transport, + } = ControlledRelay::with_opts(RelayOptions::default().reconnect(false)); + + relay.disconnect(); + relay.connect(); + + assert!( + transport + .next_dial() + .await + .send(Ok(fake_connection())) + .is_ok() + ); + + wait_for_status(&relay, RelayStatus::Connected).await; + + assert_eq!(relay.stats().attempts(), 1); + assert_eq!(relay.stats().success(), 1); + } } #[cfg(bench)] diff --git a/nostr-sdk/src/relay/mod.rs b/nostr-sdk/src/relay/mod.rs index cea7291da..eb6b47813 100644 --- a/nostr-sdk/src/relay/mod.rs +++ b/nostr-sdk/src/relay/mod.rs @@ -234,9 +234,9 @@ impl Relay { /// /// # Overview /// - /// If the relay’s status is not [`RelayStatus::Initialized`] or [`RelayStatus::Terminated`], - /// this method returns immediately without doing anything. - /// Otherwise, the connection task will be spawned, which will attempt to connect to relay. + /// Requests a connection when the relay is initialized, idle or sleeping. + /// The connection task is spawned once and retained until ban or shutdown. + /// Requests made while an old WebSocket is closing are retained. /// /// This method returns immediately and doesn't provide any information on if the connection was successful or not. /// @@ -244,39 +244,27 @@ impl Relay { /// /// By default, in case of disconnection, the connection task will automatically attempt to reconnect. /// This behavior can be disabled by changing [`RelayOptions::reconnect`] option. + #[inline] pub fn connect(&self) { - // Immediately return if can't connect - if !self.status().can_connect() { - return; - } - - // Update status - // Change it to pending to avoid issues with the health check (initialized check) - self.inner.set_status(RelayStatus::Pending, false); - - // Spawn connection task - self.inner.spawn_connection_task(None); + self.inner.request_connect(); } /// Waits for relay connection /// /// Wait for relay connection at most for the specified `timeout`. - /// The code continues when the relay is connected or the `timeout` is reached. + /// Returns when connected, when the relay becomes idle, sleeping, banned or shut down, + /// or when the timeout is reached. pub async fn wait_for_connection(&self, timeout: Duration) { + // Subscribe before reading status so a concurrent transition cannot be missed. + let mut notifications = self.inner.internal_notification_sender.subscribe(); + let status: RelayStatus = self.status(); - // Immediately returns if the relay is already connected, if it's terminated or banned. - if status.is_connected() - || status.is_terminated() - || status.is_banned() - || status.is_shutdown() - { + // Immediately returns if the relay is already connected, if it's idle, banned or shut down. + if status.is_connected() || status.is_idle() || status.is_banned() || status.is_shutdown() { return; } - // Subscribe to notifications - let mut notifications = self.inner.internal_notification_sender.subscribe(); - // Set timeout time::timeout(Some(timeout), async { while let Ok(notification) = notifications.recv().await { @@ -285,12 +273,11 @@ impl Relay { match status { // Waiting for connection RelayStatus::Initialized - | RelayStatus::Pending | RelayStatus::Connecting | RelayStatus::Disconnected => {} - // Connected or terminated/banned/sleeping/shutdown + // Connected or idle/banned/sleeping/shutdown RelayStatus::Connected - | RelayStatus::Terminated + | RelayStatus::Idle | RelayStatus::Banned | RelayStatus::Sleeping | RelayStatus::Shutdown => break, @@ -305,14 +292,18 @@ impl Relay { /// /// # Overview /// - /// If the relay’s status is not [`RelayStatus::Initialized`] or [`RelayStatus::Terminated`], - /// this method returns immediately without doing anything. - /// Otherwise, attempts to establish a connection without spawning the connection task if it fails. - /// This means that if the connection fails, no automatic retries are scheduled. - /// Use [`Relay::connect`] if you want to immediately spawn a connection task, - /// regardless of whether the initial connection succeeds. + /// Returns immediately if already connected. Otherwise, requests a connection + /// or waits for the current attempt (the next scheduled retry if disconnected). + /// Concurrent callers share that attempt, each with their own timeout. + /// + /// A newly requested attempt that fails does not schedule automatic retries. + /// Use [`Relay::connect`] to enable retries after an initial failure. + /// + /// The timeout includes waiting for an old WebSocket to close and admission policy checks. + /// Dropping the future only stops waiting; the connection attempt continues within + /// its assigned timeout. A caller joining an existing attempt does not change its timeout. /// - /// Returns an error if the connection fails or if the relay has been banned. + /// Returns an error on connection failure, timeout, disconnect, ban or shutdown. /// /// # Automatic reconnection /// @@ -323,7 +314,9 @@ impl Relay { TryConnect::new(self) } - /// Disconnect from relay and set status to [`RelayStatus::Terminated`]. + /// Disconnect from relay and set status to [`RelayStatus::Idle`]. + /// + /// The persistent task closes the WebSocket and waits for another connection request. #[inline] pub fn disconnect(&self) { self.inner.disconnect() @@ -418,13 +411,14 @@ mod tests { use std::pin::Pin; use std::sync::Arc; - use async_utility::time; use futures::StreamExt; + use tokio::time; use super::*; use crate::error::{Error, ErrorKind}; use crate::local_relay::*; use crate::policy::{AdmitPolicy, AdmitStatus}; + use crate::relay::constants::WEBSOCKET_TX_TIMEOUT; use crate::test_utils::setup_relay; #[derive(Debug)] @@ -469,7 +463,7 @@ mod tests { fn check_relay_is_sleeping(relay: &Relay) { assert_eq!(relay.status(), RelayStatus::Sleeping); assert!(relay.status().can_connect()); - assert!(!relay.inner.is_running()); + assert!(relay.inner.is_running()); } #[tokio::test] @@ -490,9 +484,19 @@ mod tests { assert_eq!(relay.status(), RelayStatus::Connected); + // Complete a protocol exchange before shutting down the server. + relay.count_events(Filter::new()).await.unwrap(); + mock.shutdown(); - time::sleep(Duration::from_millis(100)).await; + let timeout: Duration = WEBSOCKET_TX_TIMEOUT + Duration::from_secs(2); + time::timeout(timeout, async { + while relay.status() != RelayStatus::Disconnected { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); assert_eq!(relay.status(), RelayStatus::Disconnected); @@ -517,13 +521,22 @@ mod tests { assert_eq!(relay.status(), RelayStatus::Connected); - mock.shutdown(); + // Complete a protocol exchange before shutting down the server. + relay.count_events(Filter::new()).await.unwrap(); - time::sleep(Duration::from_millis(100)).await; + mock.shutdown(); - assert_eq!(relay.status(), RelayStatus::Terminated); + let timeout: Duration = WEBSOCKET_TX_TIMEOUT + Duration::from_secs(2); + time::timeout(timeout, async { + while relay.status() != RelayStatus::Idle { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); - assert!(!relay.inner.is_running()); + assert_eq!(relay.status(), RelayStatus::Idle); + assert!(relay.inner.is_running()); } #[tokio::test] @@ -548,9 +561,8 @@ mod tests { time::sleep(Duration::from_millis(100)).await; - assert_eq!(relay.status(), RelayStatus::Terminated); - - assert!(!relay.inner.is_running()); + assert_eq!(relay.status(), RelayStatus::Idle); + assert!(relay.inner.is_running()); } #[tokio::test] @@ -578,9 +590,8 @@ mod tests { time::sleep(Duration::from_millis(100)).await; - assert_eq!(relay.status(), RelayStatus::Terminated); - - assert!(!relay.inner.is_running()); + assert_eq!(relay.status(), RelayStatus::Idle); + assert!(relay.inner.is_running()); } #[tokio::test] @@ -645,9 +656,8 @@ mod tests { time::sleep(Duration::from_millis(100)).await; - assert_eq!(relay.status(), RelayStatus::Terminated); - - assert!(!relay.inner.is_running()); + assert_eq!(relay.status(), RelayStatus::Idle); + assert!(relay.inner.is_running()); } #[tokio::test] @@ -674,9 +684,8 @@ mod tests { time::sleep(Duration::from_millis(100)).await; - assert_eq!(relay.status(), RelayStatus::Terminated); - - assert!(!relay.inner.is_running()); + assert_eq!(relay.status(), RelayStatus::Idle); + assert!(relay.inner.is_running()); } #[tokio::test] @@ -705,9 +714,8 @@ mod tests { assert_eq!(err.kind(), ErrorKind::Rejected); assert_eq!(err.to_string(), "received termination request"); - assert_eq!(relay.status(), RelayStatus::Terminated); - - assert!(!relay.inner.is_running()); + assert_eq!(relay.status(), RelayStatus::Idle); + assert!(relay.inner.is_running()); } #[tokio::test] @@ -731,7 +739,13 @@ mod tests { relay.ban(); assert_eq!(relay.status(), RelayStatus::Banned); - assert!(!relay.inner.is_running()); + time::timeout(Duration::from_secs(2), async { + while relay.inner.is_running() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); // Retry to connect let res = relay.try_connect().timeout(Duration::from_secs(2)).await; @@ -897,8 +911,8 @@ mod tests { time::sleep(Duration::from_secs(2)).await; - assert_eq!(relay.status(), RelayStatus::Terminated); - assert!(!relay.inner.is_running()); + assert_eq!(relay.status(), RelayStatus::Idle); + assert!(relay.inner.is_running()); // Retry to connect let res = relay.try_connect().timeout(Duration::from_secs(2)).await; @@ -906,8 +920,8 @@ mod tests { assert_eq!(err.kind(), ErrorKind::Rejected); assert_eq!(err.to_string(), "connection rejected: reason=banned"); - assert_eq!(relay.status(), RelayStatus::Terminated); - assert!(!relay.inner.is_running()); + assert_eq!(relay.status(), RelayStatus::Idle); + assert!(relay.inner.is_running()); } #[tokio::test] diff --git a/nostr-sdk/src/relay/status.rs b/nostr-sdk/src/relay/status.rs index bf17e15bc..85ff1a173 100644 --- a/nostr-sdk/src/relay/status.rs +++ b/nostr-sdk/src/relay/status.rs @@ -1,65 +1,22 @@ use core::fmt; -use core::sync::atomic::{AtomicU8, Ordering}; - -#[derive(Debug)] -pub(super) struct AtomicRelayStatus { - value: AtomicU8, -} - -impl Default for AtomicRelayStatus { - fn default() -> Self { - Self::new(RelayStatus::Initialized) - } -} - -impl AtomicRelayStatus { - #[inline] - pub(super) fn new(status: RelayStatus) -> Self { - Self { - value: AtomicU8::new(status as u8), - } - } - - #[inline] - pub fn set(&self, status: RelayStatus) { - self.value.store(status as u8, Ordering::SeqCst); - } - - pub(super) fn load(&self) -> RelayStatus { - let val: u8 = self.value.load(Ordering::SeqCst); - match val { - 0 => RelayStatus::Initialized, - 1 => RelayStatus::Pending, - 2 => RelayStatus::Connecting, - 3 => RelayStatus::Connected, - 4 => RelayStatus::Disconnected, - 5 => RelayStatus::Terminated, - 6 => RelayStatus::Banned, - 7 => RelayStatus::Sleeping, - 8 => RelayStatus::Shutdown, - _ => unreachable!(), - } - } -} /// Relay connection status +#[non_exhaustive] #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub enum RelayStatus { /// The relay has just been created. Initialized = 0, - /// The relay will try to connect shortly. - Pending = 1, - /// Trying to connect. + /// A connection has been requested or is being established. Connecting = 2, /// Connected. Connected = 3, /// The connection failed, but another attempt will occur soon. Disconnected = 4, - /// The connection has been terminated and no retry will occur. - Terminated = 5, + /// No connection or automatic retry is scheduled. Call `connect` to reconnect. + Idle = 5, /// The relay has been banned. Banned = 6, - /// Relay is sleeping + /// The relay is sleeping and will reconnect when activity resumes. Sleeping = 7, /// The relay has been shut down and can't be used again. Shutdown = 8, @@ -69,11 +26,10 @@ impl fmt::Display for RelayStatus { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Initialized => f.write_str("Initialized"), - Self::Pending => f.write_str("Pending"), Self::Connecting => f.write_str("Connecting"), Self::Connected => f.write_str("Connected"), Self::Disconnected => f.write_str("Disconnected"), - Self::Terminated => f.write_str("Terminated"), + Self::Idle => f.write_str("Idle"), Self::Banned => f.write_str("Banned"), Self::Sleeping => f.write_str("Sleeping"), Self::Shutdown => f.write_str("Shutdown"), @@ -93,18 +49,14 @@ impl RelayStatus { matches!(self, Self::Connected) } - /// Check if is `disconnected`, `terminated`, `banned`, `sleeping` or `shutdown`. - #[inline] + /// Check if is [`RelayStatus::Disconnected`] pub(crate) fn is_disconnected(&self) -> bool { - matches!( - self, - Self::Disconnected | Self::Terminated | Self::Banned | Self::Sleeping | Self::Shutdown - ) + matches!(self, Self::Disconnected) } - /// Check if is [`RelayStatus::Terminated`] - pub(crate) fn is_terminated(&self) -> bool { - matches!(self, Self::Terminated) + /// Check if is [`RelayStatus::Idle`] + pub(crate) fn is_idle(&self) -> bool { + matches!(self, Self::Idle) } /// Check if is [`RelayStatus::Banned`] @@ -122,10 +74,25 @@ impl RelayStatus { matches!(self, Self::Shutdown) } - /// Check if relay can start a connection (status is `initialized` or `terminated`) + /// Check if relay can start a connection (initialized, idle or sleeping). #[inline] pub(crate) fn can_connect(&self) -> bool { - matches!(self, Self::Initialized | Self::Terminated | Self::Sleeping) + matches!(self, Self::Initialized | Self::Idle | Self::Sleeping) + } + + /// Check if is `disconnected`, `idle`, `banned`, `sleeping` or `shutdown`. + #[inline] + pub(crate) fn is_connection_closed(&self) -> bool { + matches!( + self, + Self::Disconnected | Self::Idle | Self::Banned | Self::Sleeping | Self::Shutdown + ) + } + + /// Check whether the relay is in a terminal state and cannot reconnect. + #[inline] + pub(crate) fn is_terminal(&self) -> bool { + matches!(self, Self::Banned | Self::Shutdown) } } @@ -133,145 +100,123 @@ impl RelayStatus { mod tests { use super::*; - #[test] - fn test_status_set() { - let relay = AtomicRelayStatus::default(); - relay.set(RelayStatus::Connected); - assert_eq!(relay.load(), RelayStatus::Connected); - } - #[test] fn test_status_initialized() { - let status = RelayStatus::Initialized; + let status: RelayStatus = RelayStatus::Initialized; assert!(status.is_initialized()); assert!(!status.is_connected()); assert!(!status.is_disconnected()); - assert!(!status.is_terminated()); + assert!(!status.is_idle()); assert!(!status.is_banned()); assert!(!status.is_sleeping()); assert!(!status.is_shutdown()); assert!(status.can_connect()); - let relay = AtomicRelayStatus::new(status); - assert_eq!(relay.load(), RelayStatus::Initialized); - } - - #[test] - fn test_status_pending() { - let status = RelayStatus::Pending; - assert!(!status.is_initialized()); - assert!(!status.is_connected()); - assert!(!status.is_disconnected()); - assert!(!status.is_terminated()); - assert!(!status.is_banned()); - assert!(!status.is_sleeping()); - assert!(!status.is_shutdown()); - assert!(!status.can_connect()); - let relay = AtomicRelayStatus::new(status); - assert_eq!(relay.load(), RelayStatus::Pending); + assert!(!status.is_connection_closed()); + assert!(!status.is_terminal()); } #[test] fn test_status_connecting() { - let status = RelayStatus::Connecting; + let status: RelayStatus = RelayStatus::Connecting; assert!(!status.is_initialized()); assert!(!status.is_connected()); assert!(!status.is_disconnected()); - assert!(!status.is_terminated()); + assert!(!status.is_idle()); assert!(!status.is_banned()); assert!(!status.is_sleeping()); assert!(!status.is_shutdown()); assert!(!status.can_connect()); - let relay = AtomicRelayStatus::new(status); - assert_eq!(relay.load(), RelayStatus::Connecting); + assert!(!status.is_connection_closed()); + assert!(!status.is_terminal()); } #[test] fn test_status_connected() { - let status = RelayStatus::Connected; + let status: RelayStatus = RelayStatus::Connected; assert!(!status.is_initialized()); assert!(status.is_connected()); assert!(!status.is_disconnected()); - assert!(!status.is_terminated()); + assert!(!status.is_idle()); assert!(!status.is_banned()); assert!(!status.is_sleeping()); assert!(!status.is_shutdown()); assert!(!status.can_connect()); - let relay = AtomicRelayStatus::new(status); - assert_eq!(relay.load(), RelayStatus::Connected); + assert!(!status.is_connection_closed()); + assert!(!status.is_terminal()); } #[test] fn test_status_disconnected() { - let status = RelayStatus::Disconnected; + let status: RelayStatus = RelayStatus::Disconnected; assert!(!status.is_initialized()); assert!(!status.is_connected()); assert!(status.is_disconnected()); - assert!(!status.is_terminated()); + assert!(!status.is_idle()); assert!(!status.is_banned()); assert!(!status.is_sleeping()); assert!(!status.is_shutdown()); assert!(!status.can_connect()); - let relay = AtomicRelayStatus::new(status); - assert_eq!(relay.load(), RelayStatus::Disconnected); + assert!(status.is_connection_closed()); + assert!(!status.is_terminal()); } #[test] - fn test_status_terminated() { - let status = RelayStatus::Terminated; + fn test_status_idle() { + let status: RelayStatus = RelayStatus::Idle; assert!(!status.is_initialized()); assert!(!status.is_connected()); - assert!(status.is_disconnected()); - assert!(status.is_terminated()); + assert!(!status.is_disconnected()); + assert!(status.is_idle()); assert!(!status.is_banned()); assert!(!status.is_sleeping()); assert!(!status.is_shutdown()); assert!(status.can_connect()); - let relay = AtomicRelayStatus::new(status); - assert_eq!(relay.load(), RelayStatus::Terminated); + assert!(status.is_connection_closed()); + assert!(!status.is_terminal()); } #[test] fn test_status_banned() { - let status = RelayStatus::Banned; + let status: RelayStatus = RelayStatus::Banned; assert!(!status.is_initialized()); assert!(!status.is_connected()); - assert!(status.is_disconnected()); - assert!(!status.is_terminated()); + assert!(!status.is_disconnected()); + assert!(!status.is_idle()); assert!(status.is_banned()); assert!(!status.is_sleeping()); assert!(!status.is_shutdown()); assert!(!status.can_connect()); - let relay = AtomicRelayStatus::new(status); - assert_eq!(relay.load(), RelayStatus::Banned); + assert!(status.is_connection_closed()); + assert!(status.is_terminal()); } #[test] fn test_status_sleeping() { - let status = RelayStatus::Sleeping; + let status: RelayStatus = RelayStatus::Sleeping; assert!(!status.is_initialized()); assert!(!status.is_connected()); - assert!(status.is_disconnected()); - assert!(!status.is_terminated()); + assert!(!status.is_disconnected()); + assert!(!status.is_idle()); assert!(!status.is_banned()); assert!(status.is_sleeping()); assert!(!status.is_shutdown()); assert!(status.can_connect()); - let relay = AtomicRelayStatus::new(status); - assert_eq!(relay.load(), RelayStatus::Sleeping); + assert!(status.is_connection_closed()); + assert!(!status.is_terminal()); } #[test] fn test_status_shutdown() { - let status = RelayStatus::Shutdown; + let status: RelayStatus = RelayStatus::Shutdown; assert!(!status.is_initialized()); assert!(!status.is_connected()); - assert!(status.is_disconnected()); - assert!(!status.is_terminated()); + assert!(!status.is_disconnected()); + assert!(!status.is_idle()); assert!(!status.is_banned()); assert!(!status.is_sleeping()); assert!(status.is_shutdown()); assert!(!status.can_connect()); - let relay = AtomicRelayStatus::new(status); - assert_eq!(relay.load(), RelayStatus::Shutdown); + assert!(status.is_connection_closed()); + assert!(status.is_terminal()); } } diff --git a/nostr-sdk/src/test_utils.rs b/nostr-sdk/src/test_utils.rs index 7caca36da..1f7053f0c 100644 --- a/nostr-sdk/src/test_utils.rs +++ b/nostr-sdk/src/test_utils.rs @@ -1,11 +1,267 @@ +use std::net::SocketAddr; +use std::pin::Pin; +use std::task::{Context, Poll}; use std::time::Duration; -use nostr::types::RelayUrl; +use async_wsocket::Message; +use futures::{Sink, SinkExt}; +use nostr::types::{RelayUrl, Url}; +use tokio::sync::mpsc::{Receiver, Sender}; +use tokio::sync::{mpsc, oneshot}; +use tokio::task::JoinHandle; +use tokio::time; use crate::authenticator::Authenticator; use crate::client::Client; +use crate::error::Error; +use crate::future::BoxedFuture; use crate::local_relay::*; -use crate::relay::Relay; +use crate::policy::{AdmitPolicy, AdmitStatus}; +use crate::relay::{Relay, RelayBuilder, RelayOptions, RelayStatus}; +use crate::transport::websocket::{WebSocketSink, WebSocketStream, WebSocketTransport}; + +pub(crate) type DialReply = oneshot::Sender>; +pub(crate) type PolicyReply = oneshot::Sender>; + +pub(crate) const TEST_TIMEOUT: Duration = Duration::from_secs(2); + +pub(crate) fn start_try_connect(relay: &Relay, timeout: Duration) -> JoinHandle> { + let relay = relay.clone(); + tokio::spawn(async move { relay.try_connect().timeout(timeout).await }) +} + +pub(crate) async fn wait_until(description: &str, mut condition: F) +where + F: FnMut() -> bool, +{ + time::timeout(TEST_TIMEOUT, async { + while !condition() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {description}")); +} + +pub(crate) async fn wait_for_status(relay: &Relay, status: RelayStatus) { + wait_until(&format!("relay status {status}"), || { + relay.status() == status + }) + .await; +} + +#[derive(Debug)] +pub(crate) struct ControlledConnectionPolicy(Sender); + +impl ControlledConnectionPolicy { + pub(crate) fn new() -> (Self, PolicyController) { + let (sender, receiver) = mpsc::channel(8); + (Self(sender), PolicyController { receiver }) + } +} + +impl AdmitPolicy for ControlledConnectionPolicy { + fn admit_connection<'a>( + &'a self, + _url: &'a RelayUrl, + ) -> BoxedFuture<'a, Result> { + Box::pin(async move { + let (reply, result) = oneshot::channel(); + self.0.send(reply).await.expect("policy controller dropped"); + result.await.expect("policy reply dropped without a result") + }) + } +} + +pub(crate) struct PolicyController { + receiver: Receiver, +} + +impl PolicyController { + pub(crate) async fn next_check(&mut self) -> PolicyReply { + time::timeout(TEST_TIMEOUT, self.receiver.recv()) + .await + .expect("timed out waiting for a connection policy check") + .expect("connection policy dropped before checking") + } + + pub(crate) fn assert_no_pending_checks(&mut self) { + assert!( + matches!( + self.receiver.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + ), + "expected no pending connection policy checks" + ); + } +} + +#[derive(Debug)] +struct ControlledWebSocketTransport(Sender); + +impl WebSocketTransport for ControlledWebSocketTransport { + fn support_ping(&self) -> bool { + false + } + + fn connect<'a>( + &'a self, + _url: &'a Url, + _proxy: Option, + ) -> BoxedFuture<'a, Result<(WebSocketSink, WebSocketStream), Error>> { + Box::pin(async move { + let (reply, result) = oneshot::channel(); + self.0.send(reply).await.expect("dial controller dropped"); + result.await.expect("dial reply dropped without a result") + }) + } +} + +pub(crate) struct DialController { + receiver: Receiver, +} + +impl DialController { + pub(crate) async fn next_dial(&mut self) -> DialReply { + time::timeout(TEST_TIMEOUT, self.receiver.recv()) + .await + .expect("timed out waiting for a transport dial") + .expect("transport dropped before dialing") + } + + pub(crate) fn assert_no_pending_dials(&mut self) { + assert!( + matches!( + self.receiver.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + ), + "expected no pending transport dials" + ); + } +} + +/// Keeps the relay separate from the controller that completes its transport attempts. +pub(crate) struct ControlledRelay { + pub(crate) relay: Relay, + pub(crate) transport: DialController, +} + +impl ControlledRelay { + pub(crate) fn new() -> Self { + Self::with_opts(RelayOptions::default()) + } + + pub(crate) fn with_opts(opts: RelayOptions) -> Self { + let builder = + Relay::builder(RelayUrl::parse("wss://relay.example.com").unwrap()).opts(opts); + Self::with_builder(builder) + } + + pub(crate) fn with_policy(policy: ControlledConnectionPolicy) -> Self { + let builder = Relay::builder(RelayUrl::parse("wss://relay.example.com").unwrap()) + .admit_policy(policy); + Self::with_builder(builder) + } + + fn with_builder(builder: RelayBuilder) -> Self { + let (requests, receiver) = mpsc::channel(8); + let relay = builder + .websocket_transport(ControlledWebSocketTransport(requests)) + .build(); + + Self { + relay, + transport: DialController { receiver }, + } + } +} + +/// Pauses WebSocket closing until the test releases it. +pub(crate) struct CloseGate { + closing: oneshot::Receiver<()>, + release: oneshot::Sender<()>, +} + +impl CloseGate { + pub(crate) async fn wait_until_closing(&mut self) { + time::timeout(TEST_TIMEOUT, &mut self.closing) + .await + .expect("timed out waiting for WebSocket closing to start") + .expect("WebSocket dropped before closing"); + } + + pub(crate) fn release(self) { + self.release + .send(()) + .expect("WebSocket close was cancelled"); + } + + pub(crate) fn is_cancelled(&self) -> bool { + self.release.is_closed() + } +} + +struct ClosingSink { + inner: WebSocketSink, + entered: Option>, + release: oneshot::Receiver<()>, +} + +impl Sink for ClosingSink { + type Error = Error; + + fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.as_mut().poll_ready(cx) + } + + fn start_send(mut self: Pin<&mut Self>, message: Message) -> Result<(), Error> { + self.inner.as_mut().start_send(message) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.inner.as_mut().poll_flush(cx) + } + + fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if let Some(entered) = self.entered.take() { + let _ = entered.send(()); + } + + if Pin::new(&mut self.release).poll(cx).is_pending() { + return Poll::Pending; + } + + self.inner.as_mut().poll_close(cx) + } +} + +/// Accepts outgoing messages without producing any incoming messages. +pub(crate) fn fake_connection() -> (WebSocketSink, WebSocketStream) { + ( + Box::pin(futures::sink::drain().sink_map_err(|error| match error {})), + Box::pin(futures::stream::pending()), + ) +} + +pub(crate) fn connection_gate_close( + connection: (WebSocketSink, WebSocketStream), +) -> ((WebSocketSink, WebSocketStream), CloseGate) { + let (entered, closing) = oneshot::channel(); + let (release, released) = oneshot::channel(); + let (inner, stream) = connection; + + ( + ( + Box::pin(ClosingSink { + inner, + entered: Some(entered), + release: released, + }), + stream, + ), + CloseGate { closing, release }, + ) +} pub(crate) async fn setup_nip42_read_local_relay() -> LocalRelay { let local = LocalRelay::builder() From c28142e162662f8242e184797ed20d4f3b23adf8 Mon Sep 17 00:00:00 2001 From: Jeff Gardner <202880+erskingardner@users.noreply.github.com> Date: Wed, 7 Oct 2026 09:41:37 +0000 Subject: [PATCH 2/2] sdk: restore subscriptions after same-second reconnections Record the connection count for each long-lived subscription so a new session restores it even when it starts in the same second. Ref https://github.com/nostrdevkit/nostr/pull/1479 Pull-Request: https://github.com/nostrdevkit/nostr/pull/1491 Signed-off-by: Yuki Kishimoto --- nostr-sdk/CHANGELOG.md | 1 + nostr-sdk/src/relay/inner.rs | 108 +++++++++++++++++++++++++++-------- nostr-sdk/src/test_utils.rs | 51 ++++++++++++++++- 3 files changed, 134 insertions(+), 26 deletions(-) diff --git a/nostr-sdk/CHANGELOG.md b/nostr-sdk/CHANGELOG.md index c7cbae359..c0149c17d 100644 --- a/nostr-sdk/CHANGELOG.md +++ b/nostr-sdk/CHANGELOG.md @@ -57,6 +57,7 @@ - Fail NIP-77 reconciliation when a download batch is rejected or closes before its requested events arrive. - Retain reconnect requests while a previous WebSocket closes (https://github.com/nostrdevkit/nostr/pull/1491). - Allow banning and shutting down idle relays (https://github.com/nostrdevkit/nostr/pull/1491). +- Restore subscriptions after reconnections within the same second (https://github.com/nostrdevkit/nostr/pull/1491). ### Deprecated diff --git a/nostr-sdk/src/relay/inner.rs b/nostr-sdk/src/relay/inner.rs index 85b5fb59d..e0a265e24 100644 --- a/nostr-sdk/src/relay/inner.rs +++ b/nostr-sdk/src/relay/inner.rs @@ -189,7 +189,7 @@ impl RelayChannels { #[derive(Debug)] struct SubscriptionData { pub filters: Vec, - pub subscribed_at: Timestamp, + pub subscribed_connection: usize, pub is_auto_closing: bool, /// Received EOSE msg pub received_eose: bool, @@ -201,10 +201,10 @@ struct SubscriptionData { impl SubscriptionData { #[inline] - fn long_lived(filters: Vec) -> Self { + fn long_lived(filters: Vec, subscribed_connection: usize) -> Self { Self { filters, - subscribed_at: Timestamp::now(), + subscribed_connection, is_auto_closing: false, received_eose: false, received_events: AtomicUsize::new(0), @@ -216,7 +216,7 @@ impl SubscriptionData { fn auto_closing(filters: Vec) -> Self { Self { filters, - subscribed_at: Timestamp::zero(), + subscribed_connection: 0, is_auto_closing: true, received_eose: false, received_events: AtomicUsize::new(0), @@ -430,7 +430,10 @@ impl InnerRelay { return Err(Error::invalid_msg("subscription ID already exists")); } - subscriptions.insert(id, SubscriptionData::long_lived(filters)); + subscriptions.insert( + id, + SubscriptionData::long_lived(filters, self.stats.success()), + ); Ok(()) } @@ -456,7 +459,7 @@ impl InnerRelay { &self, id: &SubscriptionId, filters: Vec, - update_subscribed_at: bool, + mark_subscribed: bool, ) -> Result<(), Error> { let mut subscriptions = self.atomic.subscriptions.write().await; let data = subscriptions @@ -464,8 +467,8 @@ impl InnerRelay { .ok_or_else(|| Error::not_found("subscription not found"))?; data.filters = filters; - if update_subscribed_at { - data.subscribed_at = Timestamp::now(); + if mark_subscribed { + data.subscribed_connection = self.stats.success(); data.closed = false; } @@ -508,20 +511,20 @@ impl InnerRelay { let subscriptions = self.atomic.subscriptions.read().await; match subscriptions.get(id) { Some(SubscriptionData { - subscribed_at, + subscribed_connection, closed, is_auto_closing: false, .. }) => { - // Never subscribed -> SHOULD subscribe - // Subscription closed by relay -> SHOULD subscribe - if subscribed_at.is_zero() || *closed { + if *closed { return true; } - // First connection and subscribed_at != 0 -> SHOULD NOT re-subscribe - // Many connections and subscription NOT done in current websocket session -> SHOULD re-subscribe - self.stats.connected_at() > *subscribed_at && self.stats.success() > 1 + // The successful connection count identifies the WebSocket session. + // Requests queued before the first connection are already in the outbound + // channel. Later sessions must restore subscriptions even within one second. + let current_connection: usize = self.stats.success(); + current_connection > 1 && current_connection > *subscribed_connection } // NOT subscribe if auto-closing subscription or subscription not found Some(SubscriptionData { @@ -2076,12 +2079,16 @@ mod tests { use crate::authenticator::SignerAuthenticator; use crate::error::{Error, ErrorKind}; use crate::future::BoxedFuture; + use crate::local_relay::MockRelay; use crate::policy::{AdmitPolicy, AdmitStatus}; use crate::relay::{Relay, RelayOptions}; + use crate::stream::NotificationStream; use crate::test_utils::{ ControlledConnectionPolicy, ControlledRelay, DialReply, PolicyReply, TEST_TIMEOUT, - connection_gate_close, fake_connection, start_try_connect, wait_for_status, wait_until, + connection_gate_close, fake_connection, start_try_connect, wait_for_eose, wait_for_status, + wait_for_subscription_event, wait_until, }; + use crate::transport::websocket::{DefaultWebsocketTransport, WebSocketTransport}; #[derive(Debug)] struct CountingAdmitPolicy(Arc); @@ -2560,19 +2567,34 @@ mod tests { } async fn reconnect_while_closing(opts: RelayOptions, request: ReconnectRequest) { + let local: MockRelay = MockRelay::run().await.unwrap(); + let url: RelayUrl = local.url().await; + + let websocket: DefaultWebsocketTransport = DefaultWebsocketTransport; + let ControlledRelay { relay, mut transport, } = ControlledRelay::with_opts(opts); let initial: TokioJoinHandle> = start_try_connect(&relay, TEST_TIMEOUT); + let reply: DialReply = transport.next_dial().await; - let (connection, mut close_gate) = connection_gate_close(fake_connection()); + let connection = websocket.connect((&url).into(), None).await.unwrap(); + let (connection, mut close_gate) = connection_gate_close(connection); - assert!(transport.next_dial().await.send(Ok(connection)).is_ok()); + assert!(reply.send(Ok(connection)).is_ok()); initial.await.unwrap().unwrap(); + let mut notifications: NotificationStream = relay.notifications(); + let keys: Keys = Keys::generate(); + let subscription: SubscriptionId = relay + .subscribe(Filter::new().author(keys.public_key())) + .await + .unwrap(); + wait_for_eose(&mut notifications, &subscription).await; + relay.disconnect(); close_gate.wait_until_closing().await; @@ -2600,18 +2622,23 @@ mod tests { close_gate.release(); + let reply: DialReply = transport.next_dial().await; assert!( - transport - .next_dial() - .await - .send(Ok(fake_connection())) + reply + .send(websocket.connect((&url).into(), None).await) .is_ok() ); - if let Some(pending) = pending { pending.await.unwrap(); } - wait_for_status(&relay, RelayStatus::Connected).await; + wait_for_eose(&mut notifications, &subscription).await; + + let event: Event = EventBuilder::new(Kind::TextNote, "after reconnect") + .finalize(&keys) + .unwrap(); + let event_id = event.id; + local.add_event(event).await.unwrap(); + wait_for_subscription_event(&mut notifications, &subscription, &event_id).await; assert_eq!(relay.stats().attempts(), 2); assert_eq!(relay.stats().success(), 2); @@ -3207,6 +3234,39 @@ mod tests { assert_eq!(relay.stats().attempts(), 1); assert_eq!(relay.stats().success(), 1); } + + #[tokio::test] + async fn subscription_restoration_uses_connection_counter() { + let relay: Relay = Relay::new(RelayUrl::parse("wss://relay.example.com").unwrap()); + + let id: SubscriptionId = SubscriptionId::new("generation"); + let filters: Vec = vec![Filter::new().kind(Kind::TextNote)]; + relay + .inner + .add_long_lived_subscription(id.clone(), filters.clone()) + .await + .unwrap(); + + relay.inner.stats.new_success(); + + assert!(!relay.inner.should_resubscribe(&id).await); + + relay.inner.stats.new_success(); + + assert!(relay.inner.should_resubscribe(&id).await); + + relay + .inner + .update_subscription(&id, filters, true) + .await + .unwrap(); + + assert!(!relay.inner.should_resubscribe(&id).await); + + relay.inner.stats.new_success(); + + assert!(relay.inner.should_resubscribe(&id).await); + } } #[cfg(bench)] diff --git a/nostr-sdk/src/test_utils.rs b/nostr-sdk/src/test_utils.rs index 1f7053f0c..5ab96a711 100644 --- a/nostr-sdk/src/test_utils.rs +++ b/nostr-sdk/src/test_utils.rs @@ -4,7 +4,9 @@ use std::task::{Context, Poll}; use std::time::Duration; use async_wsocket::Message; -use futures::{Sink, SinkExt}; +use futures::{Sink, SinkExt, StreamExt}; +use nostr::event::EventId; +use nostr::message::{RelayMessage, SubscriptionId}; use nostr::types::{RelayUrl, Url}; use tokio::sync::mpsc::{Receiver, Sender}; use tokio::sync::{mpsc, oneshot}; @@ -17,7 +19,8 @@ use crate::error::Error; use crate::future::BoxedFuture; use crate::local_relay::*; use crate::policy::{AdmitPolicy, AdmitStatus}; -use crate::relay::{Relay, RelayBuilder, RelayOptions, RelayStatus}; +use crate::relay::{Relay, RelayBuilder, RelayNotification, RelayOptions, RelayStatus}; +use crate::stream::NotificationStream; use crate::transport::websocket::{WebSocketSink, WebSocketStream, WebSocketTransport}; pub(crate) type DialReply = oneshot::Sender>; @@ -50,6 +53,50 @@ pub(crate) async fn wait_for_status(relay: &Relay, status: RelayStatus) { .await; } +async fn wait_for_notification( + notifications: &mut NotificationStream, + description: &str, + mut matches: F, +) where + F: FnMut(&RelayNotification) -> bool, +{ + time::timeout(TEST_TIMEOUT, async { + while let Some(notification) = notifications.next().await { + if matches(¬ification) { + return; + } + } + + panic!("notification stream closed while waiting for {description}"); + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {description}")); +} + +pub(crate) async fn wait_for_eose( + notifications: &mut NotificationStream, + subscription: &SubscriptionId, +) { + wait_for_notification(notifications, "subscription EOSE", |notification| { + matches!(notification, RelayNotification::Message { message } + if matches!(message.as_ref(), RelayMessage::EndOfStoredEvents(id) + if id.as_ref() == subscription)) + }) + .await; +} + +pub(crate) async fn wait_for_subscription_event( + notifications: &mut NotificationStream, + subscription: &SubscriptionId, + event_id: &EventId, +) { + wait_for_notification(notifications, "subscription event", |notification| { + matches!(notification, RelayNotification::Event { subscription_id, event } + if subscription_id == subscription && event.id == *event_id) + }) + .await; +} + #[derive(Debug)] pub(crate) struct ControlledConnectionPolicy(Sender);