pool: add ReqExitPolicy::WaitForEvents variant

Allow exiting from a REQ after the subscription received N events.

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

Author: yukibtc

CHANGELOG.md

@@ -96,6 +96,7 @@
 * database: add `Events::force_insert` ([Yuki Kishimoto])
 * pool: event verification cache ([Yuki Kishimoto])
 * pool: add `AdmitPolicy` trait ([Yuki Kishimoto])
+* pool: add `ReqExitPolicy::WaitForEvents` variant ([Yuki Kishimoto])
 * ffi: add Mac Catalyst support in Swift package ([Yuki Kishimoto])
 * js: add `KindStandard` enum ([Yuki Kishimoto])
 

bindings/nostr-sdk-ffi/src/relay/options.rs

@@ -175,6 +175,8 @@ impl RelayOptions {
 pub enum ReqExitPolicy {
     /// Exit on EOSE
     ExitOnEOSE,
+    /// Wait to receive N events and then exit.
+    WaitForEvents { num: u16 },
     /// After EOSE is received, keep listening for N more events that match the filter.
     WaitForEventsAfterEOSE { num: u16 },
     /// After EOSE is received, keep listening for matching events for `Duration` more time.
@@ -185,6 +187,7 @@ impl From<ReqExitPolicy> for prelude::ReqExitPolicy {
     fn from(value: ReqExitPolicy) -> Self {
         match value {
             ReqExitPolicy::ExitOnEOSE => Self::ExitOnEOSE,
+            ReqExitPolicy::WaitForEvents { num } => Self::WaitForEvents(num),
             ReqExitPolicy::WaitForEventsAfterEOSE { num } => Self::WaitForEventsAfterEOSE(num),
             ReqExitPolicy::WaitDurationAfterEOSE { duration } => {
                 Self::WaitDurationAfterEOSE(duration)

bindings/nostr-sdk-js/src/relay/options.rs

@@ -108,6 +108,14 @@ impl JsReqExitPolicy {
         }
     }
 
+    /// Wait to receive N events and then exit.
+    #[wasm_bindgen(js_name = waitForEvents)]
+    pub fn wait_for_events(num: u16) -> Self {
+        Self {
+            inner: ReqExitPolicy::WaitForEvents(num),
+        }
+    }
+
     /// After EOSE is received, keep listening for N more events that match the filter
     #[wasm_bindgen(js_name = waitForEventsAfterEOSE)]
     pub fn wait_for_events_after_eose(num: u16) -> Self {

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

@@ -1253,7 +1253,8 @@ impl InnerRelay {
         activity: &Option<Sender<SubscriptionActivity>>,
     ) -> Option<HandleAutoClosing> {
         time::timeout(opts.timeout, async move {
-            let mut counter: u16 = 0;
+            let mut wait_for_events_counter: u16 = 0;
+            let mut wait_for_events_after_eose_counter: u16 = 0;
             let mut received_eose: bool = false;
             let mut require_resubscription: bool = false;
             let mut last_event: Option<Instant> = None;
@@ -1296,14 +1297,23 @@ impl InnerRelay {
                                     last_event = Some(Instant::now());
                                 }
 
-                                if let ReqExitPolicy::WaitForEventsAfterEOSE(num) = opts.exit_policy
-                                {
-                                    if received_eose {
-                                        counter += 1;
-                                        if counter >= num {
+                                // Check exit policy
+                                match opts.exit_policy {
+                                    ReqExitPolicy::WaitForEvents(num) => {
+                                        wait_for_events_counter += 1;
+                                        if wait_for_events_counter >= num {
                                             break;
                                         }
                                     }
+                                    ReqExitPolicy::WaitForEventsAfterEOSE(num) => {
+                                        if received_eose {
+                                            wait_for_events_after_eose_counter += 1;
+                                            if wait_for_events_after_eose_counter >= num {
+                                                break;
+                                            }
+                                        }
+                                    }
+                                    _ => {}
                                 }
                             }
                         }

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

@@ -1210,6 +1210,46 @@ mod tests {
         assert_eq!(events.len(), 3);
     }
 
+    #[tokio::test]
+    async fn test_fetch_events_wait_for_events() {
+        let (relay, _mock) = setup_event_fetching_relay(5).await;
+
+        let events = relay
+            .fetch_events(
+                Filter::new().kind(Kind::TextNote),
+                Duration::from_secs(15),
+                ReqExitPolicy::WaitForEvents(2),
+            )
+            .await
+            .unwrap();
+        assert_eq!(events.len(), 2); // Requested all text notes but exit after receive 2
+
+        // Task to send additional event
+        let r = relay.clone();
+        tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_secs(2)).await;
+
+            // Signer
+            let keys = Keys::generate();
+
+            // Build and send event
+            let event = EventBuilder::metadata(&Metadata::new().name("Test"))
+                .sign_with_keys(&keys)
+                .unwrap();
+            r.send_event(&event).await.unwrap();
+        });
+
+        let events = relay
+            .fetch_events(
+                Filter::new().kind(Kind::Metadata),
+                Duration::from_secs(5),
+                ReqExitPolicy::WaitForEvents(1),
+            )
+            .await
+            .unwrap();
+        assert_eq!(events.len(), 1);
+    }
+
     #[tokio::test]
     async fn test_fetch_events_wait_for_events_after_eose() {
         let (relay, _mock) = setup_event_fetching_relay(10).await;

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

@@ -175,6 +175,8 @@ pub enum ReqExitPolicy {
     /// Exit on EOSE.
     #[default]
     ExitOnEOSE,
+    /// Wait to receive N events and then exit.
+    WaitForEvents(u16),
     /// After EOSE is received, keep listening for N more events that match the filter.
     WaitForEventsAfterEOSE(u16),
     /// After EOSE is received, keep listening for matching events for [`Duration`] more time.