sdk: retain reconnect requests across connection task teardown

nostr-sdk/CHANGELOG.md

@@ -37,6 +37,10 @@
 
 ### Fixed
 
+- Preserve reconnect requests made while a terminated connection task is still
+  releasing ownership, instead of leaving the relay pending with no task.
+- Restore long-lived subscriptions on immediate reconnect even when both
+  connections occur within one second.
 - Require a correlated COUNT response before `Relay::count_events` returns a count.
 - Preserve the broadcast lag or closure cause in event `OK` and authentication waiters instead of returning a generic premature-exit error.
 

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,84 @@ 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::test_utils::{setup_client, setup_client_with_authenticator};
+
+    #[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 = setup_client_with_authenticator(
+            url.clone(),
+            SignerAuthenticator::new(first_keys.clone()),
+        )
+        .await;
+        let second = setup_client_with_authenticator(
+            url.clone(),
+            SignerAuthenticator::new(second_keys.clone()),
+        )
+        .await;
+        let unauthenticated = 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(),
+            crate::relay::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());
+    }
 
     #[tokio::test]
     async fn unverified_event_is_rejected_before_sending() {

nostr-sdk/src/client/mod.rs

@@ -421,7 +421,12 @@ impl Client {
         let _ = self.remove_all_relays().force().await;
     }
 
-    /// Connect to a previously added relay
+    /// Connect to a previously added relay.
+    ///
+    /// This queues a connection request and returns before the relay is
+    /// connected. After [`Client::disconnect_relay`], it also works while the
+    /// previous connection task is still closing. Observe a fresh connected
+    /// status or live subscription traffic before treating the relay as ready.
     #[inline]
     pub async fn connect_relay<'a, U>(&self, url: U) -> Result<(), Error>
     where
@@ -432,7 +437,13 @@ impl Client {
         self.pool().connect_relay(&url).await
     }
 
-    /// Try to connect to a previously added relay
+    /// Try to connect to a previously added relay.
+    ///
+    /// This is a one-shot connection attempt, not a task-exit barrier. If the
+    /// previous task still owns the relay after [`Client::disconnect_relay`],
+    /// it returns a state error and queues a fresh connection. Use
+    /// [`Client::connect_relay`] for immediate reactivation without retiring
+    /// and replacing the relay object.
     #[inline]
     pub async fn try_connect_relay<'a, U>(&self, url: U, timeout: Duration) -> Result<(), Error>
     where
@@ -443,7 +454,10 @@ impl Client {
         self.pool().try_connect_relay(&url, timeout).await
     }
 
-    /// Disconnect relay
+    /// Disconnect relay.
+    ///
+    /// This requests termination and publishes `Terminated` status, but does
+    /// not wait for the connection task or WebSocket close to finish.
     #[inline]
     pub async fn disconnect_relay<'a, U>(&self, url: U) -> Result<(), Error>
     where

nostr-sdk/src/relay/api/try_connect.rs

@@ -14,6 +14,19 @@ pub struct TryConnect<'relay> {
     timeout: Duration,
 }
 
