relay-pool: register the auto-closing subscription before sending the REQ msg

Fixes #985

Tested-by: reya <[email protected]>
Signed-off-by: Yuki Kishimoto <[email protected]>

Author: yukibtc

crates/nostr-relay-pool/src/relay/inner.rs

@@ -295,7 +295,7 @@ impl InnerRelay {
         subscription.get(id).map(|d| d.filter.clone())
     }
 
-    async fn remove_subscription(&self, id: &SubscriptionId) {
+    pub(super) async fn remove_subscription(&self, id: &SubscriptionId) {
         let mut subscriptions = self.atomic.subscriptions.write().await;
         subscriptions.remove(id);
     }
@@ -1378,10 +1378,6 @@ impl InnerRelay {
     ) {
         let relay = self.clone(); // <-- FULL RELAY CLONE HERE
         task::spawn(async move {
-            relay
-                .add_auto_closing_subscription(id.clone(), filter.clone())
-                .await;
-
             // Check if CLOSE needed
             let to_close: bool = match relay
                 .handle_auto_closing(&id, &filter, opts, notifications, &activity)
@@ -1823,15 +1819,24 @@ impl InnerRelay {
         }
 
         let filter = Filter::new().ids(ids);
-        self.send_msg(ClientMessage::Req {
+        let msg: ClientMessage = ClientMessage::Req {
             subscription_id: Cow::Borrowed(down_sub_id),
             filter: Cow::Borrowed(&filter),
-        })?;
+        };
 
         // Register an auto-closing subscription
-        self.add_auto_closing_subscription(down_sub_id.clone(), filter)
+        self.add_auto_closing_subscription(down_sub_id.clone(), filter.clone())
             .await;
 
+        // Send msg
+        if let Err(e) = self.send_msg(msg) {
+            // Remove previously added subscription
+            self.remove_subscription(down_sub_id).await;
+
+            // Propagate error
+            return Err(e);
+        }
+
         *in_flight_down = true;
 
         Ok(())

crates/nostr-relay-pool/src/relay/mod.rs

@@ -490,14 +490,14 @@ impl Relay {
         filter: Filter,
         opts: SubscribeOptions,
     ) -> Result<(), Error> {
-        // Check if auto-close condition is set
+        // Check if the auto-close condition is set
         match opts.auto_close {
-            Some(opts) => self.subscribe_auto_closing(id, filter, opts, None),
+            Some(opts) => self.subscribe_auto_closing(id, filter, opts, None).await,
             None => self.subscribe_long_lived(id, filter).await,
         }
     }
 
-    fn subscribe_auto_closing(
+    async fn subscribe_auto_closing(
         &self,
         id: SubscriptionId,
         filter: Filter,
@@ -513,8 +513,19 @@ impl Relay {
         // Subscribe to notifications
         let notifications = self.inner.internal_notification_sender.subscribe();
 
+        // Register the auto-closing subscription
+        self.inner
+            .add_auto_closing_subscription(id.clone(), filter.clone())
+            .await;
+
         // Send REQ message
-        self.inner.send_msg(msg)?;
+        if let Err(e) = self.inner.send_msg(msg) {
+            // Remove previously added subscription
+            self.inner.remove_subscription(&id).await;
+
+            // Propagate error
+            return Err(e);
+        }
 
         // Spawn auto-closing handler
         self.inner
@@ -571,7 +582,8 @@ impl Relay {
 
         // Subscribe
         let id: SubscriptionId = SubscriptionId::generate();
-        self.subscribe_auto_closing(id, filter, opts, Some(tx))?;
+        self.subscribe_auto_closing(id, filter, opts, Some(tx))
+            .await?;
 
         // Handle subscription activity
         while let Some(activity) = rx.recv().await {