+struct ReservedConnection<'relay> {
+    relay: &'relay Relay,
+    transferred: bool,
+}
+
+impl Drop for ReservedConnection<'_> {
+    fn drop(&mut self) {
+        if !self.transferred {
+            self.relay.inner.finish_reserved_connection_task();
+        }
+    }
+}
+
 impl<'relay> TryConnect<'relay> {
     #[inline]
     pub(crate) fn new(relay: &'relay Relay) -> Self {
@@ -63,6 +76,16 @@ impl<'relay> IntoFuture for TryConnect<'relay> {
                 return Err(Error::connection_rejected(reason));
             }
 
+            // Reserve the relay before dialing. Otherwise a retiring task can
+            // still own it after the handshake publishes Connected.
+            if !self.relay.inner.reserve_try_connect()? {
+                return Ok(());
+            }
+            let mut reservation = ReservedConnection {
+                relay: self.relay,
+                transferred: false,
+            };
+
             // Try to connect
             // This will set the status to "terminated" if the connection fails
             let stream: (WebSocketSink, WebSocketStream) = self
@@ -71,8 +94,8 @@ impl<'relay> IntoFuture for TryConnect<'relay> {
                 ._try_connect(self.timeout, RelayStatus::Terminated)
                 .await?;
 
-            // Spawn connection task
-            self.relay.inner.spawn_connection_task(Some(stream));
+            self.relay.inner.spawn_reserved_connection_task(stream);
+            reservation.transferred = true;
 
             Ok(())
         })
@@ -81,12 +104,142 @@ impl<'relay> IntoFuture for TryConnect<'relay> {
 
 #[cfg(test)]
 mod tests {
+    use std::net::SocketAddr;
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicBool, Ordering};
+
     use async_utility::time;
-    use nostr::types::RelayUrl;
+    use nostr::types::{RelayUrl, Url};
+    use tokio::sync::Notify;
 
     use super::*;
     use crate::error::ErrorKind;
     use crate::local_relay::*;
+    use crate::policy::AdmitPolicy;
+    use crate::transport::websocket::WebSocketTransport;
+
+    #[derive(Debug, Clone, Default)]
+    struct FirstConnectionGate {
+        blocked: Arc<AtomicBool>,
+        entered: Arc<Notify>,
+        release: Arc<Notify>,
+    }
+
+    impl AdmitPolicy for FirstConnectionGate {
+        fn admit_connection<'a>(
+            &'a self,
+            _url: &'a RelayUrl,
+        ) -> BoxedFuture<'a, Result<AdmitStatus, Error>> {
+            Box::pin(async move {
+                if !self.blocked.swap(true, Ordering::SeqCst) {
+                    self.entered.notify_one();
+                    self.release.notified().await;
+                }
+                Ok(AdmitStatus::Success)
+            })
+        }
+    }
+
+    #[tokio::test]
+    async fn try_connect_rechecks_state_after_connection_policy() {
+        for status in [
+            RelayStatus::Shutdown,
+            RelayStatus::Banned,
+            RelayStatus::Connected,
+        ] {
+            let mock = MockRelay::run().await.unwrap();
+            let gate = FirstConnectionGate::default();
+            let relay = Relay::builder(mock.url().await)
+                .admit_policy(gate.clone())
+                .build();
+            let mut attempt = relay.try_connect().into_future();
+            tokio::time::timeout(Duration::from_secs(2), async {
+                tokio::select! {
+                    _ = gate.entered.notified() => (),
+                    result = &mut attempt => panic!("attempt ended before policy release: {result:?}"),
+                }
+            }).await.unwrap();
+
+            match status {
+                RelayStatus::Shutdown => relay.shutdown(),
+                RelayStatus::Banned => relay.ban(),
+                RelayStatus::Connected => relay
+                    .try_connect()
+                    .timeout(Duration::from_secs(2))
+                    .await
+                    .unwrap(),
+                _ => unreachable!(),
+            }
+            gate.release.notify_one();
+            let result = attempt.await;
+            assert_eq!(relay.status(), status);
+            if status.is_connected() {
+                result.unwrap();
+                assert_eq!(relay.stats().success(), 1);
+                assert!(relay.inner.is_running());
+            } else {
+                assert!(result.is_err());
+                assert_eq!(relay.stats().attempts(), 0);
+                assert!(!relay.inner.is_running());
+            }
+            relay.shutdown();
+            mock.shutdown();
+        }
+    }
+
+    #[derive(Debug, Clone, Default)]
+    struct PendingTransport {
+        entered: Arc<Notify>,
+    }
+
+    impl WebSocketTransport for PendingTransport {
+        fn support_ping(&self) -> bool {
+            false
+        }
+
+        fn connect<'a>(
+            &'a self,
+            _url: &'a Url,
+            _proxy: Option<SocketAddr>,
+        ) -> BoxedFuture<'a, Result<(WebSocketSink, WebSocketStream), Error>> {
+            Box::pin(async move {
+                self.entered.notify_one();
+                std::future::pending().await
+            })
+        }
+    }
+
+    #[tokio::test]
+    async fn cancelled_and_timed_out_dials_release_connection_ownership() {
+        let transport = PendingTransport::default();
+        let relay = Relay::builder("wss://relay.example.com".parse().unwrap())
+            .websocket_transport(transport.clone())
+            .build();
+        let mut attempt = relay.try_connect().into_future();
+        tokio::time::timeout(Duration::from_secs(2), async {
+            tokio::select! {
+                _ = transport.entered.notified() => (),
+                result = &mut attempt => panic!("dial ended before cancellation: {result:?}"),
+            }
+        })
+        .await
+        .unwrap();
+        assert_eq!(relay.status(), RelayStatus::Connecting);
+        assert!(relay.inner.is_running());
+        drop(attempt);
+        assert_eq!(relay.status(), RelayStatus::Terminated);
+        assert!(!relay.inner.is_running());
+
+        let error = relay
+            .try_connect()
+            .timeout(Duration::from_millis(20))
+            .await
+            .unwrap_err();
+        assert_eq!(error.kind(), ErrorKind::Timeout);
+        assert_eq!(relay.stats().attempts(), 2);
+        assert_eq!(relay.status(), RelayStatus::Terminated);
+        assert!(!relay.inner.is_running());
+    }
 
     #[tokio::test]
     async fn test_try_connect() {